From 240a2949104eec1e753589eac7255d4858c57a3e Mon Sep 17 00:00:00 2001
From: Kian Parvin <kian.parvin@canonical.com>
Date: Fri, 6 Sep 2024 11:26:45 +0200
Subject: [PATCH] PR comments

- Add further docstring
- Reorder some code
---
 internal/auth/oauth2.go         |  5 ++++-
 internal/jujuapi/streamproxy.go | 17 ++++++++++-------
 internal/rpc/streamproxy.go     |  1 +
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/internal/auth/oauth2.go b/internal/auth/oauth2.go
index 2a6b93bd1..d44f2351b 100644
--- a/internal/auth/oauth2.go
+++ b/internal/auth/oauth2.go
@@ -330,7 +330,10 @@ func (as *AuthenticationService) MintSessionToken(email string) (string, error)
 // access token JWT, returning the parsed token.
 //
 // The subject of the token contains the user's email and can be used
-// for user object creation
+// for user object creation.
+//
+// The error code returned here is used by the Juju CLI to know when to start a
+// device login flow, prompting the user to login again.
 func (as *AuthenticationService) VerifySessionToken(token string) (_ jwt.Token, err error) {
 	const op = errors.Op("auth.AuthenticationService.VerifySessionToken")
 	errorFn := func(message string) error {
diff --git a/internal/jujuapi/streamproxy.go b/internal/jujuapi/streamproxy.go
index 2ef9f270a..cbb76b6b0 100644
--- a/internal/jujuapi/streamproxy.go
+++ b/internal/jujuapi/streamproxy.go
@@ -1,4 +1,5 @@
 // Copyright 2024 Canonical.
+
 package jujuapi
 
 import (
@@ -19,7 +20,7 @@ import (
 	jimmRPC "github.com/canonical/jimm/v3/internal/rpc"
 )
 
-// A streamProxier serves all HTTP endpoints by proxying
+// A streamProxier serves the the /log endpoint by proxying
 // messages between the controller and client.
 type streamProxier struct {
 	// TODO(Kian): Refactor the apiServer to use the JIMM API rather than a concrete struct
@@ -71,13 +72,9 @@ func (s streamProxier) ServeWS(ctx context.Context, clientConn *websocket.Conn)
 		return
 	}
 
-	model, err := s.jimm.GetModel(ctx, uuid)
-	if err != nil {
-		writeError(err.Error(), errors.CodeModelNotFound)
-		return
-	}
+	modelTag := names.NewModelTag(uuid)
 
-	if ok, err := checkPermission(ctx, finalPath, user, model.ResourceTag()); err != nil {
+	if ok, err := checkPermission(ctx, finalPath, user, modelTag); err != nil {
 		writeError(err.Error(), errors.CodeUnauthorized)
 		return
 	} else if !ok {
@@ -85,6 +82,12 @@ func (s streamProxier) ServeWS(ctx context.Context, clientConn *websocket.Conn)
 		return
 	}
 
+	model, err := s.jimm.GetModel(ctx, uuid)
+	if err != nil {
+		writeError(err.Error(), errors.CodeModelNotFound)
+		return
+	}
+
 	api, err := s.jimm.Dialer.Dial(ctx, &model.Controller, model.ResourceTag(), nil)
 	if err != nil {
 		zapctx.Error(ctx, "failed to dial controller", zap.Error(err))
diff --git a/internal/rpc/streamproxy.go b/internal/rpc/streamproxy.go
index 7430a455b..5da96206f 100644
--- a/internal/rpc/streamproxy.go
+++ b/internal/rpc/streamproxy.go
@@ -1,4 +1,5 @@
 // Copyright 2024 Canonical.
+
 package rpc
 
 import (