github.com/ontio/ontology@v1.14.4/vm/neovm/utils/vm_reader.go (about)

     1  /*
     2   * Copyright (C) 2018 The ontology Authors
     3   * This file is part of The ontology library.
     4   *
     5   * The ontology is free software: you can redistribute it and/or modify
     6   * it under the terms of the GNU Lesser General Public License as published by
     7   * the Free Software Foundation, either version 3 of the License, or
     8   * (at your option) any later version.
     9   *
    10   * The ontology is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU Lesser General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU Lesser General Public License
    16   * along with The ontology.  If not, see <http://www.gnu.org/licenses/>.
    17   */
    18  
    19  package utils
    20  
    21  import (
    22  	"bytes"
    23  	"encoding/binary"
    24  	"io"
    25  )
    26  
    27  type VmReader struct {
    28  	reader     *bytes.Reader
    29  	BaseStream []byte
    30  	AllowEOF   bool // allow VmReader.ReadBytes got EOF and return 0 bytes
    31  }
    32  
    33  func NewVmReader(b []byte) *VmReader {
    34  	var vmreader VmReader
    35  	vmreader.reader = bytes.NewReader(b)
    36  	vmreader.BaseStream = b
    37  	return &vmreader
    38  }
    39  
    40  func (r *VmReader) ReadByte() (byte, error) {
    41  	byte, err := r.reader.ReadByte()
    42  	return byte, err
    43  }
    44  
    45  func (r *VmReader) ReadBytes(count int) ([]byte, error) {
    46  	// first check to avoid memory attack
    47  	amount := r.reader.Len()
    48  	if r.AllowEOF {
    49  		amount = 1 * 1024 * 1024
    50  	}
    51  	if amount < count {
    52  		return nil, io.EOF
    53  	}
    54  	b := make([]byte, count)
    55  	// when reader does not has enough bytes, it will return the actual length and error is nil
    56  	// this is a bug previously(it should use io.ReadFull), but we have to keep it to avoid hard fork
    57  	_, err := r.reader.Read(b)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	return b, nil
    62  }
    63  
    64  func (r *VmReader) ReadBytesInto(b []byte) error {
    65  	_, err := r.reader.Read(b)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	return nil
    70  }
    71  
    72  func (r *VmReader) ReadUint16() (uint16, error) {
    73  	var b [2]byte
    74  	err := r.ReadBytesInto(b[:])
    75  	if err != nil {
    76  		return 0, err
    77  	}
    78  	return binary.LittleEndian.Uint16(b[:]), nil
    79  }
    80  
    81  func (r *VmReader) ReadUint32() (uint32, error) {
    82  	var b [4]byte
    83  	err := r.ReadBytesInto(b[:])
    84  	if err != nil {
    85  		return 0, err
    86  	}
    87  	return binary.LittleEndian.Uint32(b[:]), nil
    88  }
    89  
    90  func (r *VmReader) ReadUint64() (uint64, error) {
    91  	var b [8]byte
    92  	err := r.ReadBytesInto(b[:])
    93  	if err != nil {
    94  		return 0, err
    95  	}
    96  	return binary.LittleEndian.Uint64(b[:]), nil
    97  }
    98  
    99  func (r *VmReader) ReadInt16() (int16, error) {
   100  	val, err := r.ReadUint16()
   101  	return int16(val), err
   102  }
   103  
   104  func (r *VmReader) ReadInt32() (int32, error) {
   105  	val, err := r.ReadUint32()
   106  	return int32(val), err
   107  }
   108  
   109  func (r *VmReader) Position() int {
   110  	return int(r.reader.Size()) - r.reader.Len()
   111  }
   112  
   113  func (r *VmReader) Length() int {
   114  	return r.reader.Len()
   115  }
   116  
   117  func (r *VmReader) Seek(offset int64, whence int) (int64, error) {
   118  	return r.reader.Seek(offset, whence)
   119  }
   120  
   121  func (r *VmReader) ReadVarBytes(max uint32) ([]byte, error) {
   122  	n, err := r.ReadVarInt(uint64(max))
   123  	if err != nil {
   124  		return nil, err
   125  	}
   126  	return r.ReadBytes(int(n))
   127  }
   128  
   129  func (r *VmReader) ReadVarInt(max uint64) (uint64, error) {
   130  	fb, _ := r.ReadByte()
   131  	var value uint64
   132  	var err error
   133  
   134  	switch fb {
   135  	case 0xFD:
   136  		val, e := r.ReadInt16()
   137  		value = uint64(val)
   138  		err = e
   139  	case 0xFE:
   140  		val, e := r.ReadUint32()
   141  		value = uint64(val)
   142  		err = e
   143  	case 0xFF:
   144  		val, e := r.ReadUint64()
   145  		value = uint64(val)
   146  		err = e
   147  	default:
   148  		value = uint64(fb)
   149  	}
   150  	if err != nil {
   151  		return 0, err
   152  	}
   153  	if value > max {
   154  		return 0, nil
   155  	}
   156  	return value, nil
   157  }
   158  
   159  func (r *VmReader) ReadVarString(maxlen uint32) (string, error) {
   160  	bs, err := r.ReadVarBytes(maxlen)
   161  	return string(bs), err
   162  }