Add remote_addr to context (#70)

* Add remote_addr to context

* feat: include RemoteAddress ctx wrapper

---------

Co-authored-by: Jeroen Rinzema <[email protected]>

Author: joshrendek

conn.go

@@ -2,6 +2,7 @@ package wire
 
 import (
 	"context"
+	"net"
 
 	"github.com/jackc/pgx/v5/pgtype"
 )
@@ -12,6 +13,7 @@ const (
 	ctxTypeMap ctxKey = iota
 	ctxClientMetadata
 	ctxServerMetadata
+	ctxRemoteAddr
 )
 
 // setTypeInfo constructs a new Postgres type connection info for the given value
@@ -30,6 +32,21 @@ func TypeMap(ctx context.Context) *pgtype.Map {
 	return val.(*pgtype.Map)
 }
 
+func setRemoteAddress(ctx context.Context, addr net.Addr) context.Context {
+	return context.WithValue(ctx, ctxRemoteAddr, addr)
+}
+
+// RemoteAddress returns the Postgres remote address connection info if it has been set inside
+// the given context.
+func RemoteAddress(ctx context.Context) net.Addr {
+	val := ctx.Value(ctxRemoteAddr)
+	if val == nil {
+		return nil
+	}
+
+	return val.(net.Addr)
+}
+
 // Parameters represents a parameters collection of parameter status keys and
 // their values.
 type Parameters map[ParameterStatus]string

wire.go

@@ -123,6 +123,7 @@ func (srv *Server) Serve(listener net.Listener) error {
 
 func (srv *Server) serve(ctx context.Context, conn net.Conn) error {
 	ctx = setTypeInfo(ctx, srv.types)
+	ctx = setRemoteAddress(ctx, conn.RemoteAddr())
 	defer conn.Close()
 
 	srv.logger.Debug("serving a new client connection")