From c4cc894452e996527bf6edbfa5c486fe18bc6b56 Mon Sep 17 00:00:00 2001
From: Tyler Treat <ttreat31@gmail.com>
Date: Fri, 1 Jul 2016 23:57:56 -0500
Subject: [PATCH] Expose NATS Conn

Add a NatsConn method to the Conn interface which allows retrieving the
underlying NATS Conn. Addresses #67.
---
 stan.go      | 12 ++++++++++++
 stan_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/stan.go b/stan.go
index 1ef5be1..1dc8717 100644
--- a/stan.go
+++ b/stan.go
@@ -50,6 +50,11 @@ type Conn interface {
 
 	// Close
 	Close() error
+
+	// NatsConn returns the underlying NATS conn. Use this with care. For
+	// example, closing the wrapped NATS conn will put the NATS Streaming Conn
+	// in an invalid state.
+	NatsConn() *nats.Conn
 }
 
 // Errors
@@ -271,6 +276,13 @@ func (sc *conn) Close() error {
 	return nil
 }
 
+// NatsConn returns the underlying NATS conn. Use this with care. For example,
+// closing the wrapped NATS conn will put the NATS Streaming Conn in an invalid
+// state.
+func (sc *conn) NatsConn() *nats.Conn {
+	return sc.nc
+}
+
 // Process a heartbeat from the NATS Streaming cluster
 func (sc *conn) processHeartBeat(m *nats.Msg) {
 	// No payload assumed, just reply.
diff --git a/stan_test.go b/stan_test.go
index 2a38e51..f07140a 100644
--- a/stan_test.go
+++ b/stan_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/nats-io/nats"
 	"github.com/nats-io/nats-streaming-server/server"
 	"github.com/nats-io/nats-streaming-server/test"
+	natstest "github.com/nats-io/nats/test"
 
 	natsd "github.com/nats-io/gnatsd/test"
 )
@@ -1853,3 +1854,37 @@ func TestRaceAckOnClose(t *testing.T) {
 	time.Sleep(10 * time.Millisecond)
 	sc.Close()
 }
+
+func TestNatsConn(t *testing.T) {
+	s := RunServer(clusterName)
+	defer s.Shutdown()
+	sc := NewDefaultConnection(t)
+	defer sc.Close()
+
+	// Make sure we can get the STAN-created Conn.
+	nc := sc.NatsConn()
+
+	if nc.Status() != nats.CONNECTED {
+		t.Fatal("Should have status set to CONNECTED")
+	}
+	nc.Close()
+	if nc.Status() != nats.CLOSED {
+		t.Fatal("Should have status set to CLOSED")
+	}
+
+	sc.Close()
+	if sc.NatsConn() != nil {
+		t.Fatal("Wrapped conn should be nil after close")
+	}
+
+	// Make sure we can get the Conn we provide.
+	nc = natstest.NewDefaultConnection(t)
+	sc, err := Connect(clusterName, clientName, NatsConn(nc))
+	if err != nil {
+		stackFatalf(t, "Expected to connect correctly, got err %v", err)
+	}
+	defer sc.Close()
+	if sc.NatsConn() != nc {
+		t.Fatal("Unexpected wrapped conn")
+	}
+}