obfsvpn: allow wrapping existing connections

This can be used to add support for additional networks that we don't
necessarily want to depend on in the library.

Signed-off-by: Sam Whited <sam@samwhited.com>

2 files changed, 31 insertions, 12 deletions

diff --git a/dialer.go b/dialer.go
index 444e0c2..5043aeb 100644
--- a/dialer.go
+++ b/dialer.go
@@ -120,8 +120,22 @@ func NewDialerArgs(args pt.Args) (*Dialer, error) {
 	}, nil
 }
 
+// Wrap performs the ntor handshake over an existing conection.
+// Wrap ignores the underlying Dialer config.
+func (d *Dialer) Wrap(ctx context.Context, conn net.Conn) (net.Conn, error) {
+	return d.dial(ctx, "", "", func(network, address string) (net.Conn, error) {
+		return conn, nil
+	})
+}
+
 // Dial creates an outbound net.Conn and performs the ntor handshake.
 func (d *Dialer) Dial(ctx context.Context, network, address string) (net.Conn, error) {
+	return d.dial(ctx, network, address, func(network, address string) (net.Conn, error) {
+		return d.Dialer.DialContext(ctx, network, address)
+	})
+}
+
+func (d *Dialer) dial(ctx context.Context, network, address string, f func(network, address string) (net.Conn, error)) (net.Conn, error) {
 	if d.cf == nil {
 		cf, err := (&obfs4.Transport{}).ClientFactory("")
 		if err != nil {
@@ -134,9 +148,7 @@ func (d *Dialer) Dial(ctx context.Context, network, address string) (net.Conn, e
 	if err != nil {
 		return nil, err
 	}
-	return d.cf.Dial(network, address, func(network, address string) (net.Conn, error) {
-		return d.Dialer.DialContext(ctx, network, address)
-	}, args)
+	return d.cf.Dial(network, address, f, args)
 }
 
 // Args returns the dialers options as pluggable transport arguments.
diff --git a/listener.go b/listener.go
index 4163cb9..593032f 100644
--- a/listener.go
+++ b/listener.go
@@ -38,20 +38,16 @@ func NewListenConfigCert(cert string) (*ListenConfig, error) {
 	}, nil
 }
 
-// Listen announces on the local network address.
-//
-// See func net.Dial for a description of the network and address parameters.
-func (lc *ListenConfig) Listen(ctx context.Context, network, address string) (*Listener, error) {
-	ln, err := lc.ListenConfig.Listen(ctx, network, address)
-	if err != nil {
-		return nil, err
-	}
+// Wrap takes an existing net.Listener and wraps it in a listener that is
+// configured to perform the ntor handshake.
+// Values from the inner net.ListenConfig are ignored.
+func (lc *ListenConfig) Wrap(ctx context.Context, ln net.Listener) (*Listener, error) {
 	args := make(pt.Args)
 	args.Add("node-id", lc.NodeID.Hex())
 	args.Add("private-key", lc.PrivateKey.Hex())
 	seed := ntor.KeySeed{}
 	if bytes.Equal(lc.Seed[:], seed[:]) {
-		_, err = rand.Read(seed[:])
+		_, err := rand.Read(seed[:])
 		if err != nil {
 			return nil, err
 		}
@@ -66,6 +62,17 @@ func (lc *ListenConfig) Listen(ctx context.Context, network, address string) (*L
 	return &Listener{sf: sf, ln: ln}, nil
 }
 
+// Listen announces on the local network address.
+//
+// See func net.Dial for a description of the network and address parameters.
+func (lc *ListenConfig) Listen(ctx context.Context, network, address string) (*Listener, error) {
+	ln, err := lc.ListenConfig.Listen(ctx, network, address)
+	if err != nil {
+		return nil, err
+	}
+	return lc.Wrap(ctx, ln)
+}
+
 // Listener is a network listener that accepts obfuscated connections and
 // performs the ntor handshake on them.
 type Listener struct {