From 687ae6c50e2844c70d4d80ccd706d166190ff7cd Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Fri, 21 Oct 2016 00:33:23 +0200
Subject: [PATCH] chunk writer

---
 proxy/shadowsocks/ota.go      | 25 +++++++++++++++++++++++++
 proxy/shadowsocks/ota_test.go | 13 +++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/proxy/shadowsocks/ota.go b/proxy/shadowsocks/ota.go
index 5a6633d8..fa597537 100644
--- a/proxy/shadowsocks/ota.go
+++ b/proxy/shadowsocks/ota.go
@@ -100,3 +100,28 @@ func (this *ChunkReader) Read() (*alloc.Buffer, error) {
 
 	return buffer, nil
 }
+
+type ChunkWriter struct {
+	writer io.Writer
+	auth   *Authenticator
+}
+
+func NewChunkWriter(writer io.Writer, auth *Authenticator) *ChunkWriter {
+	return &ChunkWriter{
+		writer: writer,
+		auth:   auth,
+	}
+}
+
+func (this *ChunkWriter) Release() {
+	this.writer = nil
+	this.auth = nil
+}
+
+func (this *ChunkWriter) Write(payload *alloc.Buffer) (int, error) {
+	totalLength := payload.Len()
+	authBytes := this.auth.Authenticate(nil, payload.Bytes())
+	payload.Prepend(authBytes)
+	payload.PrependUint16(uint16(totalLength))
+	return this.writer.Write(payload.Bytes())
+}
diff --git a/proxy/shadowsocks/ota_test.go b/proxy/shadowsocks/ota_test.go
index 87d8e512..eaf53fb1 100644
--- a/proxy/shadowsocks/ota_test.go
+++ b/proxy/shadowsocks/ota_test.go
@@ -22,3 +22,16 @@ func TestNormalChunkReading(t *testing.T) {
 	payload.PrependBytes(3, 4)
 	assert.Bytes(payload.Value).Equals([]byte{3, 4, 11, 12, 13, 14, 15, 16, 17, 18})
 }
+
+func TestNormalChunkWriting(t *testing.T) {
+	assert := assert.On(t)
+
+	buffer := alloc.NewBuffer().Clear()
+	writer := NewChunkWriter(buffer, NewAuthenticator(ChunkKeyGenerator(
+		[]byte{21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36})))
+
+	nBytes, err := writer.Write(alloc.NewBuffer().Clear().Append([]byte{11, 12, 13, 14, 15, 16, 17, 18}))
+	assert.Error(err).IsNil()
+	assert.Int(nBytes).Equals(buffer.Len())
+	assert.Bytes(buffer.Value).Equals([]byte{0, 8, 39, 228, 69, 96, 133, 39, 254, 26, 201, 70, 11, 12, 13, 14, 15, 16, 17, 18})
+}