CSS-7081 Add OAuth-specific methods to secrets store

internal/cmdtest/jimmsuite.go

@@ -100,7 +100,7 @@ func (s *JimmCmdSuite) SetUpTest(c *gc.C) {
 	err = s.Service.StartJWKSRotator(ctx, time.NewTicker(time.Hour).C, time.Now().UTC().AddDate(0, 3, 0))
 	c.Assert(err, gc.Equals, nil)
 
-	err = s.Service.CheckOrGenerateOAuthKey(ctx)
+	err = s.JIMM.GetCredentialStore().PutOAuthKey(ctx, []byte(jimmtest.JWTTestSecret))
 	c.Assert(err, gc.Equals, nil)
 
 	s.HTTP.StartTLS()

internal/db/secrets.go

@@ -283,6 +283,18 @@ func (d *Database) PutJWKSExpiry(ctx context.Context, expiry time.Time) error {
 	return d.UpsertSecret(ctx, &secret)
 }
 
+// CleanupJWKS removes all secrets associated with OAuth.
+func (d *Database) CleanupOAuth(ctx context.Context) error {
+	const op = errors.Op("database.CleanupOAuth")
+	secret := dbmodel.NewSecret(oauthKind, oauthKeyTag, nil)
+	err := d.DeleteSecret(ctx, &secret)
+	if err != nil {
+		zapctx.Error(ctx, "failed to cleanup OAUth key", zap.Error(err))
+		return errors.E(op, err, "failed to cleanup OAUth key")
+	}
+	return nil
+}
+
 // GetOAuthKey returns the current HS256 (symmetric) key used to sign OAuth session tokens.
 func (d *Database) GetOAuthKey(ctx context.Context) ([]byte, error) {
 	const op = errors.Op("database.GetOAuthKey")

internal/jimm/cloudcredential_test.go

@@ -1771,6 +1771,10 @@ func (s testCloudCredentialAttributeStore) CleanupJWKS(ctx context.Context) erro
 	return errors.E(errors.CodeNotImplemented)
 }
 
+func (s testCloudCredentialAttributeStore) CleanupOAuth(ctx context.Context) error {
+	return errors.E(errors.CodeNotImplemented)
+}
+
 func (s testCloudCredentialAttributeStore) GetOAuthKey(ctx context.Context) ([]byte, error) {
 	return nil, errors.E(errors.CodeNotImplemented)
 }

internal/jimm/credentials/credentials.go

@@ -52,6 +52,9 @@ type CredentialStore interface {
 	// PutJWKSPrivateKey persists the private key associated with the current JWKS within the store.
 	PutJWKSPrivateKey(ctx context.Context, pem []byte) error
 
+	// CleanupJWKS removes all secrets associated with OAuth.
+	CleanupOAuth(ctx context.Context) error
+
 	// PutJWKSExpiry sets the expiry time for the current JWKS within the store.
 	PutJWKSExpiry(ctx context.Context, expiry time.Time) error
 

internal/jimm/jimm.go

@@ -173,7 +173,7 @@ type OAuthAuthenticator interface {
 	VerifyClientCredentials(ctx context.Context, clientID string, clientSecret string) error
 }
 
-// GetCredentialStore return the OpenFGA client used by JIMM.
+// GetCredentialStore returns the credential store used by JIMM.
 func (j *JIMM) GetCredentialStore() credentials.CredentialStore {
 	return j.CredentialStore
 }

internal/jimmhttp/auth_handler_test.go

@@ -4,14 +4,12 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"math/rand"
 	"net"
 	"net/http"
 	"net/http/cookiejar"
 	"net/http/httptest"
 	"net/url"
 	"regexp"
-	"strconv"
 	"testing"
 	"time"
 
@@ -42,17 +40,14 @@ func setupDbAndSessionStore(c *qt.C) (*db.Database, *pgstore.PGStore) {
 }
 
 func setupTestServer(c *qt.C, dashboardURL string, db *db.Database, sessionStore *pgstore.PGStore) *httptest.Server {
+	// Find a random free TCP port.
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	c.Assert(err, qt.IsNil)
+	port := fmt.Sprintf("%d", listener.Addr().(*net.TCPAddr).Port)
+
 	// Create unstarted server to enable auth service
 	s := httptest.NewUnstartedServer(nil)
-	// Setup random port listener
-	minPort := 30000
-	maxPort := 50000
-
-	port := strconv.Itoa(rand.Intn(maxPort-minPort+1) + minPort)
-	l, err := net.Listen("tcp", "localhost:"+port)
-	c.Assert(err, qt.IsNil)
-	// Set the listener with a random port
-	s.Listener = l
+	s.Listener = listener
 
 	// Remember redirect url to check it matches after test server starts
 	redirectURL := "http://127.0.0.1:" + port + "/callback"

internal/jimmtest/auth.go

@@ -18,8 +18,8 @@ import (
 	"github.com/canonical/jimm/internal/openfga"
 )
 
-var (
-	jwtTestSecret = "test-secret"
+const (
+	JWTTestSecret = "test-secret"
 )
 
 // An Authenticator is an implementation of jimm.Authenticator that returns
@@ -59,7 +59,7 @@ func NewUserSessionLogin(username string) api.LoginProvider {
 		panic("failed to generate test session token")
 	}
 
-	freshToken, err := jwt.Sign(token, jwt.WithKey(jwa.HS256, []byte(jwtTestSecret)))
+	freshToken, err := jwt.Sign(token, jwt.WithKey(jwa.HS256, []byte(JWTTestSecret)))
 	if err != nil {
 		panic("failed to sign test session token")
 	}

internal/jimmtest/store.go

@@ -33,7 +33,7 @@ type InMemoryCredentialStore struct {
 // with some secrets/keys being populated.
 func NewInMemoryCredentialStore() *InMemoryCredentialStore {
 	return &InMemoryCredentialStore{
-		oauthKey: []byte("secret-oauth-key"),
+		oauthKey: []byte(JWTTestSecret),
 	}
 }
 
@@ -156,6 +156,15 @@ func (s *InMemoryCredentialStore) GetJWKSExpiry(ctx context.Context) (time.Time,
 	return s.expiry, nil
 }
 
+// CleanupJWKS removes all secrets associated with OAuth.
+func (s *InMemoryCredentialStore) CleanupOAuth(ctx context.Context) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.oauthKey = nil
+	return nil
+}
+
 // PutJWKS puts a generated RS256[4096 bit] JWKS without x5c or x5t into the credential store.
 func (s *InMemoryCredentialStore) PutJWKS(ctx context.Context, jwks jwk.Set) error {
 	s.mu.Lock()

internal/jimmtest/suite.go

@@ -84,7 +84,7 @@ func (s *JIMMSuite) SetUpTest(c *gc.C) {
 	s.cancel = cancel
 
 	// Note that the secret key here must match what is used in tests.
-	s.JIMM.OAuthAuthenticator = NewMockOAuthAuthenticator(jwtTestSecret)
+	s.JIMM.OAuthAuthenticator = NewMockOAuthAuthenticator(JWTTestSecret)
 
 	err = s.JIMM.Database.Migrate(ctx, false)
 	c.Assert(err, gc.Equals, nil)

internal/vault/vault.go

@@ -59,7 +59,7 @@ func (s *VaultStore) Get(ctx context.Context, tag names.CloudCredentialTag) (map
 		return nil, errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.path(tag))
+	secret, err := client.Logical().ReadWithContext(ctx, s.path(tag))
 	if err != nil {
 		return nil, errors.E(op, err)
 	}
@@ -96,7 +96,7 @@ func (s *VaultStore) Put(ctx context.Context, tag names.CloudCredentialTag, attr
 	for k, v := range attr {
 		data[k] = v
 	}
-	_, err = client.Logical().Write(s.path(tag), data)
+	_, err = client.Logical().WriteWithContext(ctx, s.path(tag), data)
 	if err != nil {
 		return errors.E(op, err)
 	}
@@ -112,7 +112,7 @@ func (s *VaultStore) delete(ctx context.Context, tag names.CloudCredentialTag) e
 	if err != nil {
 		return errors.E(op, err)
 	}
-	_, err = client.Logical().Delete(s.path(tag))
+	_, err = client.Logical().DeleteWithContext(ctx, s.path(tag))
 	if rerr, ok := err.(*api.ResponseError); ok && rerr.StatusCode == http.StatusNotFound {
 		// Ignore the error if attempting to delete something that isn't there.
 		err = nil
@@ -133,7 +133,7 @@ func (s *VaultStore) GetControllerCredentials(ctx context.Context, controllerNam
 		return "", "", errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.controllerCredentialsPath(controllerName))
+	secret, err := client.Logical().ReadWithContext(ctx, s.controllerCredentialsPath(controllerName))
 	if err != nil {
 		return "", "", errors.E(op, err)
 	}
@@ -169,7 +169,7 @@ func (s *VaultStore) PutControllerCredentials(ctx context.Context, controllerNam
 		usernameKey: username,
 		passwordKey: password,
 	}
-	_, err = client.Logical().Write(s.controllerCredentialsPath(controllerName), data)
+	_, err = client.Logical().WriteWithContext(ctx, s.controllerCredentialsPath(controllerName), data)
 	if err != nil {
 		return errors.E(op, err)
 	}
@@ -186,9 +186,9 @@ func (s *VaultStore) CleanupJWKS(ctx context.Context) error {
 	}
 	// Vault does not return errors on deletion requests where
 	// the secret does not exist. As such we just return the last known error.
-	client.Logical().Delete(s.getJWKSExpiryPath())
-	client.Logical().Delete(s.getJWKSPath())
-	if _, err = client.Logical().Delete(s.getJWKSPrivateKeyPath()); err != nil {
+	client.Logical().DeleteWithContext(ctx, s.getJWKSExpiryPath())
+	client.Logical().DeleteWithContext(ctx, s.getJWKSPath())
+	if _, err = client.Logical().DeleteWithContext(ctx, s.getJWKSPrivateKeyPath()); err != nil {
 		return errors.E(op, err)
 	}
 	return nil
@@ -203,7 +203,7 @@ func (s *VaultStore) GetJWKS(ctx context.Context) (jwk.Set, error) {
 		return nil, errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.getJWKSPath())
+	secret, err := client.Logical().ReadWithContext(ctx, s.getJWKSPath())
 	if err != nil {
 		return nil, errors.E(op, err)
 	}
@@ -239,7 +239,7 @@ func (s *VaultStore) GetJWKSPrivateKey(ctx context.Context) ([]byte, error) {
 		return nil, errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.getJWKSPrivateKeyPath())
+	secret, err := client.Logical().ReadWithContext(ctx, s.getJWKSPrivateKeyPath())
 	if err != nil {
 		return nil, errors.E(op, err)
 	}
@@ -269,7 +269,7 @@ func (s *VaultStore) GetJWKSExpiry(ctx context.Context) (time.Time, error) {
 		return now, errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.getJWKSExpiryPath())
+	secret, err := client.Logical().ReadWithContext(ctx, s.getJWKSExpiryPath())
 	if err != nil {
 		return now, errors.E(op, err)
 	}
@@ -310,7 +310,8 @@ func (s *VaultStore) PutJWKS(ctx context.Context, jwks jwk.Set) error {
 		return errors.E(op, err)
 	}
 
-	_, err = client.Logical().WriteBytes(
+	_, err = client.Logical().WriteBytesWithContext(
+		ctx,
 		// We persist in a similar folder to the controller credentials, but sub-route
 		// to .well-known for further extensions and mental clarity within our vault.
 		s.getJWKSPath(),
@@ -332,7 +333,8 @@ func (s *VaultStore) PutJWKSPrivateKey(ctx context.Context, pem []byte) error {
 		return errors.E(op, err)
 	}
 
-	if _, err := client.Logical().Write(
+	if _, err := client.Logical().WriteWithContext(
+		ctx,
 		// We persist in a similar folder to the controller credentials, but sub-route
 		// to .well-known for further extensions and mental clarity within our vault.
 		s.getJWKSPrivateKeyPath(),
@@ -352,7 +354,8 @@ func (s *VaultStore) PutJWKSExpiry(ctx context.Context, expiry time.Time) error
 		return errors.E(op, err)
 	}
 
-	if _, err := client.Logical().Write(
+	if _, err := client.Logical().WriteWithContext(
+		ctx,
 		s.getJWKSExpiryPath(),
 		map[string]interface{}{
 			"jwks-expiry": expiry,
@@ -385,6 +388,23 @@ func (s *VaultStore) getJWKSExpiryPath() string {
 	return path.Join(s.getWellKnownPath(), "jwks-expiry")
 }
 
+// CleanupOAuth removes all secrets associated with OAuth.
+func (s *VaultStore) CleanupOAuth(ctx context.Context) error {
+	const op = errors.Op("vault.CleanupOAuth")
+
+	client, err := s.client(ctx)
+	if err != nil {
+		return errors.E(op, err)
+	}
+
+	// Vault does not return errors on deletion requests where
+	// the secret does not exist.
+	if _, err := client.Logical().DeleteWithContext(ctx, s.GetOAuthKeyPath()); err != nil {
+		return errors.E(op, err)
+	}
+	return nil
+}
+
 // GetOAuthKey returns the current HS256 (symmetric) key used to sign OAuth session tokens.
 func (s *VaultStore) GetOAuthKey(ctx context.Context) ([]byte, error) {
 	const op = errors.Op("vault.GetOAuthKey")
@@ -394,7 +414,7 @@ func (s *VaultStore) GetOAuthKey(ctx context.Context) ([]byte, error) {
 		return nil, errors.E(op, err)
 	}
 
-	secret, err := client.Logical().Read(s.GetOAuthKeyPath())
+	secret, err := client.Logical().ReadWithContext(ctx, s.GetOAuthKeyPath())
 	if err != nil {
 		return nil, errors.E(op, err)
 	}
@@ -405,7 +425,14 @@ func (s *VaultStore) GetOAuthKey(ctx context.Context) ([]byte, error) {
 		return nil, errors.E(op, errors.CodeNotFound, msg)
 	}
 
-	keyPemB64 := secret.Data["key"].(string)
+	raw := secret.Data["key"]
+	if secret.Data["key"] == nil {
+		msg := "nil OAuth key data"
+		zapctx.Debug(ctx, msg)
+		return nil, errors.E(op, errors.CodeNotFound, msg)
+	}
+
+	keyPemB64 := raw.(string)
 
 	keyPem, err := base64.StdEncoding.DecodeString(keyPemB64)
 	if err != nil {
@@ -424,7 +451,8 @@ func (s *VaultStore) PutOAuthKey(ctx context.Context, raw []byte) error {
 		return errors.E(op, err)
 	}
 
-	if _, err := client.Logical().Write(
+	if _, err := client.Logical().WriteWithContext(
+		ctx,
 		s.GetOAuthKeyPath(),
 		map[string]interface{}{"key": raw},
 	); err != nil {
@@ -448,7 +476,7 @@ func (s *VaultStore) deleteControllerCredentials(ctx context.Context, controller
 	if err != nil {
 		return errors.E(op, err)
 	}
-	_, err = client.Logical().Delete(s.controllerCredentialsPath(controllerName))
+	_, err = client.Logical().DeleteWithContext(ctx, s.controllerCredentialsPath(controllerName))
 	if rerr, ok := err.(*api.ResponseError); ok && rerr.StatusCode == http.StatusNotFound {
 		// Ignore the error if attempting to delete something that isn't there.
 		err = nil
@@ -471,7 +499,7 @@ func (s *VaultStore) client(ctx context.Context) (*api.Client, error) {
 		return s.client_, nil
 	}
 
-	secret, err := s.Client.Logical().Write(s.AuthPath, s.AuthSecret)
+	secret, err := s.Client.Logical().WriteWithContext(ctx, s.AuthPath, s.AuthSecret)
 	if err != nil {
 		return nil, errors.E(op, err)
 	}

internal/vault/vault_test.go

@@ -200,11 +200,27 @@ func TestGetAndPutOAuthKey(t *testing.T) {
 	c.Assert(retrievedKey, qt.DeepEquals, key)
 }
 
+func TestGetOAuthKeyFailsIfDataIsNil(t *testing.T) {
+	c := qt.New(t)
+	ctx := context.Background()
+	store := newStore(c)
+
+	err := store.PutOAuthKey(ctx, nil)
+	c.Assert(err, qt.IsNil)
+
+	retrieved, err := store.GetOAuthKey(ctx)
+	c.Assert(err, qt.ErrorMatches, "nil OAuth key data")
+	c.Assert(retrieved, qt.IsNil)
+}
+
 func TestGetOAuthKeyFailsIfNotFound(t *testing.T) {
 	c := qt.New(t)
 	ctx := context.Background()
 	store := newStore(c)
 
+	err := store.CleanupOAuth(ctx)
+	c.Assert(err, qt.IsNil)
+
 	retrieved, err := store.GetOAuthKey(ctx)
 	c.Assert(err, qt.ErrorMatches, "no OAuth key exists")
 	c.Assert(retrieved, qt.IsNil)