1 files changed, 7 insertions, 11 deletions
diff --git a/pkg/server/socket/proxy.go b/pkg/server/socket/proxy.go
index c46d1c1..f6c4f72 100644
--- a/pkg/server/socket/proxy.go
+++ b/pkg/server/socket/proxy.go
@@ -3,6 +3,7 @@ package socket
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 
@@ -26,7 +27,7 @@ type proxyServer struct {
 	auth string
 	wait chan status
 	env  env.Env
-	conn Conn
+	conn *conn
 }
 
 func (sock *proxySocket) New(env env.Env) (Conn, error) {
@@ -74,15 +75,9 @@ func (s *proxyServer) Send(wq queue.Q) error {
 	return s.conn.Send(wq)
 }
 
-func (s *proxyServer) initConn(addr string) error {
-	dial := dialSocket{
-		Proto: s.sock.Proto,
-		Addr:  addr,
-	}
-
-	conn, err := dial.New(s.env)
+func (s *proxyServer) dial(addr string) error {
+	conn, err := dial(context.TODO(), s.env, s.sock.Proto, addr)
 	if err != nil {
-		dial.Close()
 		return err
 	}
 
@@ -118,17 +113,18 @@ func (s *proxyServer) Recv(rq queue.Q) error {
 		}
 	}
 
-	if err := s.initConn(req.URI); err != nil {
+	if err := s.dial(req.URI); err != nil {
 		s.wait <- status{500, "Unable to connect"}
 		return err
 	}
 
 	s.wait <- status{200, "Connection established"}
 
-	return queue.IoCopy(r, s.conn.(*conn))
+	return queue.IoCopy(r, s.conn)
 }
 
 func (s *proxyServer) Close() (err error) {
+	// TODO safe close
 	if s.conn != nil {
 		err = s.conn.Close()
 	}