CSS-9575 Add device flow integration tests

internal/auth/oauth2.go

@@ -326,6 +326,9 @@ func (as *AuthenticationService) MintSessionToken(email string) (string, error)
 // for user object creation
 func (as *AuthenticationService) VerifySessionToken(token string) (_ jwt.Token, err error) {
 	const op = errors.Op("auth.AuthenticationService.VerifySessionToken")
+	errorFn := func(message string) error {
+		return errors.E(op, message, errors.CodeUnauthorized)
+	}
 	defer func() {
 		if err != nil {
 			servermon.AuthenticationFailCount.WithLabelValues("VerifySessionToken").Inc()
@@ -335,24 +338,24 @@ func (as *AuthenticationService) VerifySessionToken(token string) (_ jwt.Token,
 	}()
 
 	if len(token) == 0 {
-		return nil, errors.E(op, "authentication failed, no token presented")
+		return nil, errorFn("no token presented")
 	}
 
 	decodedToken, err := base64.StdEncoding.DecodeString(token)
 	if err != nil {
-		return nil, errors.E(op, fmt.Sprintf("authentication failed, failed to decode token: %s", err))
+		return nil, errorFn(fmt.Sprintf("failed to decode token: %s", err))
 	}
 
 	parsedToken, err := jwt.Parse(decodedToken, jwt.WithKey(as.signingAlg, []byte(as.jwtSessionKey)))
 	if err != nil {
 		if stderrors.Is(err, jwt.ErrTokenExpired()) {
-			return nil, errors.E(op, errors.CodeUnauthorized, "JIMM session token expired")
+			return nil, errorFn("JIMM session token expired")
 		}
-		return nil, errors.E(op, err)
+		return nil, errorFn(err.Error())
 	}
 
 	if _, err = mail.ParseAddress(parsedToken.Subject()); err != nil {
-		return nil, errors.E(op, "failed to parse email")
+		return nil, errorFn("failed to parse email")
 	}
 
 	return parsedToken, nil

internal/auth/oauth2_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/canonical/jimm/internal/auth"
 	"github.com/canonical/jimm/internal/db"
 	"github.com/canonical/jimm/internal/dbmodel"
+	"github.com/canonical/jimm/internal/errors"
 	"github.com/canonical/jimm/internal/jimmtest"
 	"github.com/coreos/go-oidc/v3/oidc"
 	qt "github.com/frankban/quicktest"
@@ -56,8 +57,6 @@ func setupTestAuthSvc(ctx context.Context, c *qt.C, expiry time.Duration) (*auth
 
 // This test requires the local docker compose to be running and keycloak
 // to be available.
-//
-// TODO(ale8k): Use a mock for this and also device below, but future work???
 func TestAuthCodeURL(t *testing.T) {
 	c := qt.New(t)
 	ctx := context.Background()
@@ -205,6 +204,20 @@ func TestSessionTokenRejectsExpiredToken(t *testing.T) {
 
 	_, err = authSvc.VerifySessionToken(token)
 	c.Assert(err, qt.ErrorMatches, `JIMM session token expired`)
+	c.Assert(errors.ErrorCode(err), qt.Equals, errors.CodeUnauthorized)
+}
+
+func TestSessionTokenRejectsEmptyToken(t *testing.T) {
+	c := qt.New(t)
+
+	ctx := context.Background()
+
+	noDuration := time.Duration(0)
+	authSvc, _, _ := setupTestAuthSvc(ctx, c, noDuration)
+
+	_, err := authSvc.VerifySessionToken("")
+	c.Assert(err, qt.ErrorMatches, `no token presented`)
+	c.Assert(errors.ErrorCode(err), qt.Equals, errors.CodeUnauthorized)
 }
 
 func TestSessionTokenValidatesEmail(t *testing.T) {
@@ -220,6 +233,7 @@ func TestSessionTokenValidatesEmail(t *testing.T) {
 
 	_, err = authSvc.VerifySessionToken(token)
 	c.Assert(err, qt.ErrorMatches, "failed to parse email")
+	c.Assert(errors.ErrorCode(err), qt.Equals, errors.CodeUnauthorized)
 }
 
 func TestVerifyClientCredentials(t *testing.T) {

internal/jimmhttp/websocket_test.go

@@ -111,11 +111,12 @@ func TestWSHandlerNilServer(t *testing.T) {
 	c.Assert(err, qt.ErrorMatches, `websocket: close 1000 \(normal\)`)
 }
 
-type authFailServer struct{}
+type authFailServer struct{ c jimmtest.SimpleTester }
 
 // GetAuthenticationService returns JIMM's oauth authentication service.
 func (s authFailServer) GetAuthenticationService() jimm.OAuthAuthenticator {
-	return jimmtest.NewMockOAuthAuthenticator("")
+	authenticator := jimmtest.NewMockOAuthAuthenticator(s.c, nil)
+	return &authenticator
 }
 
 func (s authFailServer) ServeWS(ctx context.Context, conn *websocket.Conn) {}
@@ -124,7 +125,7 @@ func TestWSHandlerAuthFailsServer(t *testing.T) {
 	c := qt.New(t)
 
 	hnd := &jimmhttp.WSHandler{
-		Server: authFailServer{},
+		Server: authFailServer{c: c},
 	}
 
 	srv := httptest.NewServer(hnd)

internal/jimmtest/auth.go

@@ -21,23 +21,27 @@ import (
 
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/go-chi/chi/v5"
+	"github.com/google/uuid"
 	"github.com/gorilla/sessions"
 	"github.com/juju/juju/api"
 	jujuparams "github.com/juju/juju/rpc/params"
 	"github.com/lestrrat-go/jwx/v2/jwa"
 	"github.com/lestrrat-go/jwx/v2/jwt"
+	"golang.org/x/oauth2"
 
 	"github.com/canonical/jimm/internal/auth"
 	"github.com/canonical/jimm/internal/db"
+	jimmerrors "github.com/canonical/jimm/internal/errors"
 	"github.com/canonical/jimm/internal/jimm"
 	"github.com/canonical/jimm/internal/jimmhttp"
 	"github.com/canonical/jimm/internal/openfga"
 )
 
 const (
-	// Note that these values are deliberately different to make sure we're not
-	// reusing/misusing them.
-	JWTTestSecret      = "test-secret"
+	// Secrets used for auth in tests.
+	//
+	// Note that these values are deliberately different to make sure we're not reusing/misusing them.
+	JWTTestSecret      = "jwt-test-secret"
 	SessionStoreSecret = "another-test-secret"
 )
 
@@ -63,40 +67,107 @@ func (a Authenticator) Authenticate(_ context.Context, _ *jujuparams.LoginReques
 
 type MockOAuthAuthenticator struct {
 	jimm.OAuthAuthenticator
+	c SimpleTester
+	// PollingChan is used to simulate polling an OIDC server during the device flow.
+	// It expects a username to be received that will be used to generate the user's access token.
+	PollingChan     <-chan string
+	polledUsername  string
+	mockAccessToken string
 }
 
-func NewMockOAuthAuthenticator(secretKey string) MockOAuthAuthenticator {
-	return MockOAuthAuthenticator{}
+func NewMockOAuthAuthenticator(c SimpleTester, testChan <-chan string) MockOAuthAuthenticator {
+	return MockOAuthAuthenticator{c: c, PollingChan: testChan}
+}
+
+// Device is a mock implementation for the start of the device flow, returning dummy polling data.
+func (m *MockOAuthAuthenticator) Device(ctx context.Context) (*oauth2.DeviceAuthResponse, error) {
+	return &oauth2.DeviceAuthResponse{
+		DeviceCode:              "test-device-code",
+		UserCode:                "test-user-code",
+		VerificationURI:         "http://no-such-uri.canonical.com",
+		VerificationURIComplete: "http://no-such-uri.canonical.com",
+		Expiry:                  time.Now().Add(time.Minute),
+		Interval:                int64(time.Minute.Seconds()),
+	}, nil
+}
+
+// DeviceAccessToken is a mock implementation of the second step in the device flow where JIMM
+// polls an OIDC server for the device code.
+func (m *MockOAuthAuthenticator) DeviceAccessToken(ctx context.Context, res *oauth2.DeviceAuthResponse) (*oauth2.Token, error) {
+	select {
+	case username := <-m.PollingChan:
+		m.polledUsername = username
+		uuid, err := uuid.NewRandom()
+		if err != nil {
+			m.c.Fatalf("failed to generate UUID for device access token")
+		}
+		m.mockAccessToken = uuid.String()
+	case <-ctx.Done():
+		return &oauth2.Token{}, ctx.Err()
+	}
+	return &oauth2.Token{AccessToken: m.mockAccessToken}, nil
 }
 
 // VerifySessionToken provides the mock implementation for verifying session tokens.
 // Allowing JIMM tests to create their own session tokens that will always be accepted.
 // Notice the use of jwt.ParseInsecure to skip JWT signature verification.
-func (m MockOAuthAuthenticator) VerifySessionToken(token string) (jwt.Token, error) {
+func (m *MockOAuthAuthenticator) VerifySessionToken(token string) (jwt.Token, error) {
+	errorFn := func(err error) error {
+		return jimmerrors.E(err, jimmerrors.CodeUnauthorized)
+	}
 	decodedToken, err := base64.StdEncoding.DecodeString(token)
 	if err != nil {
-		return nil, errors.New("authentication failed, failed to decode token")
+		return nil, errorFn(errors.New("failed to decode token"))
 	}
 
 	parsedToken, err := jwt.ParseInsecure(decodedToken)
 	if err != nil {
-		return nil, err
+		return nil, errorFn(err)
 	}
 
 	if _, err = mail.ParseAddress(parsedToken.Subject()); err != nil {
-		return nil, errors.New("failed to parse email")
+		return nil, errorFn(errors.New("failed to parse email"))
 	}
 	return parsedToken, nil
 }
 
-func (m MockOAuthAuthenticator) AuthenticateBrowserSession(ctx context.Context, w http.ResponseWriter, req *http.Request) (context.Context, error) {
+// ExtractAndVerifyIDToken returns an ID token where the subject is equal to the username obtained during the device flow.
+// The auth token must match the one returned during the device flow.
+// If the polled username is empty it indicates an error that the device flow was not run prior to calling this function.
+func (m *MockOAuthAuthenticator) ExtractAndVerifyIDToken(ctx context.Context, oauth2Token *oauth2.Token) (*oidc.IDToken, error) {
+	if m.polledUsername == "" {
+		return &oidc.IDToken{}, errors.New("unknown user for mock auth login")
+	}
+	if m.mockAccessToken != oauth2Token.AccessToken {
+		return &oidc.IDToken{}, errors.New("access token does not match the generated access token")
+	}
+	return &oidc.IDToken{Subject: m.polledUsername}, nil
+}
+
+// Email returns the subject from an ID token.
+func (m *MockOAuthAuthenticator) Email(idToken *oidc.IDToken) (string, error) {
+	return idToken.Subject, nil
+}
+
+// UpdateIdentity is a no-op mock.
+func (m *MockOAuthAuthenticator) UpdateIdentity(ctx context.Context, email string, token *oauth2.Token) error {
+	return nil
+}
+
+// MintSessionToken creates an unsigned session token with the email provided.
+func (m *MockOAuthAuthenticator) MintSessionToken(email string) (string, error) {
+	return newSessionToken(m.c, email, ""), nil
+}
+
+// AuthenticateBrowserSession always returns an error.
+func (m *MockOAuthAuthenticator) AuthenticateBrowserSession(ctx context.Context, w http.ResponseWriter, req *http.Request) (context.Context, error) {
 	return ctx, errors.New("authentication failed")
 }
 
-// NewUserSessionLogin returns a login provider than be used with Juju Dial Opts
-// to define how login will take place. In this case we login using a session token
-// that the JIMM server should verify with the same test secret.
-func NewUserSessionLogin(c SimpleTester, username string) api.LoginProvider {
+// newSessionToken returns a serialised JWT that can be used in tests.
+// Tests using a mock authenticator can provide an empty signatureSecret
+// while integration tests must provide the same secret used when verifying JWTs.
+func newSessionToken(c SimpleTester, username string, signatureSecret string) string {
 	email := convertUsernameToEmail(username)
 	token, err := jwt.NewBuilder().
 		Subject(email).
@@ -105,13 +176,23 @@ func NewUserSessionLogin(c SimpleTester, username string) api.LoginProvider {
 	if err != nil {
 		c.Fatalf("failed to generate test session token")
 	}
-
-	freshToken, err := jwt.Sign(token, jwt.WithKey(jwa.HS256, []byte(JWTTestSecret)))
+	var serialisedToken []byte
+	if signatureSecret != "" {
+		serialisedToken, err = jwt.Sign(token, jwt.WithKey(jwa.HS256, []byte(JWTTestSecret)))
+	} else {
+		serialisedToken, err = jwt.NewSerializer().Serialize(token)
+	}
 	if err != nil {
-		c.Fatalf("failed to sign test session token")
+		c.Fatalf("failed to sign/serialise token")
 	}
+	return base64.StdEncoding.EncodeToString(serialisedToken)
+}
 
-	b64Token := base64.StdEncoding.EncodeToString(freshToken)
+// NewUserSessionLogin returns a login provider than be used with Juju Dial Opts
+// to define how login will take place. In this case we login using a session token
+// that the JIMM server should verify with the same test secret.
+func NewUserSessionLogin(c SimpleTester, username string) api.LoginProvider {
+	b64Token := newSessionToken(c, username, JWTTestSecret)
 	return api.NewSessionTokenLoginProvider(b64Token, nil, nil)
 }
 

internal/jimmtest/jimm.go

@@ -22,7 +22,7 @@ func NewTestJimmParams(t Tester) jimm.Params {
 			Scopes:              []string{oidc.ScopeOpenID, "profile", "email"},
 			SessionTokenExpiry:  time.Duration(time.Hour),
 			SessionCookieMaxAge: 60,
-			JWTSessionKey:       "test-secret",
+			JWTSessionKey:       JWTTestSecret,
 		},
 		DashboardFinalRedirectURL: "dashboard-url",
 	}

internal/jimmtest/suite.go

@@ -60,8 +60,9 @@ type JIMMSuite struct {
 	COFGAClient *cofga.Client
 	COFGAParams *cofga.OpenFGAParams
 
-	Server *httptest.Server
-	cancel context.CancelFunc
+	Server         *httptest.Server
+	cancel         context.CancelFunc
+	deviceFlowChan chan string
 }
 
 func (s *JIMMSuite) SetUpTest(c *gc.C) {
@@ -85,8 +86,9 @@ func (s *JIMMSuite) SetUpTest(c *gc.C) {
 	ctx, cancel := context.WithCancel(context.Background())
 	s.cancel = cancel
 
-	// Note that the secret key here must match what is used in tests.
-	s.JIMM.OAuthAuthenticator = NewMockOAuthAuthenticator(JWTTestSecret)
+	s.deviceFlowChan = make(chan string, 1)
+	authenticator := NewMockOAuthAuthenticator(c, s.deviceFlowChan)
+	s.JIMM.OAuthAuthenticator = &authenticator
 
 	err = s.JIMM.Database.Migrate(ctx, false)
 	c.Assert(err, gc.Equals, nil)
@@ -247,6 +249,14 @@ func (s *JIMMSuite) AddModel(c *gc.C, owner names.UserTag, name string, cloud na
 	return names.NewModelTag(mi.UUID)
 }
 
+// EnableDeviceFlow allows a test to use the device flow.
+// Call this non-blocking function before login to ensure the device flow won't block.
+//
+// This is necessary as the mock authenticator simulates polling an external OIDC server.
+func (s *JIMMSuite) EnableDeviceFlow(username string) {
+	s.deviceFlowChan <- username
+}
+
 // A JujuSuite is a suite that intialises a JIMM and adds the testing juju
 // controller.
 type JujuSuite struct {

internal/jujuapi/admin_test.go

@@ -255,11 +255,11 @@ func (s *adminSuite) TestDeviceLogin(c *gc.C) {
 	// Test no token present
 	var loginResult jujuparams.LoginResult
 	err = conn.APICall("Admin", 4, "", "LoginWithSessionToken", nil, &loginResult)
-	c.Assert(err, gc.ErrorMatches, "authentication failed, no token presented.*")
+	c.Assert(err, gc.ErrorMatches, "no token presented.*")
 
 	// Test token not base64 encoded
 	err = conn.APICall("Admin", 4, "", "LoginWithSessionToken", params.LoginWithSessionTokenRequest{SessionToken: string(decodedToken)}, &loginResult)
-	c.Assert(err, gc.ErrorMatches, "authentication failed, failed to decode token.*")
+	c.Assert(err, gc.ErrorMatches, "failed to decode token.*")
 
 	// Test token base64 encoded passes authentication
 	err = conn.APICall("Admin", 4, "", "LoginWithSessionToken", params.LoginWithSessionTokenRequest{SessionToken: sessionTokenResp.SessionToken}, &loginResult)