fix data steal by proxy

7 files changed, 49 insertions, 28 deletions

diff --git a/TODO b/TODO
index 2191e8f..2a4a7f9 100644
--- a/TODO
+++ b/TODO
@@ -12,3 +12,7 @@
 - docs
 
 - tproxy usage
+
+- dial with timeout
+- proxy socket with connect wait timeout
+- config split with double spaces
diff --git a/USAGE b/USAGE
index 889c477..3d4ddf1 100644
--- a/USAGE
+++ b/USAGE
@@ -6,13 +6,12 @@ tunnel add srcport,listen [dsthost:]dstport
 
 iptables -t nat -A OUTPUT -p tcp -d 198.145.29.83 -j REDIRECT --to-port 4000
 
-tunnel add 4000,listen,redirect :proxy,addr=@original.addr @proxyaddr
+tunnel add 4000,listen,redirect :proxy,addr=@original @proxyaddr
 
 - check forward usage
 
 BUGS:
-	Direct connection also possible. In this case original.addr will be equal to listen address.
-	This happens when ipv4 is used.
+	With ipv4 direct connection causes original addr will be equal to listen addr.
 
 # Transparent proxy using tproxy
 
diff --git a/pkg/http/http.go b/pkg/http/http.go
index eefe348..4180c81 100644
--- a/pkg/http/http.go
+++ b/pkg/http/http.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/textproto"
 	"regexp"
 	"strconv"
 	"strings"
@@ -27,6 +28,8 @@ type Response struct {
 	Header Header
 }
 
+type handler func(string) error
+
 var version = "HTTP/[0-9]+\\.[0-9]+"
 var requestLine = regexp.MustCompile("^([A-Z]+) +([^ ]+) +(" + version + ")$")
 var statusLine = regexp.MustCompile("^(" + version + ") +([0-9]+) +(.*)$")
@@ -52,31 +55,34 @@ func addHeader(m Header, k, v string) Header {
 	return m
 }
 
-func parse(r io.Reader, init func(string) error, f func(string) error) error {
-	scanner := bufio.NewScanner(r)
-	ok := false
+func parse(r *bufio.Reader, init handler, f handler) error {
+	proto := textproto.NewReader(r)
+	first := true
 
-	if scanner.Scan() {
-		if err := init(scanner.Text()); err != nil {
+	for {
+		line, err := proto.ReadLine()
+		if err != nil {
+			if err == io.EOF {
+				err = io.ErrUnexpectedEOF
+			}
 			return err
 		}
-	}
 
-	for scanner.Scan() {
-		if scanner.Text() == "" {
-			ok = true
-			break
+		if first {
+			if err := init(line); err != nil {
+				return err
+			}
+			first = false
+			continue
 		}
 
-		if err := f(scanner.Text()); err != nil {
-			return err
+		if line == "" {
+			return nil
 		}
-	}
 
-	if err := scanner.Err(); err != nil {
-		return err
-	} else if !ok {
-		return io.ErrUnexpectedEOF
+		if err := f(line); err != nil {
+			return err
+		}
 	}
 
 	return nil
@@ -108,7 +114,7 @@ func (r *Request) parseHeader(s string) error {
 	}
 }
 
-func ParseRequest(r io.Reader) (*Request, error) {
+func ParseRequest(r *bufio.Reader) (*Request, error) {
 	req := new(Request)
 
 	if err := parse(r, req.parseRequestLine, req.parseHeader); err != nil {
@@ -146,7 +152,7 @@ func (r *Response) parseHeader(s string) error {
 	}
 }
 
-func ParseResponse(r io.Reader) (*Response, error) {
+func ParseResponse(r *bufio.Reader) (*Response, error) {
 	resp := new(Response)
 
 	if err := parse(r, resp.parseStatusLine, resp.parseHeader); err != nil {
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) {