From 0446e823208559907ab89f6efe5ac88f2ba43edf Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano.cano@gmail.com>
Date: Wed, 27 Apr 2022 12:05:19 -0700
Subject: [PATCH] Add context methods for the authority database

---
 ca/ca.go |  9 +++++----
 db/db.go | 24 ++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/ca/ca.go b/ca/ca.go
index 2df52555..f5cf30db 100644
--- a/ca/ca.go
+++ b/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 db := a.GetAdminDatabase(); db != nil {
-		ctx = admin.NewContext(ctx, db)
+	if authDB := a.GetDatabase(); authDB != nil {
+		ctx = db.NewContext(ctx, authDB)
+	}
+	if adminDB := a.GetAdminDatabase(); adminDB != nil {
+		ctx = admin.NewContext(ctx, adminDB)
 	}
-
 	return ctx
 }
 
diff --git a/db/db.go b/db/db.go
index eccaf801..c4b1c8a7 100644
--- a/db/db.go
+++ b/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