fix data steal by proxy

4 files changed, 18 insertions, 6 deletions

diff --git a/pkg/server/hook/proxy.go b/pkg/server/hook/proxy.go
index 172b01a..a488f58 100644
--- a/pkg/server/hook/proxy.go
+++ b/pkg/server/hook/proxy.go
@@ -1,6 +1,7 @@
 package hook
 
 import (
+	"bufio"
 	"bytes"
 	"errors"
 	"fmt"
@@ -45,7 +46,9 @@ func (p *proxy) Send(rq, wq queue.Q) error {
 }
 
 func (p *proxy) Recv(rq, wq queue.Q) error {
-	resp, err := http.ParseResponse(rq.Reader())
+	r := bufio.NewReader(rq.Reader())
+
+	resp, err := http.ParseResponse(r)
 
 	if err == nil && resp.Code != http.OK {
 		err = fmt.Errorf("connect failed: %d %s", resp.Code, resp.Desc)
@@ -58,7 +61,7 @@ func (p *proxy) Recv(rq, wq queue.Q) error {
 
 	p.c <- true
 
-	return queue.Copy(rq, wq)
+	return queue.IoCopy(r, wq.Writer())
 }
 
 func (h *proxyHook) Open(env env.Env) (interface{}, error) {
diff --git a/pkg/server/queue/queue.go b/pkg/server/queue/queue.go
index 704d884..f0c1fc9 100644
--- a/pkg/server/queue/queue.go
+++ b/pkg/server/queue/queue.go
@@ -79,6 +79,10 @@ func (q Q) Dry() {
 }
 
 func (w *writer) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+
 	buf := make([]byte, len(p))
 	copy(buf, p)
 	w.q <- buf
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) {