Repository: google/exposure-notifications-server
Pull request #1341 (2021-02-04 15:38:27)

Pull logger and request_id from context in jwks service
This updates the jwks service to populate and pull critical information from the request context instead of the initial context, similar to the verifications server.

Yes, a lot of this middleware is similar to the verification server, but it's _just_ slightly different enough that it doesn't make sense to try and unify (yet).

**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
Pull logger and request_id from context in jwks 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/1341#" 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":[]} -->
sethvargo: /assign @whaught 
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/1341#" title="Author self-approved">sethvargo</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/1341#pullrequestreview-583527653" title="Approved">whaught</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,whaught]

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: @sethvargo: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | 12fed0401cd902d0c8645af2b6a48ff724fec693 | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/1341/pull-en-server-release-unit/1357352833631916032) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
sethvargo: Wut... test failure
google-oss-robot: New changes are detected. LGTM label has been removed.
sethvargo: Dunno why the go.mod got sad

Patch:
diff --git a/internal/jwks/handle_update_all.go b/internal/jwks/handle_update_all.go
new file mode 100644
index 0000000..0943d12
--- /dev/null
+++ b/internal/jwks/handle_update_all.go
@@ -0,0 +1,39 @@
+// Copyright 2021 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 (
+	"fmt"
+	"net/http"
+
+	"github.com/google/exposure-notifications-server/pkg/logging"
+)
+
+func (s *Server) handleUpdateAll() http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
+		logger := logging.FromContext(ctx).Named("handleUpdateAll")
+
+		if err := s.manager.UpdateAll(ctx); err != nil {
+			logger.Errorw("failed to update all", "error", err)
+			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+			return
+		}
+
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, http.StatusText(http.StatusOK))
+	})
+}
diff --git a/internal/jwks/jwks.go b/internal/jwks/jwks.go
index 5a32fe6..23322ad 100644
--- a/internal/jwks/jwks.go
+++ b/internal/jwks/jwks.go
@@ -36,27 +36,22 @@ import (
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/hashicorp/go-multierror"
 	"github.com/rakutentech/jwk-go/jwk"
-	"go.uber.org/zap"
 )
 
 // Manager handles updating all HealthAuthorities if they've specified a JWKS
 // URI.
 type Manager struct {
 	db         *database.DB
-	logger     *zap.SugaredLogger
 	cleanupTTL time.Duration
 }
 
 // NewManager creates a new Manager.
-func NewManager(ctx context.Context, db *database.DB, cleanupTTL time.Duration) (*Manager, error) {
-	logger := logging.FromContext(ctx).Named("jwks")
-
+func NewManager(db *database.DB, cleanupTTL time.Duration) (*Manager, error) {
 	if cleanupTTL < 0 {
 		cleanupTTL *= -1
 	}
 
 	return &Manager{
-		logger:     logger,
 		db:         db,
 		cleanupTTL: cleanupTTL,
 	}, nil
@@ -79,8 +74,7 @@ func (mgr *Manager) getKeys(ctx context.Context, ha *model.HealthAuthority) ([]b
 		return nil, fmt.Errorf("creating connection: %w", err)
 	}
 
-	var resp *http.Response
-	resp, err = http.DefaultClient.Do(req)
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("reading connection: %w", err)
 	}
@@ -169,7 +163,9 @@ func findKeyMods(ha *model.HealthAuthority, rxKeys []string) (deadKeys []int, ne
 
 // updateHA updates HealthAuthority's keys.
 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)
+	logger := logging.FromContext(ctx).Named("updateHA").
+		With("health_authority_name", ha.Name).
+		With("health_authority_id", ha.ID)
 
 	if ha.JwksURI == nil || len(*ha.JwksURI) == 0 {
 		logger.Infow("skipping jwks, no URI specified")
@@ -181,9 +177,9 @@ func (mgr *Manager) updateHA(ctx context.Context, ha *model.HealthAuthority) err
 
 	purgeBefore := time.Now().UTC().Add(-1 * mgr.cleanupTTL)
 	if count, err := haDB.PurgeHealthAuthorityKeys(ctx, ha, purgeBefore); err != nil {
-		logger.Errorw("error purging expired health authority keys", "healthAuthorityID", ha.ID, "error", err)
+		logger.Errorw("error purging expired health authority keys", "error", err)
 	} else if count > 0 {
-		logger.Infow("purged health authority keys", "healthAuthorityID", ha.ID, "count", count)
+		logger.Infow("purged health authority keys", "count", count)
 	}
 
 	// Get the keys for the health authority
@@ -242,7 +238,9 @@ func (mgr *Manager) updateHA(ctx context.Context, ha *model.HealthAuthority) err
 
 // UpdateAll reads the JWKS keys for all HealthAuthorities.
 func (mgr *Manager) UpdateAll(ctx context.Context) error {
-	mgr.logger.Info("starting jwks update")
+	logger := logging.FromContext(ctx).Named("UpdateAll")
+	logger.Debug("starting jwks update")
+	defer logger.Debug("finished jwks update")
 
 	haDB := hadb.New(mgr.db)
 	healthAuthorities, err := haDB.ListAllHealthAuthoritiesWithoutKeys(ctx)
@@ -268,8 +266,6 @@ func (mgr *Manager) UpdateAll(ctx context.Context) error {
 	}
 	wg.Wait()
 
-	mgr.logger.Info("finished jwks update")
-
 	if err := merr.ErrorOrNil(); err != nil {
 		return fmt.Errorf("failed to update all: %w", err)
 	}
diff --git a/internal/jwks/server.go b/internal/jwks/server.go
index 11bd82e..462bb47 100644
--- a/internal/jwks/server.go
+++ b/internal/jwks/server.go
@@ -17,62 +17,50 @@ package jwks
 import (
 	"context"
 	"fmt"
-	"net/http"
 
+	"github.com/google/exposure-notifications-server/internal/middleware"
 	"github.com/google/exposure-notifications-server/internal/serverenv"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/google/exposure-notifications-server/pkg/server"
+
 	"github.com/gorilla/mux"
 )
 
-// Server is the debugger server.
+// Server is the 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) {
+// NewServer makes a new server.
+func NewServer(cfg *Config, env *serverenv.ServerEnv) (*Server, error) {
 	if env.Database() == nil {
-		return nil, fmt.Errorf("expected env to have database")
+		return nil, fmt.Errorf("missing database in server env")
 	}
 
-	ctx := context.Background()
-	manager, err := NewManager(ctx, env.Database(), config.KeyCleanupTTL)
+	manager, err := NewManager(env.Database(), cfg.KeyCleanupTTL)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create manager: %w", err)
 	}
 
 	return &Server{
-		config:  config,
+		config:  cfg,
 		manager: manager,
 		env:     env,
 	}, nil
 }
 
+// Routes returns the router for this server.
 func (s *Server) Routes(ctx context.Context) *mux.Router {
+	logger := logging.FromContext(ctx).Named("jwks")
+
 	r := mux.NewRouter()
+	r.Use(middleware.PopulateRequestID())
+	r.Use(middleware.PopulateLogger(logger))
+
 	r.Handle("/health", server.HandleHealthz(ctx))
-	r.Handle("/", s.handleUpdateAll(ctx))
+	r.Handle("/", s.handleUpdateAll())
 
 	return r
 }
-
-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))
-	})
-}
diff --git a/internal/middleware/logger.go b/internal/middleware/logger.go
new file mode 100644
index 0000000..3ce974b
--- /dev/null
+++ b/internal/middleware/logger.go
@@ -0,0 +1,77 @@
+// Copyright 2021 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 middleware
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/google/exposure-notifications-server/pkg/logging"
+
+	"github.com/gorilla/mux"
+	"go.uber.org/zap"
+)
+
+const (
+	// googleCloudTraceHeader is the header with trace data.
+	googleCloudTraceHeader = "X-Cloud-Trace-Context"
+
+	// googleCloudTraceKey is the key in the structured log where trace information
+	// is expected to be present.
+	googleCloudTraceKey = "logging.googleapis.com/trace"
+)
+
+// googleCloudProjectID is the project id, populated by Terraform during service
+// deployment.
+var googleCloudProjectID = os.Getenv("PROJECT_ID")
+
+// PopulateLogger populates the logger onto the context.
+func PopulateLogger(originalLogger *zap.SugaredLogger) mux.MiddlewareFunc {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			logger := originalLogger
+
+			// Only override the logger if it's the default logger. This is only used
+			// for testing and is intentionally a strict object equality check because
+			// the default logger is a global default in the logger package.
+			if existing := logging.FromContext(ctx); existing == logging.DefaultLogger() {
+				logger = existing
+			}
+
+			// If there's a request ID, set that on the logger.
+			if id := RequestIDFromContext(ctx); id != "" {
+				logger = logger.With("request_id", id)
+			}
+
+			// On Google Cloud, extract the trace context and add it to the logger.
+			if v := r.Header.Get(googleCloudTraceHeader); v != "" && googleCloudProjectID != "" {
+				parts := strings.Split(v, "/")
+				if len(parts) > 0 && len(parts[0]) > 0 {
+					val := fmt.Sprintf("projects/%s/traces/%s", googleCloudProjectID, parts[0])
+					logger = logger.With(googleCloudTraceKey, val)
+				}
+			}
+
+			ctx = logging.WithLogger(ctx, logger)
+			r = r.Clone(ctx)
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go
new file mode 100644
index 0000000..42e4351
--- /dev/null
+++ b/internal/middleware/middleware.go
@@ -0,0 +1,20 @@
+// Copyright 2021 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 middleware defines common http middlewares.
+package middleware
+
+// contextKey is a unique type to avoid clashing with other packages that use
+// context's to pass data.
+type contextKey string
diff --git a/internal/middleware/observability.go b/internal/middleware/observability.go
new file mode 100644
index 0000000..c663941
--- /dev/null
+++ b/internal/middleware/observability.go
@@ -0,0 +1,36 @@
+// Copyright 2021 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 middleware
+
+import (
+	"net/http"
+
+	"github.com/google/exposure-notifications-server/pkg/observability"
+
+	"github.com/gorilla/mux"
+)
+
+// WithObservability sets common observability context fields.
+func WithObservability() mux.MiddlewareFunc {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+			ctx = observability.WithBuildInfo(ctx)
+
+			r = r.Clone(ctx)
+			next.ServeHTTP(w, r)
+		})
+	}
+}
diff --git a/internal/middleware/request_id.go b/internal/middleware/request_id.go
new file mode 100644
index 0000000..8c0dcea
--- /dev/null
+++ b/internal/middleware/request_id.go
@@ -0,0 +1,71 @@
+// Copyright 2021 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 middleware
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/google/uuid"
+	"github.com/gorilla/mux"
+)
+
+// contextKeyRequestID is the unique key in the context where the request ID is
+// stored.
+const contextKeyRequestID = contextKey("request_id")
+
+// PopulateRequestID populates the request context with a random UUID if one
+// does not already exist.
+func PopulateRequestID() mux.MiddlewareFunc {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			if existing := RequestIDFromContext(ctx); existing == "" {
+				u, err := uuid.NewRandom()
+				if err != nil {
+					http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+					return
+				}
+
+				ctx = withRequestID(ctx, u.String())
+				r = r.Clone(ctx)
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
+// RequestIDFromContext pulls the request ID from the context, if one was set.
+// If one was not set, it returns the empty string.
+func RequestIDFromContext(ctx context.Context) string {
+	v := ctx.Value(contextKeyRequestID)
+	if v == nil {
+		return ""
+	}
+
+	t, ok := v.(string)
+	if !ok {
+		return ""
+	}
+	return t
+}
+
+// withRequestID sets the request ID on the provided context, returning a new
+// context.
+func withRequestID(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, contextKeyRequestID, id)
+}
diff --git a/pkg/observability/observability.go b/pkg/observability/observability.go
index 498b79d..2b7d6b1 100644
--- a/pkg/observability/observability.go
+++ b/pkg/observability/observability.go
@@ -20,18 +20,35 @@ import (
 	"database/sql"
 	"fmt"
 	"io"
+	"os"
 	"sync"
 
+	"github.com/google/exposure-notifications-server/internal/buildinfo"
+
 	"contrib.go.opencensus.io/integrations/ocsql"
 	"go.opencensus.io/plugin/ocgrpc"
 	"go.opencensus.io/plugin/ochttp"
 	"go.opencensus.io/stats/view"
+	"go.opencensus.io/tag"
 )
 
 // OCSQLDriverName is the name of the SQL driver wrapped by OpenCensus
 // instrumentation code.
 const OCSQLDriverName = "ocsql"
 
+var (
+	BuildIDTagKey  = tag.MustNewKey("build_id")
+	BuildTagTagKey = tag.MustNewKey("build_tag")
+
+	KnativeServiceTagKey       = tag.MustNewKey("k_service")
+	KnativeRevisionTagKey      = tag.MustNewKey("k_revision")
+	KnativeConfigurationTagKey = tag.MustNewKey("k_configuration")
+
+	knativeService       = os.Getenv("K_SERVICE")
+	knativeRevision      = os.Getenv("K_REVISION")
+	knativeConfiguration = os.Getenv("K_CONFIGURATION")
+)
+
 func defaultViews() []*view.View {
 	var ret []*view.View
 	ret = append(ret, ochttp.DefaultClientViews...)
@@ -101,3 +118,26 @@ func NewFromEnv(config *Config) (Exporter, error) {
 		return nil, fmt.Errorf("unknown observability exporter type %v", config.ExporterType)
 	}
 }
+
+// WithBuildInfo creates a new context with the build and revision info attached
+// to the observability context.
+func WithBuildInfo(ctx context.Context) context.Context {
+	tags := make([]tag.Mutator, 0, 5)
+	tags = append(tags, tag.Upsert(BuildIDTagKey, buildinfo.BuildID))
+	tags = append(tags, tag.Upsert(BuildTagTagKey, buildinfo.BuildTag))
+
+	if knativeService != "" {
+		tags = append(tags, tag.Upsert(KnativeServiceTagKey, knativeService))
+	}
+
+	if knativeRevision != "" {
+		tags = append(tags, tag.Upsert(KnativeRevisionTagKey, knativeRevision))
+	}
+
+	if knativeConfiguration != "" {
+		tags = append(tags, tag.Upsert(KnativeConfigurationTagKey, knativeConfiguration))
+	}
+
+	ctx, _ = tag.New(ctx, tags...)
+	return ctx
+}


Test patch:
diff --git a/internal/jwks/jwks_test.go b/internal/jwks/jwks_test.go
index 094c914..bcbea21 100644
--- a/internal/jwks/jwks_test.go
+++ b/internal/jwks/jwks_test.go
@@ -139,7 +139,7 @@ func TestUpdateHA(t *testing.T) {
 			// Set up the test.
 			ctx := context.Background()
 			testDB, _ := testDatabaseInstance.NewDatabase(t)
-			mgr, err := NewManager(ctx, testDB, time.Minute)
+			mgr, err := NewManager(testDB, time.Minute)
 			if err != nil {
 				t.Fatalf("[%d] unexpected error: %v", i, err)
 			}