From 82ac6a5fff4332e5274d1efbebf0ea1dcd63b6af Mon Sep 17 00:00:00 2001
From: Boshi Lian <farmer1992@gmail.com>
Date: Sun, 11 Feb 2024 01:54:55 -0800
Subject: [PATCH] add remote signer api and e2e test of it (#320)

---
 e2e/plugin_test.go        | 33 ++++++++++++++++++++++++++++++---
 libplugin/pluginbase.go   | 11 +++++++++++
 plugin/testplugin/main.go | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/e2e/plugin_test.go b/e2e/plugin_test.go
index d3fa23d2e..4b59d511c 100644
--- a/e2e/plugin_test.go
+++ b/e2e/plugin_test.go
@@ -1,6 +1,12 @@
 package e2e_test
 
 import (
+	"bytes"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/pem"
 	"fmt"
 	"net"
 	"net/http"
@@ -110,11 +116,30 @@ func createRpcServer(r *rpcServer) net.Listener {
 
 func TestPlugin(t *testing.T) {
 
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		t.Fatalf("failed to generate private key: %v", err)
+	}
+
+	privKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
+	privKeyPem := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "RSA PRIVATE KEY",
+			Bytes: privKeyBytes,
+		},
+	)
+
+	sshkey, err := ssh.NewSignerFromKey(privateKey)
+	if err != nil {
+		t.Fatalf("failed to create ssh signer: %v", err)
+	}
+
 	sshsvr := createFakeSshServer(&ssh.ServerConfig{
-		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
-			if string(pass) != "rpcpassword" {
-				return nil, fmt.Errorf("invalid password")
+		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
+			if !bytes.Equal(key.Marshal(), sshkey.PublicKey().Marshal()) {
+				return nil, fmt.Errorf("public key mismatch")
 			}
+
 			return nil, nil
 		},
 	})
@@ -152,6 +177,8 @@ func TestPlugin(t *testing.T) {
 		sshsvr.Addr().String(),
 		"--rpcserver",
 		rpcsvr.Addr().String(),
+		"--testremotekey",
+		base64.StdEncoding.EncodeToString(privKeyPem),
 	)
 
 	if err != nil {
diff --git a/libplugin/pluginbase.go b/libplugin/pluginbase.go
index 4a9580102..b8b6dff4f 100644
--- a/libplugin/pluginbase.go
+++ b/libplugin/pluginbase.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"os"
 
+	"github.com/tg123/remotesigner/grpcsigner"
 	"github.com/tg123/sshpiper/libplugin/ioconn"
 	"google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
@@ -60,6 +61,8 @@ type SshPiperPluginConfig struct {
 	PipeStartCallback func(conn ConnMetadata)
 
 	PipeErrorCallback func(conn ConnMetadata, err error)
+
+	GrpcRemoteSignerFactory grpcsigner.SignerFactory
 }
 
 type SshPiperPlugin interface {
@@ -99,6 +102,14 @@ func NewFromGrpc(config SshPiperPluginConfig, grpc *grpc.Server, listener net.Li
 
 	RegisterSshPiperPluginServer(s.grpc, s)
 
+	if config.GrpcRemoteSignerFactory != nil {
+		gs, err := grpcsigner.NewSignerServer(config.GrpcRemoteSignerFactory)
+		if err != nil {
+			return nil, err
+		}
+		grpcsigner.RegisterSignerServer(s.grpc, gs)
+	}
+
 	return s, nil
 }
 
diff --git a/plugin/testplugin/main.go b/plugin/testplugin/main.go
index cd0419677..804e37375 100644
--- a/plugin/testplugin/main.go
+++ b/plugin/testplugin/main.go
@@ -3,10 +3,14 @@
 package main
 
 import (
+	"crypto"
+	"encoding/base64"
+	"fmt"
 	"net/rpc"
 
 	"github.com/tg123/sshpiper/libplugin"
 	"github.com/urfave/cli/v2"
+	"golang.org/x/crypto/ssh"
 )
 
 func main() {
@@ -23,6 +27,10 @@ func main() {
 				Name:     "testsshserver",
 				Required: true,
 			},
+			&cli.StringFlag{
+				Name:     "testremotekey",
+				Required: true,
+			},
 		},
 		CreateConfig: func(c *cli.Context) (*libplugin.SshPiperPluginConfig, error) {
 
@@ -36,6 +44,21 @@ func main() {
 				return nil, err
 			}
 
+			keydata, err := base64.StdEncoding.DecodeString(c.String("testremotekey"))
+			if err != nil {
+				return nil, err
+			}
+
+			key, err := ssh.ParseRawPrivateKey(keydata)
+			if err != nil {
+				return nil, err
+			}
+
+			_, ok := key.(crypto.Signer)
+			if !ok {
+				return nil, fmt.Errorf("key format not supported")
+			}
+
 			return &libplugin.SshPiperPluginConfig{
 				NewConnectionCallback: func(conn libplugin.ConnMetadata) error {
 					return rpcclient.Call("TestPlugin.NewConnection", "", nil)
@@ -56,10 +79,17 @@ func main() {
 					return &libplugin.Upstream{
 						Host:          host,
 						Port:          int32(port),
-						Auth:          libplugin.CreatePasswordAuthFromString(newpass),
+						Auth:          libplugin.CreateRemoteSignerAuth("testplugin"),
 						IgnoreHostKey: true,
 					}, nil
 				},
+				GrpcRemoteSignerFactory: func(metadata string) crypto.Signer {
+					if metadata != "testplugin" {
+						panic("metadata mismatch")
+					}
+
+					return key.(crypto.Signer)
+				},
 			}, nil
 		},
 	})