Repository: google/exposure-notifications-server
Pull request #1015 (2020-09-21 16:49:39)

Add server listener for jwks, fix context propagation
**Release Note**

<!-- Enter your extended release note in the below block. If the PR requires
additional action from users switching to the new release, include the string
"action required". If no release note is required, write "NONE". -->

```release-note
Add server listener for jwks
Fix context propagation on jwk service
```

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1015#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1015#pullrequestreview-492808675" title="Approved">jeremyfaller</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/1015#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/cmd/jwks-updater/main.go b/cmd/jwks-updater/main.go
index cb6fea5..e1c9c8b 100644
--- a/cmd/jwks-updater/main.go
+++ b/cmd/jwks-updater/main.go
@@ -27,6 +27,7 @@ import (
 	"github.com/google/exposure-notifications-server/pkg/jwks"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	_ "github.com/google/exposure-notifications-server/pkg/observability"
+	"github.com/google/exposure-notifications-server/pkg/server"
 	"github.com/sethvargo/go-signalcontext"
 )
 
@@ -49,6 +50,8 @@ func main() {
 }
 
 func realMain(ctx context.Context) error {
+	logger := logging.FromContext(ctx)
+
 	var config jwks.Config
 	env, err := setup.Setup(ctx, &config)
 	if err != nil {
@@ -56,6 +59,16 @@ func realMain(ctx context.Context) error {
 	}
 	defer env.Close(ctx)
 
-	env.JwksManager().UpdateAll()
-	return nil
+	jwksServer, err := jwks.NewServer(&config, env)
+	if err != nil {
+		return fmt.Errorf("jwks.NewServer: %w", err)
+	}
+
+	srv, err := server.New(config.Port)
+	if err != nil {
+		return fmt.Errorf("server.New: %w", err)
+	}
+	logger.Infof("listening on :%s", config.Port)
+
+	return srv.ServeHTTPHandler(ctx, jwksServer.Routes(ctx))
 }
diff --git a/internal/serverenv/env.go b/internal/serverenv/env.go
index 27541e8..caf86f8 100644
--- a/internal/serverenv/env.go
+++ b/internal/serverenv/env.go
@@ -24,7 +24,6 @@ import (
 	"github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/metrics"
 	"github.com/google/exposure-notifications-server/internal/storage"
-	"github.com/google/exposure-notifications-server/pkg/jwks"
 	"github.com/google/exposure-notifications-server/pkg/keys"
 	"github.com/google/exposure-notifications-server/pkg/observability"
 	"github.com/google/exposure-notifications-server/pkg/secrets"
@@ -42,7 +41,6 @@ type ServerEnv struct {
 	keyManager            keys.KeyManager
 	secretManager         secrets.SecretManager
 	observabilityExporter observability.Exporter
-	jwks                  jwks.Manager
 }
 
 // Option defines function types to modify the ServerEnv on creation.
@@ -120,14 +118,6 @@ func WithObservabilityExporter(oe observability.Exporter) Option {
 	}
 }
 
-// WithJwksManager creates an Option to install the JWKS manager.
-func WithJwksManager(mgr jwks.Manager) Option {
-	return func(s *ServerEnv) *ServerEnv {
-		s.jwks = mgr
-		return s
-	}
-}
-
 func (s *ServerEnv) SecretManager() secrets.SecretManager {
 	return s.secretManager
 }
@@ -156,10 +146,6 @@ func (s *ServerEnv) GetKeyManager() keys.KeyManager {
 	return s.keyManager
 }
 
-func (s *ServerEnv) JwksManager() jwks.Manager {
-	return s.jwks
-}
-
 // GetSignerForKey returns the crypto.Singer implementation to use based on the installed KeyManager.
 // If there is no KeyManager installed, this returns an error.
 func (s *ServerEnv) GetSignerForKey(ctx context.Context, keyName string) (crypto.Signer, error) {
diff --git a/internal/setup/setup.go b/internal/setup/setup.go
index 1855f5d..f623ca7 100644
--- a/internal/setup/setup.go
+++ b/internal/setup/setup.go
@@ -24,7 +24,6 @@ import (
 	"github.com/google/exposure-notifications-server/internal/metrics"
 	"github.com/google/exposure-notifications-server/internal/serverenv"
 	"github.com/google/exposure-notifications-server/internal/storage"
-	"github.com/google/exposure-notifications-server/pkg/jwks"
 	"github.com/google/exposure-notifications-server/pkg/keys"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/google/exposure-notifications-server/pkg/observability"
@@ -68,12 +67,6 @@ type SecretManagerConfigProvider interface {
 	SecretManagerConfig() *secrets.Config
 }
 
-// JWKSManagerConfigProvider signals that the config knows how to configure the
-// JWKS manager.
-type JWKSManagerConfigProvider interface {
-	JWKSManagerConfig() *jwks.Config
-}
-
 // Setup runs common initialization code for all servers. See SetupWith.
 func Setup(ctx context.Context, config interface{}) (*serverenv.ServerEnv, error) {
 	return SetupWith(ctx, config, envconfig.OsLookuper())
@@ -243,24 +236,6 @@ func SetupWith(ctx context.Context, config interface{}, l envconfig.Lookuper) (*
 
 			logger.Infow("authorizedapp", "config", aaConfig)
 		}
-
-		// Start the JWKS key downloader.
-		if provider, ok := config.(JWKSManagerConfigProvider); ok {
-			logger.Info("configuring jwks manager")
-
-			jwksConfig := provider.JWKSManagerConfig()
-			jwks, err := jwks.NewFromEnv(ctx, logger, db, jwksConfig)
-			if err != nil {
-				// Ensure the db is closed on an error.
-				defer db.Close(ctx)
-				return nil, fmt.Errorf("unable to start jwks manager: %w", err)
-			}
-
-			// Update the serverEnv setup.
-			serverEnvOpts = append(serverEnvOpts, serverenv.WithJwksManager(jwks))
-
-			logger.Infow("jwks", "config", jwksConfig)
-		}
 	}
 
 	return serverenv.New(ctx, serverEnvOpts...), nil
diff --git a/pkg/jwks/jwks.go b/pkg/jwks/jwks.go
index 7a800d5..75b20a0 100644
--- a/pkg/jwks/jwks.go
+++ b/pkg/jwks/jwks.go
@@ -30,39 +30,51 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/setup"
 	hadb "github.com/google/exposure-notifications-server/internal/verification/database"
 	"github.com/google/exposure-notifications-server/internal/verification/model"
+	"github.com/google/exposure-notifications-server/pkg/logging"
+	"github.com/hashicorp/go-multierror"
 	"github.com/rakutentech/jwk-go/jwk"
 	"go.uber.org/zap"
 )
 
+var _ setup.DatabaseConfigProvider = (*Config)(nil)
+
 type Config struct {
 	Database database.Config
+
+	Port string `env:"PORT, default=8080"`
 }
 
 func (c *Config) DatabaseConfig() *database.Config {
 	return &c.Database
 }
 
-func (c *Config) JWKSManagerConfig() *Config {
-	return c
-}
-
 // Manager handles updating all HealthAuthorities if they've specified a JWKS
 // URI.
 type Manager struct {
-	ctxt   context.Context
-	logger *zap.SugaredLogger
 	db     *database.DB
+	logger *zap.SugaredLogger
+}
+
+// NewManager creates a new Manager.
+func NewManager(ctx context.Context, db *database.DB) (*Manager, error) {
+	logger := logging.FromContext(ctx).Named("jwks")
+
+	return &Manager{
+		logger: logger,
+		db:     db,
+	}, nil
 }
 
 // getKeys reads the keys for a single HealthAuthority from its jwks server.
-func (mgr Manager) getKeys(ha *model.HealthAuthority) ([]byte, error) {
+func (mgr *Manager) getKeys(ctx context.Context, ha *model.HealthAuthority) ([]byte, error) {
 	if len(ha.JwksURI) == 0 {
 		return nil, nil
 	}
 
-	reqCtxt, done := context.WithTimeout(mgr.ctxt, 5*time.Second)
+	reqCtxt, done := context.WithTimeout(ctx, 5*time.Second)
 	defer done()
 	req, err := http.NewRequestWithContext(reqCtxt, "GET", ha.JwksURI, nil)
 	if err != nil {
@@ -157,13 +169,18 @@ func findKeyMods(ha *model.HealthAuthority, rxKeys []string) (deadKeys []int, ne
 }
 
 // updateHA updates HealthAuthority's keys.
-func (mgr *Manager) updateHA(haDB *hadb.HealthAuthorityDB, ha *model.HealthAuthority) error {
+func (mgr *Manager) updateHA(ctx context.Context, ha *model.HealthAuthority) error {
+	logger := mgr.logger.With("health_authority_name", ha.Name, "health_authority_id", ha.ID)
+
 	if len(ha.JwksURI) == 0 {
-		mgr.logger.Infof("Skipping JWKS, HealthAuthority: %q No URI specified.", ha.Name)
+		logger.Infow("skipping jwks, no URI specified")
 		return nil
 	}
 
-	resp, err := mgr.getKeys(ha)
+	// Create the hadb once to save allocations
+	haDB := hadb.New(mgr.db)
+
+	resp, err := mgr.getKeys(ctx, ha)
 	if err != nil {
 		return err
 	}
@@ -184,7 +201,7 @@ func (mgr *Manager) updateHA(haDB *hadb.HealthAuthorityDB, ha *model.HealthAutho
 	for _, i := range deadKeys {
 		hak := ha.Keys[i]
 		hak.Revoke()
-		if err := haDB.UpdateHealthAuthorityKey(mgr.ctxt, hak); err != nil {
+		if err := haDB.UpdateHealthAuthorityKey(ctx, hak); err != nil {
 			return fmt.Errorf("error updating key: %w", err)
 		}
 	}
@@ -197,53 +214,55 @@ func (mgr *Manager) updateHA(haDB *hadb.HealthAuthorityDB, ha *model.HealthAutho
 			From:         time.Now(),
 			PublicKeyPEM: key,
 		}
-		if err := haDB.AddHealthAuthorityKey(mgr.ctxt, ha, hak); err != nil {
+		if err := haDB.AddHealthAuthorityKey(ctx, ha, hak); err != nil {
 			return fmt.Errorf("error adding key: %w", err)
 		}
 		ha.Keys = append(ha.Keys, hak)
 	}
 
 	// And save the HealthAuthority.
-	haDB.UpdateHealthAuthority(mgr.ctxt, ha)
+	haDB.UpdateHealthAuthority(ctx, ha)
 
-	mgr.logger.Infof("Updated JWKS HealthAuthority: %q URI (%v): %d new, %d deleted", ha.Name, ha.JwksURI, len(newKeys), len(deadKeys))
+	logger.Infow("updated jwks",
+		"uri", ha.JwksURI,
+		"new", len(newKeys),
+		"deleted", len(deadKeys))
 
 	return nil
 }
 
 // UpdateAll reads the JWKS keys for all HealthAuthorities.
-func (mgr Manager) UpdateAll() {
-	mgr.logger.Info("Starting JWKS update")
+func (mgr *Manager) UpdateAll(ctx context.Context) error {
+	mgr.logger.Info("starting jwks update")
 
 	haDB := hadb.New(mgr.db)
-	healthAuthorities, err := haDB.ListAllHealthAuthoritiesWithoutKeys(mgr.ctxt)
+	healthAuthorities, err := haDB.ListAllHealthAuthoritiesWithoutKeys(ctx)
 	if err != nil {
-		mgr.logger.Errorf("error querying db %v", err)
-		return
+		return fmt.Errorf("failed to query db: %w", err)
 	}
 
+	var merr *multierror.Error
+	var merrLock sync.Mutex
+
 	var wg sync.WaitGroup
 	for _, ha := range healthAuthorities {
 		wg.Add(1)
 		go func(ha *model.HealthAuthority) {
 			defer wg.Done()
-			err := mgr.updateHA(haDB, ha)
+			err := mgr.updateHA(ctx, ha)
 			if err != nil {
-				mgr.logger.Errorf("error running JWKS HealthAuthority: %q error: %v", ha.Name, err)
-				return
+				merrLock.Lock()
+				merr = multierror.Append(merr, fmt.Errorf("failed to processes %v: %w", ha.Name, err))
+				merrLock.Unlock()
 			}
 		}(ha)
 	}
 	wg.Wait()
 
-	mgr.logger.Info("Done JWKS update")
-}
+	mgr.logger.Info("finished jwks update")
 
-// NewFromEnv creates a new JWKS configuration manager.
-func NewFromEnv(ctxt context.Context, logger *zap.SugaredLogger, db *database.DB, cfg *Config) (Manager, error) {
-	return Manager{
-		ctxt:   ctxt,
-		logger: logger,
-		db:     db,
-	}, nil
+	if err := merr.ErrorOrNil(); err != nil {
+		return fmt.Errorf("failed to update all: %w", err)
+	}
+	return nil
 }
diff --git a/pkg/jwks/server.go b/pkg/jwks/server.go
new file mode 100644
index 0000000..ca937f5
--- /dev/null
+++ b/pkg/jwks/server.go
@@ -0,0 +1,66 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package jwks
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	"github.com/google/exposure-notifications-server/internal/serverenv"
+	"github.com/google/exposure-notifications-server/pkg/logging"
+	"github.com/google/exposure-notifications-server/pkg/server"
+)
+
+// Server is the debugger server.
+type Server struct {
+	config  *Config
+	manager *Manager
+	env     *serverenv.ServerEnv
+}
+
+// NewServer makes a new debugger server.
+func NewServer(config *Config, env *serverenv.ServerEnv) (*Server, error) {
+	return &Server{
+		config: config,
+		env:    env,
+	}, nil
+}
+
+func (s *Server) Routes(ctx context.Context) *http.ServeMux {
+	mux := http.NewServeMux()
+	mux.Handle("/health", server.HandleHealthz(ctx))
+	mux.Handle("/", s.handleUpdateAll(ctx))
+
+	return mux
+}
+
+func (s *Server) handleUpdateAll(ctx context.Context) http.Handler {
+	logger := logging.FromContext(ctx).Named("handleUpdateAll")
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
+		if err := s.manager.UpdateAll(ctx); err != nil {
+			logger.Error(err)
+			w.WriteHeader(http.StatusInternalServerError)
+			fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
+			return
+		}
+
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, http.StatusText(http.StatusOK))
+	})
+}


Test patch:
diff --git a/pkg/jwks/jwks_test.go b/pkg/jwks/jwks_test.go
index 8f441ca..48547f5 100644
--- a/pkg/jwks/jwks_test.go
+++ b/pkg/jwks/jwks_test.go
@@ -27,7 +27,6 @@ import (
 	"github.com/google/exposure-notifications-server/internal/database"
 	hadb "github.com/google/exposure-notifications-server/internal/verification/database"
 	"github.com/google/exposure-notifications-server/internal/verification/model"
-	"github.com/google/exposure-notifications-server/pkg/logging"
 )
 
 var key1 = `{"kid":"r2v1","kty":"EC","crv":"P-256","x":"qcMMcLX1Z2afVAzypTMw1g3KN_OcdgvRDwOgpDWiswU","y":"RjK8Hc7pLLO_JADNhwZIxCXjCH95VHuWPoKVaCGkXiA"}`
@@ -127,17 +126,16 @@ func TestUpdateHA(t *testing.T) {
 			defer ts.Close()
 
 			// Set up the test.
-			cfg := &Config{}
+			ctx := context.Background()
 			testDB := database.NewTestDatabase(t)
-			logger := logging.NewLogger(true)
-			mgr, err := NewFromEnv(context.Background(), logger, testDB, cfg)
+			mgr, err := NewManager(ctx, testDB)
 			if err != nil {
 				t.Fatalf("[%d] unexpected error: %v", i, err)
 			}
 			ha := &model.HealthAuthority{JwksURI: ts.URL}
 
 			// Test networking.
-			rxKeys, err := mgr.getKeys(ha)
+			rxKeys, err := mgr.getKeys(ctx, ha)
 			if err != nil {
 				t.Fatalf("[%d] unexpected error: %v", i, err)
 			}
@@ -181,23 +179,23 @@ func TestUpdateHA(t *testing.T) {
 			keys, test.ha.Keys = test.ha.Keys, nil
 			haDB := hadb.New(mgr.db)
 			test.ha.Issuer, test.ha.Audience, test.ha.Name = "ISSUER", "AUDIENCE", "NAME"
-			if err := haDB.AddHealthAuthority(mgr.ctxt, &test.ha); err != nil {
+			if err := haDB.AddHealthAuthority(ctx, &test.ha); err != nil {
 				t.Fatalf("[%d] error adding the HealthAuthority, %v", i, err)
 			}
 			for _, key := range keys {
-				if err := haDB.AddHealthAuthorityKey(mgr.ctxt, &test.ha, key); err != nil {
+				if err := haDB.AddHealthAuthorityKey(ctx, &test.ha, key); err != nil {
 					t.Fatalf("[%d] error adding key: %v", i, err)
 				}
 			}
 			test.ha.Keys = keys
 
 			// Now, run the whole flow for a HealthAuthority.
-			if err := mgr.updateHA(haDB, &test.ha); err != nil {
+			if err := mgr.updateHA(ctx, &test.ha); err != nil {
 				t.Fatalf("[%d] error updating: %v", i, err)
 			}
 
 			// Check the DB.
-			if ha, err := haDB.GetHealthAuthorityByID(mgr.ctxt, test.ha.ID); err != nil {
+			if ha, err := haDB.GetHealthAuthorityByID(ctx, test.ha.ID); err != nil {
 				t.Fatalf("[%d] error retreiving HealthAuthority: %v", i, err)
 			} else {
 				// Check the resultant keys.