diff --git a/utils/read_full.go b/utils/read_full.go new file mode 100644 index 0000000..afb2554 --- /dev/null +++ b/utils/read_full.go @@ -0,0 +1,20 @@ +package utils + +import "io" + +const readFullBufferSize = 1024 + 1 // +1 because telegram opreates with blocks mod 4 + +func ReadFull(src io.Reader) (rv []byte, err error) { + buf := make([]byte, readFullBufferSize) + n := readFullBufferSize + + for n == len(buf) { + n, err = src.Read(buf) + if err != nil { + return nil, err + } + rv = append(rv, buf[:n]...) + } + + return rv, nil +} diff --git a/wrappers/blockcipher.go b/wrappers/blockcipher.go index 12c015f..0ae2cb3 100644 --- a/wrappers/blockcipher.go +++ b/wrappers/blockcipher.go @@ -4,7 +4,6 @@ import ( "bytes" "crypto/aes" "crypto/cipher" - "errors" "fmt" "net" "time" @@ -12,10 +11,9 @@ import ( "go.uber.org/zap" "github.com/9seconds/mtg/conntypes" + "github.com/9seconds/mtg/utils" ) -const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4 - type wrapperBlockCipher struct { buf bytes.Buffer @@ -41,35 +39,29 @@ func (w *wrapperBlockCipher) WriteTimeout(p []byte, timeout time.Duration) (int, } func (w *wrapperBlockCipher) Read(p []byte) (int, error) { - return w.read(p, readAll) - -} - -func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, error) { - return w.read(p, readAllTimeout(timeout)) -} - -func (w *wrapperBlockCipher) read(p []byte, - reader func(conntypes.StreamReadWriteCloser) ([]byte, error)) (int, error) { if w.buf.Len() > 0 { return w.flush(p) } - var buf []byte - for len(buf) == 0 || len(buf)%aes.BlockSize != 0 { - rv, err := reader(w.parent) + var currentBuffer []byte + for len(currentBuffer) == 0 || len(currentBuffer)%aes.BlockSize != 0 { + rv, err := utils.ReadFull(w.parent) if err != nil { - return 0, fmt.Errorf("cannot read from socket: %w", err) + return 0, fmt.Errorf("cannot read data: %w", err) } - buf = append(buf, rv...) + currentBuffer = append(currentBuffer, rv...) } - w.decryptor.CryptBlocks(buf, buf) - w.buf.Write(buf) + w.decryptor.CryptBlocks(currentBuffer, currentBuffer) + w.buf.Write(currentBuffer) return w.flush(p) } +func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, error) { + return w.Read(p) +} + func (w *wrapperBlockCipher) flush(p []byte) (int, error) { if w.buf.Len() > len(p) { return w.buf.Read(p) @@ -93,44 +85,6 @@ func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) { return encrypted, nil } -func readAll(src conntypes.StreamReadWriteCloser) (rv []byte, err error) { - buf := make([]byte, blockCipherReadCurrentDataBufferSize) - n := blockCipherReadCurrentDataBufferSize - - for n == len(buf) { - n, err = src.Read(buf) - if err != nil { - return nil, err - } - rv = append(rv, buf[:n]...) - } - - return rv, nil -} - -func readAllTimeout(timeout time.Duration) func(conntypes.StreamReadWriteCloser) ([]byte, error) { - return func(src conntypes.StreamReadWriteCloser) (rv []byte, err error) { - tmo := timeout - buf := make([]byte, blockCipherReadCurrentDataBufferSize) - n := blockCipherReadCurrentDataBufferSize - - for n == len(buf) { - if tmo <= 0 { - return nil, errors.New("timeout") - } - startTime := time.Now() - n, err = src.ReadTimeout(buf, tmo) - if err != nil { - return nil, err - } - rv = append(rv, buf[:n]...) - tmo -= time.Since(startTime) - } - - return rv, nil - } -} - func (w *wrapperBlockCipher) Close() error { return w.parent.Close() }