Add SetSaltGenerator

shadowsocks/client.go

@@ -69,7 +69,7 @@ func (c *ssClient) DialTCP(laddr *net.TCPAddr, raddr string) (onet.DuplexConn, e
 	if err != nil {
 		return nil, err
 	}
-	ssw := NewShadowsocksWriter(proxyConn, c.cipher, RandomSaltGenerator)
+	ssw := NewShadowsocksWriter(proxyConn, c.cipher)
 	_, err = ssw.LazyWrite(socksTargetAddr)
 	if err != nil {
 		proxyConn.Close()

shadowsocks/client_test.go

@@ -224,7 +224,7 @@ func startShadowsocksTCPEchoProxy(expectedTgtAddr string, t testing.TB) (net.Lis
 				defer running.Done()
 				defer clientConn.Close()
 				ssr := NewShadowsocksReader(clientConn, cipher)
-				ssw := NewShadowsocksWriter(clientConn, cipher, RandomSaltGenerator)
+				ssw := NewShadowsocksWriter(clientConn, cipher)
 				ssClientConn := onet.WrapConn(clientConn, ssr, ssw)
 
 				tgtAddr, err := socks.ReadAddr(ssClientConn)

shadowsocks/stream.go

@@ -52,19 +52,7 @@ var RandomSaltGenerator SaltGenerator = &randomSaltGenerator{}
 // The LazyWrite and Flush methods allow a header to be
 // added but delayed until the first write, for concatenation.
 // All methods except Flush must be called from a single thread.
-type Writer interface {
-	io.Writer
-	io.ReaderFrom
-	// LazyWrite queues p to be written, but doesn't send it until
-	// Flush() is called, a non-lazy write is made, or the buffer
-	// is filled.
-	LazyWrite(p []byte) (int, error)
-	// Flush sends the pending data, if any.  This method is
-	// thread-safe.
-	Flush() error
-}
-
-type shadowsocksWriter struct {
+type Writer struct {
 	// This type is single-threaded except when needFlush is true.
 	// mu protects needFlush, and also protects everything
 	// else while needFlush could be true.
@@ -87,16 +75,18 @@ type shadowsocksWriter struct {
 
 // NewShadowsocksWriter creates a Writer that encrypts the given Writer using
 // the shadowsocks protocol with the given shadowsocks cipher.
-func NewShadowsocksWriter(writer io.Writer, ssCipher shadowaead.Cipher, saltGenerator SaltGenerator) Writer {
-	if saltGenerator == nil {
-		saltGenerator = RandomSaltGenerator
-	}
-	return &shadowsocksWriter{writer: writer, ssCipher: ssCipher, saltGenerator: saltGenerator}
+func NewShadowsocksWriter(writer io.Writer, ssCipher shadowaead.Cipher) *Writer {
+	return &Writer{writer: writer, ssCipher: ssCipher, saltGenerator: RandomSaltGenerator}
+}
+
+// SetSaltGenerator sets the salt generator to be used. Must be called before the first write.
+func (sw *Writer) SetSaltGenerator(saltGenerator SaltGenerator) {
+	sw.saltGenerator = saltGenerator
 }
 
 // init generates a random salt, sets up the AEAD object and writes
 // the salt to the inner Writer.
-func (sw *shadowsocksWriter) init() (err error) {
+func (sw *Writer) init() (err error) {
 	if sw.aead == nil {
 		salt := make([]byte, sw.ssCipher.SaltSize())
 		if err := sw.saltGenerator.GetSalt(salt); err != nil {
@@ -120,19 +110,21 @@ func (sw *shadowsocksWriter) init() (err error) {
 
 // encryptBlock encrypts `plaintext` in-place.  The slice must have enough capacity
 // for the tag. Returns the total ciphertext length.
-func (sw *shadowsocksWriter) encryptBlock(plaintext []byte) int {
+func (sw *Writer) encryptBlock(plaintext []byte) int {
 	out := sw.aead.Seal(plaintext[:0], sw.counter, plaintext, nil)
 	increment(sw.counter)
 	return len(out)
 }
 
-func (sw *shadowsocksWriter) Write(p []byte) (int, error) {
+func (sw *Writer) Write(p []byte) (int, error) {
 	sw.byteWrapper.Reset(p)
 	n, err := sw.ReadFrom(&sw.byteWrapper)
 	return int(n), err
 }
 
-func (sw *shadowsocksWriter) LazyWrite(p []byte) (int, error) {
+// LazyWrite queues p to be written, but doesn't send it until Flush() is
+// called, a non-lazy write is made, or the buffer is filled.
+func (sw *Writer) LazyWrite(p []byte) (int, error) {
 	if err := sw.init(); err != nil {
 		return 0, err
 	}
@@ -159,7 +151,8 @@ func (sw *shadowsocksWriter) LazyWrite(p []byte) (int, error) {
 	}
 }
 
-func (sw *shadowsocksWriter) Flush() error {
+// Flush sends the pending data, if any.  This method is thread-safe.
+func (sw *Writer) Flush() error {
 	sw.mu.Lock()
 	defer sw.mu.Unlock()
 	if !sw.needFlush {
@@ -178,7 +171,7 @@ func isZero(b []byte) bool {
 }
 
 // Returns the slices of sw.buf in which to place plaintext for encryption.
-func (sw *shadowsocksWriter) buffers() (sizeBuf, payloadBuf []byte) {
+func (sw *Writer) buffers() (sizeBuf, payloadBuf []byte) {
 	// sw.buf starts with the salt.
 	saltSize := sw.ssCipher.SaltSize()
 
@@ -190,7 +183,8 @@ func (sw *shadowsocksWriter) buffers() (sizeBuf, payloadBuf []byte) {
 	return
 }
 
-func (sw *shadowsocksWriter) ReadFrom(r io.Reader) (int64, error) {
+// ReadFrom implements the io.ReaderFrom interface.
+func (sw *Writer) ReadFrom(r io.Reader) (int64, error) {
 	if err := sw.init(); err != nil {
 		return 0, err
 	}
@@ -240,15 +234,15 @@ func (sw *shadowsocksWriter) ReadFrom(r io.Reader) (int64, error) {
 
 // Adds as much of `plaintext` into the buffer as will fit, and increases
 // sw.pending accordingly.  Returns the number of bytes consumed.
-func (sw *shadowsocksWriter) enqueue(plaintext []byte) int {
+func (sw *Writer) enqueue(plaintext []byte) int {
 	_, payloadBuf := sw.buffers()
 	n := copy(payloadBuf[sw.pending:], plaintext)
 	sw.pending += n
 	return n
 }
 
 // Encrypts all pending data and writes it to the output.
-func (sw *shadowsocksWriter) flush() error {
+func (sw *Writer) flush() error {
 	if sw.pending == 0 {
 		return nil
 	}

shadowsocks/stream_test.go

@@ -157,7 +157,7 @@ func TestEndToEnd(t *testing.T) {
 	cipher := newTestCipher(t)
 
 	connReader, connWriter := io.Pipe()
-	writer := NewShadowsocksWriter(connWriter, cipher, RandomSaltGenerator)
+	writer := NewShadowsocksWriter(connWriter, cipher)
 	reader := NewShadowsocksReader(connReader, cipher)
 	expected := "Test"
 	go func() {
@@ -180,7 +180,7 @@ func TestEndToEnd(t *testing.T) {
 func TestLazyWriteFlush(t *testing.T) {
 	cipher := newTestCipher(t)
 	buf := new(bytes.Buffer)
-	writer := NewShadowsocksWriter(buf, cipher, RandomSaltGenerator)
+	writer := NewShadowsocksWriter(buf, cipher)
 	header := []byte{1, 2, 3, 4}
 	n, err := writer.LazyWrite(header)
 	if n != len(header) {
@@ -241,7 +241,7 @@ func TestLazyWriteFlush(t *testing.T) {
 func TestLazyWriteConcat(t *testing.T) {
 	cipher := newTestCipher(t)
 	buf := new(bytes.Buffer)
-	writer := NewShadowsocksWriter(buf, cipher, RandomSaltGenerator)
+	writer := NewShadowsocksWriter(buf, cipher)
 	header := []byte{1, 2, 3, 4}
 	n, err := writer.LazyWrite(header)
 	if n != len(header) {
@@ -295,7 +295,7 @@ func TestLazyWriteConcat(t *testing.T) {
 func TestLazyWriteOversize(t *testing.T) {
 	cipher := newTestCipher(t)
 	buf := new(bytes.Buffer)
-	writer := NewShadowsocksWriter(buf, cipher, RandomSaltGenerator)
+	writer := NewShadowsocksWriter(buf, cipher)
 	N := 25000 // More than one block, less than two.
 	data := make([]byte, N)
 	for i := range data {
@@ -335,7 +335,7 @@ func TestLazyWriteOversize(t *testing.T) {
 func TestLazyWriteConcurrentFlush(t *testing.T) {
 	cipher := newTestCipher(t)
 	buf := new(bytes.Buffer)
-	writer := NewShadowsocksWriter(buf, cipher, RandomSaltGenerator)
+	writer := NewShadowsocksWriter(buf, cipher)
 	header := []byte{1, 2, 3, 4}
 	n, err := writer.LazyWrite(header)
 	if n != len(header) {

shadowsocks/tcp.go

@@ -265,7 +265,8 @@ func (s *tcpService) handleConnection(listenerPort int, clientConn onet.DuplexCo
 		clientConn.SetReadDeadline(time.Time{})
 
 		ssr := NewShadowsocksReader(clientReader, cipher)
-		ssw := NewShadowsocksWriter(clientConn, cipher, &recordingSaltGenerator{saltGenerator: RandomSaltGenerator, replayCache: s.replayCache, keyID: keyID})
+		ssw := NewShadowsocksWriter(clientConn, cipher)
+		ssw.SetSaltGenerator(&recordingSaltGenerator{saltGenerator: RandomSaltGenerator, replayCache: s.replayCache, keyID: keyID})
 		clientConn = onet.WrapConn(clientConn, ssr, ssw)
 		return proxyConnection(clientConn, &proxyMetrics, s.checkAllowedIP)
 	}()

shadowsocks/tcp_test.go

@@ -144,7 +144,7 @@ func BenchmarkTCPFindCipherRepeat(b *testing.B) {
 		addr := &net.TCPAddr{IP: clientIP, Port: 54321}
 		c := conn{clientAddr: addr, reader: reader, writer: writer}
 		cipher := cipherEntries[cipherNumber].Cipher
-		go NewShadowsocksWriter(writer, cipher, RandomSaltGenerator).Write(MakeTestPayload(50))
+		go NewShadowsocksWriter(writer, cipher).Write(MakeTestPayload(50))
 		b.StartTimer()
 		_, _, _, _, _, err := findAccessKey(&c, clientIP, cipherList)
 		b.StopTimer()
@@ -207,7 +207,7 @@ func TestReplayDefense(t *testing.T) {
 	cipherEntry := snapshot[0].Value.(*CipherEntry)
 	cipher := cipherEntry.Cipher
 	reader, writer := io.Pipe()
-	go NewShadowsocksWriter(writer, cipher, RandomSaltGenerator).Write([]byte{0})
+	go NewShadowsocksWriter(writer, cipher).Write([]byte{0})
 	preamble := make([]byte, 32+2+16)
 	if _, err := io.ReadFull(reader, preamble); err != nil {
 		t.Fatal(err)