write salt/iv, addr together with content

Sending salt, address and content seperately would generate three packets
at the start of a connection. Since the size of salt is almost always 8, 16 or 32,
it leaves a strong feature.
Additionally, with TCP fastopen turned on, only the first write would be carried by
SYN packet. Writing multiple times makes it pointless.

shadowaead/stream.go

@@ -16,17 +16,19 @@ type writer struct {
 	cipher.AEAD
 	nonce []byte
 	buf   []byte
+	salt  []byte
 }
 
 // NewWriter wraps an io.Writer with AEAD encryption.
-func NewWriter(w io.Writer, aead cipher.AEAD) io.Writer { return newWriter(w, aead) }
+func NewWriter(w io.Writer, aead cipher.AEAD, salt []byte) io.Writer { return newWriter(w, aead, salt) }
 
-func newWriter(w io.Writer, aead cipher.AEAD) *writer {
+func newWriter(w io.Writer, aead cipher.AEAD, salt []byte) *writer {
 	return &writer{
 		Writer: w,
 		AEAD:   aead,
 		buf:    make([]byte, 2+aead.Overhead()+payloadSizeMask+aead.Overhead()),
 		nonce:  make([]byte, aead.NonceSize()),
+		salt:   salt,
 	}
 }
 
@@ -36,6 +38,23 @@ func (w *writer) Write(b []byte) (int, error) {
 	return int(n), err
 }
 
+// Write salt before encrypted buffer to io.Writer.
+func (w *writer) write(b []byte) (int, error) {
+	if len(w.salt) == 0 {
+		return w.Writer.Write(b)
+	}
+	buf := make([]byte, len(w.salt) + len(b))
+	copy(buf[:len(w.salt)], w.salt)
+	copy(buf[len(w.salt):], b)
+	nw, err := w.Writer.Write(buf)
+	if nw < len(w.salt) {
+		w.salt = w.salt[nw:]
+		return 0, err
+	}
+	w.salt = nil
+	return nw - len(w.salt), err
+}
+
 // ReadFrom reads from the given io.Reader until EOF or error, encrypts and
 // writes to the embedded io.Writer. Returns number of bytes read from r and
 // any error encountered.
@@ -56,7 +75,7 @@ func (w *writer) ReadFrom(r io.Reader) (n int64, err error) {
 			w.Seal(payloadBuf[:0], w.nonce, payloadBuf, nil)
 			increment(w.nonce)
 
-			_, ew := w.Writer.Write(buf)
+			_, ew := w.write(buf)
 			if ew != nil {
 				err = ew
 				break
@@ -240,11 +259,7 @@ func (c *streamConn) initWriter() error {
 	if err != nil {
 		return err
 	}
-	_, err = c.Conn.Write(salt)
-	if err != nil {
-		return err
-	}
-	c.w = newWriter(c.Conn, aead)
+	c.w = newWriter(c.Conn, aead, salt)
 	return nil
 }
 

shadowstream/stream.go

@@ -14,11 +14,12 @@ type writer struct {
 	io.Writer
 	cipher.Stream
 	buf []byte
+	iv  []byte
 }
 
 // NewWriter wraps an io.Writer with stream cipher encryption.
-func NewWriter(w io.Writer, s cipher.Stream) io.Writer {
-	return &writer{Writer: w, Stream: s, buf: make([]byte, bufSize)}
+func NewWriter(w io.Writer, s cipher.Stream, iv []byte) io.Writer {
+	return &writer{Writer: w, Stream: s, buf: make([]byte, bufSize), iv: iv}
 }
 
 func (w *writer) ReadFrom(r io.Reader) (n int64, err error) {
@@ -29,7 +30,7 @@ func (w *writer) ReadFrom(r io.Reader) (n int64, err error) {
 			n += int64(nr)
 			buf = buf[:nr]
 			w.XORKeyStream(buf, buf)
-			_, ew := w.Writer.Write(buf)
+			_, ew := w.write(buf)
 			if ew != nil {
 				err = ew
 				return
@@ -50,6 +51,23 @@ func (w *writer) Write(b []byte) (int, error) {
 	return int(n), err
 }
 
+// Write IV before encrypted buffer to io.Writer.
+func (w *writer) write(b []byte) (int, error) {
+	if len(w.iv) == 0 {
+		return w.Writer.Write(b)
+	}
+	buf := make([]byte, len(w.iv) + len(b))
+	copy(buf[:len(w.iv)], w.iv)
+	copy(buf[len(w.iv):], b)
+	nw, err := w.Writer.Write(buf)
+	if nw < len(w.iv) {
+		w.iv = w.iv[nw:]
+		return 0, err
+	}
+	w.iv = nil
+	return nw - len(w.iv), err
+}
+
 type reader struct {
 	io.Reader
 	cipher.Stream
@@ -144,10 +162,7 @@ func (c *conn) initWriter() error {
 		if _, err := io.ReadFull(rand.Reader, iv); err != nil {
 			return err
 		}
-		if _, err := c.Conn.Write(iv); err != nil {
-			return err
-		}
-		c.w = &writer{Writer: c.Conn, Stream: c.Encrypter(iv), buf: buf}
+		c.w = &writer{Writer: c.Conn, Stream: c.Encrypter(iv), buf: buf, iv: iv}
 	}
 	return nil
 }

tcp.go

@@ -73,13 +73,9 @@ func tcpLocal(addr, server string, shadow func(net.Conn) net.Conn, getAddr func(
 			rc.(*net.TCPConn).SetKeepAlive(true)
 			rc = shadow(rc)
 
-			if _, err = rc.Write(tgt); err != nil {
-				logf("failed to send target address: %v", err)
-				return
-			}
-
 			logf("proxy %s <-> %s <-> %s", c.RemoteAddr(), server, tgt)
-			_, _, err = relay(rc, c)
+			ca := &connWithAddr{Conn: c, addr: tgt}
+			_, _, err = relay(rc, ca)
 			if err != nil {
 				if err, ok := err.(net.Error); ok && err.Timeout() {
 					return // ignore i/o timeout
@@ -163,3 +159,21 @@ func relay(left, right net.Conn) (int64, int64, error) {
 	}
 	return n, rs.N, err
 }
+
+type connWithAddr struct {
+	net.Conn
+	addr socks.Addr
+}
+
+// Read reads the addr and data from the connection.
+// The format of output is aligned with shadowsocks protocol.
+func (c *connWithAddr) Read(b []byte) (n int, err error) {
+	nc := copy(b, c.addr)
+	if nc < len(c.addr) {
+		c.addr = c.addr[:nc]
+		return nc, nil
+	}
+	c.addr = nil
+	nr, err := c.Conn.Read(b[nc:])
+	return nc + nr, err
+}