From 67ddd428bf8e7f512ad63758a0e75d1e4c404312 Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Wed, 18 Aug 2021 11:31:14 +0200
Subject: [PATCH] server: add server reference to context

---
 server/server.go | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/server/server.go b/server/server.go
index ef81c2dfcc028e39317556e739170053c9a976f8..11e7b97f6d921f44ae2e2e23c4c52688dc6f81d1 100644
--- a/server/server.go
+++ b/server/server.go
@@ -46,11 +46,12 @@ var (
 )
 
 type Server struct {
+	Schema graphql.ExecutableSchema
+
 	conf    ini.File
 	db      *sql.DB
 	redis   *goRedis.Client
 	router  chi.Router
-	schema  graphql.ExecutableSchema
 	service string
 	queues  []*work.Queue
 }
@@ -76,7 +77,7 @@ func (server *Server) Router() chi.Router {
 // configured before this is called.
 func (server *Server) WithSchema(
 	schema graphql.ExecutableSchema, scopes []string) *Server {
-	server.schema = schema
+	server.Schema = schema
 
 	var (
 		complexity int
@@ -119,6 +120,7 @@ func (server *Server) WithSchema(
 	return server
 }
 
+var serverCtxKey = &contextKey{"server"}
 var remoteAddrCtxKey = &contextKey{"remoteAddr"}
 
 type contextKey struct {
@@ -186,7 +188,8 @@ func (server *Server) WithDefaultMiddleware() *Server {
 	server.router.Use(auth.Middleware(server.conf, apiconf))
 	server.router.Use(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			ctx := context.WithValue(r.Context(), remoteAddrCtxKey, r.RemoteAddr)
+			ctx := context.WithValue(r.Context(), serverCtxKey, server)
+			ctx = context.WithValue(ctx, remoteAddrCtxKey, r.RemoteAddr)
 			r = r.WithContext(ctx)
 			next.ServeHTTP(w, r)
 		})
@@ -202,6 +205,14 @@ func RemoteAddr(ctx context.Context) string {
 	return raw
 }
 
+func ForContext(ctx context.Context) *Server {
+	raw, ok := ctx.Value(serverCtxKey).(*Server)
+	if !ok {
+		panic(fmt.Errorf("Invalid server context"))
+	}
+	return raw
+}
+
 // Add user-defined middleware to the server
 func (server *Server) WithMiddleware(
 	middlewares ...func(http.Handler) http.Handler) *Server {