package shadowsocks import ( "bytes" "crypto/hmac" "crypto/sha1" "encoding/binary" "io" "v2ray.com/core/common" "v2ray.com/core/common/buf" "v2ray.com/core/common/bytespool" "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, dest []byte) { hasher := hmac.New(sha1.New, v.key()) common.Must2(hasher.Write(data)) res := hasher.Sum(nil) copy(dest, res[:AuthSize]) } 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 := uint32(0) return func() []byte { newKey := make([]byte, len(iv)+4) copy(newKey, iv) binary.BigEndian.PutUint32(newKey[len(iv):], chunkID) 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) ReadMultiBuffer() (buf.MultiBuffer, error) { size, err := serial.ReadUint16(v.reader) if err != nil { return nil, newError("failed to read size").Base(err) } size += AuthSize buffer := bytespool.Alloc(int32(size)) defer bytespool.Free(buffer) if _, err := io.ReadFull(v.reader, buffer[:size]); err != nil { return nil, err } authBytes := buffer[:AuthSize] payload := buffer[AuthSize:size] actualAuthBytes := make([]byte, AuthSize) v.auth.Authenticate(payload, actualAuthBytes) if !bytes.Equal(authBytes, actualAuthBytes) { return nil, newError("invalid auth") } var mb buf.MultiBuffer common.Must2(mb.Write(payload)) return mb, 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), } } // WriteMultiBuffer implements buf.Writer. func (w *ChunkWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { defer mb.Release() for { payloadLen, _ := mb.Read(w.buffer[2+AuthSize:]) binary.BigEndian.PutUint16(w.buffer, uint16(payloadLen)) w.auth.Authenticate(w.buffer[2+AuthSize:2+AuthSize+payloadLen], w.buffer[2:]) if err := buf.WriteAllBytes(w.writer, w.buffer[:2+AuthSize+payloadLen]); err != nil { return err } if mb.IsEmpty() { break } } return nil }