Golang gzip.NewReader函数代码示例

func (s *archiveDataSuite) SetUpTest(c *gc.C) {
	s.IsolationSuite.SetUpTest(c)

	meta, err := backups.NewMetadataJSONReader(bytes.NewBufferString(`{` +
		`"ID":"20140909-115934.asdf-zxcv-qwe",` +
		`"Checksum":"123af2cef",` +
		`"ChecksumFormat":"SHA-1, base64 encoded",` +
		`"Size":10,` +
		`"Stored":"0001-01-01T00:00:00Z",` +
		`"Started":"2014-09-09T11:59:34Z",` +
		`"Finished":"2014-09-09T12:00:34Z",` +
		`"Notes":"",` +
		`"Environment":"asdf-zxcv-qwe",` +
		`"Machine":"0",` +
		`"Hostname":"myhost",` +
		`"Version":"1.21-alpha3"` +
		`}` + "\n"))
	c.Assert(err, jc.ErrorIsNil)

	archiveFile := s.newArchiveFile(c, meta)
	compressed, err := ioutil.ReadAll(archiveFile)
	c.Assert(err, jc.ErrorIsNil)
	gzr, err := gzip.NewReader(bytes.NewBuffer(compressed))
	c.Assert(err, jc.ErrorIsNil)
	data, err := ioutil.ReadAll(gzr)
	c.Assert(err, jc.ErrorIsNil)

	s.archiveFile = bytes.NewBuffer(compressed)
	s.data = data
	s.meta = meta
}

func openStream(path string) (io.ReadCloser, error) {
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		req, err := http.NewRequest("GET", path, nil)
		if err != nil {
			return nil, err
		}
		req.Header.Set("Accept-Encoding", "gzip")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}

		gz, err := gzip.NewReader(resp.Body)
		if err != nil {
			resp.Body.Close()
			return nil, err
		}

		return gz, nil
	} else {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}

		gz, err := gzip.NewReader(f)
		if err != nil {
			f.Close()
			return nil, err
		}

		return gz, nil
	}
	panic("Unreachable")
}

func generateNewName() (string, error) {
	s1 := rand.NewSource(time.Now().UnixNano())
	r1 := rand.New(s1)
	nameIdx := r1.Intn(258000)
	surnameIdx := r1.Intn(88799)

	nameFile, err := os.Open("names.dat.gz")
	if err != nil {
		return "", err
	}

	nameReader, err := gzip.NewReader(nameFile)
	if err != nil {
		return "", err
	}
	defer nameReader.Close()

	name, err := readLine(nameReader, nameIdx)

	surnameFile, err := os.Open("names.dat.gz")
	if err != nil {
		return "", err
	}

	surnameReader, err := gzip.NewReader(surnameFile)
	if err != nil {
		return "", err
	}
	defer surnameReader.Close()

	surname, err := readLine(surnameReader, surnameIdx)

	return strings.Join([]string{name, surname}, " "), nil

}

func main() {
	flag.Parse()
	f, err := os.Open(*geoData)
	x.Check(err)

	gr, err := gzip.NewReader(f)
	x.Check(err)

	//var strBuf bytes.Buffer
	//bufReader := bufio.NewReader(gr)
	dec := json.NewDecoder(gr)

	countryToGeo := make(map[string]string)
	findFeatureArray(dec)
	for dec.More() {
		var f geojson.Feature
		err := dec.Decode(&f)
		fmt.Println(f.Properties["NAME_LONG"])
		gg, err := geojson.Marshal(f.Geometry)
		ggg := strings.Replace(string(gg), "\"", "'", -1)
		country, ok := f.Properties["NAME_LONG"].(string)
		if ok {
			countryToGeo[country] = ggg
		}
		//fmt.Printf("\"%s\"", ggg)
		if err != nil {
			fmt.Println(err)
		}
	}
	gr.Close()
	f.Close()

	f, err = os.Open(*rdf)
	x.Check(err)

	gr, err = gzip.NewReader(f)
	x.Check(err)

	scanner := bufio.NewScanner(gr)

	out, err := os.Create("countryGeoData")
	x.Check(err)
	defer out.Close()
	count1, count2 := 0, 0
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "@en") {
			items := strings.Split(line, "\t")
			country := strings.Trim(strings.Split(items[2], "@")[0], "\"")
			fmt.Println(country)
			if geoD, ok := countryToGeo[country]; ok {
				count1++
				out.WriteString(fmt.Sprintf("%s <loc> \"%s\"^^<geo:geojson> .\n", items[0], geoD))
			} else {
				count2++
			}
		}
	}
	fmt.Println(count1, count2)
}

func extractTarGz(body []byte, location string) (string, error) {
	bodyCopy := make([]byte, len(body))
	copy(bodyCopy, body)
	tarFile, _ := gzip.NewReader(bytes.NewReader(body))
	tarReader := tar.NewReader(tarFile)

	var dirList []string

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		dirList = append(dirList, header.Name)
	}

	basedir := findBaseDir(dirList)

	tarFile, _ = gzip.NewReader(bytes.NewReader(bodyCopy))
	tarReader = tar.NewReader(tarFile)

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			//return location, err
		}

		path := filepath.Join(location, strings.Replace(header.Name, basedir, "", -1))
		info := header.FileInfo()

		if info.IsDir() {
			if err = os.MkdirAll(path, info.Mode()); err != nil {
				return location, err
			}
			continue
		}

		if header.Typeflag == tar.TypeSymlink {
			err = os.Symlink(header.Linkname, path)
			continue
		}

		file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, info.Mode())
		if err == nil {
			defer file.Close()
		}
		_, err = io.Copy(file, tarReader)
		if err != nil {
			//return location, err
		}
	}
	return location, nil
}

func indexTarball(tarPath string, logChan chan *store.Record) bool {
	currentTar.Set(tarPath)
	handle, err := os.Open(tarPath)
	if err != nil {
		log.Printf("Error reading %s: %s\n", tarPath, err)
		tarsFailed.Add(1)
		return false
	}
	defer handle.Close()
	fileinfo, err := handle.Stat()
	if err != nil {
		log.Printf("Error stating %s: %s\n", tarPath, err)
		tarsFailed.Add(1)
		return false
	}
	tarBytesRead.Add(fileinfo.Size())
	unzippedHandle, err := gzip.NewReader(handle)
	if err != nil {
		log.Printf("Error unzipping tarball %s: %s\n", tarPath, err)
		tarsFailed.Add(1)
		return false
	}
	parentReader := tar.NewReader(unzippedHandle)
	for {
		parentHeader, err := parentReader.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			tarsFailed.Add(1)
			log.Printf("Error indexing %v: %v", tarPath, err)
			break
		}
		if parentHeader.Typeflag != tar.TypeReg && parentHeader.Typeflag != tar.TypeRegA {
			continue
		}
		if filepath.Ext(parentHeader.Name) != ".gz" {
			continue
		}
		parentGzipHandle, err := gzip.NewReader(parentReader)
		if err != nil {
			nestedTarsFailed.Add(1)
			log.Printf("Error gunzipping trace %s/%s: %v", tarPath, parentHeader.Name, err)
			continue
		}
		if err := indexNestedTarball(parentGzipHandle, logChan); err != nil {
			nestedTarsFailed.Add(1)
			continue
		}
		nestedTarsIndexed.Add(1)
	}
	tarsIndexed.Add(1)
	return true
}

func Fuzz(data []byte) int {
	fr, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return 0
	}
	if len(fr.Comment) > 1<<20 || len(fr.Name) > 1<<20 || len(fr.Extra) > 1<<20 {
		panic("huge header")
	}
	uncomp := make([]byte, 64<<10)
	n, err := fr.Read(uncomp)
	if err != nil && err != io.EOF {
		return 0
	}
	if n == len(uncomp) {
		return 0 // too large
	}
	uncomp = uncomp[:n]
	for c := 0; c <= 9; c++ {
		buf := new(bytes.Buffer)
		gw, err := gzip.NewWriterLevel(buf, c)
		if err != nil {
			panic(err)
		}
		gw.Header = fr.Header
		n, err := gw.Write(uncomp)
		if err != nil {
			panic(err)
		}
		if n != len(uncomp) {
			panic("short write")
		}
		if err := gw.Close(); err != nil {
			panic(err)
		}
		fr1, err := gzip.NewReader(buf)
		if err != nil {
			panic(err)
		}
		uncomp1, err := ioutil.ReadAll(fr1)
		if err != nil {
			panic(err)
		}
		if !bytes.Equal(uncomp, uncomp1) {
			panic("data differs")
		}
	}
	return 1
}

func (m *Model) loadIndex(repo string, dir string) []protocol.FileInfo {
	id := fmt.Sprintf("%x", sha1.Sum([]byte(m.repoDirs[repo])))
	name := id + ".idx.gz"
	name = filepath.Join(dir, name)

	idxf, err := os.Open(name)
	if err != nil {
		return nil
	}
	defer idxf.Close()

	gzr, err := gzip.NewReader(idxf)
	if err != nil {
		return nil
	}
	defer gzr.Close()

	var im protocol.IndexMessage
	err = im.DecodeXDR(gzr)
	if err != nil || im.Repository != repo {
		return nil
	}

	return im.Files
}

func checkToolsContent(c *gc.C, data []byte, uploaded string) {
	zr, err := gzip.NewReader(bytes.NewReader(data))
	c.Check(err, gc.IsNil)
	defer zr.Close()
	tr := tar.NewReader(zr)
	found := false
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		c.Check(err, gc.IsNil)
		if strings.ContainsAny(hdr.Name, "/\\") {
			c.Fail()
		}
		if hdr.Typeflag != tar.TypeReg {
			c.Fail()
		}
		content, err := ioutil.ReadAll(tr)
		c.Check(err, gc.IsNil)
		c.Check(string(content), gc.Equals, uploaded)
		found = true
	}
	c.Check(found, jc.IsTrue)
}

// BenchmarkParseProtoGzip is like BenchmarkParseProto above, but parses gzipped
// protobuf format.
func BenchmarkParseProtoGzip(b *testing.B) {
	b.StopTimer()
	data, err := ioutil.ReadFile("testdata/protobuf.gz")
	if err != nil {
		b.Fatal(err)
	}
	b.StartTimer()

	for i := 0; i < b.N; i++ {
		family := &dto.MetricFamily{}
		in, err := gzip.NewReader(bytes.NewReader(data))
		if err != nil {
			b.Fatal(err)
		}
		for {
			family.Reset()
			if _, err := pbutil.ReadDelimited(in, family); err != nil {
				if err == io.EOF {
					break
				}
				b.Fatal(err)
			}
		}
	}
}

// doRequest runs a request with our client
func (c *Client) doRequest(r *request) (time.Duration, *http.Response, error) {
	req, err := r.toHTTP()
	if err != nil {
		return 0, nil, err
	}
	start := time.Now()
	resp, err := c.config.HttpClient.Do(req)
	diff := time.Now().Sub(start)

	// If the response is compressed, we swap the body's reader.
	if resp != nil && resp.Header != nil {
		var reader io.ReadCloser
		switch resp.Header.Get("Content-Encoding") {
		case "gzip":
			greader, err := gzip.NewReader(resp.Body)
			if err != nil {
				return 0, nil, err
			}

			// The gzip reader doesn't close the wrapped reader so we use
			// multiCloser.
			reader = &multiCloser{
				reader:       greader,
				inorderClose: []io.Closer{greader, resp.Body},
			}
		default:
			reader = resp.Body
		}
		resp.Body = reader
	}

	return diff, resp, err
}

func (t *TarInfo) Load(file io.ReadSeeker) {
	var reader *tar.Reader
	file.Seek(0, 0)
	gzipReader, err := gzip.NewReader(file)
	if err != nil {
		// likely not a gzip compressed file
		file.Seek(0, 0)
		reader = tar.NewReader(file)
	} else {
		reader = tar.NewReader(gzipReader)
	}
	for {
		header, err := reader.Next()
		if err == io.EOF {
			// end of tar file
			break
		} else if err != nil {
			// error occured
			logger.Debug("[TarInfoLoad] Error when reading tar stream tarsum. Disabling TarSum, TarFilesInfo. Error: %s", err.Error())
			t.Error = TarError(err.Error())
			return
		}
		t.TarSum.Append(header, reader)
		t.TarFilesInfo.Append(header)
	}
}

func TestGzipOK(t *testing.T) {
	c := makeEnv()

	w := httptest.NewRecorder()

	h := GzipMiddleWare(&c, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("super things"))
	}))

	r, _ := http.NewRequest("GET", "/", nil)
	r.Header.Set("Accept-Encoding", "gzip")

	h.ServeHTTP(w, r)

	if w.Code != http.StatusOK {
		t.Errorf("expect 200, get %d, %s", w.Code, w.Body.String())
	}

	gr, err := gzip.NewReader(w.Body)
	if err != nil {
		t.Errorf("couldn't create a gzip raader, %v", err)
	}
	body, err := ioutil.ReadAll(gr)
	if err != nil {
		t.Errorf("Couldn't read gzip content. %v", err)
	}

	if string(body) != "super things" {
		t.Errorf("body not as expected.  have \"%s\"", body)
	}

	if w.HeaderMap.Get("Content-Encoding") != "gzip" {
		t.Errorf("expected gzip encoding - have %v", w.HeaderMap)
	}
}

func TestGzip(t *testing.T) {
	req, _ := http.NewRequest(vodka.GET, "/", nil)
	rec := httptest.NewRecorder()
	c := vodka.NewContext(req, vodka.NewResponse(rec), vodka.New())
	h := func(c *vodka.Context) error {
		c.Response().Write([]byte("test")) // For Content-Type sniffing
		return nil
	}

	// Skip if no Accept-Encoding header
	Gzip()(h)(c)
	assert.Equal(t, http.StatusOK, rec.Code)
	assert.Equal(t, "test", rec.Body.String())

	req, _ = http.NewRequest(vodka.GET, "/", nil)
	req.Header.Set(vodka.AcceptEncoding, "gzip")
	rec = httptest.NewRecorder()
	c = vodka.NewContext(req, vodka.NewResponse(rec), vodka.New())

	// Gzip
	Gzip()(h)(c)
	assert.Equal(t, http.StatusOK, rec.Code)
	assert.Equal(t, "gzip", rec.Header().Get(vodka.ContentEncoding))
	assert.Contains(t, rec.Header().Get(vodka.ContentType), vodka.TextPlain)
	r, err := gzip.NewReader(rec.Body)
	defer r.Close()
	if assert.NoError(t, err) {
		buf := new(bytes.Buffer)
		buf.ReadFrom(r)
		assert.Equal(t, "test", buf.String())
	}
}

// TestAcceptEncoding hits the health endpoint while explicitly
// disabling decompression on a custom client's Transport and setting
// it conditionally via the request's Accept-Encoding headers.
func TestAcceptEncoding(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := StartTestServer(t)
	defer s.Stop()
	client, err := testContext.GetHTTPClient()
	if err != nil {
		t.Fatal(err)
	}

	testData := []struct {
		acceptEncoding string
		newReader      func(io.Reader) io.Reader
	}{
		{"",
			func(b io.Reader) io.Reader {
				return b
			},
		},
		{util.GzipEncoding,
			func(b io.Reader) io.Reader {
				r, err := gzip.NewReader(b)
				if err != nil {
					t.Fatalf("could not create new gzip reader: %s", err)
				}
				return r
			},
		},
		{util.SnappyEncoding,
			func(b io.Reader) io.Reader {
				return snappy.NewReader(b)
			},
		},
	}
	for _, d := range testData {
		req, err := http.NewRequest("GET", testContext.HTTPRequestScheme()+"://"+s.ServingAddr()+healthPath, nil)
		if err != nil {
			t.Fatalf("could not create request: %s", err)
		}
		if d.acceptEncoding != "" {
			req.Header.Set(util.AcceptEncodingHeader, d.acceptEncoding)
		}
		resp, err := client.Do(req)
		if err != nil {
			t.Fatalf("could not make request to %s: %s", req.URL, err)
		}
		defer resp.Body.Close()
		if ce := resp.Header.Get(util.ContentEncodingHeader); ce != d.acceptEncoding {
			t.Fatalf("unexpected content encoding: '%s' != '%s'", ce, d.acceptEncoding)
		}
		r := d.newReader(resp.Body)
		b, err := ioutil.ReadAll(r)
		if err != nil {
			t.Fatalf("could not read '%s' response body: %s", d.acceptEncoding, err)
		}
		expected := "ok"
		if !strings.Contains(string(b), expected) {
			t.Errorf("expected body to contain %q, got %q", expected, b)
		}
	}
}