package shadowsocks import ( "bytes" "crypto/hmac" "crypto/sha1" "io" "v2ray.com/core/common/buf" "v2ray.com/core/common/serial" ) const ( // AuthSize is the number of extra bytes for Shadowsocks OTA. AuthSize = 10 ) type KeyGenerator func() []byte type Authenticator struct { key KeyGenerator } func NewAuthenticator(keygen KeyGenerator) *Authenticator { return &Authenticator{ key: keygen, } } func (v *Authenticator) Authenticate(data []byte) buf.Supplier { hasher := hmac.New(sha1.New, v.key()) hasher.Write(data) res := hasher.Sum(nil) return func(b []byte) (int, error) { return copy(b, res[:AuthSize]), nil } } func HeaderKeyGenerator(key []byte, iv []byte) func() []byte { return func() []byte { newKey := make([]byte, 0, len(key)+len(iv)) newKey = append(newKey, iv...) newKey = append(newKey, key...) return newKey } } func ChunkKeyGenerator(iv []byte) func() []byte { chunkID := 0 return func() []byte { newKey := make([]byte, 0, len(iv)+4) newKey = append(newKey, iv...) newKey = serial.IntToBytes(chunkID, newKey) chunkID++ return newKey } } type ChunkReader struct { reader io.Reader auth *Authenticator } func NewChunkReader(reader io.Reader, auth *Authenticator) *ChunkReader { return &ChunkReader{ reader: reader, auth: auth, } } func (v *ChunkReader) Read() (buf.MultiBuffer, error) { buffer := buf.New() if err := buffer.AppendSupplier(buf.ReadFullFrom(v.reader, 2)); err != nil { buffer.Release() return nil, err } // There is a potential buffer overflow here. Large buffer is 64K bytes, // while uin16 + 10 will be more than that length := serial.BytesToUint16(buffer.BytesTo(2)) + AuthSize if length > buf.Size { // Theoretically the size of a chunk is 64K, but most Shadowsocks implementations used <4K buffer. buffer.Release() buffer = buf.NewLocal(int(length) + 128) } buffer.Clear() if err := buffer.AppendSupplier(buf.ReadFullFrom(v.reader, int(length))); err != nil { buffer.Release() return nil, err } authBytes := buffer.BytesTo(AuthSize) payload := buffer.BytesFrom(AuthSize) actualAuthBytes := make([]byte, AuthSize) v.auth.Authenticate(payload)(actualAuthBytes) if !bytes.Equal(authBytes, actualAuthBytes) { buffer.Release() return nil, newError("invalid auth") } buffer.SliceFrom(AuthSize) return buf.NewMultiBufferValue(buffer), nil } type ChunkWriter struct { writer io.Writer auth *Authenticator buffer []byte } func NewChunkWriter(writer io.Writer, auth *Authenticator) *ChunkWriter { return &ChunkWriter{ writer: writer, auth: auth, buffer: make([]byte, 32*1024), } } func (w *ChunkWriter) Write(mb buf.MultiBuffer) error { defer mb.Release() for { payloadLen, err := mb.Read(w.buffer[2+AuthSize:]) if err != nil { return err } serial.Uint16ToBytes(uint16(payloadLen), w.buffer[:0]) w.auth.Authenticate(w.buffer[2+AuthSize : 2+AuthSize+payloadLen])(w.buffer[2:]) if _, err := w.writer.Write(w.buffer[:2+AuthSize+payloadLen]); err != nil { return err } if mb.IsEmpty() { break } } return nil }