feat: intercept ssh keyManager facade in model proxy

internal/jimm/jimm.go

@@ -491,6 +491,12 @@ func (j *JIMM) ServiceAccountManager() ServiceAccountManager {
 	return j.serviceAccountManager
 }
 
+// SSHKeyManager returns a manager that enables operations
+// related to ssh keys.
+func (j *JIMM) SSHKeyManager() SSHKeyManager {
+	return j.sshKeyManager
+}
+
 type permission struct {
 	resource string
 	relation string

internal/jujuapi/websocket.go

@@ -185,6 +185,7 @@ func (s apiProxier) ServeWS(ctx context.Context, clientConn *websocket.Conn) {
 		AuditLog:                auditLogger,
 		LoginService:            s.jimm.LoginManager(),
 		AuthenticatedIdentityID: auth.SessionIdentityFromContext(ctx),
+		SSHKeyManager:           s.jimm.SSHKeyManager(),
 	}
 	if err := rpcproxy.ProxySockets(ctx, proxyHelpers); err != nil {
 		zapctx.Error(ctx, "failed to start jimm model proxy", zap.Error(err))

internal/rpcproxy/export_test.go

@@ -2,4 +2,7 @@
 
 package rpcproxy
 
-type Message = message
+type (
+	Message          = message
+	KeyManagerFacade = keyManagerFacade
+)

internal/rpcproxy/keymanager.go

@@ -0,0 +1,132 @@
+// Copyright 2025 Canonical.
+
+package rpcproxy
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"fmt"
+	"regexp"
+
+	jujuparams "github.com/juju/juju/rpc/params"
+	"github.com/juju/utils/v3/ssh"
+	gossh "golang.org/x/crypto/ssh"
+
+	"github.com/canonical/jimm/v3/internal/errors"
+	"github.com/canonical/jimm/v3/internal/jimm/sshkeys"
+	"github.com/canonical/jimm/v3/internal/openfga"
+)
+
+var isFingerprintRegexp = regexp.MustCompile("^[0-9a-f]{2}(:[0-9a-f]{2}){15}$")
+
+// keyManagerFacade is intended to be a temporary struct used to emulate logic
+// that will eventually live in jujuapi. This struct contains all
+// the api layer logic for SSH key management methods that are currently
+// used by the rpcProxy.
+type keyManagerFacade struct {
+	SSHKeyManager
+	user *openfga.User
+}
+
+// ListKeys lists the authenticated user's SSH keys
+// in the format defined in args.
+func (s *keyManagerFacade) ListKeys(ctx context.Context, args jujuparams.ListSSHKeys) (jujuparams.StringsResults, error) {
+	keys, err := s.ListUserPublicKeys(ctx, s.user)
+	if err != nil {
+		return jujuparams.StringsResults{}, err
+	}
+
+	var formatter func(key sshkeys.PublicKey) string
+	switch args.Mode {
+	case ssh.FullKeys:
+		formatter = marshalAuthorizedKeyWithComment
+	case ssh.Fingerprints:
+		formatter = fingerprintWithComment
+	default:
+		return jujuparams.StringsResults{}, fmt.Errorf("unknown mode (%v)", args.Mode)
+	}
+
+	// The Juju CLI takes the first element from the slice of stringsResults.
+	res := jujuparams.StringsResult{}
+	for _, key := range keys {
+		res.Result = append(res.Result, formatter(key))
+	}
+	return jujuparams.StringsResults{Results: []jujuparams.StringsResult{res}}, nil
+}
+
+// AddKeys saves the SSH keys defined in args and associates them
+// with the authenticated user.
+func (s *keyManagerFacade) AddKeys(ctx context.Context, args jujuparams.ModifyUserSSHKeys) (jujuparams.ErrorResults, error) {
+	var res []jujuparams.ErrorResult
+	errF := func(err error, msg string) jujuparams.ErrorResult {
+		return jujuparams.ErrorResult{Error: &jujuparams.Error{
+			Code:    string(errors.ErrorCode(err)),
+			Message: fmt.Sprintf("%s: %s", msg, err.Error()),
+		}}
+	}
+
+	for i, key := range args.Keys {
+		out, comment, _, _, err := gossh.ParseAuthorizedKey([]byte(key))
+		if err != nil {
+			res = append(res, errF(err, fmt.Sprintf("Failed to parse key (entry %d)", i)))
+		}
+		jimmKey := sshkeys.PublicKey{
+			PublicKey: out,
+			Comment:   comment,
+		}
+		if err := s.AddUserPublicKey(ctx, s.user, jimmKey); err != nil {
+			res = append(res, errF(err, fmt.Sprintf("Failed to add key (comment %s)", comment)))
+		}
+	}
+
+	return jujuparams.ErrorResults{Results: res}, nil
+}
+
+// DeleteKeys removes saved keys associated with the authenticated user
+// and finds keys to remove either by comment or fingerprint.
+func (s *keyManagerFacade) DeleteKeys(ctx context.Context, args jujuparams.ModifyUserSSHKeys) (jujuparams.ErrorResults, error) {
+	var res []jujuparams.ErrorResult
+	errF := func(err error, msg string) jujuparams.ErrorResult {
+		return jujuparams.ErrorResult{Error: &jujuparams.Error{
+			Code:    string(errors.ErrorCode(err)),
+			Message: fmt.Sprintf("%s: %s", msg, err.Error()),
+		}}
+	}
+
+	for _, key := range args.Keys {
+		if isFingerprintRegexp.MatchString(key) {
+			err := s.RemoveUserKeyByFingerprint(ctx, s.user, key)
+			if err != nil {
+				res = append(res, errF(err, fmt.Sprintf("Failed to remove key by fingerprint (%s)", key)))
+			}
+		} else {
+			err := s.RemoveUserKeyByComment(ctx, s.user, key)
+			if err != nil {
+				res = append(res, errF(err, fmt.Sprintf("Failed to remove key by comment (%s)", key)))
+			}
+		}
+	}
+
+	return jujuparams.ErrorResults{Results: res}, nil
+}
+
+func marshalAuthorizedKeyWithComment(key sshkeys.PublicKey) string {
+	// Copied from gossh.MarshalAuthorizedKey with an addition for the comment.
+	// Errors from the buffer's Write..() methods are always nil.
+	b := &bytes.Buffer{}
+	b.WriteString(key.Type())
+	b.WriteByte(' ')
+	e := base64.NewEncoder(base64.StdEncoding, b)
+	_, _ = e.Write(key.Marshal())
+	e.Close()
+	b.WriteByte(' ')
+	b.WriteString(key.Comment)
+	b.WriteByte('\n')
+	return b.String()
+}
+
+func fingerprintWithComment(key sshkeys.PublicKey) string {
+	fingerprint := gossh.FingerprintLegacyMD5(key)
+	return fmt.Sprintf("%s (%s)", fingerprint, key.Comment)
+}

internal/rpcproxy/keymanager_test.go

@@ -0,0 +1,148 @@
+// Copyright 2025 Canonical.
+
+package rpcproxy_test
+
+import (
+	"bytes"
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"encoding/base64"
+	"fmt"
+	"regexp"
+	"testing"
+
+	qt "github.com/frankban/quicktest"
+	"github.com/frankban/quicktest/qtsuite"
+	"github.com/juju/juju/rpc/params"
+	"github.com/juju/utils/v3/ssh"
+	gossh "golang.org/x/crypto/ssh"
+
+	"github.com/canonical/jimm/v3/internal/jimm/sshkeys"
+	"github.com/canonical/jimm/v3/internal/openfga"
+	"github.com/canonical/jimm/v3/internal/rpcproxy"
+	"github.com/canonical/jimm/v3/internal/testutils/jimmtest/mocks"
+)
+
+type keyManagerFacadeSuite struct {
+	keyManagerFacade        rpcproxy.KeyManagerFacade
+	key1                    sshkeys.PublicKey
+	key2                    sshkeys.PublicKey
+	addKeyF                 func(ctx context.Context, user *openfga.User, publicKey sshkeys.PublicKey) error
+	removeKeyByCommentF     func(ctx context.Context, user *openfga.User, comment string) error
+	removeKeyByFingerprintF func(ctx context.Context, user *openfga.User, fingerprint string) error
+}
+
+func (k *keyManagerFacadeSuite) Init(c *qt.C) {
+	pk1, err := rsa.GenerateKey(rand.Reader, 2048)
+	c.Assert(err, qt.IsNil)
+	pk2, err := rsa.GenerateKey(rand.Reader, 2048)
+	c.Assert(err, qt.IsNil)
+
+	pubKey1, err := gossh.NewPublicKey(&pk1.PublicKey)
+	c.Assert(err, qt.IsNil)
+
+	pubKey2, err := gossh.NewPublicKey(&pk2.PublicKey)
+	c.Assert(err, qt.IsNil)
+
+	k.key1 = sshkeys.PublicKey{
+		PublicKey: pubKey1,
+		Comment:   "comment-1",
+	}
+
+	k.key2 = sshkeys.PublicKey{
+		PublicKey: pubKey2,
+		Comment:   "comment-2",
+	}
+
+	keyManager := mocks.SSHKeyManager{
+		ListUserPublicKeys_: func(ctx context.Context, user *openfga.User) ([]sshkeys.PublicKey, error) {
+			return []sshkeys.PublicKey{k.key1, k.key2}, nil
+		},
+		AddUserPublicKey_: func(ctx context.Context, user *openfga.User, publicKey sshkeys.PublicKey) error {
+			return k.addKeyF(ctx, user, publicKey)
+		},
+		RemoveUserKeyByComment_: func(ctx context.Context, user *openfga.User, comment string) error {
+			return k.removeKeyByCommentF(ctx, user, comment)
+		},
+		RemoveUserKeyByFingerprint_: func(ctx context.Context, user *openfga.User, fingerprint string) error {
+			return k.removeKeyByFingerprintF(ctx, user, fingerprint)
+		},
+	}
+	k.keyManagerFacade = rpcproxy.KeyManagerFacade{
+		SSHKeyManager: &keyManager,
+	}
+}
+
+var isFingerprintRegex = regexp.MustCompile(`[0-9a-f]{2}(:[0-9a-f]{2}){15}`)
+
+func (k *keyManagerFacadeSuite) TestListKeysShort(c *qt.C) {
+	c.Parallel()
+	ctx := context.Background()
+
+	res, err := k.keyManagerFacade.ListKeys(ctx, params.ListSSHKeys{Mode: ssh.Fingerprints})
+	c.Assert(err, qt.IsNil)
+
+	c.Assert(res.Results[0].Result, qt.HasLen, 2)
+	c.Assert(res.Results[0].Result[0], qt.Matches, `.+ \(comment-1\)`)
+	c.Assert(isFingerprintRegex.MatchString(res.Results[0].Result[0]), qt.IsTrue)
+	c.Assert(res.Results[0].Result[1], qt.Matches, `.+ \(comment-2\)`)
+	c.Assert(isFingerprintRegex.MatchString(res.Results[0].Result[1]), qt.IsTrue)
+}
+
+func (k *keyManagerFacadeSuite) TestListKeysFull(c *qt.C) {
+	c.Parallel()
+	ctx := context.Background()
+
+	res, err := k.keyManagerFacade.ListKeys(ctx, params.ListSSHKeys{Mode: ssh.FullKeys})
+	c.Assert(err, qt.IsNil)
+
+	c.Assert(res.Results[0].Result, qt.HasLen, 2)
+	c.Assert(res.Results[0].Result[0], qt.Matches, `ssh-rsa .+ comment-1\n`)
+	c.Assert(res.Results[0].Result[1], qt.Matches, `ssh-rsa .+ comment-2\n`)
+}
+
+func (k *keyManagerFacadeSuite) TestAddKeys(c *qt.C) {
+	c.Parallel()
+	ctx := context.Background()
+
+	k.addKeyF = func(ctx context.Context, user *openfga.User, publicKey sshkeys.PublicKey) error {
+		c.Check(publicKey.Marshal(), qt.DeepEquals, k.key1.Marshal())
+		c.Check(publicKey.Comment, qt.Equals, k.key1.Comment)
+		return nil
+	}
+	var b bytes.Buffer
+	e := base64.NewEncoder(base64.StdEncoding, &b)
+	_, _ = e.Write(k.key1.Marshal()) // Writes to a bytes buffer always return a nil error.
+	e.Close()
+	authorizedKey := fmt.Sprintf("%s %s %s", k.key1.Type(), &b, k.key1.Comment)
+	_, _ = k.keyManagerFacade.AddKeys(ctx, params.ModifyUserSSHKeys{Keys: []string{authorizedKey}})
+}
+
+func (k *keyManagerFacadeSuite) TestDeleteKeysByComment(c *qt.C) {
+	c.Parallel()
+	ctx := context.Background()
+
+	k.removeKeyByCommentF = func(ctx context.Context, user *openfga.User, comment string) error {
+		c.Check(comment, qt.Equals, "comment-1")
+		return nil
+	}
+	_, _ = k.keyManagerFacade.DeleteKeys(ctx, params.ModifyUserSSHKeys{Keys: []string{"comment-1"}})
+}
+
+func (k *keyManagerFacadeSuite) TestDeleteKeysByFingerprint(c *qt.C) {
+	c.Parallel()
+	ctx := context.Background()
+
+	fp := gossh.FingerprintLegacyMD5(k.key1)
+
+	k.removeKeyByFingerprintF = func(ctx context.Context, user *openfga.User, fingerprint string) error {
+		c.Check(fingerprint, qt.Equals, fp)
+		return nil
+	}
+	_, _ = k.keyManagerFacade.DeleteKeys(ctx, params.ModifyUserSSHKeys{Keys: []string{fp}})
+}
+
+func TestKeyManagerFacade(t *testing.T) {
+	qtsuite.Run(qt.New(t), &keyManagerFacadeSuite{})
+}