Golang io.LimitReader函数代码示例

func (r *ZipScannerImpl) Reader() (io.Reader, error) {
	switch r.fh.Method {
	case zip.Deflate:
		if r.Debug {
			fmt.Println("inflating...")
		}

		r.fr = flate.NewReader(r.reader)

	case zip.Store:
		if r.Debug {
			fmt.Println("reading...")
		}

		if r.fh.UncompressedSize > 0 {
			r.fr = io.LimitReader(r.reader, int64(r.fh.UncompressedSize))
		} else if r.fh.UncompressedSize == 0 && (r.fh.Flags&hasDataDescriptor) == 0 {
			// file of 0 bytes or directory ?
			r.fr = io.LimitReader(r.reader, 0)
		} else {
			return r.readError(NoUncompressedSize)
		}

	default:
		return r.readError(UnsupportedCompression)
	}

	r.err = nil
	return r.fr, r.err
}

func (s *BufferTest) TestMBuf(t *C) {
	pool := NewBufferPool(1000*1024*1024, 200*1024*1024)
	h := pool.NewPoolHandle()

	n := uint64(2 * BUF_SIZE)
	mb := MBuf{}.Init(h, n)
	t.Assert(len(mb.buffers), Equals, 2)

	r := io.LimitReader(&SeqReader{}, int64(n))

	for {
		nread, err := mb.WriteFrom(r)
		t.Assert(err, IsNil)
		if nread == 0 {
			break
		}
	}
	t.Assert(mb.wbuf, Equals, 1)
	t.Assert(mb.wp, Equals, BUF_SIZE)

	diff, err := CompareReader(mb, io.LimitReader(&SeqReader{}, int64(n)))
	t.Assert(err, IsNil)
	t.Assert(diff, Equals, -1)

	t.Assert(mb.rbuf, Equals, 1)
	t.Assert(mb.rp, Equals, BUF_SIZE)

	t.Assert(h.inUseBuffers, Equals, int64(2))
	mb.Free()
	t.Assert(h.inUseBuffers, Equals, int64(0))
}

func (cOff chunkOffTs) readPreChunk(r io.ReadSeeker) (*preChunk, error) {
	pc := preChunk{ts: cOff.ts}

	if _, err := r.Seek(cOff.offset, 0); err != nil {
		return nil, err
	}

	lr := io.LimitReader(r, cOff.size)

	var length uint32
	if err := binary.Read(lr, binary.BigEndian, &length); err != nil {
		return nil, err
	}
	lr = io.LimitReader(lr, int64(length))

	compType, err := kagus.ReadByte(lr)
	if err != nil {
		return nil, err
	}
	pc.compression = compType

	buf := new(bytes.Buffer)
	if _, err := io.Copy(buf, lr); err != nil {
		return nil, err
	}
	pc.data = buf.Bytes()

	return &pc, err

}

func (s *BufferTest) TestBufferWrite(t *C) {
	h := NewBufferPool(1000 * 1024 * 1024)

	n := uint64(2 * BUF_SIZE)
	mb := MBuf{}.Init(h, n, true)
	t.Assert(len(mb.buffers), Equals, 2)

	nwritten, err := io.Copy(mb, io.LimitReader(&SeqReader{}, int64(n)))
	t.Assert(nwritten, Equals, int64(n))
	t.Assert(err, IsNil)

	diff, err := CompareReader(mb, io.LimitReader(&SeqReader{}, int64(n)))
	t.Assert(err, IsNil)
	t.Assert(diff, Equals, -1)

	cur, err := mb.Seek(0, 1)
	t.Assert(err, IsNil)
	t.Assert(cur, Equals, int64(n))

	cur, err = mb.Seek(0, 2)
	t.Assert(err, IsNil)
	t.Assert(cur, Equals, int64(n))

	cur, err = mb.Seek(0, 0)
	t.Assert(err, IsNil)
	t.Assert(cur, Equals, int64(0))
	t.Assert(mb.rbuf, Equals, 0)
	t.Assert(mb.rp, Equals, 0)

	diff, err = CompareReader(mb, io.LimitReader(&SeqReader{}, int64(n)))
	t.Assert(err, IsNil)
	t.Assert(diff, Equals, -1)
}

// LimitReader returns a Reader that reads from r but stops with EOF after n
// bytes. The underlying implementation is a *io.LimitedReader unless n is
// greater than the maximum value representable as an int64, in which case it
// uses io.MultiReader with a pair of *io.LimitReader.
//
func LimitReader(r io.Reader, n uint64) io.Reader {
	if n > uint64(math.MaxInt64) {
		return io.MultiReader(io.LimitReader(r, math.MaxInt64), io.LimitReader(r, int64(n-uint64(math.MaxInt64))))
	} else {
		return io.LimitReader(r, int64(n))
	}
}

// archiveFileVisitor is called for each file in an archive. It may set
// tempFile and signature.
func archiveFileVisitor(dir string, tempFile *string, signature *[]byte, archivePath string, filedata io.Reader) error {
	var err error
	filename := path.Base(archivePath)
	archiveDir := path.Dir(archivePath)
	l.Debugf("considering file %s", archivePath)
	switch filename {
	case "syncthing", "syncthing.exe":
		archiveDirs := strings.Split(archiveDir, "/")
		if len(archiveDirs) > 1 {
			// Don't consider "syncthing" files found too deeply, as they may be
			// other things.
			return nil
		}
		l.Debugf("found upgrade binary %s", archivePath)
		*tempFile, err = writeBinary(dir, io.LimitReader(filedata, maxBinarySize))
		if err != nil {
			return err
		}

	case "release.sig":
		l.Debugf("found signature %s", archivePath)
		*signature, err = ioutil.ReadAll(io.LimitReader(filedata, maxSignatureSize))
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *CompressionSnappyDecoder) readHeader() (int, error) {
	header := make([]byte, 4, 4)
	_, err := c.source.Read(header[:3])
	if err != nil {
		return 0, err
	}
	headerVal := binary.LittleEndian.Uint32(header)
	c.isOriginal = headerVal%2 == 1
	c.chunkLength = int(headerVal / 2)
	if !c.isOriginal {
		// ORC does not use snappy's framing as implemented in the
		// github.com/golang/snappy Reader implementation. As a result
		// we have to read and decompress the entire chunk.
		// TODO: find reader implementation with optional framing.
		r := io.LimitReader(c.source, int64(c.chunkLength))
		src, err := ioutil.ReadAll(r)
		if err != nil {
			return 0, err
		}
		decodedBytes, err := snappy.Decode(nil, src)
		if err != nil {
			return 0, err
		}
		c.decoded = bytes.NewReader(decodedBytes)
	} else {
		c.decoded = io.LimitReader(c.source, int64(c.chunkLength))
	}
	return 0, nil
}

func (c *Connection) readReply() (*replyMsg, os.Error) {
	size_bits, _ := ioutil.ReadAll(io.LimitReader(c.conn, 4))
	size := binary.LittleEndian.Uint32(size_bits)
	rest, _ := ioutil.ReadAll(io.LimitReader(c.conn, int64(size)-4))
	reply := parseReply(rest)
	return reply, nil
}

// Extract form fields and file data from a HTTP POST Policy
func extractPostPolicyFormValues(reader *multipart.Reader) (filePart io.Reader, fileName string, formValues map[string]string, err error) {
	/// HTML Form values
	formValues = make(map[string]string)
	fileName = ""
	for err == nil {
		var part *multipart.Part
		part, err = reader.NextPart()
		if part != nil {
			canonicalFormName := http.CanonicalHeaderKey(part.FormName())
			if canonicalFormName != "File" {
				var buffer []byte
				limitReader := io.LimitReader(part, maxFormFieldSize+1)
				buffer, err = ioutil.ReadAll(limitReader)
				if err != nil {
					return nil, "", nil, err
				}
				if int64(len(buffer)) > maxFormFieldSize {
					return nil, "", nil, errSizeUnexpected
				}
				formValues[canonicalFormName] = string(buffer)
			} else {
				filePart = io.LimitReader(part, maxObjectSize)
				fileName = part.FileName()
				// As described in S3 spec, we expect file to be the last form field
				break
			}
		}
	}
	return filePart, fileName, formValues, nil
}

/* Gets the message of reply from database. */
func (self *Connection) readReply() (*opReply, os.Error) {
	size_bits, _ := ioutil.ReadAll(io.LimitReader(self.conn, 4))
	size := pack.Uint32(size_bits)
	rest, _ := ioutil.ReadAll(io.LimitReader(self.conn, int64(size)-4))
	reply := parseReply(rest)

	return reply, nil
}

func TestVolumeCreateBadDispersionValues(t *testing.T) {
	tmpfile := tests.Tempfile()
	defer os.Remove(tmpfile)

	// Create the app
	app := NewTestApp(tmpfile)
	defer app.Close()
	router := mux.NewRouter()
	app.SetRoutes(router)

	// Setup the server
	ts := httptest.NewServer(router)
	defer ts.Close()

	// VolumeCreate JSON Request
	request := []byte(`{
        "size" : 100,
        "durability": {
        	"type": "disperse",
        	"disperse": {
            	"data" : 8,
            	"redundancy" : 1
        	}
    	}
    }`)

	// Send request
	r, err := http.Post(ts.URL+"/volumes", "application/json", bytes.NewBuffer(request))
	tests.Assert(t, err == nil)
	tests.Assert(t, r.StatusCode == http.StatusBadRequest)
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, r.ContentLength))
	tests.Assert(t, err == nil)
	r.Body.Close()
	tests.Assert(t, strings.Contains(string(body), "Invalid dispersion combination"))

	// VolumeCreate JSON Request
	request = []byte(`{
        "size" : 100,
        "durability": {
        	"type": "disperse",
        	"disperse": {
            	"data" : 4,
            	"redundancy" : 3
        	}
    	}
    }`)

	// Send request
	r, err = http.Post(ts.URL+"/volumes", "application/json", bytes.NewBuffer(request))
	tests.Assert(t, err == nil)
	tests.Assert(t, r.StatusCode == http.StatusBadRequest)
	body, err = ioutil.ReadAll(io.LimitReader(r.Body, r.ContentLength))
	tests.Assert(t, err == nil)
	r.Body.Close()
	tests.Assert(t, strings.Contains(string(body), "Invalid dispersion combination"))
}

func (decoder *packetDecoder) Decode(msg interface{}) (err error) {
	decoder.buffer.Reset()
	if _, err = decoder.buffer.ReadFrom(io.LimitReader(decoder.reader, int64(decoder.n))); err != nil {
		return err
	}
	n := decoder.decodeHead(decoder.buffer.Next(decoder.n))
	if _, err = decoder.buffer.ReadFrom(io.LimitReader(decoder.reader, int64(n))); err != nil {
		return err
	}
	return decoder.base.Decode(msg)
}

func main() {
	// START OMIT
	var r io.Reader = ByteReader('A')
	r = io.LimitReader(r, 1e6)
	r = LogReader{r}
	io.Copy(ioutil.Discard, r)
	// STOP OMIT

	return
	io.Copy(ioutil.Discard, LogReader{io.LimitReader(ByteReader('A'), 1e6)})
}

func TestPrependedBlocks(t *testing.T) {
	const BLOCKSIZE = 100
	const BLOCK_COUNT = 20
	checksum := NewFileChecksumGenerator(BLOCKSIZE)

	file1 := io.LimitReader(
		readers.NewNonRepeatingSequence(0),
		BLOCKSIZE*BLOCK_COUNT,
	)

	file2 := io.LimitReader(
		io.MultiReader(
			readers.OneReader(BLOCKSIZE), // Off by one block
			readers.NewNonRepeatingSequence(0),
		),
		BLOCKSIZE*BLOCK_COUNT,
	)

	output1 := bytes.NewBuffer(nil)
	chksum1, _ := checksum.GenerateChecksums(file1, output1)

	output2 := bytes.NewBuffer(nil)
	chksum2, _ := checksum.GenerateChecksums(file2, output2)

	if bytes.Compare(chksum1, chksum2) == 0 {
		t.Fatal("Checksums should be different")
	}

	weaksize, strongSize := checksum.GetChecksumSizes()
	sums1, _ := chunks.LoadChecksumsFromReader(output1, weaksize, strongSize)
	sums2, _ := chunks.LoadChecksumsFromReader(output2, weaksize, strongSize)

	if len(sums1) != len(sums2) {
		t.Fatalf("Checksum lengths differ %v vs %v", len(sums1), len(sums2))
	}

	if sums1[0].Match(sums2[0]) {
		t.Error("Chunk sums1[0] should differ from sums2[0]")
	}

	for i, _ := range sums2 {
		if i == 0 {
			continue
		}

		if !sums1[i-1].Match(sums2[i]) {
			t.Errorf("Chunk sums1[%v] equal sums2[%v]", i-1, i)
		}

	}
}

func passThru(w io.Writer, req *http.Request) error {
	var body bytes.Buffer

	_, err := io.Copy(&body, io.LimitReader(req.Body, maxSnippetSize+1))
	req.Body.Close()

	if err != nil {
		return fmt.Errorf("Error reading body: %q", err)
	}
	if body.Len() > maxSnippetSize {
		return fmt.Errorf("Snippet is too large")
	}

	snip := &Snippet{Body: body.Bytes()}
	id := snip.Id()
	key := []byte(id)

	var output bytes.Buffer

	if err = db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketCache)
		data := b.Get(key)
		if data == nil || *flagDisableCache {
			client := http.Client{}
			r, err := client.Post(*flagCompileURL, req.Header.Get("Content-Type"), &body)
			if err != nil {
				return err
			}
			defer r.Body.Close()

			data, err = ioutil.ReadAll(io.LimitReader(r.Body, maxSnippetSize+1))
			if len(data) > maxSnippetSize {
				return fmt.Errorf("Output is too large.")
			}
			if err = b.Put(key, data); err != nil {
				return err
			}
		}
		output.Write(data)
		return nil
	}); err != nil {
		return err
	}

	if _, err := io.Copy(w, &output); err != nil {
		return err
	}

	return nil
}