<commit_before>// Package osmpbf decodes OpenStreetMap pbf files.
// Use this package by creating a NewDecoder and passing it a PBF file. Use
// Decode to return Node, Way and Relation structs.
package osmpbf

import (
	"bytes"
	"code.google.com/p/goprotobuf/proto"
	"compress/zlib"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/qedus/osmpbf/OSMPBF"
	"io"
	"io/ioutil"
)

const (
	maxBlobHeaderSize = 64 * 1024
	maxBlobSize       = 32 * 1024 * 1024
)

var (
	parseCapabilities = map[string]bool{
		"OsmSchema-V0.6": true,
		"DenseNodes":     true,
	}
)

type Node struct {
	ID   int64
	Lat  float64
	Lon  float64
	Tags map[string]string

	// TODO: Add DenseInfo
}

type Way struct {
	ID      int64
	Tags    map[string]string
	NodeIDs []int64

	// TODO: Add Info
}

type Relation struct {
	ID      int64
	Tags    map[string]string
	Members []Member

	// TODO: Add Info
	// TODO: Add roles_sid
}

type MemberType int

const (
	NodeType MemberType = iota
	WayType
	RelationType
)

type Member struct {
	ID   int64
	Type MemberType
	Role string
}

// A Decoder reads and decodes OpenStreetMap PBF data from an input stream.
type Decoder struct {
	r io.Reader

	objectQueue []interface{}
	objectIndex int
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r, make([]interface{}, 0), 0}
}

// Decode reads the next object from the input stream and returns either a
// Node, Way or Relation struct representing the underlying OpenStreetMap PBF
// data.
//
// The end of the input stream is reported by an io.EOF error.
func (dec *Decoder) Decode() (interface{}, error) {
	if dec.objectIndex >= len(dec.objectQueue) {
		dec.objectQueue = dec.objectQueue[:0]
		dec.objectIndex = 0
		if err := dec.readNextPrimitiveBlock(); err != nil {
			return nil, err
		}
	}

	dec.objectIndex++
	return dec.objectQueue[dec.objectIndex-1], nil
}

func (dec *Decoder) readNextPrimitiveBlock() error {
	for {
		blobHeaderSize, err := dec.readBlobHeaderSize()
		if err != nil {
			return err
		}

		blobHeader, err := dec.readBlobHeader(blobHeaderSize)
		if err != nil {
			return err
		}

		blob, err := dec.readBlob(blobHeader)
		if err != nil {
			return err
		}

		switch blobHeader.GetType() {
		case "OSMHeader":
			if err := dec.readOSMHeader(blob); err != nil {
				return err
			}
		case "OSMData":
			return dec.readOSMData(blob)
		default:
			// Skip over unknown type
		}
	}
}

func (dec *Decoder) readBlobHeaderSize() (uint32, error) {
	buf := make([]byte, 4)
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return 0, err
	}
	size := binary.BigEndian.Uint32(buf)

	if size >= maxBlobHeaderSize {
		return 0, errors.New("BlobHeader size >= 64Kb")
	}
	return size, nil
}

func (dec *Decoder) readBlobHeader(size uint32) (*OSMPBF.BlobHeader, error) {
	buf := make([]byte, size)
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return nil, err
	}

	blobHeader := &OSMPBF.BlobHeader{}
	if err := proto.Unmarshal(buf, blobHeader); err != nil {
		return nil, err
	}

	if blobHeader.GetDatasize() >= maxBlobSize {
		return nil, errors.New("Blob size >= 32Mb")
	}
	return blobHeader, nil
}

func (dec *Decoder) readBlob(blobHeader *OSMPBF.BlobHeader) (*OSMPBF.Blob, error) {
	buf := make([]byte, blobHeader.GetDatasize())
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return nil, err
	}

	blob := &OSMPBF.Blob{}
	if err := proto.Unmarshal(buf, blob); err != nil {
		return nil, err
	}
	return blob, nil
}

func (dec *Decoder) readOSMHeader(blob *OSMPBF.Blob) error {
	data, err := getData(blob)
	if err != nil {
		return err
	}
	headerBlock := &OSMPBF.HeaderBlock{}
	if err := proto.Unmarshal(data, headerBlock); err != nil {
		return err
	}

	// Check we have the parse capabilities
	requiredFeatures := headerBlock.GetRequiredFeatures()
	for _, feature := range requiredFeatures {
		if !parseCapabilities[feature] {
			return fmt.Errorf("parser does not have %s capability",
				feature)
		}
	}
	return nil
}

func (dec *Decoder) readOSMData(blob *OSMPBF.Blob) error {
	data, err := getData(blob)
	if err != nil {
		return err
	}

	primitiveBlock := &OSMPBF.PrimitiveBlock{}
	if err := proto.Unmarshal(data, primitiveBlock); err != nil {
		return err
	}
	dec.parsePrimitiveBlock(primitiveBlock)
	return nil
}

func getData(blob *OSMPBF.Blob) ([]byte, error) {
	switch {
	case blob.Raw != nil:
		return blob.GetRaw(), nil
	case blob.ZlibData != nil:
		compressedData := bytes.NewBuffer(blob.GetZlibData())
		r, err := zlib.NewReader(compressedData)
		if err != nil {
			return nil, err
		}
		data, err := ioutil.ReadAll(r)
		if err != nil {
			return nil, err
		}
		if len(data) != int(blob.GetRawSize()) {
			return nil, fmt.Errorf(
				"raw blob data size %d but expected %d",
				len(data), blob.GetRawSize())
		}
		return data, nil
	}
	return nil, errors.New("unknown blob data")
}

func (dec *Decoder) parsePrimitiveBlock(pb *OSMPBF.PrimitiveBlock) {
	primitiveGroups := pb.GetPrimitivegroup()
	for _, pg := range primitiveGroups {
		dec.parsePrimitiveGroup(pb, pg)
	}
}

func (dec *Decoder) parsePrimitiveGroup(pb *OSMPBF.PrimitiveBlock, pg *OSMPBF.PrimitiveGroup) {
	dec.parseNodes(pb, pg.GetNodes())
	dec.parseDenseNodes(pb, pg.GetDense())
	dec.parseWays(pb, pg.GetWays())
	dec.parseRelations(pb, pg.GetRelations())
}

func (dec *Decoder) parseNodes(pb *OSMPBF.PrimitiveBlock, nodes []*OSMPBF.Node) {
	st := pb.GetStringtable().GetS()
	granularity := int64(pb.GetGranularity())
	latOffset := pb.GetLatOffset()
	lonOffset := pb.GetLonOffset()

	for _, node := range nodes {
		id := node.GetId()
		lat := node.GetLat()
		lon := node.GetLon()

		latitude := 1e-9 * float64((latOffset + (granularity * lat)))
		longitude := 1e-9 * float64((lonOffset + (granularity * lon)))

		tags := extractTags(st, node.GetKeys(), node.GetVals())
		dec.objectQueue = append(dec.objectQueue,
			&Node{id, latitude, longitude, tags})
		panic("Please test this first")
	}
}

func (dec *Decoder) parseDenseNodes(pb *OSMPBF.PrimitiveBlock, dn *OSMPBF.DenseNodes) {
	st := pb.GetStringtable().GetS()
	granularity := int64(pb.GetGranularity())
	latOffset := pb.GetLatOffset()
	lonOffset := pb.GetLonOffset()
	ids := dn.GetId()
	lats := dn.GetLat()
	lons := dn.GetLon()
	tu := tagUnpacker{st, dn.GetKeysVals(), 0}
	id, lat, lon := int64(0), int64(0), int64(0)
	for index := range ids {
		id = ids[index] + id
		lat = lats[index] + lat
		lon = lons[index] + lon
		latitude := 1e-9 * float64((latOffset + (granularity * lat)))
		longitude := 1e-9 * float64((lonOffset + (granularity * lon)))
		tags := tu.next()
		dec.objectQueue = append(dec.objectQueue,
			&Node{id, latitude, longitude, tags})
	}
}

type tagUnpacker struct {
	stringTable [][]byte
	keysVals    []int32
	index       int
}

func (tu *tagUnpacker) next() map[string]string {
	tags := make(map[string]string)
	for ; tu.index < len(tu.keysVals); tu.index++ {

		keyID := tu.keysVals[tu.index]
		tu.index++
		if keyID == 0 {
			break
		}
		valID := tu.keysVals[tu.index]
		key := string(tu.stringTable[keyID])
		val := string(tu.stringTable[valID])
		tags[key] = val
	}
	return tags
}

func extractTags(stringTable [][]byte, keyIDs, valueIDs []uint32) map[string]string {
	tags := make(map[string]string)
	for index := range keyIDs {
		key := string(stringTable[keyIDs[index]])
		val := string(stringTable[valueIDs[index]])
		tags[key] = val
	}
	return tags
}

func (dec *Decoder) parseWays(pb *OSMPBF.PrimitiveBlock, ways []*OSMPBF.Way) {
	st := pb.GetStringtable().GetS()
	for _, way := range ways {
		id := way.GetId()
		tags := extractTags(st, way.GetKeys(), way.GetVals())

		// NodeIDs
		refs := way.GetRefs()
		nodeID := int64(0)
		nodeIDs := make([]int64, 0, len(refs))
		for index := range refs {
			nodeID = refs[index] + nodeID
			nodeIDs = append(nodeIDs, nodeID)
		}
		dec.objectQueue = append(dec.objectQueue,
			&Way{id, tags, nodeIDs})
	}
}

func extractMembers(stringTable [][]byte, rel *OSMPBF.Relation) []Member {
	memIDs := rel.GetMemids()
	types := rel.GetTypes()
	roleIDs := rel.GetRolesSid()

	memID := int64(0)
	members := make([]Member, 0, len(memIDs))
	for index := range memIDs {
		memID = memIDs[index] + memID
		memType := NodeType
		switch types[index] {
		case OSMPBF.Relation_WAY:
			memType = WayType
		case OSMPBF.Relation_RELATION:
			memType = RelationType
		}
		role := string(stringTable[roleIDs[index]])
		members = append(members, Member{memID, memType, role})
	}
	return members
}

func (dec *Decoder) parseRelations(pb *OSMPBF.PrimitiveBlock, relations []*OSMPBF.Relation) {
	st := pb.GetStringtable().GetS()
	for _, rel := range relations {
		id := rel.GetId()
		tags := extractTags(st, rel.GetKeys(), rel.GetVals())
		members := extractMembers(st, rel)

		dec.objectQueue = append(dec.objectQueue,
			&Relation{id, tags, members})
	}
}
<commit_msg>Typical fileblock contains 8000 entries.<commit_after>// Package osmpbf decodes OpenStreetMap pbf files.
// Use this package by creating a NewDecoder and passing it a PBF file. Use
// Decode to return Node, Way and Relation structs.
package osmpbf

import (
	"bytes"
	"code.google.com/p/goprotobuf/proto"
	"compress/zlib"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/qedus/osmpbf/OSMPBF"
	"io"
	"io/ioutil"
)

const (
	maxBlobHeaderSize = 64 * 1024
	maxBlobSize       = 32 * 1024 * 1024
)

var (
	parseCapabilities = map[string]bool{
		"OsmSchema-V0.6": true,
		"DenseNodes":     true,
	}
)

type Node struct {
	ID   int64
	Lat  float64
	Lon  float64
	Tags map[string]string

	// TODO: Add DenseInfo
}

type Way struct {
	ID      int64
	Tags    map[string]string
	NodeIDs []int64

	// TODO: Add Info
}

type Relation struct {
	ID      int64
	Tags    map[string]string
	Members []Member

	// TODO: Add Info
	// TODO: Add roles_sid
}

type MemberType int

const (
	NodeType MemberType = iota
	WayType
	RelationType
)

type Member struct {
	ID   int64
	Type MemberType
	Role string
}

// A Decoder reads and decodes OpenStreetMap PBF data from an input stream.
type Decoder struct {
	r io.Reader

	objectQueue []interface{}
	objectIndex int
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r, make([]interface{}, 0, 8000), 0}
}

// Decode reads the next object from the input stream and returns either a
// Node, Way or Relation struct representing the underlying OpenStreetMap PBF
// data.
//
// The end of the input stream is reported by an io.EOF error.
func (dec *Decoder) Decode() (interface{}, error) {
	if dec.objectIndex >= len(dec.objectQueue) {
		dec.objectQueue = dec.objectQueue[:0]
		dec.objectIndex = 0
		if err := dec.readNextPrimitiveBlock(); err != nil {
			return nil, err
		}
	}

	dec.objectIndex++
	return dec.objectQueue[dec.objectIndex-1], nil
}

func (dec *Decoder) readNextPrimitiveBlock() error {
	for {
		blobHeaderSize, err := dec.readBlobHeaderSize()
		if err != nil {
			return err
		}

		blobHeader, err := dec.readBlobHeader(blobHeaderSize)
		if err != nil {
			return err
		}

		blob, err := dec.readBlob(blobHeader)
		if err != nil {
			return err
		}

		switch blobHeader.GetType() {
		case "OSMHeader":
			if err := dec.readOSMHeader(blob); err != nil {
				return err
			}
		case "OSMData":
			return dec.readOSMData(blob)
		default:
			// Skip over unknown type
		}
	}
}

func (dec *Decoder) readBlobHeaderSize() (uint32, error) {
	buf := make([]byte, 4)
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return 0, err
	}
	size := binary.BigEndian.Uint32(buf)

	if size >= maxBlobHeaderSize {
		return 0, errors.New("BlobHeader size >= 64Kb")
	}
	return size, nil
}

func (dec *Decoder) readBlobHeader(size uint32) (*OSMPBF.BlobHeader, error) {
	buf := make([]byte, size)
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return nil, err
	}

	blobHeader := &OSMPBF.BlobHeader{}
	if err := proto.Unmarshal(buf, blobHeader); err != nil {
		return nil, err
	}

	if blobHeader.GetDatasize() >= maxBlobSize {
		return nil, errors.New("Blob size >= 32Mb")
	}
	return blobHeader, nil
}

func (dec *Decoder) readBlob(blobHeader *OSMPBF.BlobHeader) (*OSMPBF.Blob, error) {
	buf := make([]byte, blobHeader.GetDatasize())
	if _, err := io.ReadFull(dec.r, buf); err != nil {
		return nil, err
	}

	blob := &OSMPBF.Blob{}
	if err := proto.Unmarshal(buf, blob); err != nil {
		return nil, err
	}
	return blob, nil
}

func (dec *Decoder) readOSMHeader(blob *OSMPBF.Blob) error {
	data, err := getData(blob)
	if err != nil {
		return err
	}
	headerBlock := &OSMPBF.HeaderBlock{}
	if err := proto.Unmarshal(data, headerBlock); err != nil {
		return err
	}

	// Check we have the parse capabilities
	requiredFeatures := headerBlock.GetRequiredFeatures()
	for _, feature := range requiredFeatures {
		if !parseCapabilities[feature] {
			return fmt.Errorf("parser does not have %s capability",
				feature)
		}
	}
	return nil
}

func (dec *Decoder) readOSMData(blob *OSMPBF.Blob) error {
	data, err := getData(blob)
	if err != nil {
		return err
	}

	primitiveBlock := &OSMPBF.PrimitiveBlock{}
	if err := proto.Unmarshal(data, primitiveBlock); err != nil {
		return err
	}
	dec.parsePrimitiveBlock(primitiveBlock)
	return nil
}

func getData(blob *OSMPBF.Blob) ([]byte, error) {
	switch {
	case blob.Raw != nil:
		return blob.GetRaw(), nil
	case blob.ZlibData != nil:
		compressedData := bytes.NewBuffer(blob.GetZlibData())
		r, err := zlib.NewReader(compressedData)
		if err != nil {
			return nil, err
		}
		data, err := ioutil.ReadAll(r)
		if err != nil {
			return nil, err
		}
		if len(data) != int(blob.GetRawSize()) {
			return nil, fmt.Errorf(
				"raw blob data size %d but expected %d",
				len(data), blob.GetRawSize())
		}
		return data, nil
	}
	return nil, errors.New("unknown blob data")
}

func (dec *Decoder) parsePrimitiveBlock(pb *OSMPBF.PrimitiveBlock) {
	primitiveGroups := pb.GetPrimitivegroup()
	for _, pg := range primitiveGroups {
		dec.parsePrimitiveGroup(pb, pg)
	}
}

func (dec *Decoder) parsePrimitiveGroup(pb *OSMPBF.PrimitiveBlock, pg *OSMPBF.PrimitiveGroup) {
	dec.parseNodes(pb, pg.GetNodes())
	dec.parseDenseNodes(pb, pg.GetDense())
	dec.parseWays(pb, pg.GetWays())
	dec.parseRelations(pb, pg.GetRelations())
}

func (dec *Decoder) parseNodes(pb *OSMPBF.PrimitiveBlock, nodes []*OSMPBF.Node) {
	st := pb.GetStringtable().GetS()
	granularity := int64(pb.GetGranularity())
	latOffset := pb.GetLatOffset()
	lonOffset := pb.GetLonOffset()

	for _, node := range nodes {
		id := node.GetId()
		lat := node.GetLat()
		lon := node.GetLon()

		latitude := 1e-9 * float64((latOffset + (granularity * lat)))
		longitude := 1e-9 * float64((lonOffset + (granularity * lon)))

		tags := extractTags(st, node.GetKeys(), node.GetVals())
		dec.objectQueue = append(dec.objectQueue,
			&Node{id, latitude, longitude, tags})
		panic("Please test this first")
	}
}

func (dec *Decoder) parseDenseNodes(pb *OSMPBF.PrimitiveBlock, dn *OSMPBF.DenseNodes) {
	st := pb.GetStringtable().GetS()
	granularity := int64(pb.GetGranularity())
	latOffset := pb.GetLatOffset()
	lonOffset := pb.GetLonOffset()
	ids := dn.GetId()
	lats := dn.GetLat()
	lons := dn.GetLon()
	tu := tagUnpacker{st, dn.GetKeysVals(), 0}
	id, lat, lon := int64(0), int64(0), int64(0)
	for index := range ids {
		id = ids[index] + id
		lat = lats[index] + lat
		lon = lons[index] + lon
		latitude := 1e-9 * float64((latOffset + (granularity * lat)))
		longitude := 1e-9 * float64((lonOffset + (granularity * lon)))
		tags := tu.next()
		dec.objectQueue = append(dec.objectQueue,
			&Node{id, latitude, longitude, tags})
	}
}

type tagUnpacker struct {
	stringTable [][]byte
	keysVals    []int32
	index       int
}

func (tu *tagUnpacker) next() map[string]string {
	tags := make(map[string]string)
	for ; tu.index < len(tu.keysVals); tu.index++ {

		keyID := tu.keysVals[tu.index]
		tu.index++
		if keyID == 0 {
			break
		}
		valID := tu.keysVals[tu.index]
		key := string(tu.stringTable[keyID])
		val := string(tu.stringTable[valID])
		tags[key] = val
	}
	return tags
}

func extractTags(stringTable [][]byte, keyIDs, valueIDs []uint32) map[string]string {
	tags := make(map[string]string)
	for index := range keyIDs {
		key := string(stringTable[keyIDs[index]])
		val := string(stringTable[valueIDs[index]])
		tags[key] = val
	}
	return tags
}

func (dec *Decoder) parseWays(pb *OSMPBF.PrimitiveBlock, ways []*OSMPBF.Way) {
	st := pb.GetStringtable().GetS()
	for _, way := range ways {
		id := way.GetId()
		tags := extractTags(st, way.GetKeys(), way.GetVals())

		// NodeIDs
		refs := way.GetRefs()
		nodeID := int64(0)
		nodeIDs := make([]int64, 0, len(refs))
		for index := range refs {
			nodeID = refs[index] + nodeID
			nodeIDs = append(nodeIDs, nodeID)
		}
		dec.objectQueue = append(dec.objectQueue,
			&Way{id, tags, nodeIDs})
	}
}

func extractMembers(stringTable [][]byte, rel *OSMPBF.Relation) []Member {
	memIDs := rel.GetMemids()
	types := rel.GetTypes()
	roleIDs := rel.GetRolesSid()

	memID := int64(0)
	members := make([]Member, 0, len(memIDs))
	for index := range memIDs {
		memID = memIDs[index] + memID
		memType := NodeType
		switch types[index] {
		case OSMPBF.Relation_WAY:
			memType = WayType
		case OSMPBF.Relation_RELATION:
			memType = RelationType
		}
		role := string(stringTable[roleIDs[index]])
		members = append(members, Member{memID, memType, role})
	}
	return members
}

func (dec *Decoder) parseRelations(pb *OSMPBF.PrimitiveBlock, relations []*OSMPBF.Relation) {
	st := pb.GetStringtable().GetS()
	for _, rel := range relations {
		id := rel.GetId()
		tags := extractTags(st, rel.GetKeys(), rel.GetVals())
		members := extractMembers(st, rel)

		dec.objectQueue = append(dec.objectQueue,
			&Relation{id, tags, members})
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Andreas Louca. All rights reserved.
// Use of this source code is goverend by a BSD-style
// license that can be found in the LICENSE file.

package gosnmp

import (
	"encoding/asn1"
	"fmt"
)

type Asn1BER byte

const (
	Integer          Asn1BER = 0x02
	BitString                = 0x03
	OctetString              = 0x04
	Null                     = 0x05
	ObjectIdentifier         = 0x06
	Counter32                = 0x41
	Gauge32                  = 0x42
	TimeTicks                = 0x43
	Opaque                   = 0x44
	NsapAddress              = 0x45
	Counter64                = 0x46
	Uinteger32               = 0x47
)

// Different packet structure is needed during decode, to trick encoding/asn1 to decode the SNMP packet

type Variable struct {
	Name  asn1.ObjectIdentifier
	Type  Asn1BER
	Value interface{}
}

type VarBind struct {
	Name  asn1.ObjectIdentifier
	Value asn1.RawValue
}

type PDU struct {
	RequestId   int32
	ErrorStatus int
	ErrorIndex  int
	VarBindList []VarBind
}
type PDUResponse struct {
	RequestId   int32
	ErrorStatus int
	ErrorIndex  int
	VarBindList []*Variable
}

type Message struct {
	Version   int
	Community []uint8
	Data      asn1.RawValue
}

func decode(data []byte) (*PDUResponse, error) {
	m := Message{}
	_, err := asn1.Unmarshal(data, &m)
	if err != nil {
		return nil, err
	}
	choice := m.Data.FullBytes[0]
	switch choice {
	// SNMP Response
	case 0xa0, 0xa1, 0xa2:

		pdu := new(PDU)

		// hack ANY -> IMPLICIT SEQUENCE
		m.Data.FullBytes[0] = 0x30
		_, err = asn1.Unmarshal(m.Data.FullBytes, pdu)
		if err != nil {
			return nil, fmt.Errorf("Error decoding pdu: %#v, %#v, %s", m.Data.FullBytes, pdu, err)
		}

		// make response pdu
		resp := new(PDUResponse)
		// Copy values from parsed pdu
		resp.RequestId = pdu.RequestId
		resp.ErrorIndex = pdu.ErrorIndex
		resp.ErrorStatus = pdu.ErrorStatus

		resp.VarBindList = make([]*Variable, len(pdu.VarBindList))

		// Decode all vars
		for c, v := range pdu.VarBindList {

			val, err := decodeValue(v.Value)
			if err != nil {
				return nil, err
			} else {
				val.Name = v.Name
				resp.VarBindList[c] = val
			}
		}

		return resp, nil
	default:
		return nil, fmt.Errorf("Unable to decode type: %#v\n", choice)
	}
	return nil, fmt.Errorf("Unknown CHOICE: %x", choice)
}

func decodeValue(data asn1.RawValue) (retVal *Variable, err error) {
	retVal = new(Variable)

	switch Asn1BER(data.FullBytes[0]) {

	// Integer
	case Integer:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Integer
		retVal.Value = ret
	// Octet
	case OctetString:
		retVal.Type = OctetString
		retVal.Value = string(data.FullBytes[2:])
	// Counter32
	case Counter32:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Counter32
		retVal.Value = ret
	case TimeTicks:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = TimeTicks
		retVal.Value = ret
	// Gauge32
	case Gauge32:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Gauge32
		retVal.Value = ret
	case Counter64:
		ret, err := parseInt64(data.FullBytes[2:])

		// Decode it
		if err != nil {
			break
		}

		retVal.Type = Counter64
		retVal.Value = ret
	default:
		err = fmt.Errorf("Unable to decode %x - not implemented", data.FullBytes[0])
	}

	return
}
<commit_msg>Implemented NoSuchObject/NoSuchInstance<commit_after>// Copyright 2012 Andreas Louca. All rights reserved.
// Use of this source code is goverend by a BSD-style
// license that can be found in the LICENSE file.

package gosnmp

import (
	"encoding/asn1"
	"fmt"
)

type Asn1BER byte

const (
	Integer          Asn1BER = 0x02
	BitString                = 0x03
	OctetString              = 0x04
	Null                     = 0x05
	ObjectIdentifier         = 0x06
	Counter32                = 0x41
	Gauge32                  = 0x42
	TimeTicks                = 0x43
	Opaque                   = 0x44
	NsapAddress              = 0x45
	Counter64                = 0x46
	Uinteger32               = 0x47
	NoSuchObject             = 0x80
	NoSuchInstance           = 0x81
)

// Different packet structure is needed during decode, to trick encoding/asn1 to decode the SNMP packet

type Variable struct {
	Name  asn1.ObjectIdentifier
	Type  Asn1BER
	Value interface{}
}

type VarBind struct {
	Name  asn1.ObjectIdentifier
	Value asn1.RawValue
}

type PDU struct {
	RequestId   int32
	ErrorStatus int
	ErrorIndex  int
	VarBindList []VarBind
}
type PDUResponse struct {
	RequestId   int32
	ErrorStatus int
	ErrorIndex  int
	VarBindList []*Variable
}

type Message struct {
	Version   int
	Community []uint8
	Data      asn1.RawValue
}

func decode(data []byte) (*PDUResponse, error) {
	m := Message{}
	_, err := asn1.Unmarshal(data, &m)
	if err != nil {
		return nil, err
	}
	choice := m.Data.FullBytes[0]
	switch choice {
	// SNMP Response
	case 0xa0, 0xa1, 0xa2:

		pdu := new(PDU)

		// hack ANY -> IMPLICIT SEQUENCE
		m.Data.FullBytes[0] = 0x30
		_, err = asn1.Unmarshal(m.Data.FullBytes, pdu)
		if err != nil {
			return nil, fmt.Errorf("Error decoding pdu: %#v, %#v, %s", m.Data.FullBytes, pdu, err)
		}

		// make response pdu
		resp := new(PDUResponse)
		// Copy values from parsed pdu
		resp.RequestId = pdu.RequestId
		resp.ErrorIndex = pdu.ErrorIndex
		resp.ErrorStatus = pdu.ErrorStatus

		resp.VarBindList = make([]*Variable, len(pdu.VarBindList))

		// Decode all vars
		for c, v := range pdu.VarBindList {

			val, err := decodeValue(v.Value)
			if err != nil {
				return nil, err
			} else {
				val.Name = v.Name
				resp.VarBindList[c] = val
			}
		}

		return resp, nil
	default:
		return nil, fmt.Errorf("Unable to decode type: %#v\n", choice)
	}
	return nil, fmt.Errorf("Unknown CHOICE: %x", choice)
}

func decodeValue(data asn1.RawValue) (retVal *Variable, err error) {
	retVal = new(Variable)

	switch Asn1BER(data.FullBytes[0]) {

	// Integer
	case Integer:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Integer
		retVal.Value = ret
	// Octet
	case OctetString:
		retVal.Type = OctetString
		retVal.Value = string(data.FullBytes[2:])
	// Counter32
	case Counter32:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Counter32
		retVal.Value = ret
	case TimeTicks:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = TimeTicks
		retVal.Value = ret
	// Gauge32
	case Gauge32:
		ret, err := parseInt(data.FullBytes[2:])
		if err != nil {
			break
		}
		retVal.Type = Gauge32
		retVal.Value = ret
	case Counter64:
		ret, err := parseInt64(data.FullBytes[2:])

		// Decode it
		if err != nil {
			break
		}

		retVal.Type = Counter64
		retVal.Value = ret
	case NoSuchInstance:
		return nil, fmt.Errorf("No such instace")
	case NoSuchObject:
		return nil, fmt.Errorf("No such object")
	default:
		err = fmt.Errorf("Unable to decode %x - not implemented", data.FullBytes[0])
	}

	return
}
<|endoftext|>


<commit_before>// khan
// https://github.com/topfreegames/khan
//
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package util

// VERSION identifies Khan's current version
var VERSION = "3.0.1"
<commit_msg>Release v3.1.0<commit_after>// khan
// https://github.com/topfreegames/khan
//
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package util

// VERSION identifies Khan's current version
var VERSION = "3.1.0"
<|endoftext|>


<commit_before>package main

import (
	"C"

	"github.com/blueoceans/amp-cache-url"
)

func init() {
}

func main() {
}

//export isCacheURL
func isCacheURL(rawurl string) bool {
	return ampCacheURL.IsCacheURL(rawurl)
}

//export getCacheURL
func getCacheURL(rawurl string) (string, error) {
	return ampCacheURL.GetCacheURL(rawurl)
}

//export getOriginURL
func getOriginURL(rawurl string) (string, error) {
	return ampCacheURL.GetOriginURL(rawurl)
}
<commit_msg>same as godoc<commit_after>package main

import (
	"C"

	"github.com/blueoceans/amp-cache-url"
)

func init() {
}

func main() {
}

//export IsCacheURL
func IsCacheURL(rawurl string) bool {
	return ampCacheURL.IsCacheURL(rawurl)
}

//export GetCacheURL
func GetCacheURL(rawurl string) (string, error) {
	return ampCacheURL.GetCacheURL(rawurl)
}

//export GetOriginURL
func GetOriginURL(rawurl string) (string, error) {
	return ampCacheURL.GetOriginURL(rawurl)
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

// msgReader is a helper that reads values from a PostgreSQL message.
type msgReader struct {
	reader            *bufio.Reader
	buf               [128]byte
	msgBytesRemaining int32
	err               error
}

// Err returns any error that the msgReader has experienced
func (r *msgReader) Err() error {
	return r.err
}

// fatal tells r that a Fatal error has occurred
func (r *msgReader) fatal(err error) {
	r.err = err
}

// rxMsg reads the type and size of the next message.
func (r *msgReader) rxMsg() (t byte, err error) {
	if r.err != nil {
		return 0, err
	}
	if r.msgBytesRemaining > 0 {
		io.CopyN(ioutil.Discard, r.reader, int64(r.msgBytesRemaining))
	}
	fmt.Println("1")
	b := r.buf[0:5]
	fmt.Println("2")
	fmt.Printf("%#v\n", b)
	_, err = io.ReadFull(r.reader, b)
	fmt.Println(err)
	if r.err != nil {
		fmt.Println(err)
	}
	fmt.Println("3")
	r.msgBytesRemaining = int32(binary.BigEndian.Uint32(b[1:])) - 4
	fmt.Println("4")
	return b[0], err
}

func (r *msgReader) readByte() byte {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 1
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b, err := r.reader.ReadByte()
	if err != nil {
		r.fatal(err)
		return 0
	}

	return b
}

func (r *msgReader) readInt16() int16 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 2
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:2]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int16(binary.BigEndian.Uint16(b))
}

func (r *msgReader) readInt32() int32 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 4
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:4]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int32(binary.BigEndian.Uint32(b))
}

func (r *msgReader) readInt64() int64 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 8
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:8]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int64(binary.BigEndian.Uint64(b))
}

func (r *msgReader) readOid() Oid {
	return Oid(r.readInt32())
}

// readCString reads a null terminated string
func (r *msgReader) readCString() string {
	if r.err != nil {
		return ""
	}

	b, err := r.reader.ReadBytes(0)
	if err != nil {
		r.fatal(err)
		return ""
	}

	r.msgBytesRemaining -= int32(len(b))
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return ""
	}

	return string(b[0 : len(b)-1])
}

// readString reads count bytes and returns as string
func (r *msgReader) readString(count int32) string {
	if r.err != nil {
		return ""
	}

	r.msgBytesRemaining -= count
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return ""
	}

	var b []byte
	if count <= int32(len(r.buf)) {
		b = r.buf[0:int(count)]
	} else {
		b = make([]byte, int(count))
	}

	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return ""
	}

	return string(b)
}

// readBytes reads count bytes and returns as []byte
func (r *msgReader) readBytes(count int32) []byte {
	if r.err != nil {
		return nil
	}

	r.msgBytesRemaining -= count
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return nil
	}

	b := make([]byte, int(count))

	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return nil
	}

	return b
}
<commit_msg>Update<commit_after>package proxy

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

// msgReader is a helper that reads values from a PostgreSQL message.
type msgReader struct {
	reader            *bufio.Reader
	buf               [128]byte
	msgBytesRemaining int32
	err               error
}

// Err returns any error that the msgReader has experienced
func (r *msgReader) Err() error {
	return r.err
}

// fatal tells r that a Fatal error has occurred
func (r *msgReader) fatal(err error) {
	r.err = err
}

// rxMsg reads the type and size of the next message.
func (r *msgReader) rxMsg() (t byte, err error) {
	if r.err != nil {
		return 0, err
	}
	if r.msgBytesRemaining > 0 {
		io.CopyN(ioutil.Discard, r.reader, int64(r.msgBytesRemaining))
	}
	fmt.Println("1")
	b := r.buf[0:5]
	fmt.Println("2")
	_, err = io.ReadAtLeast(r.reader, b, len(b))
	// ReadFull(r.reader, b)
	fmt.Println(err)
	if r.err != nil {
		fmt.Println(err)
	}
	fmt.Println("3")
	r.msgBytesRemaining = int32(binary.BigEndian.Uint32(b[1:])) - 4
	fmt.Println("4")
	return b[0], err
}

func (r *msgReader) readByte() byte {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 1
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b, err := r.reader.ReadByte()
	if err != nil {
		r.fatal(err)
		return 0
	}

	return b
}

func (r *msgReader) readInt16() int16 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 2
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:2]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int16(binary.BigEndian.Uint16(b))
}

func (r *msgReader) readInt32() int32 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 4
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:4]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int32(binary.BigEndian.Uint32(b))
}

func (r *msgReader) readInt64() int64 {
	if r.err != nil {
		return 0
	}

	r.msgBytesRemaining -= 8
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return 0
	}

	b := r.buf[0:8]
	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return 0
	}

	return int64(binary.BigEndian.Uint64(b))
}

func (r *msgReader) readOid() Oid {
	return Oid(r.readInt32())
}

// readCString reads a null terminated string
func (r *msgReader) readCString() string {
	if r.err != nil {
		return ""
	}

	b, err := r.reader.ReadBytes(0)
	if err != nil {
		r.fatal(err)
		return ""
	}

	r.msgBytesRemaining -= int32(len(b))
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return ""
	}

	return string(b[0 : len(b)-1])
}

// readString reads count bytes and returns as string
func (r *msgReader) readString(count int32) string {
	if r.err != nil {
		return ""
	}

	r.msgBytesRemaining -= count
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return ""
	}

	var b []byte
	if count <= int32(len(r.buf)) {
		b = r.buf[0:int(count)]
	} else {
		b = make([]byte, int(count))
	}

	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return ""
	}

	return string(b)
}

// readBytes reads count bytes and returns as []byte
func (r *msgReader) readBytes(count int32) []byte {
	if r.err != nil {
		return nil
	}

	r.msgBytesRemaining -= count
	if r.msgBytesRemaining < 0 {
		r.fatal(errors.New("read past end of message"))
		return nil
	}

	b := make([]byte, int(count))

	_, err := io.ReadFull(r.reader, b)
	if err != nil {
		r.fatal(err)
		return nil
	}

	return b
}
<|endoftext|>


<commit_before>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v1alpha3

import (
	xdsapi "github.com/envoyproxy/go-control-plane/envoy/api/v2"

	"istio.io/istio/pilot/pkg/model"
	"istio.io/istio/pilot/pkg/networking/plugin"
)

type ConfigGeneratorImpl struct {
	// List of plugins that modify code generated by this config generator
	Plugins []plugin.Plugin
	// List of outbound clusters keyed by configNamespace & locality. For use by gateways only.
	// Must be rebuilt for each push epoch
	PrecomputedOutboundClustersForGateways map[string]map[string][]*xdsapi.Cluster
	// TODO: add others in future
}

func NewConfigGenerator(plugins []plugin.Plugin) *ConfigGeneratorImpl {
	return &ConfigGeneratorImpl{
		Plugins: plugins,
	}
}

func (configgen *ConfigGeneratorImpl) CanUsePrecomputedCDS(proxy *model.Proxy) bool {
	networkView := model.GetNetworkView(proxy)
	// If we have only more than one network view for the proxy, then recompute CDS.
	// Because, by default, we cache the CDS output for proxies in the UnnamedNetwork only.
	if len(networkView) > 1 {
		return false
	}

	return networkView[model.UnnamedNetwork]
}
<commit_msg>post cds cache remove cleanup (#14416)<commit_after>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v1alpha3

import (
	"istio.io/istio/pilot/pkg/networking/plugin"
)

type ConfigGeneratorImpl struct {
	// List of plugins that modify code generated by this config generator
	Plugins []plugin.Plugin
}

func NewConfigGenerator(plugins []plugin.Plugin) *ConfigGeneratorImpl {
	return &ConfigGeneratorImpl{
		Plugins: plugins,
	}
}
<|endoftext|>


<commit_before>package pastebin

import (
	// Go Builtin Packages
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"strings"
	"time"

	// Google Appengine Packages
	"appengine"
	"appengine/datastore"
	"appengine/user"

	// The Gorilla Web Toolkit
	"github.com/gorilla/csrf"
	"github.com/gorilla/mux"

	// Local Packages
	"pastebin/auth"
	counter "pastebin/counter"
	"pastebin/models"
	"pastebin/utils"
)

// FIXME: load auth key from elsewhere!
const csrf_auth_key string = "u<:>ZoTv3d<45!Bn?ionCt4*4&t;SpV;"

func init() {
	r := mux.NewRouter().StrictSlash(true)

	r.HandleFunc("/pastebin/about", utils.ExtraSugar(about)).Methods("GET").Name("about")
	r.HandleFunc("/pastebin/login", auth.Login).Methods("GET").Name("login")
	r.HandleFunc("/pastebin/logout", auth.Logout).Methods("GET").Name("logout")

	r.HandleFunc("/pastebin/", utils.ExtraSugar(pastebin)).Methods("GET", "POST").Name("pastebin")
	r.HandleFunc("/pastebin/clean", clean).Methods("GET").Name("pastecleanr") // Order is important! :o
	r.HandleFunc("/pastebin/search/", search).Methods("GET").Name("pastesearch")
	r.HandleFunc("/pastebin/{paste_id}", utils.ExtraSugar(pasteframe)).Methods("GET").Name("pasteframe")
	r.HandleFunc("/pastebin/{paste_id}/content", utils.ExtraSugar(pastecontent)).Methods("GET").Name("pastecontent")
	r.HandleFunc("/pastebin/{paste_id}/download", utils.ExtraSugar(pastecontent)).Methods("GET").Name("pastedownload")
	r.HandleFunc("/pastebin/{paste_id}/delete", utils.ExtraSugar(pastecontent)).Methods("POST").Name("pastedelete")

	r.NotFoundHandler = http.HandlerFunc(Http404)

	CSRF := csrf.Protect([]byte(csrf_auth_key), csrf.Secure(!appengine.IsDevAppServer()))
	http.Handle("/pastebin/", CSRF(r))
}

func Http404(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	var tmpl = template.Must(template.ParseFiles("templates/404.tmpl"))
	if err := tmpl.Execute(w, nil); err != nil {
		log.Panic(err)
	}
}

func about(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/about.tmpl"))
	if err := tmpl.Execute(w, map[string]interface{}{
		"user": user.Current(c),
	}); err != nil {
		log.Panic(err)
	}
}

func pastebin(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	if r.Method == "GET" {
		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl"))

		// http://www.gorillatoolkit.org/pkg/csrf
		if err := tmpl.Execute(w, map[string]interface{}{
			csrf.TemplateTag: csrf.TemplateField(r),
			"user":           user.Current(c),
		}); err != nil {
			log.Panic(err)
		}
	} else if r.Method == "POST" {
		if err := r.ParseForm(); err != nil {
			log.Panic(err)
		}

		paste_id, err := models.NewPaste(c, r)
		if err != nil {
			if _, ok := err.(models.ValidationError); !ok {
				http.Error(w, err.Error(), 400)
				return
			} else {
				http.Error(w, "BARF!@ Something's broken!@%", 500)
				return
			}
		}

		if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
			w.Write([]byte("/pastebin/" + paste_id))
		} else {
			// http://tools.ietf.org/html/rfc2616#section-10.3.4 / Http 303
			http.Redirect(w, r, "/pastebin/"+paste_id, http.StatusSeeOther)
		}
	}
}

func pasteframe(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	usr := user.Current(c)
	v := mux.Vars(r)
	paste_id := v["paste_id"]

	if paste, err := models.GetPaste(c, paste_id); err == datastore.ErrNoSuchEntity {
		Http404(w, r)
		return
	} else if err != nil {
		log.Panic(err)
	} else {
		showDeleteBtn := false
		if usr != nil {
			if paste.UserID == usr.ID || user.IsAdmin(c) {
				showDeleteBtn = true
			}
		}

		counter.Increment(c, paste_id)
		p_count, _ := counter.Count(c, paste_id)

		var p_content string
		if paste.Format == "plain" && paste.Zlib == true {
			var _b_content bytes.Buffer
			_p_content := bufio.NewWriter(&_b_content)
			paste.ZContent(_p_content)
			p_content = _b_content.String()
		} else {
			p_content = string(paste.Content)
		}

		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/paste.tmpl"))
		if err := tmpl.Execute(w, map[string]interface{}{
			csrf.TemplateTag: csrf.TemplateField(r),
			"paste_id":       paste_id,
			"paste":          paste,
			"p_content":      p_content,
			"p_count":        p_count,
			"user":           usr,
			"deleteBtn":      showDeleteBtn,
		}); err != nil {
			log.Panic(err)
		}
	}
}

func pastecontent(w http.ResponseWriter, r *http.Request) {
	// This is what keeps people from abusing our pastebin ^_^
	w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'self' 'unsafe-inline'; img-src *")

	c := appengine.NewContext(r)
	v := mux.Vars(r)
	paste_id := v["paste_id"]

	if paste, err := models.GetPaste(c, paste_id); err == datastore.ErrNoSuchEntity {
		Http404(w, r)
		return
	} else if err != nil {
		log.Panic(err)
	} else {
		// Add a Content-Disposition header on the /download route
		if dl := strings.Split(r.URL.Path, "/"); dl[len(dl)-1] == "download" {
			var p_title, p_extn, dl_disposition string

			if len(paste.Title) > 0 {
				p_title = paste.Title
			} else {
				p_title = paste_id
			}

			if paste.Format == "html" {
				p_extn = "html"
			} else {
				p_extn = "txt"
			}

			dl_disposition = fmt.Sprintf("attachment; filename=\"%s.%s\"", p_title, p_extn)
			w.Header().Set("Content-Disposition", dl_disposition)
		}

		// Check ownership and expire paste on the /delete route
		if dl := strings.Split(r.URL.Path, "/"); dl[len(dl)-1] == "delete" {
			if usr := user.Current(c); usr != nil {
				if paste.UserID == usr.ID || user.IsAdmin(c) {
					paste.Delete(c, paste_id)
					if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
						w.Write([]byte("/pastebin/"))
					} else {
						// http://tools.ietf.org/html/rfc2616#section-10.3.4 / Http 303
						http.Redirect(w, r, "/pastebin/", http.StatusSeeOther)
					}
					return
				}
			}
		}

		paste.ZContent(w)
	}
}

func clean(w http.ResponseWriter, r *http.Request) {
	// https://cloud.google.com/appengine/docs/go/config/cron#securing_urls_for_cron
	if r.Header.Get("X-Appengine-Cron") != "true" {
		http.Error(w, "The /clean route is avaiable to cron job only.", http.StatusForbidden)
	}

	c := appengine.NewContext(r)
	threemonthsago := time.Now().AddDate(0, 0, -90) // 3 months/90 days ago

	old_stuff := datastore.NewQuery(models.PasteDSKind).
		Filter("date_published <", threemonthsago).
		KeysOnly().Limit(150) // Find up to 150 old pastes
	old_keys, err := old_stuff.GetAll(c, nil)
	if err != nil {
		log.Panic(err)
	}

	var paste_ids []*datastore.Key
	for _, old_key := range old_keys {
		paste_id := old_key.StringID()
		if last, _ := counter.Last(c, paste_id); threemonthsago.After(last) == true {
			paste_ids = append(paste_ids, old_key)
		}
	}

	log.Printf("About to delete the following pastes: %s", paste_ids)
	if err := datastore.DeleteMulti(c, paste_ids); err != nil {
		log.Panic(err)
	}

	// Clear counter shards here
	var shardc_dkeys []*datastore.Key
	for _, paste_id := range paste_ids {
		c_key := datastore.NewKey(c, "GeneralCounterShardConfig", paste_id.StringID(), 0, nil)
		shardc_dkeys = append(shardc_dkeys, c_key)

		shard_keys := datastore.NewQuery("GeneralCounterShard").Filter("Name =", paste_id.StringID()).KeysOnly()
		if shard_dkeys, err := shard_keys.GetAll(c, nil); err == nil {
			if derr := datastore.DeleteMulti(c, shard_dkeys); derr != nil {
				log.Panic(derr)
			}
		} else {
			log.Panic(err)
		}
	}

	if err := datastore.DeleteMulti(c, shardc_dkeys); err != nil {
		log.Panic(err)
	}
}

func search(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	usr := user.Current(c)

	if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
		cursor := r.URL.Query().Get("c")
		// Let's abuse an empty Paste object to validate/clean tags
		p := new(models.Paste)
		p.Tags = strings.Split(r.URL.Query().Get("tags"), " ")
		p.Validate()

		q := datastore.NewQuery(models.PasteDSKind)
		// q = q.Project("title", "date_published", "tags") // <-- That's not allowed for '='' filter queries O_o
		for _, tag := range p.Tags {
			q = q.Filter("tags =", tag)
		}
		q = q.Order("-date_published").Limit(10)
		if len(cursor) > 0 {
			if cursor, err := datastore.DecodeCursor(cursor); err != nil {
				http.Error(w, "Oops, invalid cursor supplied.", http.StatusBadRequest)
			} else {
				q = q.Start(cursor)
			}
		}

		var results []interface{}
		t := q.Run(c)
		for {
			q := models.Paste{}
			key, err := t.Next(&q)
			if err == datastore.Done {
				break
			}
			if err != nil {
				log.Panic(c, "Running query: ", err)
				break
			}

			results = append(results, map[string]interface{}{
				"paste_id": key.StringID(),
				"title":    q.Title,
				"tags":     q.Tags,
				"date":     q.Date.Format(time.ANSIC),
			})
		}

		q_result := map[string]interface{}{
			"paste": map[string]interface{}{
				"results": results,
				"tags":    p.Tags,
			},
		}

		if cursor, err := t.Cursor(); err == nil {
			q_result["cursor"] = cursor.String()
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(q_result)
	} else {
		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/search.tmpl"))
		if err := tmpl.Execute(w, map[string]interface{}{
			"user": usr,
		}); err != nil {
			log.Panic(err)
		}
	}
}
<commit_msg>Ha! Take that!@<commit_after>package pastebin

import (
	// Go Builtin Packages
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"strings"
	"time"

	// Google Appengine Packages
	"appengine"
	"appengine/datastore"
	"appengine/user"

	// The Gorilla Web Toolkit
	"github.com/gorilla/csrf"
	"github.com/gorilla/mux"

	// Local Packages
	"pastebin/auth"
	counter "pastebin/counter"
	"pastebin/models"
	"pastebin/utils"
)

// FIXME: load auth key from elsewhere!
const csrf_auth_key string = "u<:>ZoTv3d<45!Bn?ionCt4*4&t;SpV;"

func init() {
	r := mux.NewRouter().StrictSlash(true)

	r.HandleFunc("/pastebin/about", utils.ExtraSugar(about)).Methods("GET").Name("about")
	r.HandleFunc("/pastebin/login", auth.Login).Methods("GET").Name("login")
	r.HandleFunc("/pastebin/logout", auth.Logout).Methods("GET").Name("logout")

	r.HandleFunc("/pastebin/", utils.ExtraSugar(pastebin)).Methods("GET", "POST").Name("pastebin")
	r.HandleFunc("/pastebin/clean", clean).Methods("GET").Name("pastecleanr") // Order is important! :o
	r.HandleFunc("/pastebin/search/", search).Methods("GET").Name("pastesearch")
	r.HandleFunc("/pastebin/{paste_id}", utils.ExtraSugar(pasteframe)).Methods("GET").Name("pasteframe")
	r.HandleFunc("/pastebin/{paste_id}/content", utils.ExtraSugar(pastecontent)).Methods("GET").Name("pastecontent")
	r.HandleFunc("/pastebin/{paste_id}/download", utils.ExtraSugar(pastecontent)).Methods("GET").Name("pastedownload")
	r.HandleFunc("/pastebin/{paste_id}/delete", utils.ExtraSugar(pastecontent)).Methods("POST").Name("pastedelete")

	r.NotFoundHandler = http.HandlerFunc(Http404)

	CSRF := csrf.Protect([]byte(csrf_auth_key), csrf.Secure(!appengine.IsDevAppServer()))
	http.Handle("/pastebin/", CSRF(r))
}

func Http404(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	var tmpl = template.Must(template.ParseFiles("templates/404.tmpl"))
	if err := tmpl.Execute(w, nil); err != nil {
		log.Panic(err)
	}
}

func about(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/about.tmpl"))
	if err := tmpl.Execute(w, map[string]interface{}{
		"user": user.Current(c),
	}); err != nil {
		log.Panic(err)
	}
}

func pastebin(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	if r.Method == "GET" {
		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl"))

		// http://www.gorillatoolkit.org/pkg/csrf
		if err := tmpl.Execute(w, map[string]interface{}{
			csrf.TemplateTag: csrf.TemplateField(r),
			"user":           user.Current(c),
		}); err != nil {
			log.Panic(err)
		}
	} else if r.Method == "POST" {
		if err := r.ParseForm(); err != nil {
			log.Panic(err)
		}

		paste_id, err := models.NewPaste(c, r)
		if err != nil {
			if _, ok := err.(models.ValidationError); !ok {
				http.Error(w, err.Error(), 400)
				return
			} else {
				http.Error(w, "BARF!@ Something's broken!@%", 500)
				return
			}
		}

		if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
			w.Write([]byte("/pastebin/" + paste_id))
		} else {
			// http://tools.ietf.org/html/rfc2616#section-10.3.4 / Http 303
			http.Redirect(w, r, "/pastebin/"+paste_id, http.StatusSeeOther)
		}
	}
}

func pasteframe(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	usr := user.Current(c)
	v := mux.Vars(r)
	paste_id := v["paste_id"]

	if paste, err := models.GetPaste(c, paste_id); err == datastore.ErrNoSuchEntity {
		Http404(w, r)
		return
	} else if err != nil {
		log.Panic(err)
	} else {
		showDeleteBtn := false
		if usr != nil {
			if paste.UserID == usr.ID || user.IsAdmin(c) {
				showDeleteBtn = true
			}
		}

		counter.Increment(c, paste_id)
		p_count, _ := counter.Count(c, paste_id)

		var p_content string
		if paste.Format == "plain" && paste.Zlib == true {
			var _b_content bytes.Buffer
			_p_content := bufio.NewWriter(&_b_content)
			paste.ZContent(_p_content)
			p_content = _b_content.String()
		} else {
			p_content = string(paste.Content)
		}

		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/paste.tmpl"))
		if err := tmpl.Execute(w, map[string]interface{}{
			csrf.TemplateTag: csrf.TemplateField(r),
			"paste_id":       paste_id,
			"paste":          paste,
			"p_content":      p_content,
			"p_count":        p_count,
			"user":           usr,
			"deleteBtn":      showDeleteBtn,
		}); err != nil {
			log.Panic(err)
		}
	}
}

func pastecontent(w http.ResponseWriter, r *http.Request) {
	// This is what keeps people from abusing our pastebin ^_^
	w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'self' 'unsafe-inline'; img-src *")

	c := appengine.NewContext(r)
	v := mux.Vars(r)
	paste_id := v["paste_id"]

	if paste, err := models.GetPaste(c, paste_id); err == datastore.ErrNoSuchEntity {
		Http404(w, r)
		return
	} else if err != nil {
		log.Panic(err)
	} else {
		// Add a Content-Disposition header on the /download route
		if dl := strings.Split(r.URL.Path, "/"); dl[len(dl)-1] == "download" {
			var p_title, p_extn, dl_disposition string

			if len(paste.Title) > 0 {
				p_title = paste.Title
			} else {
				p_title = paste_id
			}

			if paste.Format == "html" {
				p_extn = "html"
			} else {
				p_extn = "txt"
			}

			dl_disposition = fmt.Sprintf("attachment; filename=\"%s.%s\"", p_title, p_extn)
			w.Header().Set("Content-Disposition", dl_disposition)
		}

		// Check ownership and expire paste on the /delete route
		if dl := strings.Split(r.URL.Path, "/"); dl[len(dl)-1] == "delete" {
			if usr := user.Current(c); usr != nil {
				if paste.UserID == usr.ID || user.IsAdmin(c) {
					paste.Delete(c, paste_id)
					if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
						w.Write([]byte("/pastebin/"))
					} else {
						// http://tools.ietf.org/html/rfc2616#section-10.3.4 / Http 303
						http.Redirect(w, r, "/pastebin/", http.StatusSeeOther)
					}
					return
				}
			}
		}

		paste.ZContent(w)
	}
}

func clean(w http.ResponseWriter, r *http.Request) {
	// https://cloud.google.com/appengine/docs/go/config/cron#securing_urls_for_cron
	if r.Header.Get("X-Appengine-Cron") != "true" {
		http.Error(w, "The /clean route is avaiable to cron job only.", http.StatusForbidden)
	}

	c := appengine.NewContext(r)
	threemonthsago := time.Now().AddDate(0, 0, -90) // 3 months/90 days ago

	old_stuff := datastore.NewQuery(models.PasteDSKind).
		Filter("date_published <", threemonthsago).
		KeysOnly().Limit(150) // Find up to 150 old pastes
	old_keys, err := old_stuff.GetAll(c, nil)
	if err != nil {
		log.Panic(err)
	}

	var paste_ids []*datastore.Key
	for _, old_key := range old_keys {
		paste_id := old_key.StringID()
		if last, _ := counter.Last(c, paste_id); threemonthsago.After(last) == true {
			paste_ids = append(paste_ids, old_key)
		}
	}

	log.Printf("About to delete the following pastes: %s", paste_ids)
	if err := datastore.DeleteMulti(c, paste_ids); err != nil {
		log.Panic(err)
	}

	// Clear counter shards here
	var shardc_dkeys []*datastore.Key
	for _, paste_id := range paste_ids {
		c_key := datastore.NewKey(c, "GeneralCounterShardConfig", paste_id.StringID(), 0, nil)
		shardc_dkeys = append(shardc_dkeys, c_key)

		shard_keys := datastore.NewQuery("GeneralCounterShard").Filter("Name =", paste_id.StringID()).KeysOnly()
		if shard_dkeys, err := shard_keys.GetAll(c, nil); err == nil {
			if derr := datastore.DeleteMulti(c, shard_dkeys); derr != nil {
				log.Panic(derr)
			}
		} else {
			log.Panic(err)
		}
	}

	if err := datastore.DeleteMulti(c, shardc_dkeys); err != nil {
		log.Panic(err)
	}
}

func search(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	usr := user.Current(c)

	if r.Header.Get("X-Requested-With") == "XMLHttpRequest" { // AJAX
		cursor := r.URL.Query().Get("c")
		// Let's abuse an empty Paste object to validate/clean tags
		p := new(models.Paste)
		p.Tags = strings.Split(r.URL.Query().Get("tags"), " ")
		p.Validate()

		q := datastore.NewQuery(models.PasteDSKind)
		// q = q.Project("title", "date_published", "tags") // <-- That's not allowed for '='' filter queries O_o
		for _, tag := range p.Tags {
			q = q.Filter("tags =", tag)
		}
		q = q.Order("-date_published").Limit(10)
		if len(cursor) > 0 {
			if cursor, err := datastore.DecodeCursor(cursor); err != nil {
				http.Error(w, "Oops, invalid cursor supplied.", http.StatusBadRequest)
			} else {
				q = q.Start(cursor)
			}
		}

		var results []interface{}
		t := q.Run(c)
		for {
			q := models.Paste{}
			key, err := t.Next(&q)
			if err == datastore.Done {
				break
			}
			if err != nil {
				log.Panic(c, "Running query: ", err)
				break
			}

			results = append(results, map[string]interface{}{
				"paste_id": key.StringID(),
				"title":    template.HTMLEscapeString(q.Title),
				"tags":     q.Tags,
				"date":     q.Date.Format(time.ANSIC),
			})
		}

		q_result := map[string]interface{}{
			"paste": map[string]interface{}{
				"results": results,
				"tags":    p.Tags,
			},
		}

		if cursor, err := t.Cursor(); err == nil {
			q_result["cursor"] = cursor.String()
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(q_result)
	} else {
		var tmpl = template.Must(template.ParseFiles("templates/base.tmpl", "pastebin/templates/pastebin.tmpl", "pastebin/templates/search.tmpl"))
		if err := tmpl.Execute(w, map[string]interface{}{
			"user": usr,
		}); err != nil {
			log.Panic(err)
		}
	}
}
<|endoftext|>


<commit_before>// Main package wraps sprite_sass tool for use with the command line
// See -h for list of available options
package main

/*
#cgo LDFLAGS: -Llibsass -lsass -lstdc++
#cgo CFLAGS: -Ilibsass

#include <stdlib.h>
#include <sass_interface.h>
*/
import "C"

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	sprite "github.com/drewwells/sprite_sass"
)

var (
	Dir, Input, Output, Includes, Style string
	Comments                            bool
)

func init() {
	flag.StringVar(&Output, "output", "", "Output file")
	flag.StringVar(&Output, "o", "", "Output file")
	flag.StringVar(&Includes, "p", "", "SASS import path")
	flag.StringVar(&Dir, "dir", ".", "Image directory")
	flag.StringVar(&Dir, "d", ".", "Image directory")
	flag.StringVar(&Style, "style", "nested", "CSS nested style")
	flag.StringVar(&Style, "s", "nested", "CSS nested style")
	flag.BoolVar(&Comments, "comment", true, "Turn on source comments")
	flag.BoolVar(&Comments, "c", true, "Turn on source comments")
}

func main() {
	flag.Parse()
	for _, v := range flag.Args() {
		if strings.HasPrefix(v, "-") {
			log.Fatalf("Please specify flags before other arguments: %s", v)
		}
	}

	if len(flag.Args()) == 0 {
		fmt.Println("Please specify input and output filepaths.")
		fmt.Println("\nAvailable options:")
		flag.PrintDefaults()
		return
	} else if len(flag.Args()) > 1 {
		log.Fatal("Only one input file is supported at this time")
	}

	style, ok := sprite.Style[Style]

	if !ok {
		style = sprite.NESTED_STYLE
	}

	ctx := sprite.Context{
		OutputStyle: style,
		ImageDir:    Dir,
		Comments:    Comments,
	}

	fRead, err := os.Open(flag.Arg(0))
	if err != nil {
		panic(err)
	}

	if Includes != "" {
		ctx.IncludePaths = strings.Split(Includes, ",")
	}

	var output io.WriteCloser

	if Output == "" {
		output = os.Stdout
	} else {
		output, err = os.Create(Output)
		if err != nil {
			panic(err)
		}
	}

	err = ctx.Run(fRead, output, filepath.Dir(Input))
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Add the passed file as a working directory for import lookups<commit_after>// Main package wraps sprite_sass tool for use with the command line
// See -h for list of available options
package main

/*
#cgo LDFLAGS: -Llibsass -lsass -lstdc++
#cgo CFLAGS: -Ilibsass

#include <stdlib.h>
#include <sass_interface.h>
*/
import "C"

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	sprite "github.com/drewwells/sprite_sass"
)

var (
	Dir, Input, Output, Includes, Style string
	Comments                            bool
)

func init() {
	flag.StringVar(&Output, "output", "", "Output file")
	flag.StringVar(&Output, "o", "", "Output file")
	flag.StringVar(&Includes, "p", "", "SASS import path")
	flag.StringVar(&Dir, "dir", ".", "Image directory")
	flag.StringVar(&Dir, "d", ".", "Image directory")
	flag.StringVar(&Style, "style", "nested", "CSS nested style")
	flag.StringVar(&Style, "s", "nested", "CSS nested style")
	flag.BoolVar(&Comments, "comment", true, "Turn on source comments")
	flag.BoolVar(&Comments, "c", true, "Turn on source comments")
}

func main() {
	flag.Parse()
	for _, v := range flag.Args() {
		if strings.HasPrefix(v, "-") {
			log.Fatalf("Please specify flags before other arguments: %s", v)
		}
	}

	if len(flag.Args()) == 0 {
		fmt.Println("Please specify input and output filepaths.")
		fmt.Println("\nAvailable options:")
		flag.PrintDefaults()
		return
	} else if len(flag.Args()) > 1 {
		log.Fatal("Only one input file is supported at this time")
	}

	style, ok := sprite.Style[Style]

	if !ok {
		style = sprite.NESTED_STYLE
	}

	ctx := sprite.Context{
		OutputStyle:  style,
		ImageDir:     Dir,
		Comments:     Comments,
		IncludePaths: []string{filepath.Dir(flag.Arg(0))},
	}

	fRead, err := os.Open(flag.Arg(0))
	if err != nil {
		panic(err)
	}

	if Includes != "" {
		ctx.IncludePaths = strings.Split(Includes, ",")
	}

	var output io.WriteCloser

	if Output == "" {
		output = os.Stdout
	} else {
		output, err = os.Create(Output)
		if err != nil {
			panic(err)
		}
	}

	err = ctx.Run(fRead, output, filepath.Dir(Input))
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package artifact

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/jobmanagement_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// TestRetrieve tests that we can successfully retrieve fresh files.
func TestRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idA")
	keys := []string{"foo", "bar", "baz/baz/baz"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	for _, a := range artifacts {
		filename := makeFilename(dst, a.Name)
		if err := Retrieve(ctx, client, a, rt, dst); err != nil {
			t.Errorf("failed to retrieve %v: %v", a.Name, err)
			continue
		}
		verifySHA256(t, filename, a.Sha256)
	}
}

// TestMultiRetrieve tests that we can successfully retrieve fresh files
// concurrently.
func TestMultiRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idB")
	keys := []string{"1", "2", "3", "4", "a/5", "a/6", "a/7", "a/8", "a/a/9", "a/a/10", "a/b/11", "a/b/12"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	if err := LegacyMultiRetrieve(ctx, client, 10, artifacts, rt, dst); err != nil {
		t.Errorf("failed to retrieve: %v", err)
	}

	for _, a := range artifacts {
		verifySHA256(t, makeFilename(dst, a.Name), a.Sha256)
	}
}

// populate stages a set of artifacts with the given keys, each with
// slightly different sizes and chucksizes.
func populate(ctx context.Context, cc *grpc.ClientConn, t *testing.T, keys []string, size int, st string) (string, []*pb.ArtifactMetadata) {
	scl := pb.NewLegacyArtifactStagingServiceClient(cc)

	var artifacts []*pb.ArtifactMetadata
	for i, key := range keys {
		a := stage(ctx, scl, t, key, size+7*i, 97+i, st)
		artifacts = append(artifacts, a)
	}
	token, err := Commit(ctx, scl, artifacts, st)
	if err != nil {
		t.Fatalf("failed to commit manifest: %v", err)
		return "", nil
	}
	return token, artifacts
}

// stage stages an artifact with the given key, size and chuck size. The content is
// always 'z's.
func stage(ctx context.Context, scl pb.LegacyArtifactStagingServiceClient, t *testing.T, key string, size, chunkSize int, st string) *pb.ArtifactMetadata {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	hash := hex.EncodeToString(sha256W.Sum(nil))
	md := makeArtifact(key, hash)
	pmd := &pb.PutArtifactMetadata{
		Metadata:            md,
		StagingSessionToken: st,
	}

	stream, err := scl.PutArtifact(ctx)
	if err != nil {
		t.Fatalf("put failed: %v", err)
	}
	header := &pb.PutArtifactRequest{
		Content: &pb.PutArtifactRequest_Metadata{
			Metadata: pmd,
		},
	}
	if err := stream.Send(header); err != nil {
		t.Fatalf("send header failed: %v", err)
	}

	for i := 0; i < size; i += chunkSize {
		end := i + chunkSize
		if size < end {
			end = size
		}

		chunk := &pb.PutArtifactRequest{
			Content: &pb.PutArtifactRequest_Data{
				Data: &pb.ArtifactChunk{
					Data: data[i:end],
				},
			},
		}
		if err := stream.Send(chunk); err != nil {
			t.Fatalf("send chunk[%v:%v] failed: %v", i, end, err)
		}
	}
	if _, err := stream.CloseAndRecv(); err != nil {
		t.Fatalf("close failed: %v", err)
	}
	return md
}

// Test for new artifact retrieval.

func TestNewRetrieveWithManyFiles(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.resolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func TestNewRetrieveWithResolution(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.unresolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func checkStagedFiles(mds []*pb.ArtifactMetadata, dest string, expected map[string]string, t *testing.T) {
	if len(mds) != len(expected) {
		t.Errorf("wrong number of artifacts staged %v vs %v", len(mds), len(expected))
	}
	for _, md := range mds {
		filename := filepath.Join(dest, filepath.FromSlash(md.Name))
		fd, err := os.Open(filename)
		if err != nil {
			t.Errorf("error opening file %v", err)
		}
		defer fd.Close()

		data := make([]byte, 1<<20)
		n, err := fd.Read(data)
		if err != nil {
			t.Errorf("error reading file %v", err)
		}

		if string(data[:n]) != expected[md.Name] {
			t.Errorf("missmatched contents for %v: '%s' vs '%s'", md.Name, string(data[:n]), expected[md.Name])
		}
	}
}

type fakeRetrievalService struct {
	artifacts map[string]string // name -> content
}

func (fake *fakeRetrievalService) resolvedArtifacts() []*pipeline_v1.ArtifactInformation {
	var artifacts []*pipeline_v1.ArtifactInformation
	for name, contents := range fake.artifacts {
		payload, _ := proto.Marshal(&pipeline_v1.ArtifactStagingToRolePayload{
			StagedName: name})
		artifacts = append(artifacts, &pipeline_v1.ArtifactInformation{
			TypeUrn:     "resolved",
			TypePayload: []byte(contents),
			RoleUrn:     URNStagingTo,
			RolePayload: payload,
		})
	}
	return artifacts
}

func (fake *fakeRetrievalService) unresolvedArtifacts() []*pipeline_v1.ArtifactInformation {
	return []*pipeline_v1.ArtifactInformation{
		&pipeline_v1.ArtifactInformation{
			TypeUrn: "unresolved",
		},
	}
}

func (fake *fakeRetrievalService) ResolveArtifacts(ctx context.Context, request *pb.ResolveArtifactsRequest, opts ...grpc.CallOption) (*pb.ResolveArtifactsResponse, error) {
	response := pb.ResolveArtifactsResponse{}
	for _, dep := range request.Artifacts {
		if dep.TypeUrn == "unresolved" {
			response.Replacements = append(response.Replacements, fake.resolvedArtifacts()...)
		} else {
			response.Replacements = append(response.Replacements, dep)
		}
	}
	return &response, nil
}

func (fake *fakeRetrievalService) GetArtifact(ctx context.Context, request *pb.GetArtifactRequest, opts ...grpc.CallOption) (pb.ArtifactRetrievalService_GetArtifactClient, error) {
	if request.Artifact.TypeUrn == "resolved" {
		return &fakeGetArtifactResponseStream{data: request.Artifact.TypePayload}, nil
	}
	return nil, errors.Errorf("Unsupported artifact %v", request.Artifact)
}

type fakeGetArtifactResponseStream struct {
	data  []byte
	index int
}

func (fake *fakeGetArtifactResponseStream) Recv() (*pb.GetArtifactResponse, error) {
	if fake.index < len(fake.data) {
		fake.index += 1
		return &pb.GetArtifactResponse{Data: fake.data[fake.index-1 : fake.index]}, nil
	}
	return nil, io.EOF
}

func (fake *fakeGetArtifactResponseStream) RecvMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) SendMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Header() (metadata.MD, error) {
	return nil, nil
}

func (fake *fakeGetArtifactResponseStream) Trailer() metadata.MD {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Context() context.Context {
	return nil
}

func (fake *fakeGetArtifactResponseStream) CloseSend() error {
	return nil
}

func verifySHA256(t *testing.T, filename, hash string) {
	actual, err := computeSHA256(filename)
	if err != nil {
		t.Errorf("failed to compute hash for %v: %v", filename, err)
		return
	}
	if actual != hash {
		t.Errorf("file %v has bad SHA256: %v, want %v", filename, actual, hash)
	}
}

func makeTempDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "artifact_test_")
	if err != nil {
		t.Errorf("Test failure: cannot create temporary directory: %+v", err)
	}
	return dir
}

func makeTempFiles(t *testing.T, dir string, keys []string, size int) []string {
	var sha256s []string
	for i, key := range keys {
		hash := makeTempFile(t, makeFilename(dir, key), size+i)
		sha256s = append(sha256s, hash)
	}
	return sha256s
}

func makeTempFile(t *testing.T, filename string, size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	if err := os.MkdirAll(filepath.Dir(filename), 0755); err != nil {
		t.Fatalf("cannot create directory for %s: %v", filename, err)
	}
	if err := ioutil.WriteFile(filename, data, 0644); err != nil {
		t.Fatalf("cannot create file %s: %v", filename, err)
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	return hex.EncodeToString(sha256W.Sum(nil))
}

func makeArtifact(key, hash string) *pb.ArtifactMetadata {
	return &pb.ArtifactMetadata{
		Name:        key,
		Sha256:      hash,
		Permissions: 0644,
	}
}

func makeFilename(dir, key string) string {
	return filepath.Join(dir, filepath.FromSlash(key))
}
<commit_msg>Update materialize_test.go<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package artifact

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/jobmanagement_v1"
	ppb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// TestRetrieve tests that we can successfully retrieve fresh files.
func TestRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idA")
	keys := []string{"foo", "bar", "baz/baz/baz"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	for _, a := range artifacts {
		filename := makeFilename(dst, a.Name)
		if err := Retrieve(ctx, client, a, rt, dst); err != nil {
			t.Errorf("failed to retrieve %v: %v", a.Name, err)
			continue
		}
		verifySHA256(t, filename, a.Sha256)
	}
}

// TestMultiRetrieve tests that we can successfully retrieve fresh files
// concurrently.
func TestMultiRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idB")
	keys := []string{"1", "2", "3", "4", "a/5", "a/6", "a/7", "a/8", "a/a/9", "a/a/10", "a/b/11", "a/b/12"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	if err := LegacyMultiRetrieve(ctx, client, 10, artifacts, rt, dst); err != nil {
		t.Errorf("failed to retrieve: %v", err)
	}

	for _, a := range artifacts {
		verifySHA256(t, makeFilename(dst, a.Name), a.Sha256)
	}
}

// populate stages a set of artifacts with the given keys, each with
// slightly different sizes and chucksizes.
func populate(ctx context.Context, cc *grpc.ClientConn, t *testing.T, keys []string, size int, st string) (string, []*pb.ArtifactMetadata) {
	scl := pb.NewLegacyArtifactStagingServiceClient(cc)

	var artifacts []*pb.ArtifactMetadata
	for i, key := range keys {
		a := stage(ctx, scl, t, key, size+7*i, 97+i, st)
		artifacts = append(artifacts, a)
	}
	token, err := Commit(ctx, scl, artifacts, st)
	if err != nil {
		t.Fatalf("failed to commit manifest: %v", err)
		return "", nil
	}
	return token, artifacts
}

// stage stages an artifact with the given key, size and chuck size. The content is
// always 'z's.
func stage(ctx context.Context, scl pb.LegacyArtifactStagingServiceClient, t *testing.T, key string, size, chunkSize int, st string) *pb.ArtifactMetadata {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	hash := hex.EncodeToString(sha256W.Sum(nil))
	md := makeArtifact(key, hash)
	pmd := &pb.PutArtifactMetadata{
		Metadata:            md,
		StagingSessionToken: st,
	}

	stream, err := scl.PutArtifact(ctx)
	if err != nil {
		t.Fatalf("put failed: %v", err)
	}
	header := &pb.PutArtifactRequest{
		Content: &pb.PutArtifactRequest_Metadata{
			Metadata: pmd,
		},
	}
	if err := stream.Send(header); err != nil {
		t.Fatalf("send header failed: %v", err)
	}

	for i := 0; i < size; i += chunkSize {
		end := i + chunkSize
		if size < end {
			end = size
		}

		chunk := &pb.PutArtifactRequest{
			Content: &pb.PutArtifactRequest_Data{
				Data: &pb.ArtifactChunk{
					Data: data[i:end],
				},
			},
		}
		if err := stream.Send(chunk); err != nil {
			t.Fatalf("send chunk[%v:%v] failed: %v", i, end, err)
		}
	}
	if _, err := stream.CloseAndRecv(); err != nil {
		t.Fatalf("close failed: %v", err)
	}
	return md
}

// Test for new artifact retrieval.

func TestNewRetrieveWithManyFiles(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.resolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func TestNewRetrieveWithResolution(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.unresolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func checkStagedFiles(mds []*pb.ArtifactMetadata, dest string, expected map[string]string, t *testing.T) {
	if len(mds) != len(expected) {
		t.Errorf("wrong number of artifacts staged %v vs %v", len(mds), len(expected))
	}
	for _, md := range mds {
		filename := filepath.Join(dest, filepath.FromSlash(md.Name))
		fd, err := os.Open(filename)
		if err != nil {
			t.Errorf("error opening file %v", err)
		}
		defer fd.Close()

		data := make([]byte, 1<<20)
		n, err := fd.Read(data)
		if err != nil {
			t.Errorf("error reading file %v", err)
		}

		if string(data[:n]) != expected[md.Name] {
			t.Errorf("missmatched contents for %v: '%s' vs '%s'", md.Name, string(data[:n]), expected[md.Name])
		}
	}
}

type fakeRetrievalService struct {
	artifacts map[string]string // name -> content
}

func (fake *fakeRetrievalService) resolvedArtifacts() []*ppb.ArtifactInformation {
	var artifacts []*ppb.ArtifactInformation
	for name, contents := range fake.artifacts {
		payload, _ := proto.Marshal(&ppb.ArtifactStagingToRolePayload{
			StagedName: name})
		artifacts = append(artifacts, &ppb.ArtifactInformation{
			TypeUrn:     "resolved",
			TypePayload: []byte(contents),
			RoleUrn:     URNStagingTo,
			RolePayload: payload,
		})
	}
	return artifacts
}

func (fake *fakeRetrievalService) unresolvedArtifacts() []*ppb.ArtifactInformation {
	return []*ppb.ArtifactInformation{
		&ppb.ArtifactInformation{
			TypeUrn: "unresolved",
		},
	}
}

func (fake *fakeRetrievalService) ResolveArtifacts(ctx context.Context, request *pb.ResolveArtifactsRequest, opts ...grpc.CallOption) (*pb.ResolveArtifactsResponse, error) {
	response := pb.ResolveArtifactsResponse{}
	for _, dep := range request.Artifacts {
		if dep.TypeUrn == "unresolved" {
			response.Replacements = append(response.Replacements, fake.resolvedArtifacts()...)
		} else {
			response.Replacements = append(response.Replacements, dep)
		}
	}
	return &response, nil
}

func (fake *fakeRetrievalService) GetArtifact(ctx context.Context, request *pb.GetArtifactRequest, opts ...grpc.CallOption) (pb.ArtifactRetrievalService_GetArtifactClient, error) {
	if request.Artifact.TypeUrn == "resolved" {
		return &fakeGetArtifactResponseStream{data: request.Artifact.TypePayload}, nil
	}
	return nil, errors.Errorf("Unsupported artifact %v", request.Artifact)
}

type fakeGetArtifactResponseStream struct {
	data  []byte
	index int
}

func (fake *fakeGetArtifactResponseStream) Recv() (*pb.GetArtifactResponse, error) {
	if fake.index < len(fake.data) {
		fake.index++
		return &pb.GetArtifactResponse{Data: fake.data[fake.index-1 : fake.index]}, nil
	}
	return nil, io.EOF
}

func (fake *fakeGetArtifactResponseStream) RecvMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) SendMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Header() (metadata.MD, error) {
	return nil, nil
}

func (fake *fakeGetArtifactResponseStream) Trailer() metadata.MD {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Context() context.Context {
	return context.Background()
}

func (fake *fakeGetArtifactResponseStream) CloseSend() error {
	return nil
}

func verifySHA256(t *testing.T, filename, hash string) {
	actual, err := computeSHA256(filename)
	if err != nil {
		t.Errorf("failed to compute hash for %v: %v", filename, err)
		return
	}
	if actual != hash {
		t.Errorf("file %v has bad SHA256: %v, want %v", filename, actual, hash)
	}
}

func makeTempDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "artifact_test_")
	if err != nil {
		t.Errorf("Test failure: cannot create temporary directory: %+v", err)
	}
	return dir
}

func makeTempFiles(t *testing.T, dir string, keys []string, size int) []string {
	var sha256s []string
	for i, key := range keys {
		hash := makeTempFile(t, makeFilename(dir, key), size+i)
		sha256s = append(sha256s, hash)
	}
	return sha256s
}

func makeTempFile(t *testing.T, filename string, size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	if err := os.MkdirAll(filepath.Dir(filename), 0755); err != nil {
		t.Fatalf("cannot create directory for %s: %v", filename, err)
	}
	if err := ioutil.WriteFile(filename, data, 0644); err != nil {
		t.Fatalf("cannot create file %s: %v", filename, err)
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	return hex.EncodeToString(sha256W.Sum(nil))
}

func makeArtifact(key, hash string) *pb.ArtifactMetadata {
	return &pb.ArtifactMetadata{
		Name:        key,
		Sha256:      hash,
		Permissions: 0644,
	}
}

func makeFilename(dir, key string) string {
	return filepath.Join(dir, filepath.FromSlash(key))
}
<|endoftext|>