Add context methods for the authority database

ca/ca.go

@@ -308,11 +308,12 @@ func (ca *CA) Init(cfg *config.Config) (*CA, error) {
 // buildContext builds the server base context.
 func buildContext(a *authority.Authority) context.Context {
 	ctx := authority.NewContext(context.Background(), a)
-
+	if authDB := a.GetDatabase(); authDB != nil {
-	if db := a.GetAdminDatabase(); db != nil {
+		ctx = db.NewContext(ctx, authDB)
-		ctx = admin.NewContext(ctx, db)
+	}
+	if adminDB := a.GetAdminDatabase(); adminDB != nil {
+		ctx = admin.NewContext(ctx, adminDB)
 	}
-
 	return ctx
 }
 

db/db.go

@@ -1,6 +1,7 @@
 package db
 
 import (
+	"context"
 	"crypto/x509"
 	"encoding/json"
 	"strconv"
@@ -58,6 +59,29 @@ type AuthDB interface {
 	Shutdown() error
 }
 
+type dbKey struct{}
+
+// NewContext adds the given authority database to the context.
+func NewContext(ctx context.Context, db AuthDB) context.Context {
+	return context.WithValue(ctx, dbKey{}, db)
+}
+
+// FromContext returns the current authority database from the given context.
+func FromContext(ctx context.Context) (db AuthDB, ok bool) {
+	db, ok = ctx.Value(dbKey{}).(AuthDB)
+	return
+}
+
+// MustFromContext returns the current database from the given context. It
+// will panic if it's not in the context.
+func MustFromContext(ctx context.Context) AuthDB {
+	if db, ok := FromContext(ctx); !ok {
+		panic("authority database is not in the context")
+	} else {
+		return db
+	}
+}
+
 // DB is a wrapper over the nosql.DB interface.
 type DB struct {
 	nosql.DB