Chore: embed shadowsocks2

transport/shadowsocks/shadowstream/stream.go

@@ -0,0 +1,197 @@
+package shadowstream
+
+import (
+	"crypto/cipher"
+	"crypto/rand"
+	"io"
+	"net"
+)
+
+const bufSize = 2048
+
+type Writer struct {
+	io.Writer
+	cipher.Stream
+	buf [bufSize]byte
+}
+
+// NewWriter wraps an io.Writer with stream cipher encryption.
+func NewWriter(w io.Writer, s cipher.Stream) *Writer { return &Writer{Writer: w, Stream: s} }
+
+func (w *Writer) Write(p []byte) (n int, err error) {
+	buf := w.buf[:]
+	for nw := 0; n < len(p) && err == nil; n += nw {
+		end := n + len(buf)
+		if end > len(p) {
+			end = len(p)
+		}
+		w.XORKeyStream(buf, p[n:end])
+		nw, err = w.Writer.Write(buf[:end-n])
+	}
+	return
+}
+
+func (w *Writer) ReadFrom(r io.Reader) (n int64, err error) {
+	buf := w.buf[:]
+	for {
+		nr, er := r.Read(buf)
+		n += int64(nr)
+		b := buf[:nr]
+		w.XORKeyStream(b, b)
+		if _, err = w.Writer.Write(b); err != nil {
+			return
+		}
+		if er != nil {
+			if er != io.EOF { // ignore EOF as per io.ReaderFrom contract
+				err = er
+			}
+			return
+		}
+	}
+}
+
+type Reader struct {
+	io.Reader
+	cipher.Stream
+	buf [bufSize]byte
+}
+
+// NewReader wraps an io.Reader with stream cipher decryption.
+func NewReader(r io.Reader, s cipher.Stream) *Reader { return &Reader{Reader: r, Stream: s} }
+
+func (r *Reader) Read(p []byte) (n int, err error) {
+	n, err = r.Reader.Read(p)
+	if err != nil {
+		return 0, err
+	}
+	r.XORKeyStream(p, p[:n])
+	return
+}
+
+func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
+	buf := r.buf[:]
+	for {
+		nr, er := r.Reader.Read(buf)
+		if nr > 0 {
+			r.XORKeyStream(buf, buf[:nr])
+			nw, ew := w.Write(buf[:nr])
+			n += int64(nw)
+			if ew != nil {
+				err = ew
+				return
+			}
+		}
+		if er != nil {
+			if er != io.EOF { // ignore EOF as per io.Copy contract (using src.WriteTo shortcut)
+				err = er
+			}
+			return
+		}
+	}
+}
+
+// A Conn represents a Shadowsocks connection. It implements the net.Conn interface.
+type Conn struct {
+	net.Conn
+	Cipher
+	r       *Reader
+	w       *Writer
+	readIV  []byte
+	writeIV []byte
+}
+
+// NewConn wraps a stream-oriented net.Conn with stream cipher encryption/decryption.
+func NewConn(c net.Conn, ciph Cipher) *Conn { return &Conn{Conn: c, Cipher: ciph} }
+
+func (c *Conn) initReader() error {
+	if c.r == nil {
+		iv, err := c.ObtainReadIV()
+		if err != nil {
+			return err
+		}
+		c.r = NewReader(c.Conn, c.Decrypter(iv))
+	}
+	return nil
+}
+
+func (c *Conn) Read(b []byte) (int, error) {
+	if c.r == nil {
+		if err := c.initReader(); err != nil {
+			return 0, err
+		}
+	}
+	return c.r.Read(b)
+}
+
+func (c *Conn) WriteTo(w io.Writer) (int64, error) {
+	if c.r == nil {
+		if err := c.initReader(); err != nil {
+			return 0, err
+		}
+	}
+	return c.r.WriteTo(w)
+}
+
+func (c *Conn) initWriter() error {
+	if c.w == nil {
+		iv, err := c.ObtainWriteIV()
+		if err != nil {
+			return err
+		}
+		if _, err := c.Conn.Write(iv); err != nil {
+			return err
+		}
+		c.w = NewWriter(c.Conn, c.Encrypter(iv))
+	}
+	return nil
+}
+
+func (c *Conn) Write(b []byte) (int, error) {
+	if c.w == nil {
+		if err := c.initWriter(); err != nil {
+			return 0, err
+		}
+	}
+	return c.w.Write(b)
+}
+
+func (c *Conn) ReadFrom(r io.Reader) (int64, error) {
+	if c.w == nil {
+		if err := c.initWriter(); err != nil {
+			return 0, err
+		}
+	}
+	return c.w.ReadFrom(r)
+}
+
+func (c *Conn) ObtainWriteIV() ([]byte, error) {
+	if len(c.writeIV) == c.IVSize() {
+		return c.writeIV, nil
+	}
+
+	iv := make([]byte, c.IVSize())
+
+	if _, err := rand.Read(iv); err != nil {
+		return nil, err
+	}
+
+	c.writeIV = iv
+
+	return iv, nil
+}
+
+func (c *Conn) ObtainReadIV() ([]byte, error) {
+	if len(c.readIV) == c.IVSize() {
+		return c.readIV, nil
+	}
+
+	iv := make([]byte, c.IVSize())
+
+	if _, err := io.ReadFull(c.Conn, iv); err != nil {
+		return nil, err
+	}
+
+	c.readIV = iv
+
+	return iv, nil
+}