From ca02d9b60976b5a541e381399c00bb9f59e2866c Mon Sep 17 00:00:00 2001
From: Elara6331 <elara@elara.ws>
Date: Sun, 4 Aug 2024 19:49:43 -0700
Subject: [PATCH] Fix known hosts handling in proxy backend

---
 internal/backends/proxy.go | 38 +++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/internal/backends/proxy.go b/internal/backends/proxy.go
index c9bb21f..65bfb2a 100644
--- a/internal/backends/proxy.go
+++ b/internal/backends/proxy.go
@@ -42,6 +42,7 @@ import (
 // proxySettings represents settings for the proxy backend.
 type proxySettings struct {
 	Server      string     `cty:"server"`
+	Port        *uint       `cty:"port"`
 	User        *string    `cty:"user"`
 	PrivkeyPath *string    `cty:"privkey"`
 	UserMap     *cty.Value `cty:"user_map"`
@@ -77,6 +78,11 @@ func Proxy(route config.Route) router.Handler {
 				opts.User = &user.Name
 			}
 		}
+		
+		if opts.Port == nil {
+			port := uint(22)
+			opts.Port = &port
+		}
 
 		auth := goph.Auth{
 			gossh.PasswordCallback(requestPassword(opts, sess)),
@@ -96,25 +102,27 @@ func Proxy(route config.Route) router.Handler {
 			auth = append(goph.Auth{gossh.PublicKeys(pk)}, auth...)
 		}
 
-		c, err := goph.New(*opts.User, opts.Server, auth)
+		c, err := goph.NewConn(&goph.Config{
+			Auth: auth,
+			User: *opts.User,
+			Addr: opts.Server,
+			Port: *opts.Port,
+			Callback: func(host string, remote net.Addr, key gossh.PublicKey) error {
+				found, err := goph.CheckKnownHost(host, remote, key, "")
+				if !found {
+					if err = goph.AddKnownHost(host, remote, key, ""); err != nil {
+						return err
+					}
+				} else if err != nil {
+					return err
+				}
+				return nil
+			},
+		})
 		if err != nil {
 			return err
 		}
 
-		knownHostHandler, err := goph.DefaultKnownHosts()
-		if err != nil {
-			return err
-		}
-
-		c.Config.Callback = func(host string, remote net.Addr, key gossh.PublicKey) error {
-			println("hi")
-			err = goph.AddKnownHost(host, remote, key, "")
-			if err != nil {
-				return err
-			}
-			return knownHostHandler(host, remote, key)
-		}
-
 		baseCmd := sess.Command()
 
 		var userCmd string