diff --git a/Makefile b/Makefile index 5baa35f..f6c61b7 100644 --- a/Makefile +++ b/Makefile @@ -24,7 +24,7 @@ kill: sudo pkill -9 sower || true client: build kill - sudo $(PWD)/sower -f '' -s 127.0.0.1:5533 -H "127.0.0.1:8080" + sudo $(PWD)/sower -s 127.0.0.1:5533 -H "127.0.0.1:8080" server: build $(PWD)/sower -f '' diff --git a/proxy/parser/parse.go b/proxy/parser/parse.go index 2e81f13..c8e8d52 100644 --- a/proxy/parser/parse.go +++ b/proxy/parser/parse.go @@ -32,12 +32,14 @@ func NewOtherConn(c net.Conn, domain, port string) net.Conn { typ: OTHER, domain: domain, port: port, + init: true, Conn: c, } } func NewHttpConn(c net.Conn) net.Conn { return &conn{ typ: HTTP, + init: true, Conn: c, } } @@ -45,36 +47,43 @@ func NewHttpsConn(c net.Conn, port string) net.Conn { return &conn{ typ: HTTPS, port: port, + init: true, Conn: c, } } func (c *conn) Write(b []byte) (n int, err error) { - if !c.init { + if c.init { var pkg []byte + var prefixLen int switch c.typ { case OTHER: // type + domain + ':' + port + data - pkg = make([]byte, 0, 1+len(c.domain)+1+len(c.port)+len(b)) + prefixLen = 1 + len(c.domain) + 1 + len(c.port) + pkg = make([]byte, 0, prefixLen+len(b)) pkg = append(pkg, OTHER) pkg = append(pkg, byte(len(c.domain)+1+len(c.port))) pkg = append(pkg, []byte(c.domain+":"+c.port)...) case HTTP: // type + data - pkg = make([]byte, 0, 1+len(b)) + prefixLen = 1 + pkg = make([]byte, 0, prefixLen+len(b)) pkg = append(pkg, HTTP) case HTTPS: // type + port + data - pkg = make([]byte, 0, 1+2+len(b)) + prefixLen = 1 + 2 + pkg = make([]byte, 0, prefixLen+len(b)) pkg = append(pkg, HTTPS) port, _ := strconv.Atoi(c.port) pkg = append(pkg, byte(port>>8), byte(port)) } - c.init = true - return c.Conn.Write(append(pkg, b...)) + c.init = false + n, err := c.Conn.Write(append(pkg, b...)) + // n should larger than prefix length, if not, err is not nil + return n - prefixLen, err } return c.Conn.Write(b) @@ -111,9 +120,10 @@ func ParseAddr(conn net.Conn) (net.Conn, string, string, error) { if _, err := io.ReadFull(conn, buf); err != nil { return conn, "", "", err } + port := strconv.Itoa(int(buf[0])<<8 + int(buf[1])) conn, domain, err := ParseHttpsHost(conn) - return conn, domain, strconv.Itoa(int(buf[0])<<8 + int(buf[1])), err + return conn, domain, port, err default: return conn, "", "", errors.Errorf("not supported type (%v)", buf[0]) diff --git a/proxy/parser/parse_test.go b/proxy/parser/parse_test.go index 917b33c..add722a 100644 --- a/proxy/parser/parse_test.go +++ b/proxy/parser/parse_test.go @@ -8,6 +8,7 @@ import ( "net/http" "testing" + "github.com/wweir/sower/proxy/shadow" "github.com/wweir/sower/util" ) @@ -66,3 +67,31 @@ func TestParseAddr3(t *testing.T) { t.Error(err, host, port) } } + +func TestParseAddr4(t *testing.T) { + c1, c2 := net.Pipe() + + go func() { + c1 = shadow.Shadow(c1, "AES_128_GCM", "12345678") + c1 = NewHttpConn(c1) + req, _ := http.NewRequest("GET", "http://wweir.cc", bytes.NewReader([]byte{1, 2, 3})) + req.Write(c1) + }() + + c2 = shadow.Shadow(c2, "AES_128_GCM", "12345678") + c2, host, port, err := ParseAddr(c2) + + if err != nil || host != "wweir.cc" || port != "80" { + t.Error(err, host, port) + } + + req, err := http.ReadRequest(bufio.NewReader(c2)) + if err != nil { + t.Error(err) + } + + data, err := ioutil.ReadAll(req.Body) + if err != nil || len(data) != 3 || data[0] != 1 { + t.Error(err, data) + } +}