Fix code

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 ''

proxy/parser/parse.go

@@ -32,49 +32,58 @@ 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,
 	}
 }
 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])

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)
+	}
+}