fix data steal by proxy

2 files changed, 9 insertions, 4 deletions

diff --git a/pkg/server/socket/listen.go b/pkg/server/socket/listen.go
index 94fb85c..910e5de 100644
--- a/pkg/server/socket/listen.go
+++ b/pkg/server/socket/listen.go
@@ -80,12 +80,12 @@ func (s *listenSocket) Open(env env.Env) (Conn, error) {
 			conn.Close()
 			return nil, err
 		} else {
-			env.Set("original.addr", original)
+			env.Set("original", original)
 		}
 	}
 
 	if s.tproxy {
-		env.Set("original.addr", la.String())
+		env.Set("original", la.String())
 	}
 
 	if original == "" {
diff --git a/pkg/server/socket/proxy.go b/pkg/server/socket/proxy.go
index ef14f48..47f86c9 100644
--- a/pkg/server/socket/proxy.go
+++ b/pkg/server/socket/proxy.go
@@ -1,6 +1,7 @@
 package socket
 
 import (
+	"bufio"
 	"bytes"
 	"errors"
 	"fmt"
@@ -94,12 +95,16 @@ func (s *proxyServer) initConn(addr string) error {
 }
 
 func (s *proxyServer) Recv(rq queue.Q) error {
-	req, err := http.ParseRequest(rq.Reader())
+	r := bufio.NewReader(rq.Reader())
+
+	req, err := http.ParseRequest(r)
 	if err != nil {
 		s.wait <- status{400, "Bad Request"}
 		return err
 	}
 
+	// TODO check if extra data is available in reader
+
 	if req.Method != "CONNECT" {
 		s.wait <- status{400, "Bad Request"}
 		return errors.New("bad method")
@@ -122,7 +127,7 @@ func (s *proxyServer) Recv(rq queue.Q) error {
 
 	s.wait <- status{200, "Connection established"}
 
-	return s.conn.Recv(rq)
+	return queue.IoCopy(r, s.conn.(*conn))
 }
 
 func (s *proxyServer) Close() (err error) {