From 5de4cd8edf4f5de56c8c6910d7e00012dfe3818e Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Tue, 17 Nov 2020 11:14:01 -0500
Subject: [PATCH] Stash remote address in context

---
 server/server.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/server/server.go b/server/server.go
index 79c1d5a74c945ccde07539ca6d364ce6a94c7013..b8728131f91fcd24e037a33d2f30b9baa65d01cd 100644
--- a/server/server.go
+++ b/server/server.go
@@ -118,6 +118,12 @@ func (server *Server) WithSchema(
 	return server
 }
 
+var remoteAddrCtxKey = &contextKey{"remoteAddr"}
+
+type contextKey struct {
+	name string
+}
+
 // Adds the default middleware to this server, including:
 //
 // - Configuration middleware
@@ -177,9 +183,24 @@ func (server *Server) WithDefaultMiddleware() *Server {
 	server.router.Use(middleware.Logger)
 	server.router.Use(middleware.Timeout(timeout))
 	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)
+			r = r.WithContext(ctx)
+			next.ServeHTTP(w, r)
+		})
+	})
 	return server
 }
 
+func RemoteAddr(ctx context.Context) string {
+	raw, ok := ctx.Value(remoteAddrCtxKey).(string)
+	if !ok {
+		panic(fmt.Errorf("Invalid authentication context"))
+	}
+	return raw
+}
+
 // Add user-defined middleware to the server
 func (server *Server) WithMiddleware(
 	middlewares ...func(http.Handler) http.Handler) *Server {