Golang io.ReadFull函数代码示例

// Extracts a usable sized symmetric key and IV for the stream cipher from the huge master key, and
// creates a CTR stream cipher.
func (s *Session) makeCipher() (cipher.Stream, error) {
	// Create the key derivation function
	hasher := func() hash.Hash { return s.hash.New() }
	hkdf := hkdf.New(hasher, s.secret.Bytes(), hkdfSalt, hkdfInfo)

	// Extract the symmetric key
	key := make([]byte, s.keybits/8)
	n, err := io.ReadFull(hkdf, key)
	if n != len(key) || err != nil {
		return nil, err
	}
	// Create the block cipher
	block, err := s.crypter(key)
	if err != nil {
		return nil, err
	}
	// Extract the IV for the counter mode
	iv := make([]byte, block.BlockSize())
	n, err = io.ReadFull(hkdf, iv)
	if n != len(iv) || err != nil {
		return nil, err
	}
	// Create the stream cipher
	return cipher.NewCTR(block, iv), nil
}

// ServeAgent serves the agent protocol on the given connection. It
// returns when an I/O error occurs.
func ServeAgent(agent Agent, c io.ReadWriter) error {
	s := &server{agent}

	var length [4]byte
	for {
		if _, err := io.ReadFull(c, length[:]); err != nil {
			return err
		}
		l := binary.BigEndian.Uint32(length[:])
		if l > maxAgentResponseBytes {
			// We also cap requests.
			return fmt.Errorf("agent: request too large: %d", l)
		}

		req := make([]byte, l)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}

		repData := s.processRequestBytes(req)
		if len(repData) > maxAgentResponseBytes {
			return fmt.Errorf("agent: reply too large: %d bytes", len(repData))
		}

		binary.BigEndian.PutUint32(length[:], uint32(len(repData)))
		if _, err := c.Write(length[:]); err != nil {
			return err
		}
		if _, err := c.Write(repData); err != nil {
			return err
		}
	}
}

func (s *Server) Handle(serial int, RpcAddr string) {
	conn := s.conns[serial]
	head := make([]byte, 2)
	defer s.Close(serial)
	for {
		_, err := io.ReadFull(conn, head)
		if err != nil {
			fmt.Println(err)
			break
		}
		var info_len uint16
		buf := bytes.NewReader(head)
		if err := binary.Read(buf, binary.BigEndian, &info_len); err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		body := make([]byte, info_len)
		body_len, err := io.ReadFull(conn, body)
		if err != nil {
			fmt.Println(err)
			break
		}
		reply := CallRpc(RpcAddr, "Arith.Multiply", int(info_len), int(body_len))
		fmt.Printf("reply:%d\n", reply)
		fmt.Println("len", info_len, body_len, string(body))
		s.WriteTo(serial, body)
	}
}

// recvMessage receives any message from the backend, or returns an error if
// a problem occurred while reading the message.
func (cn *conn) recvMessage() (byte, *readBuf, error) {
	// workaround for a QueryRow bug, see exec
	if cn.saveMessageType != 0 {
		t, r := cn.saveMessageType, cn.saveMessageBuffer
		cn.saveMessageType = 0
		cn.saveMessageBuffer = nil
		return t, r, nil
	}

	x := cn.scratch[:5]
	_, err := io.ReadFull(cn.buf, x)
	if err != nil {
		return 0, nil, err
	}
	t := x[0]

	b := readBuf(x[1:])
	n := b.int32() - 4
	var y []byte
	if n <= len(cn.scratch) {
		y = cn.scratch[:n]
	} else {
		y = make([]byte, n)
	}
	_, err = io.ReadFull(cn.buf, y)
	if err != nil {
		return 0, nil, err
	}

	return t, (*readBuf)(&y), nil
}

func TestPartialRead(t *testing.T) {
	f, err := os.Open("testdata/gnu.tar")
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	defer f.Close()

	tr := NewReader(f)

	// Read the first four bytes; Next() should skip the last byte.
	hdr, err := tr.Next()
	if err != nil || hdr == nil {
		t.Fatalf("Didn't get first file: %v", err)
	}
	buf := make([]byte, 4)
	if _, err := io.ReadFull(tr, buf); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if expected := []byte("Kilt"); !bytes.Equal(buf, expected) {
		t.Errorf("Contents = %v, want %v", buf, expected)
	}

	// Second file
	hdr, err = tr.Next()
	if err != nil || hdr == nil {
		t.Fatalf("Didn't get second file: %v", err)
	}
	buf = make([]byte, 6)
	if _, err := io.ReadFull(tr, buf); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if expected := []byte("Google"); !bytes.Equal(buf, expected) {
		t.Errorf("Contents = %v, want %v", buf, expected)
	}
}

func (self *Session) reader() {
	log.Printf("session[%d] start reader...", self.id)
	defer func() {
		log.Println("reader quit...")
		self.ctrl <- true
		if self.autoReconnect {
			self.reconnect <- true
		} else {
			self.events <- newEvent(EVENT_DISCONNECT, self, nil)
		}
	}()
	header := make([]byte, 2)
	for {
		if _, err := io.ReadFull(self.conn, header); err != nil {
			break
		}
		size := binary.BigEndian.Uint16(header)
		log.Printf("size=%d", size)
		data := make([]byte, size)
		if _, err := io.ReadFull(self.conn, data); err != nil {
			self.events <- newEvent(EVENT_RECV_ERROR, self, err)
			break
		}
		log.Printf("len(data)=%d", len(data))
		log.Printf("payload: %v", data)
		api, payload, err := self.proto.Decode(data)
		log.Printf("api=%v, payload=%v, err=%v", api, payload, err)
		if err != nil {
			self.events <- newEvent(EVENT_RECV_ERROR, self, err)
			break
		}
		msg := NewMessage(api, payload)
		self.events <- newEvent(EVENT_MESSAGE, self, msg)
	}
}

func readN(r io.Reader, b []byte, n int) ([]byte, error) {
	if n == 0 && b == nil {
		return make([]byte, 0), nil
	}

	if cap(b) >= n {
		b = b[:n]
		_, err := io.ReadFull(r, b)
		return b, err
	}
	b = b[:cap(b)]

	pos := 0
	for len(b) < n {
		diff := n - len(b)
		if diff > bytesAllocLimit {
			diff = bytesAllocLimit
		}
		b = append(b, make([]byte, diff)...)

		_, err := io.ReadFull(r, b[pos:])
		if err != nil {
			return nil, err
		}

		pos = len(b)
	}

	return b, nil
}

// readDirectoryHeader attempts to read a directory header from r.
// It returns io.ErrUnexpectedEOF if it cannot read a complete header,
// and ErrFormat if it doesn't find a valid header signature.
func readDirectoryHeader(f *File, r io.Reader) error {
	var buf [directoryHeaderLen]byte
	if _, err := io.ReadFull(r, buf[:]); err != nil {
		return err
	}
	b := readBuf(buf[:])
	if sig := b.uint32(); sig != directoryHeaderSignature {
		return ErrFormat
	}
	f.CreatorVersion = b.uint16()
	f.ReaderVersion = b.uint16()
	f.Flags = b.uint16()
	f.Method = b.uint16()
	f.ModifiedTime = b.uint16()
	f.ModifiedDate = b.uint16()
	f.CRC32 = b.uint32()
	f.CompressedSize = b.uint32()
	f.UncompressedSize = b.uint32()
	filenameLen := int(b.uint16())
	extraLen := int(b.uint16())
	commentLen := int(b.uint16())
	b = b[4:] // skipped start disk number and internal attributes (2x uint16)
	f.ExternalAttrs = b.uint32()
	f.headerOffset = int64(b.uint32())
	d := make([]byte, filenameLen+extraLen+commentLen)
	if _, err := io.ReadFull(r, d); err != nil {
		return err
	}
	f.Name = string(d[:filenameLen])
	f.Extra = d[filenameLen : filenameLen+extraLen]
	f.Comment = string(d[filenameLen+extraLen:])
	return nil
}

func TestLargeMessage(t *testing.T) {
	server, ct := setUp(t, 0, math.MaxUint32, normal)
	callHdr := &CallHdr{
		Host:   "localhost",
		Method: "foo.Large",
	}
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s, err := ct.NewStream(context.Background(), callHdr)
			if err != nil {
				t.Errorf("failed to open stream: %v", err)
			}
			if err := ct.Write(s, expectedRequestLarge, &Options{Last: true, Delay: false}); err != nil {
				t.Errorf("failed to send data: %v", err)
			}
			p := make([]byte, len(expectedResponseLarge))
			_, recvErr := io.ReadFull(s, p)
			if recvErr != nil || !bytes.Equal(p, expectedResponseLarge) {
				t.Errorf("Error: %v, want <nil>; Result len: %d, want len %d", recvErr, len(p), len(expectedResponseLarge))
			}
			_, recvErr = io.ReadFull(s, p)
			if recvErr != io.EOF {
				t.Errorf("Error: %v; want <EOF>", recvErr)
			}
		}()
	}
	wg.Wait()
	ct.Close()
	server.stop()
}

func DecryptIO(reader io.Reader, writer io.Writer, passphrase string) (mac []byte, e error) {
	key := PassphraseToKey(passphrase)
	hashMac := hmac.New(sha256.New, key)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	iv := make([]byte, aes.BlockSize)
	if _, err = io.ReadFull(reader, iv); err != nil {
		return nil, err
	}
	stream := cipher.NewOFB(block, iv)
	cipherReader := &cipher.StreamReader{S: stream, R: reader}
	pp := make([]byte, len([]byte(passphrase)))
	if _, err = io.ReadFull(cipherReader, pp); err != nil {
		return nil, err
	}
	if passphrase != string(pp) {
		return nil, errors.New("Incorrect passphrase")
	}
	mw := io.MultiWriter(writer, hashMac)
	if _, err = io.Copy(mw, cipherReader); err != nil {
		return nil, err
	}
	return hashMac.Sum(nil), nil
}

// The initial response from the datanode:
// +-----------------------------------------------------------+
// |  varint length + BlockOpResponseProto                     |
// +-----------------------------------------------------------+
func (br *BlockReader) readBlockReadResponse(r io.Reader) (*hdfs.BlockOpResponseProto, error) {
	varintBytes := make([]byte, binary.MaxVarintLen32)
	_, err := io.ReadFull(r, varintBytes)
	if err != nil {
		return nil, err
	}

	respLength, varintLength := binary.Uvarint(varintBytes)
	if varintLength < 1 {
		return nil, io.ErrUnexpectedEOF
	}

	// We may have grabbed too many bytes when reading the varint.
	respBytes := make([]byte, respLength)
	extraLength := copy(respBytes, varintBytes[varintLength:])
	_, err = io.ReadFull(r, respBytes[extraLength:])
	if err != nil {
		return nil, err
	}

	resp := &hdfs.BlockOpResponseProto{}
	err = proto.Unmarshal(respBytes, resp)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

// FromReader extracts data from a serialized DeleteOp into its concrete
// structure.
func (op *DeleteOp) FromReader(r io.Reader) error {
	var b [4]byte
	_, err := io.ReadFull(r, b[:]) //skip ZERO
	if err != nil {
		return err
	}
	name, err := readCStringFromReader(r)
	if err != nil {
		return err
	}
	op.Collection = string(name)
	_, err = io.ReadFull(r, b[:]) //Grab the flags
	if err != nil {
		return err
	}
	op.Flags = uint32(getInt32(b[:], 0))

	selectorAsSlice, err := ReadDocument(r)
	if err != nil {
		return err
	}
	op.Selector = &bson.D{}
	err = bson.Unmarshal(selectorAsSlice, op.Selector)

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

// Unmarshal parses an identity.
func Unmarshal(in []byte) (*Identity, error) {
	buf := sbuf.NewBufferFrom(in)
	id := &Identity{
		private: new([ed25519.PrivateKeySize]byte),
		public:  new([ed25519.PublicKeySize]byte),
	}

	_, err := io.ReadFull(buf, id.private[:])
	if err != nil {
		return nil, ErrInvalidIdentity
	}

	_, err = io.ReadFull(buf, id.public[:])
	if err != nil {
		return nil, ErrInvalidIdentity
	}

	if (buf.Len() % ed25519.PublicKeySize) != 0 {
		return nil, ErrInvalidIdentity
	}

	for {
		if buf.Len() == 0 {
			break
		}

		peer := new([ed25519.PublicKeySize]byte)
		io.ReadFull(buf, peer[:])
		id.peers = append(id.peers, peer)
	}

	return id, nil
}

// NewReaderDict is like NewReader but uses a preset dictionary.
// NewReaderDict ignores the dictionary if the compressed data does not refer to it.
func NewReaderDict(r io.Reader, dict []byte) (io.ReadCloser, os.Error) {
	z := new(reader)
	if fr, ok := r.(flate.Reader); ok {
		z.r = fr
	} else {
		z.r = bufio.NewReader(r)
	}
	_, err := io.ReadFull(z.r, z.scratch[0:2])
	if err != nil {
		return nil, err
	}
	h := uint(z.scratch[0])<<8 | uint(z.scratch[1])
	if (z.scratch[0]&0x0f != zlibDeflate) || (h%31 != 0) {
		return nil, HeaderError
	}
	if z.scratch[1]&0x20 != 0 {
		_, err = io.ReadFull(z.r, z.scratch[0:4])
		if err != nil {
			return nil, err
		}
		checksum := uint32(z.scratch[0])<<24 | uint32(z.scratch[1])<<16 | uint32(z.scratch[2])<<8 | uint32(z.scratch[3])
		if checksum != adler32.Checksum(dict) {
			return nil, DictionaryError
		}
		z.decompressor = flate.NewReaderDict(z.r, dict)
	} else {
		z.decompressor = flate.NewReader(z.r)
	}
	z.digest = adler32.New()
	return z, nil
}

func performOneRPC(ct ClientTransport) {
	callHdr := &CallHdr{
		Host:   "localhost",
		Method: "foo.Small",
	}
	s, err := ct.NewStream(context.Background(), callHdr)
	if err != nil {
		return
	}
	opts := Options{
		Last:  true,
		Delay: false,
	}
	if err := ct.Write(s, expectedRequest, &opts); err == nil || err == io.EOF {
		time.Sleep(5 * time.Millisecond)
		// The following s.Recv()'s could error out because the
		// underlying transport is gone.
		//
		// Read response
		p := make([]byte, len(expectedResponse))
		io.ReadFull(s, p)
		// Read io.EOF
		io.ReadFull(s, p)
	}
}