Golang io.CopyN函數代碼示例

func TestConcurrent(t *testing.T) {
	var grp sync.WaitGroup
	buf := New()

	var rs []io.ReadCloser
	for i := 0; i < 1000; i++ {
		rs = append(rs, buf.NextReader())
	}

	testData := bytes.NewBuffer(nil)
	io.CopyN(testData, rand.Reader, 32*1024*10)

	for _, r := range rs {
		grp.Add(1)
		go func(r io.ReadCloser) {
			defer grp.Done()
			defer r.Close()
			data, err := ioutil.ReadAll(r)
			if err != nil {
				t.Error(err)
			}
			if !bytes.Equal(testData.Bytes(), data) {
				t.Error("unexpected result...", testData.Len(), len(data))
			}
		}(r)
	}

	r := bytes.NewReader(testData.Bytes())
	for r.Len() > 0 {
		io.CopyN(buf, r, 32*1024*2)
		<-time.After(100 * time.Millisecond)
	}
	buf.Close()
	grp.Wait()
}

func TestMixReader(t *testing.T) {
	r := NewMixReader(nil, nil)
	if _, err := r.Read(make([]byte, 1)); err == nil {
		t.Error("nil source should failed.")
	}

	r = NewMixReader(bytes.NewBuffer([]byte{0x00}), bytes.NewReader([]byte{0x00}))
	if _, err := io.CopyN(NewBytesWriter(make([]byte, 2)), r, 2); err != nil {
		t.Error("should not be nil")
	}
	if _, err := r.Read(make([]byte, 1)); err == nil {
		t.Error("should dry")
	}

	r = NewMixReader(bytes.NewBuffer([]byte{0x00}), bytes.NewReader([]byte{0x00}))
	if _, err := io.CopyN(NewBytesWriter(make([]byte, 1)), r, 1); err != nil {
		t.Error("should not be nil")
	}
	if _, err := io.CopyN(NewBytesWriter(make([]byte, 1)), r, 1); err != nil {
		t.Error("should not be nil")
	}
	if _, err := r.Read(make([]byte, 1)); err == nil {
		t.Error("should dry")
	}
}

func ExampleReadTwoOpenedUncompressedFiles() {
	var fs http.FileSystem = assets

	f0, err := fs.Open("/not-worth-compressing-file.txt")
	if err != nil {
		panic(err)
	}
	defer f0.Close()
	_ = f0.(notWorthGzipCompressing)
	f1, err := fs.Open("/not-worth-compressing-file.txt")
	if err != nil {
		panic(err)
	}
	defer f1.Close()
	_ = f1.(notWorthGzipCompressing)

	_, err = io.CopyN(os.Stdout, f0, 9)
	if err != nil {
		panic(err)
	}
	_, err = io.CopyN(os.Stdout, f1, 9)
	if err != nil {
		panic(err)
	}

	// Output:
	// Its normaIts norma
}

func ExampleBytes() {
	buf := newWriter(make([]byte, 0, 10))
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, io.NewSectionReader(*&buf, 0, 100))

	io.WriteString(buf, "Hello ")
	r := io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 5)
	io.WriteString(buf, "World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 6)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)

	io.WriteString(buf, "Hello World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 4)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)
	//Output:
	// HelloHello World WorldabcdefgHelloHello Worldabcdefg
}

// decryptRemoteState is used to help decrypt the remote state
func (m *Memberlist) decryptRemoteState(bufConn io.Reader) ([]byte, error) {
	// Read in enough to determine message length
	cipherText := bytes.NewBuffer(nil)
	cipherText.WriteByte(byte(encryptMsg))
	_, err := io.CopyN(cipherText, bufConn, 4)
	if err != nil {
		return nil, err
	}

	// Ensure we aren't asked to download too much. This is to guard against
	// an attack vector where a huge amount of state is sent
	moreBytes := binary.BigEndian.Uint32(cipherText.Bytes()[1:5])
	if moreBytes > maxPushStateBytes {
		return nil, fmt.Errorf("Remote node state is larger than limit (%d)", moreBytes)
	}

	// Read in the rest of the payload
	_, err = io.CopyN(cipherText, bufConn, int64(moreBytes))
	if err != nil {
		return nil, err
	}

	// Decrypt the cipherText
	dataBytes := cipherText.Bytes()[:5]
	cipherBytes := cipherText.Bytes()[5:]

	// Decrypt the payload
	keys := m.config.Keyring.GetKeys()
	return decryptPayload(keys, cipherBytes, dataBytes)
}

func ExampleReadTwoOpenedCompressedFiles() {
	var fs http.FileSystem = assets

	f0, err := fs.Open("/sample-file.txt")
	if err != nil {
		panic(err)
	}
	defer f0.Close()
	_ = f0.(gzipByter)
	f1, err := fs.Open("/sample-file.txt")
	if err != nil {
		panic(err)
	}
	defer f1.Close()
	_ = f1.(gzipByter)

	_, err = io.CopyN(os.Stdout, f0, 9)
	if err != nil {
		panic(err)
	}
	_, err = io.CopyN(os.Stdout, f1, 9)
	if err != nil {
		panic(err)
	}

	// Output:
	// This fileThis file
}

// Copy writes the entire database to a writer.
// A reader transaction is maintained during the copy so it is safe to continue
// using the database while a copy is in progress.
// Copy will write exactly tx.Size() bytes into the writer.
func (tx *Tx) Copy(w io.Writer) error {
	// Open reader on the database.
	f, err := os.Open(tx.db.path)
	if err != nil {
		_ = tx.Rollback()
		return err
	}

	// Copy the meta pages.
	tx.db.metalock.Lock()
	_, err = io.CopyN(w, f, int64(tx.db.pageSize*2))
	tx.db.metalock.Unlock()
	if err != nil {
		_ = tx.Rollback()
		_ = f.Close()
		return fmt.Errorf("meta copy: %s", err)
	}

	// Copy data pages.
	if _, err := io.CopyN(w, f, tx.Size()-int64(tx.db.pageSize*2)); err != nil {
		_ = tx.Rollback()
		_ = f.Close()
		return err
	}

	return f.Close()
}

// Write out the changes to the file in place.
func writeNew(f *os.File, v string, offset, length int64) error {
	_, err := f.Seek(0, 0) // reset reader
	if err != nil {
		return err
	}
	// Read from old file, write changes out to buffer.
	var buf bytes.Buffer
	_, err = io.CopyN(&buf, f, offset) // write up to change
	if err != nil {
		return err
	}
	_, err = io.CopyN(ioutil.Discard, f, length) // scrap the old
	if err != nil {
		return err
	}
	_, err = io.WriteString(&buf, v) // write the new
	if err != nil {
		return err
	}
	_, err = io.Copy(&buf, f) // write the rest
	if err != nil {
		return err
	}

	stat, err := f.Stat()
	if err != nil {
		return err
	}

	return ioutil.WriteFile(f.Name(), buf.Bytes(), stat.Mode())
}

func roundtrip(ctx *ClientContext, conn net.Conn, requestSize, requestBytes []byte) ([]byte, error) {
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Time{}
	}

	var err error
	if err = conn.SetDeadline(deadline); err != nil {
		return nil, makeClientErr(err.Error())
	}
	if _, err = conn.Write(requestSize); err != nil {
		return nil, makeClientErrf("Failed to write 4 bytes for request size: %s", err)
	}
	if _, err = conn.Write(requestBytes); err != nil {
		return nil, makeClientErrf("Failed to write %d bytes for request: %s", err)
	}
	buf := bytes.NewBuffer(make([]byte, 0, 4))
	if _, err = io.CopyN(buf, conn, 4); err != nil {
		return nil, makeClientErrf(
			"Failed to read 4 bytes for response size from '%s': %s",
			conn.RemoteAddr().String(), err)
	}
	responseSize := binary.BigEndian.Uint32(buf.Bytes())
	if _, err = io.CopyN(buf, conn, int64(responseSize)); err != nil {
		return nil, makeClientErrf(
			"Failed to read %d bytes for response from '%s': %s",
			conn.RemoteAddr().String(), err)
	}
	return buf.Bytes(), nil
}

// WriteTo writes the entire database to a writer.
// If err == nil then exactly tx.Size() bytes will be written into the writer.
func (tx *Tx) WriteTo(w io.Writer) (n int64, err error) {
	// Attempt to open reader with WriteFlag
	f, err := os.OpenFile(tx.db.path, os.O_RDONLY|tx.WriteFlag, 0)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	// Copy the meta pages.
	tx.db.metalock.Lock()
	n, err = io.CopyN(w, f, int64(tx.db.pageSize*2))
	tx.db.metalock.Unlock()
	if err != nil {
		return n, fmt.Errorf("meta copy: %s", err)
	}

	// Copy data pages.
	wn, err := io.CopyN(w, f, tx.Size()-int64(tx.db.pageSize*2))
	n += wn
	if err != nil {
		return n, err
	}

	return n, f.Close()
}

// WriteTo writes the entire database to a writer.
// If err == nil then exactly tx.Size() bytes will be written into the writer.
func (tx *Tx) WriteTo(w io.Writer) (n int64, err error) {
	// Attempt to open reader directly.
	var f *os.File
	if f, err = os.OpenFile(tx.db.path, os.O_RDONLY|odirect, 0); err != nil {
		// Fallback to a regular open if that doesn't work.
		if f, err = os.OpenFile(tx.db.path, os.O_RDONLY, 0); err != nil {
			return 0, err
		}
	}

	// Copy the meta pages.
	tx.db.metalock.Lock()
	n, err = io.CopyN(w, f, int64(tx.db.pageSize*2))
	tx.db.metalock.Unlock()
	if err != nil {
		_ = f.Close()
		return n, fmt.Errorf("meta copy: %s", err)
	}

	// Copy data pages.
	wn, err := io.CopyN(w, f, tx.Size()-int64(tx.db.pageSize*2))
	n += wn
	if err != nil {
		_ = f.Close()
		return n, err
	}

	return n, f.Close()
}

// ReadRecord reads a whole record into memory (up to 32 KB), otherwise the record is discarded.
func ReadRecord(r io.Reader) (*bytes.Buffer, error) {

	var buf bytes.Buffer

	for {
		size, last, err := ReadRecordMarker(r)

		if err != nil {
			return nil, err
		}

		if size < 1 {
			return nil, errors.New("A TCP record must be at least one byte in size")
		}

		if size >= maxRecordSize {
			io.CopyN(ioutil.Discard, r, int64(size))

			return nil, fmt.Errorf("Discarded record exceeding maximum size of %v bytes", maxRecordSize)
		}

		if n, err := io.CopyN(&buf, r, int64(size)); err != nil {
			return nil, fmt.Errorf("Unable to read entire record. Read %v, expected %v", n, size)
		}

		if last {
			break
		}
	}

	return &buf, nil
}

// readBlockHeader reads the block header.
func readBlockHeader(r io.Reader) (h *blockHeader, n int, err error) {
	var buf bytes.Buffer
	buf.Grow(20)

	// block header size
	z, err := io.CopyN(&buf, r, 1)
	n = int(z)
	if err != nil {
		return nil, n, err
	}
	s := buf.Bytes()[0]
	if s == 0 {
		return nil, n, errIndexIndicator
	}

	// read complete header
	headerLen := (int(s) + 1) * 4
	buf.Grow(headerLen - 1)
	z, err = io.CopyN(&buf, r, int64(headerLen-1))
	n += int(z)
	if err != nil {
		return nil, n, err
	}

	// unmarshal block header
	h = new(blockHeader)
	if err = h.UnmarshalBinary(buf.Bytes()); err != nil {
		return nil, n, err
	}

	return h, n, nil
}

// Reads transfer-encoding: chunked payloads from the connection reader.
func (c *Connection) readChunkedData() error {
	var err error
	var line []byte
	var size uint64
	var start time.Time

	start = time.Now()
	writer := &nonEmptyWriter{os.Stdout}

	var buffer *bytes.Buffer
	var decompressor *gzip.Reader
	var zipReader *bufio.Reader
	var data []byte

	if c.conf.GZip == true {
		buffer = bytes.NewBufferString("")
	}

	for err == nil {
		line, _, err = c.reader.ReadLine()
		if err != nil {
			return err
		}
		size, err = decodeHexString(line)
		if err != nil {
			str := fmt.Sprintf("Expected hex, got %v", string(line))
			return errors.New(str)
		}
		if c.conf.GZip == false {
			_, err = io.CopyN(writer, c.reader, int64(size))
		} else {
			_, err = io.CopyN(buffer, c.reader, int64(size))
			if err != nil {
				return err
			}
			if decompressor == nil {
				decompressor, err = gzip.NewReader(buffer)
				defer decompressor.Close()
				if err != nil {
					return err
				}
				zipReader = bufio.NewReader(decompressor)
			}
			data = make([]byte, 512, 512)
			_, err = zipReader.Read(data)
			if err != nil {
				return err
			}
			strBuffer := bytes.NewBuffer(data)
			io.CopyN(writer, strBuffer, int64(len(data)))
		}
		if c.conf.TTL > 0 {
			if time.Now().Sub(start).Nanoseconds() > c.conf.TTL {
				return nil
			}
		}
	}
	return err
}

// NewDecryptReader creates an io.ReadCloser wrapping an io.Reader.
// It has to read the entire io.Reader to disk using a temp file so that it can
// hash the contents to verify that it is safe to decrypt.
// If the file is athenticated, the DecryptReader will be returned and
// the resulting bytes will be the plaintext.
func NewDecryptReader(r io.Reader, pass []byte) (d io.ReadCloser, err error) {
	mac := make([]byte, hmacSize)
	aesKey, hmacKey, iv, header, err := decodeHeader(r, pass)
	h := hmac.New(hashFunc, hmacKey)
	h.Write(header)
	if err != nil {
		return nil, err
	}
	dst, err := tmpfile.New(&tmpfile.Context{
		Dir:    os.TempDir(),
		Suffix: "drive-encrypted-",
	})
	if err != nil {
		return nil, err
	}
	// If there is an error, try to delete the temp file.
	defer func() {
		if err != nil {
			dst.Done()
		}
	}()
	b, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	d = &decryptReader{
		tmpFile: dst,
		sReader: &cipher.StreamReader{R: dst, S: cipher.NewCTR(b, iv)},
	}
	w := io.MultiWriter(h, dst)
	buf := bufio.NewReaderSize(r, _16KB)
	for {
		b, err := buf.Peek(_16KB)
		if err != nil && err != io.EOF {
			return nil, err
		}
		if err == io.EOF {
			left := buf.Buffered()
			if left < hmacSize {
				return nil, DecryptErr
			}
			copy(mac, b[left-hmacSize:left])
			_, err = io.CopyN(w, buf, int64(left-hmacSize))
			if err != nil {
				return nil, err
			}
			break
		}
		_, err = io.CopyN(w, buf, _16KB-hmacSize)
		if err != nil {
			return nil, err
		}
	}
	if !hmac.Equal(mac, h.Sum(nil)) {
		return nil, DecryptErr
	}
	dst.Seek(0, 0)
	return d, nil
}