Golang io.ReadAtLeast函数代码示例

func decodeTransliterations() {
	r, err := zlib.NewReader(strings.NewReader(tableData))
	if err != nil {
		panic(err)
	}
	defer r.Close()
	tmp1 := make([]byte, 2)
	tmp2 := tmp1[:1]
	for {
		if _, err := io.ReadAtLeast(r, tmp1, 2); err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chr := getUint16(tmp1)
		if _, err := io.ReadAtLeast(r, tmp2, 1); err != nil {
			panic(err)
		}
		b := make([]byte, int(tmp2[0]))
		if _, err := io.ReadFull(r, b); err != nil {
			panic(err)
		}
		transliterations[int(chr)] = []rune(string(b))
	}
}

// readerContains reports whether subslice is within r.
func readerContains(r io.Reader, subslice []byte) bool {

	if r == nil || len(subslice) == 0 {
		return false
	}

	bufflen := len(subslice) * 4
	halflen := bufflen / 2
	buff := make([]byte, bufflen)
	var err error
	var n, i int

	for {
		i++
		if i == 1 {
			n, err = io.ReadAtLeast(r, buff[:halflen], halflen)
		} else {
			if i != 2 {
				// shift left to catch overlapping matches
				copy(buff[:], buff[halflen:])
			}
			n, err = io.ReadAtLeast(r, buff[halflen:], halflen)
		}

		if n > 0 && bytes.Contains(buff, subslice) {
			return true
		}

		if err != nil {
			break
		}
	}
	return false
}

func testReadAtLeast(fn string) {
	fmt.Println("testReadAtLeast---------------------------------------")
	r, err := os.Open(fn)
	if err != nil {
		fmt.Println("open file error:", fn, err)
		return
	}
	buf := make([]byte, 20)
	defer r.Close()
	_, rerr := io.ReadAtLeast(r, buf, 1)
	for {
		if rerr == nil {
			// fmt.Println(buf)
			fmt.Println(string(buf))
			fmt.Println("######", len(buf))
			_, rerr = io.ReadAtLeast(r, buf, 1)
		} else {
			if rerr == io.EOF {
				fmt.Println(" read at the end!!!")
			} else {
				fmt.Println("read error : ", rerr)
			}
			break
		}
	}
}

func (s *Socks5) serverVerifyUserPass(conn net.Conn) error {
	var req [512]byte
	un, err := io.ReadAtLeast(conn, req[:], 2)
	if err != nil {
		return err
	}
	if req[0] != USER_PASS_VERIFY_VER {
		conn.Write([]byte{USER_PASS_VERIFY_VER, USER_PASS_VERIFY_FAILED})
		return ErrNoProxy
	}
	userLen := int(req[2])
	pn, err := io.ReadAtLeast(conn, req[un:], userLen+1)
	if err != nil {
		return err
	}
	passLen := int(req[userLen+2])
	_, err = io.ReadFull(conn, req[un+pn:3+userLen+passLen])
	if err != nil {
		return err
	}

	user := req[2 : 2+userLen]
	pass := req[3+userLen : 3+userLen+passLen]
	verified := s.userPass.Verify(string(user), string(pass))
	if verified {
		conn.Write([]byte{USER_PASS_VERIFY_VER, USER_PASS_VERIFY_SUCCESS})
		return nil
	}
	conn.Write([]byte{USER_PASS_VERIFY_VER, USER_PASS_VERIFY_FAILED})
	return ErrNoProxy
}

func (h *Hilite) Render(lang string, text []byte) ([]byte, error) {
	h.m.Lock()
	defer h.m.Unlock()
	size := make([]byte, 4)
	binary.BigEndian.PutUint32(size, uint32(len(lang)))
	h.w.Write(size)
	h.w.Write([]byte(lang))
	binary.BigEndian.PutUint32(size, uint32(len(text)))
	h.w.Write(size)
	h.w.Write(text)
	_, err := io.ReadAtLeast(h.r, size, 4)
	if err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(size)
	if length == 0 {
		return nil, io.EOF
	}
	out := make([]byte, length)
	_, err = io.ReadAtLeast(h.r, out, int(length))
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (t *PrepareReply) Unmarshal(rr io.Reader) error {
	var wire byteReader
	var ok bool
	if wire, ok = rr.(byteReader); !ok {
		wire = bufio.NewReader(rr)
	}
	var b [10]byte
	var bs []byte
	bs = b[:9]
	if _, err := io.ReadAtLeast(wire, bs, 9); err != nil {
		return err
	}
	t.Balnum = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	t.OK = uint8(bs[4])
	t.Ballot = int32((uint32(bs[5]) | (uint32(bs[6]) << 8) | (uint32(bs[7]) << 16) | (uint32(bs[8]) << 24)))
	alen1, err := binary.ReadVarint(wire)
	if err != nil {
		return err
	}
	t.Cstruct = make([]int32, alen1)
	for i := int64(0); i < alen1; i++ {
		bs = b[:4]
		if _, err := io.ReadAtLeast(wire, bs, 4); err != nil {
			return err
		}
		t.Cstruct[i] = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	}
	return nil
}

func streamsEqual(fa, fb io.Reader) bool {
	bufa := make([]byte, bufsz)
	bufb := make([]byte, bufsz)
	for {
		ra, erra := io.ReadAtLeast(fa, bufa, bufsz)
		rb, errb := io.ReadAtLeast(fb, bufb, bufsz)
		if erra == io.EOF && errb == io.EOF {
			return true
		}
		if erra != nil || errb != nil {
			// if both files finished in the middle of a
			// ReadAtLeast, (returning io.ErrUnexpectedEOF), then we
			// still need to check what was read to know whether
			// they're equal.  Otherwise, we know they're not equal
			// (because we count any read error as a being non-equal
			// also).
			tailMightBeEqual := erra == io.ErrUnexpectedEOF && errb == io.ErrUnexpectedEOF
			if !tailMightBeEqual {
				return false
			}
		}
		if !bytes.Equal(bufa[:ra], bufb[:rb]) {
			return false
		}
	}
}

func ExampleReadAtLeast() {
	r := strings.NewReader("some io.Reader stream to be read\n")

	buf := make([]byte, 33)
	if _, err := io.ReadAtLeast(r, buf, 4); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf)

	// buffer smaller than minimal read size.
	shortBuf := make([]byte, 3)
	if _, err := io.ReadAtLeast(r, shortBuf, 4); err != nil {
		fmt.Println("error:", err)
	}

	// minimal read size bigger than io.Reader stream
	longBuf := make([]byte, 64)
	if _, err := io.ReadAtLeast(r, longBuf, 64); err != nil {
		fmt.Println("error:", err)
	}

	// Output:
	// some io.Reader stream to be read
	//
	// error: short buffer
	// error: EOF
}

func GetProto(conn net.Conn) (*comm.OhHai, error) {
	// Read the first 8 bytes.
	buf := make([]byte, 8)
	_, err := io.ReadAtLeast(conn, buf, 8)
	if err != nil {
		return nil, err
	}

	// Convert first 8 bytes to int
	lenbuf := bytes.NewBuffer(buf)
	var length int64
	err = binary.Read(lenbuf, binary.LittleEndian, &length)
	if err != nil {
		panic(err)
	}

	// Get the protobuf from the connection
	buf = make([]byte, length)
	_, err = io.ReadAtLeast(conn, buf, int(length))
	if err != nil {
		return nil, err
	}

	// Unmarshal the protobuf
	message := &comm.OhHai{}
	err = proto.Unmarshal(buf, message)
	if err != nil {
		return nil, err
	}

	return message, nil
}

func HandshakeSample(c net.Conn, br *bufio.Reader, bw *bufio.Writer, timeout time.Duration) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = r.(error)
		}
	}()
	// Send C0+C1
	err = bw.WriteByte(0x03)
	c1 := CreateRandomBlock(RTMP_SIG_SIZE)
	for i := 0; i < 8; i++ {
		c1[i] = 0
	}
	bw.Write(c1)
	err = bw.Flush()
	CheckError(err, "Handshake() Flush C0+C1")
	// Read S0+S1+S2
	s0, err := br.ReadByte()
	CheckError(err, "Handshake() Read S0")
	if s0 != 0x03 {
		return errors.New(fmt.Sprintf("Handshake() Got S0: %x", s0))
	}
	s1 := make([]byte, RTMP_SIG_SIZE)
	_, err = io.ReadAtLeast(br, s1, RTMP_SIG_SIZE)
	CheckError(err, "Handshake() Read S1")
	bw.Write(s1)
	err = bw.Flush()
	CheckError(err, "Handshake() Flush C2")
	_, err = io.ReadAtLeast(br, s1, RTMP_SIG_SIZE)
	CheckError(err, "Handshake() Read S2")
	return
}

func (t *Commit) Unmarshal(rr io.Reader) error {
	var wire byteReader
	var ok bool
	if wire, ok = rr.(byteReader); !ok {
		wire = bufio.NewReader(rr)
	}
	var b [24]byte
	var bs []byte
	bs = b[:12]
	if _, err := io.ReadAtLeast(wire, bs, 12); err != nil {
		return err
	}
	t.LeaderId = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	t.Replica = int32((uint32(bs[4]) | (uint32(bs[5]) << 8) | (uint32(bs[6]) << 16) | (uint32(bs[7]) << 24)))
	t.Instance = int32((uint32(bs[8]) | (uint32(bs[9]) << 8) | (uint32(bs[10]) << 16) | (uint32(bs[11]) << 24)))
	alen1, err := binary.ReadVarint(wire)
	if err != nil {
		return err
	}
	t.Command = make([]state.Command, alen1)
	for i := int64(0); i < alen1; i++ {
		t.Command[i].Unmarshal(wire)
	}
	bs = b[:24]
	if _, err := io.ReadAtLeast(wire, bs, 24); err != nil {
		return err
	}
	t.Seq = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	t.Deps[0] = int32((uint32(bs[4]) | (uint32(bs[5]) << 8) | (uint32(bs[6]) << 16) | (uint32(bs[7]) << 24)))
	t.Deps[1] = int32((uint32(bs[8]) | (uint32(bs[9]) << 8) | (uint32(bs[10]) << 16) | (uint32(bs[11]) << 24)))
	t.Deps[2] = int32((uint32(bs[12]) | (uint32(bs[13]) << 8) | (uint32(bs[14]) << 16) | (uint32(bs[15]) << 24)))
	t.Deps[3] = int32((uint32(bs[16]) | (uint32(bs[17]) << 8) | (uint32(bs[18]) << 16) | (uint32(bs[19]) << 24)))
	t.Deps[4] = int32((uint32(bs[20]) | (uint32(bs[21]) << 8) | (uint32(bs[22]) << 16) | (uint32(bs[23]) << 24)))
	return nil
}

// Read attempts to read an entire frame into data. If that is not possible, io.ErrShortBuffer
// is returned and subsequent calls will attempt to read the last frame. A frame is complete when
// err is nil.
func (r *lengthDelimitedFrameReader) Read(data []byte) (int, error) {
	if r.remaining <= 0 {
		header := [4]byte{}
		n, err := io.ReadAtLeast(r.r, header[:4], 4)
		if err != nil {
			return 0, err
		}
		if n != 4 {
			return 0, io.ErrUnexpectedEOF
		}
		frameLength := int(binary.BigEndian.Uint32(header[:]))
		r.remaining = frameLength
	}

	expect := r.remaining
	max := expect
	if max > len(data) {
		max = len(data)
	}
	n, err := io.ReadAtLeast(r.r, data[:max], int(max))
	r.remaining -= n
	if err == io.ErrShortBuffer || r.remaining > 0 {
		return n, io.ErrShortBuffer
	}
	if err != nil {
		return n, err
	}
	if n != expect {
		return n, io.ErrUnexpectedEOF
	}

	return n, nil
}

func (t *Accept) Unmarshal(rr io.Reader) error {
	var wire byteReader
	var ok bool
	if wire, ok = rr.(byteReader); !ok {
		wire = bufio.NewReader(rr)
	}
	var b [12]byte
	var bs []byte
	bs = b[:12]
	if _, err := io.ReadAtLeast(wire, bs, 12); err != nil {
		return err
	}
	t.LeaderId = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	t.Instance = int32((uint32(bs[4]) | (uint32(bs[5]) << 8) | (uint32(bs[6]) << 16) | (uint32(bs[7]) << 24)))
	t.Ballot = int32((uint32(bs[8]) | (uint32(bs[9]) << 8) | (uint32(bs[10]) << 16) | (uint32(bs[11]) << 24)))
	alen1, err := binary.ReadVarint(wire)
	if err != nil {
		return err
	}
	t.Command = make([]state.Command, alen1)
	for i := int64(0); i < alen1; i++ {
		t.Command[i].Unmarshal(wire)
	}
	if _, err := io.ReadAtLeast(wire, bs, 12); err != nil {
		return err
	}
	t.LeaseInstance = int32((uint32(bs[0]) | (uint32(bs[1]) << 8) | (uint32(bs[2]) << 16) | (uint32(bs[3]) << 24)))
	t.OriginReplica = int32((uint32(bs[4]) | (uint32(bs[5]) << 8) | (uint32(bs[6]) << 16) | (uint32(bs[7]) << 24)))
	t.PropId = int32((uint32(bs[8]) | (uint32(bs[9]) << 8) | (uint32(bs[10]) << 16) | (uint32(bs[11]) << 24)))
	return nil
}

// loadUtxoView returns a utxo view loaded from a file.
func loadUtxoView(filename string) (*blockchain.UtxoViewpoint, error) {
	// The utxostore file format is:
	// <tx hash><serialized utxo len><serialized utxo>
	//
	// The serialized utxo len is a little endian uint32 and the serialized
	// utxo uses the format described in chainio.go.

	filename = filepath.Join("testdata", filename)
	fi, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	// Choose read based on whether the file is compressed or not.
	var r io.Reader
	if strings.HasSuffix(filename, ".bz2") {
		r = bzip2.NewReader(fi)
	} else {
		r = fi
	}
	defer fi.Close()

	view := blockchain.NewUtxoViewpoint()
	for {
		// Hash of the utxo entry.
		var hash wire.ShaHash
		_, err := io.ReadAtLeast(r, hash[:], len(hash[:]))
		if err != nil {
			// Expected EOF at the right offset.
			if err == io.EOF {
				break
			}
			return nil, err
		}

		// Num of serialize utxo entry bytes.
		var numBytes uint32
		err = binary.Read(r, binary.LittleEndian, &numBytes)
		if err != nil {
			return nil, err
		}

		// Serialized utxo entry.
		serialized := make([]byte, numBytes)
		_, err = io.ReadAtLeast(r, serialized, int(numBytes))
		if err != nil {
			return nil, err
		}

		// Deserialize it and add it to the view.
		utxoEntry, err := blockchain.TstDeserializeUtxoEntry(serialized)
		if err != nil {
			return nil, err
		}
		view.Entries()[hash] = utxoEntry
	}

	return view, nil
}

// LoadNext reads and returns the next BSON document in the stream. If the
// BSONSource was created with NewBSONSource then each returned []byte will be
// a slice of a single reused I/O buffer. If the BSONSource was created with
// NewBufferlessBSONSource then each returend []byte will be individually
// allocated
func (bs *BSONSource) LoadNext() []byte {
	var into []byte
	if bs.reusableBuf == nil {
		into = make([]byte, 4)
	} else {
		into = bs.reusableBuf
	}
	// read the bson object size (a 4 byte integer)
	_, err := io.ReadAtLeast(bs.Stream, into[0:4], 4)
	if err != nil {
		if err != io.EOF {
			bs.err = err
			return nil
		}
		// we hit EOF right away, so we're at the end of the stream.
		bs.err = nil
		return nil
	}

	bsonSize := int32(
		(uint32(into[0]) << 0) |
			(uint32(into[1]) << 8) |
			(uint32(into[2]) << 16) |
			(uint32(into[3]) << 24),
	)

	// Verify that the size of the BSON object we are about to read can
	// actually fit into the buffer that was provided. If not, either the BSON is
	// invalid, or the buffer passed in is too small.
	// Verify that we do not have an invalid BSON document with size < 5.
	if bsonSize > MaxBSONSize || bsonSize < 5 {
		bs.err = fmt.Errorf("invalid BSONSize: %v bytes", bsonSize)
		return nil
	}
	if int(bsonSize) > cap(into) {
		bigInto := make([]byte, bsonSize)
		copy(bigInto, into)
		into = bigInto
		if bs.reusableBuf != nil {
			bs.reusableBuf = bigInto
		}
	}
	into = into[:int(bsonSize)]
	_, err = io.ReadAtLeast(bs.Stream, into[4:], int(bsonSize-4))
	if err != nil {
		if err != io.EOF {
			bs.err = err
			return nil
		}
		// this case means we hit EOF but read a partial document,
		// so there's a broken doc in the stream. Treat this as error.
		bs.err = fmt.Errorf("invalid bson: %v", err)
		return nil
	}

	bs.err = nil
	return into
}