implement auth in jump server

internal/jimm/ssh/ssh.go

@@ -0,0 +1,48 @@
+// Copyright 2025 Canonical.
+
+package ssh
+
+import (
+	"context"
+
+	"github.com/canonical/jimm/v3/internal/dbmodel"
+	"github.com/canonical/jimm/v3/internal/errors"
+	"github.com/canonical/jimm/v3/internal/openfga"
+)
+
+type identityManager interface {
+	FetchIdentity(ctx context.Context, id string) (*openfga.User, error)
+}
+
+type modelManager interface {
+	GetModel(ctx context.Context, uuid string) (dbmodel.Model, error)
+}
+
+type sshKeyManager interface {
+	VerifyPublicKey(ctx context.Context, user *openfga.User, fingerprint string) (bool, error)
+}
+
+// sshManager provides a means to manage ssh server within JIMM.
+type sshManager struct {
+	modelManager
+	identityManager
+	sshKeyManager
+}
+
+// NewSSHManager returns a new SSHManager that offers jimm functionality to the SSHJumpServer.
+func NewSSHManager(identityManager identityManager, modelManager modelManager, sshKeyManager sshKeyManager) (*sshManager, error) {
+	if identityManager == nil {
+		return nil, errors.E("identityManager cannot be nil")
+	}
+	if modelManager == nil {
+		return nil, errors.E("modelManager cannot be nil")
+	}
+	if sshKeyManager == nil {
+		return nil, errors.E("sshManager cannot be nil")
+	}
+	return &sshManager{
+		modelManager:    modelManager,
+		identityManager: identityManager,
+		sshKeyManager:   sshKeyManager,
+	}, nil
+}

internal/jimm/ssh/ssh_test.go

@@ -0,0 +1,18 @@
+// Copyright 2025 Canonical.
+
+package ssh_test
+
+import (
+	"testing"
+
+	qt "github.com/frankban/quicktest"
+
+	"github.com/canonical/jimm/v3/internal/jimm/ssh"
+	"github.com/canonical/jimm/v3/internal/testutils/jimmtest/mocks"
+)
+
+func TestSSHManagerCreation(t *testing.T) {
+	c := qt.New(t)
+	_, err := ssh.NewSSHManager(&mocks.IdentityManager{}, &mocks.ModelManager{}, &mocks.SSHKeyManager{})
+	c.Assert(err, qt.IsNil)
+}

internal/ssh/ssh.go

@@ -9,6 +9,7 @@ import (
 	"net"
 
 	"github.com/gliderlabs/ssh"
+	"github.com/juju/names/v5"
 	"github.com/juju/zaputil/zapctx"
 	"go.uber.org/zap"
 	gossh "golang.org/x/crypto/ssh"
@@ -19,10 +20,18 @@ import (
 // juju_ssh_default_port is the default port we expect the juju controllers to respond on.
 const juju_ssh_default_port = 17022
 
-// Resolver is the interface with the methods needed by the ssh jump server to route request.
-type Resolver interface {
+type publicKeySSHUserKey struct{}
+
+// SSHManager is the interface with the methods needed by the ssh jump server to route request.
+type SSHManager interface {
 	// AddrFromModelUUID is the method to resolve the address of the controller to contact given the model UUID.
-	AddrFromModelUUID(ctx context.Context, user openfga.User, modelUUID string) (string, error)
+	AddrFromModelUUID(ctx context.Context, user *openfga.User, modelTag names.ModelTag) (string, error)
+
+	// FetchIdentity
+	FetchIdentity(ctx context.Context, id string) (*openfga.User, error)
+
+	// VerifyPublicKey verifies the identityName is
+	VerifyPublicKey(ctx context.Context, user *openfga.User, fingerprint string) (bool, error)
 }
 
 // forwardMessage is the struct holding the information about the jump message received by the ssh client.
@@ -33,11 +42,11 @@ type forwardMessage struct {
 	SrcPort  uint32
 }
 
-// Server is the custom struct to embed the gliderlabs.ssh server and a resolver.
+// Server is the custom struct to embed the gliderlabs.ssh server and a sshManager.
 type Server struct {
 	*ssh.Server
 
-	resolver Resolver
+	sshManager SSHManager
 }
 
 // Config is the struct holding the configuration for the jump server.
@@ -48,23 +57,33 @@ type Config struct {
 }
 
 // NewJumpServer creates the jump server struct.
-func NewJumpServer(ctx context.Context, config Config, resolver Resolver) (Server, error) {
+func NewJumpServer(ctx context.Context, config Config, sshManager SSHManager) (Server, error) {
 	zapctx.Info(ctx, "NewJumpServer")
 
-	if resolver == nil {
+	if sshManager == nil {
 		return Server{}, fmt.Errorf("Cannot create JumpSSHServer with a nil resolver.")
 	}
 	server := Server{
 		Server: &ssh.Server{
 			Addr: fmt.Sprintf(":%s", config.Port),
 			ChannelHandlers: map[string]ssh.ChannelHandler{
-				"direct-tcpip": directTCPIPHandler(resolver),
+				"direct-tcpip": directTCPIPHandler(sshManager),
 			},
 			PublicKeyHandler: func(ctx ssh.Context, key ssh.PublicKey) bool {
+				user, err := sshManager.FetchIdentity(ctx, ctx.User())
+				if err != nil {
+					zapctx.Info(ctx, fmt.Sprintf("cannot find user %s", ctx.User()))
+					return false
+				}
+				if ok, err := sshManager.VerifyPublicKey(ctx, user, gossh.FingerprintSHA256(key)); !ok || err != nil {
+					zapctx.Info(ctx, fmt.Sprintf("cannot verify key for user %s", ctx.User()), zap.Error(err))
+					return false
+				}
+				ctx.SetValue(publicKeySSHUserKey{}, user)
 				return true
 			},
 		},
-		resolver: resolver,
+		sshManager: sshManager,
 	}
 	s, err := gossh.ParsePrivateKey([]byte(config.HostKey))
 	if err != nil {
@@ -75,22 +94,31 @@ func NewJumpServer(ctx context.Context, config Config, resolver Resolver) (Serve
 	return server, nil
 }
 
-func directTCPIPHandler(resolver Resolver) func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
+func directTCPIPHandler(sshManager SSHManager) func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
 	return func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
 		d := forwardMessage{}
-
 		k := newChan.ExtraData()
 
 		if err := gossh.Unmarshal(k, &d); err != nil {
-			rejectConnectionAndLogError(ctx, newChan, "Failed to parse channel data", err)
+			rejectConnectionAndLogError(ctx, newChan, "failed to parse channel data", err)
 			return
 		}
 		if d.DestPort == 0 {
 			d.DestPort = juju_ssh_default_port
 		}
-		addr, err := resolver.AddrFromModelUUID(ctx, openfga.User{}, d.DestAddr)
+		if !names.IsValidModel(d.DestAddr) {
+			rejectConnectionAndLogError(ctx, newChan, "invalid model uuid", nil)
+			return
+		}
+		modelTag := names.NewModelTag(d.DestAddr)
+		user, err := fetchAndVerifySSHUser(ctx, modelTag)
 		if err != nil {
-			rejectConnectionAndLogError(ctx, newChan, "Failed to resolve address from model uuid", err)
+			rejectConnectionAndLogError(ctx, newChan, err.Error(), err)
+			return
+		}
+		addr, err := sshManager.AddrFromModelUUID(ctx, user, modelTag)
+		if err != nil {
+			rejectConnectionAndLogError(ctx, newChan, "failed to resolve address from model uuid", err)
 			return
 		}
 		dest := net.JoinHostPort(addr, fmt.Sprint(d.DestPort))
@@ -105,13 +133,13 @@ func directTCPIPHandler(resolver Resolver) func(srv *ssh.Server, conn *gossh.Ser
 			},
 		})
 		if err != nil {
-			rejectConnectionAndLogError(ctx, newChan, fmt.Sprintf("Failed to connect to %s: %v", dest, err), err)
+			rejectConnectionAndLogError(ctx, newChan, fmt.Sprintf("failed to connect to %s: %v", dest, err), err)
 			return
 		}
 
 		dstChan, reqs, err := client.OpenChannel("direct-tcpip", gossh.Marshal(d))
 		if err != nil {
-			rejectConnectionAndLogError(ctx, newChan, "Failed to open destination channel", err)
+			rejectConnectionAndLogError(ctx, newChan, "failed to open destination channel", err)
 			return
 		}
 		// gossh.Request are requests sent outside of the normal stream of data (ex. pty-req for an interactive session).
@@ -132,21 +160,38 @@ func directTCPIPHandler(resolver Resolver) func(srv *ssh.Server, conn *gossh.Ser
 			defer dstChan.Close()
 			_, err := io.Copy(srcDest, dstChan)
 			if err != nil {
-				rejectConnectionAndLogError(ctx, newChan, "Failed to copy data from src to dts", err)
+				rejectConnectionAndLogError(ctx, newChan, "failed to copy data from src to dts", err)
 			}
 		}()
 		go func() {
 			defer srcDest.Close()
 			defer dstChan.Close()
 			_, err := io.Copy(dstChan, srcDest)
 			if err != nil {
-				rejectConnectionAndLogError(ctx, newChan, "Failed to copy data from dst to src", err)
+				rejectConnectionAndLogError(ctx, newChan, "failed to copy data from dst to src", err)
 			}
 		}()
 		zapctx.Info(ctx, fmt.Sprintf("Proxying connection from %s:%d to %s:%d \n", d.SrcAddr, d.SrcPort, d.DestAddr, d.DestPort))
 	}
 }
 
+// fetchAndVerifySSHUser extracts the user from the context and checks the user has permission to ssh.
+func fetchAndVerifySSHUser(ctx ssh.Context, modelTag names.ModelTag) (*openfga.User, error) {
+	user, ok := ctx.Value(publicKeySSHUserKey{}).(*openfga.User)
+	if !ok {
+		return nil, fmt.Errorf("fo user in the context")
+	}
+	ok, err := user.IsModelWriter(ctx, modelTag)
+	if err != nil {
+		return nil, fmt.Errorf("failed to resolve address from model uuid")
+	}
+	if !ok {
+		return nil, fmt.Errorf("user doesn't have permission")
+	}
+	return user, nil
+}
+
+// rejectConnectionAndLogError logs the error and rejects the channel with a message.
 func rejectConnectionAndLogError(ctx context.Context, newChan gossh.NewChannel, msg string, err error) {
 	zapctx.Error(ctx, msg, zap.Error(err))
 	err = newChan.Reject(gossh.ConnectionFailed, msg)