<commit_before>package taggolib

import (
	"bytes"
	"encoding/binary"
	"io"
	"io/ioutil"
	"strconv"
	"strings"
	"time"

	"github.com/eaburns/bit"
)

var (
	// oggMagicNumber is the magic number used to identify a OGG audio stream
	oggMagicNumber = []byte("OggS")
	// oggVorbisWord is used to denote the beginning of a Vorbis information block
	oggVorbisWord = []byte("vorbis")
)

// oggParser represents a OGG audio metadata tag parser
type oggParser struct {
	duration time.Duration
	durChan  chan struct{}
	encoder  string
	idHeader *oggIDHeader
	reader   io.ReadSeeker
	tags     map[string]string

	// Shared buffer and unsigned integers stored as fields to prevent unneeded allocations
	buffer []byte
	ui8    uint8
	ui32   uint32
	ui64   uint64
}

// Album returns the Album tag for this stream
func (o oggParser) Album() string {
	return o.tags[tagAlbum]
}

// AlbumArtist returns the AlbumArtist tag for this stream
func (o oggParser) AlbumArtist() string {
	return o.tags[tagAlbumArtist]
}

// Artist returns the Artist tag for this stream
func (o oggParser) Artist() string {
	return o.tags[tagArtist]
}

// BitDepth returns the bits-per-sample of this stream
func (o oggParser) BitDepth() int {
	// Ogg Vorbis should always provide 16 bit depth
	return 16
}

// Bitrate calculates the audio bitrate for this stream
func (o oggParser) Bitrate() int {
	// TODO: see how max/min bitrate play into calculations
	return int(o.idHeader.NomBitrate) / 1000
}

// Channels returns the number of channels for this stream
func (o oggParser) Channels() int {
	return int(o.idHeader.ChannelCount)
}

// Comment returns the Comment tag for this stream
func (o oggParser) Comment() string {
	return o.tags[tagComment]
}

// Date returns the Date tag for this stream
func (o oggParser) Date() string {
	return o.tags[tagDate]
}

// DiscNumber returns the DiscNumber tag for this stream
func (o oggParser) DiscNumber() int {
	disc, err := strconv.Atoi(o.tags[tagDiscNumber])
	if err != nil {
		return 0
	}

	return disc
}

// Duration returns the time duration for this stream
func (o oggParser) Duration() time.Duration {
	return o.duration
}

// Encoder returns the encoder for this stream
func (o oggParser) Encoder() string {
	return o.encoder
}

// Format returns the name of the OGG format
func (o oggParser) Format() string {
	return "OGG"
}

// Genre returns the Genre tag for this stream
func (o oggParser) Genre() string {
	return o.tags[tagGenre]
}

// SampleRate returns the sample rate in Hertz for this stream
func (o oggParser) SampleRate() int {
	return int(o.idHeader.SampleRate)
}

// Tag attempts to return the raw, unprocessed tag with the specified name for this stream
func (o oggParser) Tag(name string) string {
	return o.tags[name]
}

// Title returns the Title tag for this stream
func (o oggParser) Title() string {
	return o.tags[tagTitle]
}

// TrackNumber returns the TrackNumber tag for this stream
func (o oggParser) TrackNumber() int {
	// Check for a /, such as 2/8
	track, err := strconv.Atoi(strings.Split(o.tags[tagTrackNumber], "/")[0])
	if err != nil {
		return 0
	}

	return track
}

// newOGGParser creates a parser for OGG audio streams
func newOGGParser(reader io.ReadSeeker) (*oggParser, error) {
	// Create OGG parser
	parser := &oggParser{
		buffer: make([]byte, 128),
		reader: reader,
	}

	// Parse the required ID header
	if err := parser.parseOGGIDHeader(); err != nil {
		return nil, err
	}

	// Parse the required comment header
	if err := parser.parseOGGCommentHeader(); err != nil {
		return nil, err
	}

	// Parse the file's duration
	if err := parser.parseOGGDuration(); err != nil {
		return nil, err
	}

	// Return parser
	return parser, nil
}

// oggPageHeader represents the information contained in an Ogg Page header
type oggPageHeader struct {
	CapturePattern  []byte
	Version         uint8
	HeaderType      uint8
	GranulePosition uint64
	BitstreamSerial uint32
	PageSequence    uint32
	Checksum        []byte
	PageSegments    uint8
}

// parseOGGPageHeader parses an Ogg page header
func (o *oggParser) parseOGGPageHeader(skipMagicNumber bool) (*oggPageHeader, error) {
	// Create page header
	pageHeader := new(oggPageHeader)

	// Unless skip is specified, check for capture pattern
	if !skipMagicNumber {
		if _, err := o.reader.Read(o.buffer[:4]); err != nil {
			return nil, err
		}
		pageHeader.CapturePattern = o.buffer[:4]

		// Verify proper capture pattern
		if !bytes.Equal(pageHeader.CapturePattern, oggMagicNumber) {
			return nil, ErrInvalidStream
		}
	} else {
		// If skipped, assume capture pattern is correct magic number
		pageHeader.CapturePattern = oggMagicNumber
	}

	// Version (must always be 0)
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.Version = o.ui8

	// Verify mandated version 0
	if pageHeader.Version != 0 {
		return nil, ErrInvalidStream
	}

	// Header type
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.HeaderType = o.ui8

	// Granule position
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui64); err != nil {
		return nil, err
	}
	pageHeader.GranulePosition = o.ui64

	// Bitstream serial number
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return nil, err
	}
	pageHeader.BitstreamSerial = o.ui32

	// Page sequence number
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return nil, err
	}
	pageHeader.PageSequence = o.ui32

	// Checksum
	if _, err := o.reader.Read(o.buffer[:4]); err != nil {
		return nil, err
	}
	pageHeader.Checksum = o.buffer[:4]

	// Page segments
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.PageSegments = o.ui8

	// Segment table is next, but we won't need it for tag parsing, so seek ahead
	// size of uint8 (1 byte) multiplied by number of page segments
	if _, err := o.reader.Seek(int64(pageHeader.PageSegments), 1); err != nil {
		return nil, err

	}
	return pageHeader, nil
}

// parseOGGCommonHeader parses information common to all Ogg Vorbis headers
func (o *oggParser) parseOGGCommonHeader() (byte, error) {
	// Read the first byte to get header type
	if _, err := o.reader.Read(o.buffer[:1]); err != nil {
		return 0, err
	}

	// Store first byte at end of buffer so we can return it later without more allocations
	o.buffer[len(o.buffer)-1] = o.buffer[0]

	// Read for 'vorbis' identification word
	if _, err := o.reader.Read(o.buffer[:len(oggVorbisWord)]); err != nil {
		return 0, err
	}

	// Ensure 'vorbis' identification word is present
	if !bytes.Equal(o.buffer[:len(oggVorbisWord)], oggVorbisWord) {
		return 0, ErrInvalidStream
	}

	// Return header type from end of buffer
	return o.buffer[len(o.buffer)-1], nil
}

// oggIDHeader represents the information contained in an Ogg Vorbis identification header
type oggIDHeader struct {
	VorbisVersion uint32
	ChannelCount  uint8
	SampleRate    uint32
	MaxBitrate    uint32
	NomBitrate    uint32
	MinBitrate    uint32
	Blocksize0    uint8
	Blocksize1    uint8
	Framing       bool
}

// parseOGGIDHeader parses the required identification header for an Ogg Vorbis stream
func (o *oggParser) parseOGGIDHeader() error {
	// Read OGG page header, skipping the capture pattern because New() already verified
	// the magic number for us
	if _, err := o.parseOGGPageHeader(true); err != nil {
		return err
	}

	// Check for valid common header
	headerType, err := o.parseOGGCommonHeader()
	if err != nil {
		return err
	}

	// Ensure header type 1: identification header
	if headerType != byte(1) {
		return ErrInvalidStream
	}

	// Read fields found in identification header
	header := new(oggIDHeader)

	// Vorbis version
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return err
	}
	header.VorbisVersion = o.ui32

	// Ensure Vorbis version is 0, per specification
	if header.VorbisVersion != 0 {
		return ErrInvalidStream
	}

	// Channel count
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return err
	}
	header.ChannelCount = o.ui8

	// uint32 x 4: sample rate, maximum bitrate, nominal bitrate, minimum bitrate
	uint32Slice := make([]uint32, 4)
	for i := 0; i < 4; i++ {
		// Read in one uint32
		if err := binary.Read(o.reader, binary.LittleEndian, &uint32Slice[i]); err != nil {
			return err
		}
	}

	// Copy out slice values
	header.SampleRate = uint32Slice[0]
	header.MaxBitrate = uint32Slice[1]
	header.NomBitrate = uint32Slice[2]
	header.MinBitrate = uint32Slice[3]

	// Create and use a bit reader to parse the following fields
	//    4 - Blocksize 0
	//    4 - Blocksize 1
	//    1 - Framing flag
	fields, err := bit.NewReader(o.reader).ReadFields(4, 4, 1)
	if err != nil {
		return err
	}

	header.Blocksize0 = uint8(fields[0])
	header.Blocksize1 = uint8(fields[1])
	header.Framing = fields[2] == 1

	// Store ID header
	o.idHeader = header
	return nil
}

// parseOGGCommentHeader parses the Vorbis Comment tags in an Ogg Vorbis file
func (o *oggParser) parseOGGCommentHeader() error {
	// Read OGG page header, specifying false to check the capture pattern
	if _, err := o.parseOGGPageHeader(false); err != nil {
		return err
	}

	// Parse common header
	headerType, err := o.parseOGGCommonHeader()
	if err != nil {
		return err
	}

	// Verify header type (3: Vorbis Comment)
	if headerType != byte(3) {
		return ErrInvalidStream
	}

	// Read vendor length
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return err
	}

	// Read vendor string, store as encoder
	if _, err := o.reader.Read(o.buffer[:o.ui32]); err != nil {
		return err
	}
	o.encoder = string(o.buffer[:o.ui32])

	// Read comment length (new allocation for use with loop counter)
	var commentLength uint32
	if err := binary.Read(o.reader, binary.LittleEndian, &commentLength); err != nil {
		return err
	}

	// Begin iterating tags, and building tag map
	tagMap := map[string]string{}
	for i := 0; i < int(commentLength); i++ {
		// Read tag string length
		if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
			return err
		}

		// Read tag string
		n, err := o.reader.Read(o.buffer[:o.ui32])
		if err != nil {
			return err
		}

		// Split tag name and data, store in map
		pair := strings.Split(string(o.buffer[:n]), "=")
		tagMap[strings.ToUpper(pair[0])] = pair[1]
	}

	// Seek one byte forward to prepare for the setup header
	if _, err := o.reader.Seek(1, 1); err != nil {
		return err
	}

	// Store tags
	o.tags = tagMap
	return nil
}

// parseOGGDuration reads out the rest of the file to find the last OGG page header, which
// contains information needed to parse the file duration
func (o *oggParser) parseOGGDuration() error {
	// Seek as far forward as sanely possible so we don't need to read tons of excess data
	// For now, a value of 4096 bytes before the end appears to work, and should give a bit
	// of wiggle-room without causing us to read the entire file
	if _, err := o.reader.Seek(-4096, 2); err != nil {
		return err
	}

	// Read the rest of the file to find the last page header
	vorbisFile, err := ioutil.ReadAll(o.reader)
	if err != nil {
		return err
	}

	// Find the index of the last OGG page header
	index := bytes.LastIndex(vorbisFile, oggMagicNumber)
	if index == -1 {
		return ErrInvalidStream
	}

	// Read using the in-memory bytes to grab the last page header information
	o.reader = bytes.NewReader(vorbisFile[index:])
	pageHeader, err := o.parseOGGPageHeader(false)
	if err != nil {
		return nil
	}

	// Calculate duration using last granule position divided by sample rate
	o.duration = time.Duration(pageHeader.GranulePosition/uint64(o.idHeader.SampleRate)) * time.Second
	return nil
}
<commit_msg>Remove unused oggParser.durChan field<commit_after>package taggolib

import (
	"bytes"
	"encoding/binary"
	"io"
	"io/ioutil"
	"strconv"
	"strings"
	"time"

	"github.com/eaburns/bit"
)

var (
	// oggMagicNumber is the magic number used to identify a OGG audio stream
	oggMagicNumber = []byte("OggS")
	// oggVorbisWord is used to denote the beginning of a Vorbis information block
	oggVorbisWord = []byte("vorbis")
)

// oggParser represents a OGG audio metadata tag parser
type oggParser struct {
	duration time.Duration
	encoder  string
	idHeader *oggIDHeader
	reader   io.ReadSeeker
	tags     map[string]string

	// Shared buffer and unsigned integers stored as fields to prevent unneeded allocations
	buffer []byte
	ui8    uint8
	ui32   uint32
	ui64   uint64
}

// Album returns the Album tag for this stream
func (o oggParser) Album() string {
	return o.tags[tagAlbum]
}

// AlbumArtist returns the AlbumArtist tag for this stream
func (o oggParser) AlbumArtist() string {
	return o.tags[tagAlbumArtist]
}

// Artist returns the Artist tag for this stream
func (o oggParser) Artist() string {
	return o.tags[tagArtist]
}

// BitDepth returns the bits-per-sample of this stream
func (o oggParser) BitDepth() int {
	// Ogg Vorbis should always provide 16 bit depth
	return 16
}

// Bitrate calculates the audio bitrate for this stream
func (o oggParser) Bitrate() int {
	// TODO: see how max/min bitrate play into calculations
	return int(o.idHeader.NomBitrate) / 1000
}

// Channels returns the number of channels for this stream
func (o oggParser) Channels() int {
	return int(o.idHeader.ChannelCount)
}

// Comment returns the Comment tag for this stream
func (o oggParser) Comment() string {
	return o.tags[tagComment]
}

// Date returns the Date tag for this stream
func (o oggParser) Date() string {
	return o.tags[tagDate]
}

// DiscNumber returns the DiscNumber tag for this stream
func (o oggParser) DiscNumber() int {
	disc, err := strconv.Atoi(o.tags[tagDiscNumber])
	if err != nil {
		return 0
	}

	return disc
}

// Duration returns the time duration for this stream
func (o oggParser) Duration() time.Duration {
	return o.duration
}

// Encoder returns the encoder for this stream
func (o oggParser) Encoder() string {
	return o.encoder
}

// Format returns the name of the OGG format
func (o oggParser) Format() string {
	return "OGG"
}

// Genre returns the Genre tag for this stream
func (o oggParser) Genre() string {
	return o.tags[tagGenre]
}

// SampleRate returns the sample rate in Hertz for this stream
func (o oggParser) SampleRate() int {
	return int(o.idHeader.SampleRate)
}

// Tag attempts to return the raw, unprocessed tag with the specified name for this stream
func (o oggParser) Tag(name string) string {
	return o.tags[name]
}

// Title returns the Title tag for this stream
func (o oggParser) Title() string {
	return o.tags[tagTitle]
}

// TrackNumber returns the TrackNumber tag for this stream
func (o oggParser) TrackNumber() int {
	// Check for a /, such as 2/8
	track, err := strconv.Atoi(strings.Split(o.tags[tagTrackNumber], "/")[0])
	if err != nil {
		return 0
	}

	return track
}

// newOGGParser creates a parser for OGG audio streams
func newOGGParser(reader io.ReadSeeker) (*oggParser, error) {
	// Create OGG parser
	parser := &oggParser{
		buffer: make([]byte, 128),
		reader: reader,
	}

	// Parse the required ID header
	if err := parser.parseOGGIDHeader(); err != nil {
		return nil, err
	}

	// Parse the required comment header
	if err := parser.parseOGGCommentHeader(); err != nil {
		return nil, err
	}

	// Parse the file's duration
	if err := parser.parseOGGDuration(); err != nil {
		return nil, err
	}

	// Return parser
	return parser, nil
}

// oggPageHeader represents the information contained in an Ogg Page header
type oggPageHeader struct {
	CapturePattern  []byte
	Version         uint8
	HeaderType      uint8
	GranulePosition uint64
	BitstreamSerial uint32
	PageSequence    uint32
	Checksum        []byte
	PageSegments    uint8
}

// parseOGGPageHeader parses an Ogg page header
func (o *oggParser) parseOGGPageHeader(skipMagicNumber bool) (*oggPageHeader, error) {
	// Create page header
	pageHeader := new(oggPageHeader)

	// Unless skip is specified, check for capture pattern
	if !skipMagicNumber {
		if _, err := o.reader.Read(o.buffer[:4]); err != nil {
			return nil, err
		}
		pageHeader.CapturePattern = o.buffer[:4]

		// Verify proper capture pattern
		if !bytes.Equal(pageHeader.CapturePattern, oggMagicNumber) {
			return nil, ErrInvalidStream
		}
	} else {
		// If skipped, assume capture pattern is correct magic number
		pageHeader.CapturePattern = oggMagicNumber
	}

	// Version (must always be 0)
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.Version = o.ui8

	// Verify mandated version 0
	if pageHeader.Version != 0 {
		return nil, ErrInvalidStream
	}

	// Header type
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.HeaderType = o.ui8

	// Granule position
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui64); err != nil {
		return nil, err
	}
	pageHeader.GranulePosition = o.ui64

	// Bitstream serial number
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return nil, err
	}
	pageHeader.BitstreamSerial = o.ui32

	// Page sequence number
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return nil, err
	}
	pageHeader.PageSequence = o.ui32

	// Checksum
	if _, err := o.reader.Read(o.buffer[:4]); err != nil {
		return nil, err
	}
	pageHeader.Checksum = o.buffer[:4]

	// Page segments
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return nil, err
	}
	pageHeader.PageSegments = o.ui8

	// Segment table is next, but we won't need it for tag parsing, so seek ahead
	// size of uint8 (1 byte) multiplied by number of page segments
	if _, err := o.reader.Seek(int64(pageHeader.PageSegments), 1); err != nil {
		return nil, err

	}
	return pageHeader, nil
}

// parseOGGCommonHeader parses information common to all Ogg Vorbis headers
func (o *oggParser) parseOGGCommonHeader() (byte, error) {
	// Read the first byte to get header type
	if _, err := o.reader.Read(o.buffer[:1]); err != nil {
		return 0, err
	}

	// Store first byte at end of buffer so we can return it later without more allocations
	o.buffer[len(o.buffer)-1] = o.buffer[0]

	// Read for 'vorbis' identification word
	if _, err := o.reader.Read(o.buffer[:len(oggVorbisWord)]); err != nil {
		return 0, err
	}

	// Ensure 'vorbis' identification word is present
	if !bytes.Equal(o.buffer[:len(oggVorbisWord)], oggVorbisWord) {
		return 0, ErrInvalidStream
	}

	// Return header type from end of buffer
	return o.buffer[len(o.buffer)-1], nil
}

// oggIDHeader represents the information contained in an Ogg Vorbis identification header
type oggIDHeader struct {
	VorbisVersion uint32
	ChannelCount  uint8
	SampleRate    uint32
	MaxBitrate    uint32
	NomBitrate    uint32
	MinBitrate    uint32
	Blocksize0    uint8
	Blocksize1    uint8
	Framing       bool
}

// parseOGGIDHeader parses the required identification header for an Ogg Vorbis stream
func (o *oggParser) parseOGGIDHeader() error {
	// Read OGG page header, skipping the capture pattern because New() already verified
	// the magic number for us
	if _, err := o.parseOGGPageHeader(true); err != nil {
		return err
	}

	// Check for valid common header
	headerType, err := o.parseOGGCommonHeader()
	if err != nil {
		return err
	}

	// Ensure header type 1: identification header
	if headerType != byte(1) {
		return ErrInvalidStream
	}

	// Read fields found in identification header
	header := new(oggIDHeader)

	// Vorbis version
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return err
	}
	header.VorbisVersion = o.ui32

	// Ensure Vorbis version is 0, per specification
	if header.VorbisVersion != 0 {
		return ErrInvalidStream
	}

	// Channel count
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui8); err != nil {
		return err
	}
	header.ChannelCount = o.ui8

	// uint32 x 4: sample rate, maximum bitrate, nominal bitrate, minimum bitrate
	uint32Slice := make([]uint32, 4)
	for i := 0; i < 4; i++ {
		// Read in one uint32
		if err := binary.Read(o.reader, binary.LittleEndian, &uint32Slice[i]); err != nil {
			return err
		}
	}

	// Copy out slice values
	header.SampleRate = uint32Slice[0]
	header.MaxBitrate = uint32Slice[1]
	header.NomBitrate = uint32Slice[2]
	header.MinBitrate = uint32Slice[3]

	// Create and use a bit reader to parse the following fields
	//    4 - Blocksize 0
	//    4 - Blocksize 1
	//    1 - Framing flag
	fields, err := bit.NewReader(o.reader).ReadFields(4, 4, 1)
	if err != nil {
		return err
	}

	header.Blocksize0 = uint8(fields[0])
	header.Blocksize1 = uint8(fields[1])
	header.Framing = fields[2] == 1

	// Store ID header
	o.idHeader = header
	return nil
}

// parseOGGCommentHeader parses the Vorbis Comment tags in an Ogg Vorbis file
func (o *oggParser) parseOGGCommentHeader() error {
	// Read OGG page header, specifying false to check the capture pattern
	if _, err := o.parseOGGPageHeader(false); err != nil {
		return err
	}

	// Parse common header
	headerType, err := o.parseOGGCommonHeader()
	if err != nil {
		return err
	}

	// Verify header type (3: Vorbis Comment)
	if headerType != byte(3) {
		return ErrInvalidStream
	}

	// Read vendor length
	if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
		return err
	}

	// Read vendor string, store as encoder
	if _, err := o.reader.Read(o.buffer[:o.ui32]); err != nil {
		return err
	}
	o.encoder = string(o.buffer[:o.ui32])

	// Read comment length (new allocation for use with loop counter)
	var commentLength uint32
	if err := binary.Read(o.reader, binary.LittleEndian, &commentLength); err != nil {
		return err
	}

	// Begin iterating tags, and building tag map
	tagMap := map[string]string{}
	for i := 0; i < int(commentLength); i++ {
		// Read tag string length
		if err := binary.Read(o.reader, binary.LittleEndian, &o.ui32); err != nil {
			return err
		}

		// Read tag string
		n, err := o.reader.Read(o.buffer[:o.ui32])
		if err != nil {
			return err
		}

		// Split tag name and data, store in map
		pair := strings.Split(string(o.buffer[:n]), "=")
		tagMap[strings.ToUpper(pair[0])] = pair[1]
	}

	// Seek one byte forward to prepare for the setup header
	if _, err := o.reader.Seek(1, 1); err != nil {
		return err
	}

	// Store tags
	o.tags = tagMap
	return nil
}

// parseOGGDuration reads out the rest of the file to find the last OGG page header, which
// contains information needed to parse the file duration
func (o *oggParser) parseOGGDuration() error {
	// Seek as far forward as sanely possible so we don't need to read tons of excess data
	// For now, a value of 4096 bytes before the end appears to work, and should give a bit
	// of wiggle-room without causing us to read the entire file
	if _, err := o.reader.Seek(-4096, 2); err != nil {
		return err
	}

	// Read the rest of the file to find the last page header
	vorbisFile, err := ioutil.ReadAll(o.reader)
	if err != nil {
		return err
	}

	// Find the index of the last OGG page header
	index := bytes.LastIndex(vorbisFile, oggMagicNumber)
	if index == -1 {
		return ErrInvalidStream
	}

	// Read using the in-memory bytes to grab the last page header information
	o.reader = bytes.NewReader(vorbisFile[index:])
	pageHeader, err := o.parseOGGPageHeader(false)
	if err != nil {
		return nil
	}

	// Calculate duration using last granule position divided by sample rate
	o.duration = time.Duration(pageHeader.GranulePosition/uint64(o.idHeader.SampleRate)) * time.Second
	return nil
}
<|endoftext|>


<commit_before>package generator

import (
	"fmt"
	"go/build"
	"go/types"
	"log"
	"path/filepath"
	"reflect"
	"strings"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/imports"
)

func (f *Fake) loadPackages(c Cacher, workingDir string) error {
	log.Println("loading packages...")
	p, ok := c.Load(f.TargetPackage)
	if ok {
		f.Packages = p
		log.Printf("loaded %v packages from cache\n", len(f.Packages))
		return nil
	}
	importPath := f.TargetPackage
	if !filepath.IsAbs(importPath) {
		ctx := getBuildContext(workingDir)
		bp, err := ctx.Import(f.TargetPackage, workingDir, build.FindOnly)
		if err != nil {
			return err
		}
		importPath = bp.ImportPath
	}
	p, err := packages.Load(&packages.Config{
		Mode:  packages.NeedName | packages.NeedFiles | packages.NeedImports | packages.NeedDeps | packages.NeedTypes,
		Dir:   workingDir,
		Tests: true,
	}, importPath)
	if err != nil {
		return err
	}
	for i := range p {
		if len(p[i].Errors) > 0 {
			if i == 0 {
				err = p[i].Errors[0]
			}
			for j := range p[i].Errors {
				log.Printf("error loading packages: %v", strings.TrimPrefix(fmt.Sprintf("%v", p[i].Errors[j]), "-: "))
			}
		}
	}
	if err != nil {
		return err
	}
	f.Packages = p
	c.Store(f.TargetPackage, p)
	log.Printf("loaded %v packages\n", len(f.Packages))
	return nil
}

func (f *Fake) findPackage() error {
	var target *types.TypeName
	var pkg *packages.Package
	for i := range f.Packages {
		if f.Packages[i].Types == nil || f.Packages[i].Types.Scope() == nil {
			continue
		}
		pkg = f.Packages[i]
		if f.Mode == Package {
			break
		}

		raw := pkg.Types.Scope().Lookup(f.TargetName)
		if raw != nil {
			if typeName, ok := raw.(*types.TypeName); ok {
				target = typeName
				break
			}
		}
		pkg = nil
	}
	if pkg == nil {
		switch f.Mode {
		case Package:
			return fmt.Errorf("cannot find package with name: %s", f.TargetPackage)
		case InterfaceOrFunction:
			return fmt.Errorf("cannot find package with target: %s", f.TargetName)
		}
	}
	f.Target = target
	f.Package = pkg
	f.TargetPackage = imports.VendorlessPath(pkg.PkgPath)
	t := f.Imports.Add(pkg.Name, f.TargetPackage)
	f.TargetAlias = t.Alias
	if f.Mode != Package {
		f.TargetName = target.Name()
	}

	if f.Mode == InterfaceOrFunction {
		if !f.IsInterface() && !f.IsFunction() {
			return fmt.Errorf("cannot generate an fake for %s because it is not an interface or function", f.TargetName)
		}
	}

	if f.IsInterface() {
		log.Printf("Found interface with name: [%s]\n", f.TargetName)
	}
	if f.IsFunction() {
		log.Printf("Found function with name: [%s]\n", f.TargetName)
	}
	if f.Mode == Package {
		log.Printf("Found package with name: [%s]\n", f.TargetPackage)
	}
	return nil
}

// addImportsFor inspects the given type and adds imports to the fake if importable
// types are found.
func (f *Fake) addImportsFor(typ types.Type) {
	if typ == nil {
		return
	}

	switch t := typ.(type) {
	case *types.Basic:
		return
	case *types.Pointer:
		f.addImportsFor(t.Elem())
	case *types.Map:
		f.addImportsFor(t.Key())
		f.addImportsFor(t.Elem())
	case *types.Chan:
		f.addImportsFor(t.Elem())
	case *types.Named:
		if t.Obj() != nil && t.Obj().Pkg() != nil {
			f.Imports.Add(t.Obj().Pkg().Name(), t.Obj().Pkg().Path())
		}
	case *types.Slice:
		f.addImportsFor(t.Elem())
	case *types.Array:
		f.addImportsFor(t.Elem())
	case *types.Interface:
		return
	case *types.Signature:
		f.addTypesForMethod(t)
	case *types.Struct:
		for i := 0; i < t.NumFields(); i++ {
			f.addImportsFor(t.Field(i).Type())
		}
	default:
		log.Printf("!!! WARNING: Missing case for type %s\n", reflect.TypeOf(typ).String())
	}
}
<commit_msg>load types info<commit_after>package generator

import (
	"fmt"
	"go/build"
	"go/types"
	"log"
	"path/filepath"
	"reflect"
	"strings"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/imports"
)

func (f *Fake) loadPackages(c Cacher, workingDir string) error {
	log.Println("loading packages...")
	p, ok := c.Load(f.TargetPackage)
	if ok {
		f.Packages = p
		log.Printf("loaded %v packages from cache\n", len(f.Packages))
		return nil
	}
	importPath := f.TargetPackage
	if !filepath.IsAbs(importPath) {
		ctx := getBuildContext(workingDir)
		bp, err := ctx.Import(f.TargetPackage, workingDir, build.FindOnly)
		if err != nil {
			return err
		}
		importPath = bp.ImportPath
	}
	p, err := packages.Load(&packages.Config{
		Mode:  packages.NeedName | packages.NeedFiles | packages.NeedImports | packages.NeedDeps | packages.NeedTypes | packages.NeedTypesInfo,
		Dir:   workingDir,
		Tests: true,
	}, importPath)
	if err != nil {
		return err
	}
	for i := range p {
		if len(p[i].Errors) > 0 {
			if i == 0 {
				err = p[i].Errors[0]
			}
			for j := range p[i].Errors {
				log.Printf("error loading packages: %v", strings.TrimPrefix(fmt.Sprintf("%v", p[i].Errors[j]), "-: "))
			}
		}
	}
	if err != nil {
		return err
	}
	f.Packages = p
	c.Store(f.TargetPackage, p)
	log.Printf("loaded %v packages\n", len(f.Packages))
	return nil
}

func (f *Fake) findPackage() error {
	var target *types.TypeName
	var pkg *packages.Package
	for i := range f.Packages {
		if f.Packages[i].Types == nil || f.Packages[i].Types.Scope() == nil {
			continue
		}
		pkg = f.Packages[i]
		if f.Mode == Package {
			break
		}

		raw := pkg.Types.Scope().Lookup(f.TargetName)
		if raw != nil {
			if typeName, ok := raw.(*types.TypeName); ok {
				target = typeName
				break
			}
		}
		pkg = nil
	}
	if pkg == nil {
		switch f.Mode {
		case Package:
			return fmt.Errorf("cannot find package with name: %s", f.TargetPackage)
		case InterfaceOrFunction:
			return fmt.Errorf("cannot find package with target: %s", f.TargetName)
		}
	}
	f.Target = target
	f.Package = pkg
	f.TargetPackage = imports.VendorlessPath(pkg.PkgPath)
	t := f.Imports.Add(pkg.Name, f.TargetPackage)
	f.TargetAlias = t.Alias
	if f.Mode != Package {
		f.TargetName = target.Name()
	}

	if f.Mode == InterfaceOrFunction {
		if !f.IsInterface() && !f.IsFunction() {
			return fmt.Errorf("cannot generate an fake for %s because it is not an interface or function", f.TargetName)
		}
	}

	if f.IsInterface() {
		log.Printf("Found interface with name: [%s]\n", f.TargetName)
	}
	if f.IsFunction() {
		log.Printf("Found function with name: [%s]\n", f.TargetName)
	}
	if f.Mode == Package {
		log.Printf("Found package with name: [%s]\n", f.TargetPackage)
	}
	return nil
}

// addImportsFor inspects the given type and adds imports to the fake if importable
// types are found.
func (f *Fake) addImportsFor(typ types.Type) {
	if typ == nil {
		return
	}

	switch t := typ.(type) {
	case *types.Basic:
		return
	case *types.Pointer:
		f.addImportsFor(t.Elem())
	case *types.Map:
		f.addImportsFor(t.Key())
		f.addImportsFor(t.Elem())
	case *types.Chan:
		f.addImportsFor(t.Elem())
	case *types.Named:
		if t.Obj() != nil && t.Obj().Pkg() != nil {
			f.Imports.Add(t.Obj().Pkg().Name(), t.Obj().Pkg().Path())
		}
	case *types.Slice:
		f.addImportsFor(t.Elem())
	case *types.Array:
		f.addImportsFor(t.Elem())
	case *types.Interface:
		return
	case *types.Signature:
		f.addTypesForMethod(t)
	case *types.Struct:
		for i := 0; i < t.NumFields(); i++ {
			f.addImportsFor(t.Field(i).Type())
		}
	default:
		log.Printf("!!! WARNING: Missing case for type %s\n", reflect.TypeOf(typ).String())
	}
}
<|endoftext|>


<commit_before>package parser

import (
  "bytes"
  "rubex"
  "strconv"
  "fmt"
)

// Type tags so we know what kind of token we have
type Lexeme int
const (
  LPAREN = iota
  RPAREN
  LBRACE
  RBRACE
  COMMA
  DOT
  EQUAL
  STRING
  REGEXP
  POS
  GVAR
  LVAR
  KWD
  ID
  FUNC
  TYPE
  PATH
  IMPORT
  READ
  EOF
  ERROR
)

var lexemeName [21]string
var matcher [21]*rubex.Regexp
var symbolLexeme map[string]Lexeme
var symbolPattern *rubex.Regexp
var numberPattern *rubex.Regexp

func init() {
  // Is there a more elegant way to do this?
  lexemeName[LPAREN] = "LPAREN"
  lexemeName[RPAREN] = "RPAREN"
  lexemeName[LBRACE] = "LBRACE"
  lexemeName[RBRACE] = "RBRACE"
  lexemeName[COMMA]  = "COMMA"
  lexemeName[DOT]    = "DOT"
  lexemeName[EQUAL]  = "EQUAL"
  lexemeName[STRING] = "STRING"
  lexemeName[REGEXP] = "REGEXP"
  lexemeName[POS]    = "POS"
  lexemeName[GVAR]   = "GVAR"
  lexemeName[LVAR]   = "LVAR"
  lexemeName[KWD]    = "KWD"
  lexemeName[ID]     = "ID"
  lexemeName[FUNC]   = "FUNC"
  lexemeName[TYPE]   = "TYPE"
  lexemeName[PATH]   = "PATH"
  lexemeName[IMPORT] = "IMPORT"
  lexemeName[READ]   = "READ"
  lexemeName[EOF]    = "EOF"
  lexemeName[ERROR]  = "ERROR"

  matcher[STRING], _ = rubex.Compile(`\A"(\\.|[^"\\])*"|\A'(\\.|[^'\\])*'`)
  // the pattern and options of the regexp are in captures 1 and 3
  matcher[REGEXP], _ = rubex.Compile(`\A\/((\\.|[^\/\\])*)\/([imxouesn]*)`)
  matcher[POS], _    = rubex.Compile(`\A(top|bottom|before|after)`)
  matcher[GVAR], _   = rubex.Compile(`\A\$\w+`)
  matcher[LVAR], _   = rubex.Compile(`\A%\w+`)
  matcher[KWD], _    = rubex.Compile(`\A[a-zA-Z_:][-\w:.]*:`)
  matcher[ID], _     = rubex.Compile(`\A\$|^[_a-z][\w\$]*`)
  matcher[TYPE], _   = rubex.Compile(`\A[A-Z]\w*`)
  matcher[PATH], _   = rubex.Compile(`\A[-+.*?:\/\w]+`)

  // Map parens, braces, etc to their lexemes
  symbolLexeme = make(map[string]Lexeme, 7)
  symbolLexeme["("] = LPAREN
  symbolLexeme[")"] = RPAREN
  symbolLexeme["{"] = LBRACE
  symbolLexeme["}"] = RBRACE
  symbolLexeme[","] = COMMA
  symbolLexeme["."] = DOT
  symbolLexeme["="] = EQUAL
  symbolPattern, _ = rubex.Compile(`\A[\(\)\{\}\,\.=]`)

  numberPattern, _ = rubex.Compile(`\A\d+`)
}

// A token has a type (aka lexeme), a value, and a line number
type Token struct {
  Lexeme
  Value string
  ExtraValue string
  LineNum int
}

func (t *Token) Inspect() string {
  return fmt.Sprintf("[%s: %s, %s, %d]", lexemeName[t.Lexeme], t.Value, t.ExtraValue, t.LineNum)
}

/*
  Represent a tokenizer with a struct containing the remaining source text and
  the line number. Easier than using a stateless tokenizing function that
  returns them as extra values and requires the parser to keep track of them.
*/
type Tokenizer struct {
  Source []byte
  LineNum int
  Lookahead *Token
  unterminatedComment bool
}

func (t *Tokenizer) hasPrefix(s string) bool {
  return bytes.HasPrefix(t.Source, []byte(s))
}

// Discard leading spaces (excluding newlines) in the source text.
func (t *Tokenizer) discardSpaces() {
  t.Source = bytes.TrimLeft(t.Source, " \t")
}

// Discard leading text until a newline (or EOF) is found.
func(t *Tokenizer) discardLine() {
  if i := bytes.IndexByte(t.Source, '\n'); i >= 0 {
    t.Source = t.Source[i:]
  } else {
    t.Source = t.Source[len(t.Source):]
  }
}

// Discard the leading comment in the source text.
func (t *Tokenizer) discardComment() {
  if t.hasPrefix("#") || t.hasPrefix("//") {
    t.discardLine()
  } else if t.hasPrefix("/*") {
    t.discardBlockComment()
  }
}

// Helper for discarding block comments.
// TO DO: ERROR HANDLING FOR UNTERMINATED COMMENTS
func (t *Tokenizer) discardBlockComment() {
  depth, i, length := 1, 2, len(t.Source)
  error := false
  for depth > 0 {
    if i >= length {
      error = true
      break
    }
    switch t.Source[i] {
    case '\n':
      t.LineNum++
    case '/':
      i++
      if i >= length {
        error = true
        break
      }
      if t.Source[i] == '*' {
        depth++
      }
    case '*':
      i++
      if i >= length {
        error = true
        break
      }
      if t.Source[i] == '/' {
        depth--
      }
    }
    i++
  }
  t.Source = t.Source[i:]
  if error {
    t.Lookahead = &Token{ Lexeme: ERROR, Value: "unterminated comment", ExtraValue: "", LineNum: t.LineNum }
    t.unterminatedComment = true
  }
}

// Discard all leading whitespace and comments from the source text. Need to
// tally up the newlines to keep LineNum up to date.
func (t *Tokenizer) discardWhitespaceAndComments() {
  for len(t.Source) > 0 {
    switch {
    case t.hasPrefix("\n"):
      t.LineNum++
      t.Source = t.Source[1:]
    case t.hasPrefix(" ") || t.hasPrefix("\t"):
      t.discardSpaces()
    case t.hasPrefix("#") || t.hasPrefix("//") || t.hasPrefix("/*"):
      t.discardComment()
    default:
      return
    }
  }
}

func (t *Tokenizer) popToken(lexeme Lexeme, value string, length int) *Token {
  val := &Token { Lexeme: lexeme, Value: value, ExtraValue: "", LineNum: t.LineNum }
  t.Source = t.Source[length:]
  return val
}

func (t *Tokenizer) popError(message string) *Token {
  val := &Token { Lexeme: ERROR, Value: message, ExtraValue: "", LineNum: t.LineNum }
  t.discardLine()
  return val
}

// Helper for unquoting strings. The main difficulty is that Go strings are
// exclusively double-quoted, so single-quoted strings need to be converted
// before being passed to strconv.Unquote(...).
func unquote(chars []byte) string {
  var converted []byte
  if chars[0] == '\'' {
    converted = bytes.Replace(chars, []byte(`\'`), []byte(`'`), -1)
    converted = bytes.Replace(chars, []byte(`"`), []byte(`\"`), -1)
    converted[0] = '"'
    converted[len(converted)-1] = '"'
  } else {
    converted = chars
  }
  val, _ := strconv.Unquote(string(converted))
  return val
}

// The heart of the tokenizer. This function tries to munch off a token from
// the head of the source text.
func (t *Tokenizer) munch() *Token {
  src := t.Source
  if len(src) == 0 {
    return t.popToken(EOF, "", 0)
  } else if t.hasPrefix("*/") {
    return t.popError("unmatched comment terminator")
  } else if c := string(symbolPattern.Find(src)); len(c) > 0 {
    return t.popToken(symbolLexeme[c], c, 1)
  } else if c := string(numberPattern.Find(src)); len(c) > 0 {
    return t.popToken(STRING, c, len(c))
  } else if t.hasPrefix("'") || t.hasPrefix("\"") {
    if c := matcher[STRING].Find(src); len(c) > 0 {
      unquoted := unquote(c)
      return t.popToken(STRING, unquoted, len(c))
    } else {
      return t.popError("unterminated string literal")
    }
  } else if t.hasPrefix("/") {
    if cs := matcher[REGEXP].FindSubmatch(src); len(cs) > 0 {
      pattern := cs[1]
      options := cs[3]
      val := t.popToken(REGEXP, string(pattern), len(cs[0]))
      val.ExtraValue = string(options)
      return val
    } else {
      return t.popError("unterminated regular expression literal")
    }
  } else if c := matcher[KWD].Find(src); len(c) > 0 {
    return t.popToken(KWD, string(c[:len(c)-1]), len(c))
  } else if c := matcher[GVAR].Find(src); len(c) > 0 {
    return t.popToken(GVAR, string(c[1:]), len(c))
  } else if c := matcher[LVAR].Find(src); len(c) > 0 {
    return t.popToken(LVAR, string(c[1:]), len(c))
  } else if c := string(matcher[ID].Find(src)); len(c) > 0 {
    if matcher[POS].MatchString(c) {
      return t.popToken(POS, c, len(c))
    } else if c == "read" {
      return t.popToken(READ, "", len(c))
    } else {
      return t.popToken(ID, c, len(c))
    }
  } else if c := string(matcher[TYPE].Find(src)); len(c) > 0 {
    return t.popToken(TYPE, c, len(c))
  } else if t.hasPrefix("@import") {
    tok := t.popToken(IMPORT, "", 7)
    t.discardWhitespaceAndComments()
    if c := string(matcher[PATH].Find(t.Source)); len(c) > 0 {
      tok.Value = c
      t.Source = t.Source[len(c):]
    } else if c := matcher[STRING].Find(t.Source); len(c) > 0 {
      tok.Value = unquote(c)
      t.Source = t.Source[len(c):]
    } else {
      tok = t.popError("malformed import")
    }
    return tok
  } else if t.hasPrefix("@func") {
    return t.popToken(FUNC, "", 5)
  } else {
    return t.popError("unrecognized token")
  }
  return t.popError("unrecognized token")
}

func (t *Tokenizer) Peek() *Token {
  return t.Lookahead
}

func (t *Tokenizer) Pop() *Token {
  val := t.Lookahead
  t.discardWhitespaceAndComments()
  if !t.unterminatedComment {
    t.Lookahead = t.munch()
  } else {
    t.unterminatedComment = false
  }
  return val
}

func MakeTokenizer(src []byte) *Tokenizer {
  t := Tokenizer { Source: src, Lookahead: nil, LineNum: 1, unterminatedComment: false }
  t.Pop()
  return &t
}
<commit_msg>Updating some comments.<commit_after>package parser

import (
  "bytes"
  "rubex"
  "strconv"
  "fmt"
)

// Type tags so we know what kind of token we have
type Lexeme int
const (
  LPAREN = iota
  RPAREN
  LBRACE
  RBRACE
  COMMA
  DOT
  EQUAL
  STRING
  REGEXP
  POS
  GVAR
  LVAR
  KWD
  ID
  FUNC
  TYPE
  PATH
  IMPORT
  READ
  EOF
  ERROR
)

var lexemeName [21]string
var matcher [21]*rubex.Regexp
var symbolLexeme map[string]Lexeme
var symbolPattern *rubex.Regexp
var numberPattern *rubex.Regexp

func init() {
  // Is there a more elegant way to do this?
  lexemeName[LPAREN] = "LPAREN"
  lexemeName[RPAREN] = "RPAREN"
  lexemeName[LBRACE] = "LBRACE"
  lexemeName[RBRACE] = "RBRACE"
  lexemeName[COMMA]  = "COMMA"
  lexemeName[DOT]    = "DOT"
  lexemeName[EQUAL]  = "EQUAL"
  lexemeName[STRING] = "STRING"
  lexemeName[REGEXP] = "REGEXP"
  lexemeName[POS]    = "POS"
  lexemeName[GVAR]   = "GVAR"
  lexemeName[LVAR]   = "LVAR"
  lexemeName[KWD]    = "KWD"
  lexemeName[ID]     = "ID"
  lexemeName[FUNC]   = "FUNC"
  lexemeName[TYPE]   = "TYPE"
  lexemeName[PATH]   = "PATH"
  lexemeName[IMPORT] = "IMPORT"
  lexemeName[READ]   = "READ"
  lexemeName[EOF]    = "EOF"
  lexemeName[ERROR]  = "ERROR"

  matcher[STRING], _ = rubex.Compile(`\A"(\\.|[^"\\])*"|\A'(\\.|[^'\\])*'`)
  // the pattern and options of the regexp are in captures 1 and 3
  matcher[REGEXP], _ = rubex.Compile(`\A\/((\\.|[^\/\\])*)\/([imxouesn]*)`)
  matcher[POS], _    = rubex.Compile(`\A(top|bottom|before|after)`)
  matcher[GVAR], _   = rubex.Compile(`\A\$\w+`)
  matcher[LVAR], _   = rubex.Compile(`\A%\w+`)
  matcher[KWD], _    = rubex.Compile(`\A[a-zA-Z_:][-\w:.]*:`)
  matcher[ID], _     = rubex.Compile(`\A\$|^[_a-z][\w\$]*`)
  matcher[TYPE], _   = rubex.Compile(`\A[A-Z]\w*`)
  matcher[PATH], _   = rubex.Compile(`\A[-+.*?:\/\w]+`)

  // Map parens, braces, etc to their lexemes
  symbolLexeme = make(map[string]Lexeme, 7)
  symbolLexeme["("] = LPAREN
  symbolLexeme[")"] = RPAREN
  symbolLexeme["{"] = LBRACE
  symbolLexeme["}"] = RBRACE
  symbolLexeme[","] = COMMA
  symbolLexeme["."] = DOT
  symbolLexeme["="] = EQUAL
  symbolPattern, _ = rubex.Compile(`\A[\(\)\{\}\,\.=]`)

  numberPattern, _ = rubex.Compile(`\A\d+`)
}

// A token has a type (aka lexeme), a value, and a line number
type Token struct {
  Lexeme
  Value string
  ExtraValue string
  LineNum int
}

func (t *Token) Inspect() string {
  return fmt.Sprintf("[%s: %s, %s, %d]", lexemeName[t.Lexeme], t.Value, t.ExtraValue, t.LineNum)
}

/*
  Represent a tokenizer with a struct containing the remaining source text and
  the line number. Easier than using a stateless tokenizing function that
  returns them as extra values and requires the parser to keep track of them.
*/
type Tokenizer struct {
  Source []byte
  LineNum int
  Lookahead *Token
  unterminatedComment bool
}

func (t *Tokenizer) hasPrefix(s string) bool {
  return bytes.HasPrefix(t.Source, []byte(s))
}

// Discard leading spaces (excluding newlines) in the source text.
func (t *Tokenizer) discardSpaces() {
  t.Source = bytes.TrimLeft(t.Source, " \t")
}

// Discard leading text until a newline (or EOF) is found.
func(t *Tokenizer) discardLine() {
  if i := bytes.IndexByte(t.Source, '\n'); i >= 0 {
    t.Source = t.Source[i:]
  } else {
    t.Source = t.Source[len(t.Source):]
  }
}

// Discard the leading comment in the source text.
func (t *Tokenizer) discardComment() {
  if t.hasPrefix("#") || t.hasPrefix("//") {
    t.discardLine()
  } else if t.hasPrefix("/*") {
    t.discardBlockComment()
  }
}

// Helper for discarding block comments.
// TO DO: ERROR HANDLING FOR UNTERMINATED COMMENTS
func (t *Tokenizer) discardBlockComment() {
  depth, i, length := 1, 2, len(t.Source)
  error := false
  for depth > 0 {
    if i >= length {
      error = true
      break
    }
    switch t.Source[i] {
    case '\n':
      t.LineNum++
    case '/':
      i++
      if i >= length {
        error = true
        break
      }
      if t.Source[i] == '*' {
        depth++
      }
    case '*':
      i++
      if i >= length {
        error = true
        break
      }
      if t.Source[i] == '/' {
        depth--
      }
    }
    i++
  }
  t.Source = t.Source[i:]
  if error {
    t.Lookahead = &Token{ Lexeme: ERROR, Value: "unterminated comment", ExtraValue: "", LineNum: t.LineNum }
    t.unterminatedComment = true
  }
}

// Discard all leading whitespace and comments from the source text. Need to
// tally up the newlines to keep LineNum up to date.
func (t *Tokenizer) discardWhitespaceAndComments() {
  for len(t.Source) > 0 {
    switch {
    case t.hasPrefix("\n"):
      t.LineNum++
      t.Source = t.Source[1:]
    case t.hasPrefix(" ") || t.hasPrefix("\t"):
      t.discardSpaces()
    case t.hasPrefix("#") || t.hasPrefix("//") || t.hasPrefix("/*"):
      t.discardComment()
    default:
      return
    }
  }
}

// Returns the next token and simultaneously discards the specified number of
// characters from the source text.
func (t *Tokenizer) popToken(lexeme Lexeme, value string, length int) *Token {
  val := &Token { Lexeme: lexeme, Value: value, ExtraValue: "", LineNum: t.LineNum }
  t.Source = t.Source[length:]
  return val
}

// Returns an error token and discards the rest of the line.
func (t *Tokenizer) popError(message string) *Token {
  val := &Token { Lexeme: ERROR, Value: message, ExtraValue: "", LineNum: t.LineNum }
  t.discardLine()
  return val
}

// Helper for unquoting strings. The main difficulty is that Go strings are
// exclusively double-quoted, so single-quoted strings need to be converted
// before being passed to strconv.Unquote(...).
func unquote(chars []byte) string {
  var converted []byte
  if chars[0] == '\'' {
    converted = bytes.Replace(chars, []byte(`\'`), []byte(`'`), -1)
    converted = bytes.Replace(chars, []byte(`"`), []byte(`\"`), -1)
    converted[0] = '"'
    converted[len(converted)-1] = '"'
  } else {
    converted = chars
  }
  val, _ := strconv.Unquote(string(converted))
  return val
}

// The heart of the tokenizer. This function tries to munch off a token from
// the head of the source text.
func (t *Tokenizer) munch() *Token {
  src := t.Source
  if len(src) == 0 {
    return t.popToken(EOF, "", 0)
  } else if t.hasPrefix("*/") {
    return t.popError("unmatched comment terminator")
  } else if c := string(symbolPattern.Find(src)); len(c) > 0 {
    return t.popToken(symbolLexeme[c], c, 1)
  } else if c := string(numberPattern.Find(src)); len(c) > 0 {
    return t.popToken(STRING, c, len(c))
  } else if t.hasPrefix("'") || t.hasPrefix("\"") {
    if c := matcher[STRING].Find(src); len(c) > 0 {
      unquoted := unquote(c)
      return t.popToken(STRING, unquoted, len(c))
    } else {
      return t.popError("unterminated string literal")
    }
  } else if t.hasPrefix("/") {
    if cs := matcher[REGEXP].FindSubmatch(src); len(cs) > 0 {
      pattern := cs[1]
      options := cs[3]
      val := t.popToken(REGEXP, string(pattern), len(cs[0]))
      val.ExtraValue = string(options)
      return val
    } else {
      return t.popError("unterminated regular expression literal")
    }
  } else if c := matcher[KWD].Find(src); len(c) > 0 {
    return t.popToken(KWD, string(c[:len(c)-1]), len(c))
  } else if c := matcher[GVAR].Find(src); len(c) > 0 {
    return t.popToken(GVAR, string(c[1:]), len(c))
  } else if c := matcher[LVAR].Find(src); len(c) > 0 {
    return t.popToken(LVAR, string(c[1:]), len(c))
  } else if c := string(matcher[ID].Find(src)); len(c) > 0 {
    if matcher[POS].MatchString(c) {
      return t.popToken(POS, c, len(c))
    } else if c == "read" {
      return t.popToken(READ, "", len(c))
    } else {
      return t.popToken(ID, c, len(c))
    }
  } else if c := string(matcher[TYPE].Find(src)); len(c) > 0 {
    return t.popToken(TYPE, c, len(c))
  } else if t.hasPrefix("@import") {
    tok := t.popToken(IMPORT, "", 7)
    t.discardWhitespaceAndComments()
    if c := string(matcher[PATH].Find(t.Source)); len(c) > 0 {
      tok.Value = c
      t.Source = t.Source[len(c):]
    } else if c := matcher[STRING].Find(t.Source); len(c) > 0 {
      tok.Value = unquote(c)
      t.Source = t.Source[len(c):]
    } else {
      tok = t.popError("malformed import")
    }
    return tok
  } else if t.hasPrefix("@func") {
    return t.popToken(FUNC, "", 5)
  } else {
    return t.popError("unrecognized token")
  }
  return t.popError("unrecognized token")
}

/*
  The following three functions constitute the API for the tokenizer.
*/

func MakeTokenizer(src []byte) *Tokenizer {
  t := Tokenizer { Source: src, Lookahead: nil, LineNum: 1, unterminatedComment: false }
  t.Pop()
  return &t
}

func (t *Tokenizer) Peek() *Token {
  return t.Lookahead
}

func (t *Tokenizer) Pop() *Token {
  val := t.Lookahead
  t.discardWhitespaceAndComments()
  if !t.unterminatedComment {
    t.Lookahead = t.munch()
  } else {
    t.unterminatedComment = false
  }
  return val
}<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"text/template"

	"github.com/Sirupsen/logrus"
	"github.com/justone/pmb/api"
	"github.com/kardianos/osext"
)

func handleOSXCommand(bus *pmb.PMB, command string, subcommand string, arguments []string) error {

	var err error

	logrus.Debugf("Handling %s with args of %s", command, arguments)

	// launch agent name
	agentName := fmt.Sprintf("org.endot.pmb.%s", subcommand)
	logrus.Debugf("Name of launchagent: %s", agentName)

	// figure out launch agent config path
	launchAgentFile := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), agentName)
	logrus.Debugf("launchagent file: %s", launchAgentFile)

	var homeDir string
	if homeDir = os.Getenv("HOME"); len(homeDir) == 0 {
		homeDir = "/tmp"
	}

	// create launch data
	executable, err := osext.Executable()
	if err != nil {
		return err
	}

	launchData := struct {
		Name, Executable, Broker, HomeDir string
		Args                              []string
	}{
		agentName, executable, bus.BrokerURI(), homeDir, arguments,
	}

	switch command {
	case "list":
		fmt.Printf(`
Available commands for running '%s' as a background process (agent):

start - Starts agent via launchctl.
stop - Stops agent via launchctl.
restart - Restarts agent via launchctl.
configure - This will configure the agent, but not start it.
unconfigure - This will remove the agent configuration.

`, fmt.Sprintf("pmb %s", subcommand))

	case "restart":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = stop(launchAgentFile, agentName)
		if err != nil {
			return err
		}

		err = start(launchAgentFile, agentName)
		if err != nil {
			return err
		}
	case "stop":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = stop(launchAgentFile, agentName)
		if err != nil {
			return err
		}

		err = unconfigure(launchAgentFile)
		if err != nil {
			return err
		}
	case "start":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = start(launchAgentFile, agentName)
		if err != nil {
			return err
		}
	case "configure":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}
	case "unconfigure":
		err = unconfigure(launchAgentFile)
		if err != nil {
			return err
		}
	}

	return nil
}

func start(launchAgentFile string, agentName string) error {
	listCmd := exec.Command("/bin/launchctl", "list", agentName)
	err := listCmd.Run()

	if _, ok := err.(*exec.ExitError); ok {
		// launch agent wasn't loaded yet, so load to start
		startCmd := exec.Command("/bin/launchctl", "load", launchAgentFile)
		startErr := startCmd.Run()
		if startErr != nil {
			return startErr
		}
	} else if err != nil {
		// some error running the list command
		return err
	} else {
		// launch agent was already loaded
		logrus.Infof("Already running")
	}

	return nil
}

func stop(launchAgentFile string, agentName string) error {
	listCmd := exec.Command("/bin/launchctl", "list", agentName)
	err := listCmd.Run()

	if err == nil {
		// launch agent was loaded, so unload to stop
		stopCmd := exec.Command("/bin/launchctl", "unload", launchAgentFile)
		stopErr := stopCmd.Run()
		if stopErr != nil {
			return stopErr
		}
	} else if _, ok := err.(*exec.ExitError); ok {
		// launch agent wasn't already loaded
		logrus.Infof("Already stopped")
	} else {
		// some error running the list command
		return err
	}

	return nil
}

func configure(launchAgentFile string, config string) error {

	err := ioutil.WriteFile(launchAgentFile, []byte(config), 0644)
	if err != nil {
		return err
	}

	logrus.Debugf("Created %s: %s", launchAgentFile, config)

	return nil
}

func generateLaunchConfig(launchData interface{}) string {
	configureTemplate := `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple Computer//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
	<dict>
		<key>Label</key>
		<string>{{ .Name }}</string>
		<key>KeepAlive</key>
		<true/>
		<key>EnvironmentVariables</key>
		<dict>
			<key>PATH</key>
			<string>/usr/bin:/bin:/usr/sbin:/sbin:/usr/local/bin</string>
			<key>PMB_PRIMARY_URI</key>
			<string>{{ .Broker }}</string>
		</dict>
		<key>StandardOutPath</key>
		<string>{{ .HomeDir }}/.pmb.log</string>
		<key>StandardErrorPath</key>
		<string>{{ .HomeDir }}/.pmb.log</string>
		<key>ProgramArguments</key>
		<array>
			<string>{{ .Executable }}</string>
{{ range .Args }}
			<string>{{ . }}</string>
{{ end }}
		</array>
	</dict>
</plist>`

	tmpl := template.Must(template.New("configure").Parse(configureTemplate))
	var output bytes.Buffer

	err := tmpl.Execute(&output, launchData)
	if err != nil {
		return ""
	}

	return output.String()
}

func unconfigure(launchAgentFile string) error {
	logrus.Debugf("Removing %s", launchAgentFile)
	return os.Remove(launchAgentFile)
}
<commit_msg>fixing an errant "primary", replacing with broker<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"text/template"

	"github.com/Sirupsen/logrus"
	"github.com/justone/pmb/api"
	"github.com/kardianos/osext"
)

func handleOSXCommand(bus *pmb.PMB, command string, subcommand string, arguments []string) error {

	var err error

	logrus.Debugf("Handling %s with args of %s", command, arguments)

	// launch agent name
	agentName := fmt.Sprintf("org.endot.pmb.%s", subcommand)
	logrus.Debugf("Name of launchagent: %s", agentName)

	// figure out launch agent config path
	launchAgentFile := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), agentName)
	logrus.Debugf("launchagent file: %s", launchAgentFile)

	var homeDir string
	if homeDir = os.Getenv("HOME"); len(homeDir) == 0 {
		homeDir = "/tmp"
	}

	// create launch data
	executable, err := osext.Executable()
	if err != nil {
		return err
	}

	launchData := struct {
		Name, Executable, Broker, HomeDir string
		Args                              []string
	}{
		agentName, executable, bus.BrokerURI(), homeDir, arguments,
	}

	switch command {
	case "list":
		fmt.Printf(`
Available commands for running '%s' as a background process (agent):

start - Starts agent via launchctl.
stop - Stops agent via launchctl.
restart - Restarts agent via launchctl.
configure - This will configure the agent, but not start it.
unconfigure - This will remove the agent configuration.

`, fmt.Sprintf("pmb %s", subcommand))

	case "restart":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = stop(launchAgentFile, agentName)
		if err != nil {
			return err
		}

		err = start(launchAgentFile, agentName)
		if err != nil {
			return err
		}
	case "stop":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = stop(launchAgentFile, agentName)
		if err != nil {
			return err
		}

		err = unconfigure(launchAgentFile)
		if err != nil {
			return err
		}
	case "start":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}

		err = start(launchAgentFile, agentName)
		if err != nil {
			return err
		}
	case "configure":
		err = configure(launchAgentFile, generateLaunchConfig(launchData))
		if err != nil {
			return err
		}
	case "unconfigure":
		err = unconfigure(launchAgentFile)
		if err != nil {
			return err
		}
	}

	return nil
}

func start(launchAgentFile string, agentName string) error {
	listCmd := exec.Command("/bin/launchctl", "list", agentName)
	err := listCmd.Run()

	if _, ok := err.(*exec.ExitError); ok {
		// launch agent wasn't loaded yet, so load to start
		startCmd := exec.Command("/bin/launchctl", "load", launchAgentFile)
		startErr := startCmd.Run()
		if startErr != nil {
			return startErr
		}
	} else if err != nil {
		// some error running the list command
		return err
	} else {
		// launch agent was already loaded
		logrus.Infof("Already running")
	}

	return nil
}

func stop(launchAgentFile string, agentName string) error {
	listCmd := exec.Command("/bin/launchctl", "list", agentName)
	err := listCmd.Run()

	if err == nil {
		// launch agent was loaded, so unload to stop
		stopCmd := exec.Command("/bin/launchctl", "unload", launchAgentFile)
		stopErr := stopCmd.Run()
		if stopErr != nil {
			return stopErr
		}
	} else if _, ok := err.(*exec.ExitError); ok {
		// launch agent wasn't already loaded
		logrus.Infof("Already stopped")
	} else {
		// some error running the list command
		return err
	}

	return nil
}

func configure(launchAgentFile string, config string) error {

	err := ioutil.WriteFile(launchAgentFile, []byte(config), 0644)
	if err != nil {
		return err
	}

	logrus.Debugf("Created %s: %s", launchAgentFile, config)

	return nil
}

func generateLaunchConfig(launchData interface{}) string {
	configureTemplate := `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple Computer//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
	<dict>
		<key>Label</key>
		<string>{{ .Name }}</string>
		<key>KeepAlive</key>
		<true/>
		<key>EnvironmentVariables</key>
		<dict>
			<key>PATH</key>
			<string>/usr/bin:/bin:/usr/sbin:/sbin:/usr/local/bin</string>
			<key>PMB_BROKER_URI</key>
			<string>{{ .Broker }}</string>
		</dict>
		<key>StandardOutPath</key>
		<string>{{ .HomeDir }}/.pmb.log</string>
		<key>StandardErrorPath</key>
		<string>{{ .HomeDir }}/.pmb.log</string>
		<key>ProgramArguments</key>
		<array>
			<string>{{ .Executable }}</string>
{{ range .Args }}
			<string>{{ . }}</string>
{{ end }}
		</array>
	</dict>
</plist>`

	tmpl := template.Must(template.New("configure").Parse(configureTemplate))
	var output bytes.Buffer

	err := tmpl.Execute(&output, launchData)
	if err != nil {
		return ""
	}

	return output.String()
}

func unconfigure(launchAgentFile string) error {
	logrus.Debugf("Removing %s", launchAgentFile)
	return os.Remove(launchAgentFile)
}
<|endoftext|>


<commit_before>package persistence

import (
	"github.com/premkit/healthcheck/log"

	"github.com/boltdb/bolt"
)

// DB is the lazy-loaded reference to the BoltDB instance.  Use the GetDB() function to obtain this.
var DB *bolt.DB

// GetDB returns the singleton instance of the BoltDB connection.  This is not a threadsafe object,
// but transactions are.  Any caller using this object should use a transaction.
func GetDB() (*bolt.DB, error) {
	if DB != nil {
		return DB, nil
	}

	conn, err := bolt.Open("/data/healthcheck.db", 0600, nil)
	if err != nil {
		log.Error(err)
		return nil, err
	}

	if err := initializeDatabase(conn); err != nil {
		return nil, err
	}

	DB = conn
	return DB, nil
}

func initializeDatabase(conn *bolt.DB) error {
	// Perform some initialization
	err := conn.Update(func(tx *bolt.Tx) error {
		// Create the default buckets
		_, err := tx.CreateBucketIfNotExists([]byte("Index"))
		if err != nil {
			log.Error(err)
			return err
		}

		return nil
	})

	return err
}
<commit_msg>Make data directory, and read the location from an envvar<commit_after>package persistence

import (
	"os"
	"path"

	"github.com/premkit/healthcheck/log"

	"github.com/boltdb/bolt"
)

// DB is the lazy-loaded reference to the BoltDB instance.  Use the GetDB() function to obtain this.
var DB *bolt.DB

// GetDB returns the singleton instance of the BoltDB connection.  This is not a threadsafe object,
// but transactions are.  Any caller using this object should use a transaction.
func GetDB() (*bolt.DB, error) {
	if DB != nil {
		return DB, nil
	}

	// Use the environment variable for the datadir, if set
	dataDirectory := os.Getenv("DATA_DIRECTORY")
	if dataDirectory == "" {
		dataDirectory = "/data"
	}

	if err := os.MkdirAll(dataDirectory, 0600); err != nil {
		log.Error(err)
		return nil, err
	}

	path := path.Join(dataDirectory, "healthcheck.db")
	conn, err := bolt.Open(path, 0600, nil)
	if err != nil {
		log.Error(err)
		return nil, err
	}

	if err := initializeDatabase(conn); err != nil {
		return nil, err
	}

	DB = conn
	return DB, nil
}

func initializeDatabase(conn *bolt.DB) error {
	// Perform some initialization
	err := conn.Update(func(tx *bolt.Tx) error {
		// Create the default buckets
		_, err := tx.CreateBucketIfNotExists([]byte("Index"))
		if err != nil {
			log.Error(err)
			return err
		}

		return nil
	})

	return err
}
<|endoftext|>


<commit_before>package pixelgl

import (
	"fmt"
	"image/color"

	"github.com/faiface/glhf"
	"github.com/faiface/mainthread"
	"github.com/faiface/pixel"
	"github.com/go-gl/mathgl/mgl32"
	"github.com/pkg/errors"
)

// Canvas is an off-screen rectangular BasicTarget and Picture at the same time, that you can draw
// onto.
//
// It supports TrianglesPosition, TrianglesColor, TrianglesPicture and PictureColor.
type Canvas struct {
	gf     *GLFrame
	shader *glhf.Shader

	cmp    pixel.ComposeMethod
	mat    mgl32.Mat3
	col    mgl32.Vec4
	smooth bool

	sprite *pixel.Sprite
}

var _ pixel.ComposeTarget = (*Canvas)(nil)

// NewCanvas creates a new empty, fully transparent Canvas with given bounds.
func NewCanvas(bounds pixel.Rect) *Canvas {
	c := &Canvas{
		gf:  NewGLFrame(bounds),
		mat: mgl32.Ident3(),
		col: mgl32.Vec4{1, 1, 1, 1},
	}

	c.SetBounds(bounds)

	var shader *glhf.Shader
	mainthread.Call(func() {
		var err error
		shader, err = glhf.NewShader(
			canvasVertexFormat,
			canvasUniformFormat,
			canvasVertexShader,
			canvasFragmentShader,
		)
		if err != nil {
			panic(errors.Wrap(err, "failed to create Canvas, there's a bug in the shader"))
		}
	})
	c.shader = shader

	return c
}

// MakeTriangles creates a specialized copy of the supplied Triangles that draws onto this Canvas.
//
// TrianglesPosition, TrianglesColor and TrianglesPicture are supported.
func (c *Canvas) MakeTriangles(t pixel.Triangles) pixel.TargetTriangles {
	return &canvasTriangles{
		GLTriangles: NewGLTriangles(c.shader, t),
		dst:         c,
	}
}

// MakePicture create a specialized copy of the supplied Picture that draws onto this Canvas.
//
// PictureColor is supported.
func (c *Canvas) MakePicture(p pixel.Picture) pixel.TargetPicture {
	if cp, ok := p.(*canvasPicture); ok {
		return &canvasPicture{
			GLPicture: cp.GLPicture,
			dst:       c,
		}
	}
	if gp, ok := p.(GLPicture); ok {
		return &canvasPicture{
			GLPicture: gp,
			dst:       c,
		}
	}
	return &canvasPicture{
		GLPicture: NewGLPicture(p),
		dst:       c,
	}
}

// SetMatrix sets a Matrix that every point will be projected by.
func (c *Canvas) SetMatrix(m pixel.Matrix) {
	for i := range m {
		c.mat[i] = float32(m[i])
	}
}

// SetColorMask sets a color that every color in triangles or a picture will be multiplied by.
func (c *Canvas) SetColorMask(col color.Color) {
	rgba := pixel.Alpha(1)
	if col != nil {
		rgba = pixel.ToRGBA(col)
	}
	c.col = mgl32.Vec4{
		float32(rgba.R),
		float32(rgba.G),
		float32(rgba.B),
		float32(rgba.A),
	}
}

// SetComposeMethod sets a Porter-Duff composition method to be used in the following draws onto
// this Canvas.
func (c *Canvas) SetComposeMethod(cmp pixel.ComposeMethod) {
	c.cmp = cmp
}

// SetBounds resizes the Canvas to the new bounds. Old content will be preserved.
func (c *Canvas) SetBounds(bounds pixel.Rect) {
	c.gf.SetBounds(bounds)
	if c.sprite == nil {
		c.sprite = pixel.NewSprite(nil, pixel.Rect{})
	}
	c.sprite.Set(c, c.Bounds())
	c.sprite.SetMatrix(pixel.IM.Moved(c.Bounds().Center()))
}

// Bounds returns the rectangular bounds of the Canvas.
func (c *Canvas) Bounds() pixel.Rect {
	return c.gf.Bounds()
}

// SetSmooth sets whether stretched Pictures drawn onto this Canvas should be drawn smooth or
// pixely.
func (c *Canvas) SetSmooth(smooth bool) {
	c.smooth = smooth
}

// Smooth returns whether stretched Pictures drawn onto this Canvas are set to be drawn smooth or
// pixely.
func (c *Canvas) Smooth() bool {
	return c.smooth
}

// must be manually called inside mainthread
func (c *Canvas) setGlhfBounds() {
	bx, by, bw, bh := intBounds(c.gf.Bounds())
	glhf.Bounds(bx, by, bw, bh)
}

// must be manually called inside mainthread
func setBlendFunc(cmp pixel.ComposeMethod) {
	switch cmp {
	case pixel.ComposeOver:
		glhf.BlendFunc(glhf.One, glhf.OneMinusSrcAlpha)
	case pixel.ComposeIn:
		glhf.BlendFunc(glhf.DstAlpha, glhf.Zero)
	case pixel.ComposeOut:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.Zero)
	case pixel.ComposeAtop:
		glhf.BlendFunc(glhf.DstAlpha, glhf.OneMinusSrcAlpha)
	case pixel.ComposeRover:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.One)
	case pixel.ComposeRin:
		glhf.BlendFunc(glhf.Zero, glhf.SrcAlpha)
	case pixel.ComposeRout:
		glhf.BlendFunc(glhf.Zero, glhf.OneMinusSrcAlpha)
	case pixel.ComposeRatop:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.SrcAlpha)
	case pixel.ComposeXor:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.OneMinusSrcAlpha)
	case pixel.ComposePlus:
		glhf.BlendFunc(glhf.One, glhf.One)
	case pixel.ComposeCopy:
		glhf.BlendFunc(glhf.One, glhf.Zero)
	default:
		panic(errors.New("Canvas: invalid compose method"))
	}
}

// Clear fills the whole Canvas with a single color.
func (c *Canvas) Clear(color color.Color) {
	c.gf.Dirty()

	rgba := pixel.ToRGBA(color)

	// color masking
	rgba = rgba.Mul(pixel.RGBA{
		R: float64(c.col[0]),
		G: float64(c.col[1]),
		B: float64(c.col[2]),
		A: float64(c.col[3]),
	})

	mainthread.CallNonBlock(func() {
		c.setGlhfBounds()
		c.gf.Frame().Begin()
		glhf.Clear(
			float32(rgba.R),
			float32(rgba.G),
			float32(rgba.B),
			float32(rgba.A),
		)
		c.gf.Frame().End()
	})
}

// Color returns the color of the pixel over the given position inside the Canvas.
func (c *Canvas) Color(at pixel.Vec) pixel.RGBA {
	return c.gf.Color(at)
}

// Texture returns the underlying OpenGL Texture of this Canvas.
//
// Implements GLPicture interface.
func (c *Canvas) Texture() *glhf.Texture {
	return c.gf.Texture()
}

// Draw draws a rectangle equal to Canvas's Bounds containing the Canvas's content to another
// Target.
//
// Note, that the matrix and the color mask of this Canvas have no effect here.
func (c *Canvas) Draw(t pixel.Target) {
	c.sprite.Draw(t)
}

type canvasTriangles struct {
	*GLTriangles
	dst *Canvas
}

func (ct *canvasTriangles) draw(tex *glhf.Texture, bounds pixel.Rect) {
	ct.dst.gf.Dirty()

	// save the current state vars to avoid race condition
	cmp := ct.dst.cmp
	mat := ct.dst.mat
	col := ct.dst.col
	smt := ct.dst.smooth

	mainthread.CallNonBlock(func() {
		ct.dst.setGlhfBounds()
		setBlendFunc(cmp)

		frame := ct.dst.gf.Frame()
		shader := ct.dst.shader

		frame.Begin()
		shader.Begin()

		dstBounds := ct.dst.Bounds()
		shader.SetUniformAttr(canvasBounds, mgl32.Vec4{
			float32(dstBounds.Min.X()),
			float32(dstBounds.Min.Y()),
			float32(dstBounds.W()),
			float32(dstBounds.H()),
		})
		shader.SetUniformAttr(canvasTransform, mat)
		shader.SetUniformAttr(canvasColorMask, col)

		if tex == nil {
			ct.vs.Begin()
			ct.vs.Draw()
			ct.vs.End()
		} else {
			tex.Begin()

			bx, by, bw, bh := intBounds(bounds)
			shader.SetUniformAttr(canvasTexBounds, mgl32.Vec4{
				float32(bx),
				float32(by),
				float32(bw),
				float32(bh),
			})

			if tex.Smooth() != smt {
				tex.SetSmooth(smt)
			}

			ct.vs.Begin()
			ct.vs.Draw()
			ct.vs.End()

			tex.End()
		}

		shader.End()
		frame.End()
	})
}

func (ct *canvasTriangles) Draw() {
	ct.draw(nil, pixel.Rect{})
}

type canvasPicture struct {
	GLPicture
	dst *Canvas
}

func (cp *canvasPicture) Draw(t pixel.TargetTriangles) {
	ct := t.(*canvasTriangles)
	if cp.dst != ct.dst {
		panic(fmt.Errorf("(%T).Draw: TargetTriangles generated by different Canvas", cp))
	}
	ct.draw(cp.GLPicture.Texture(), cp.GLPicture.Bounds())
}

const (
	canvasPosition int = iota
	canvasColor
	canvasTexture
	canvasIntensity
)

var canvasVertexFormat = glhf.AttrFormat{
	canvasPosition:  {Name: "position", Type: glhf.Vec2},
	canvasColor:     {Name: "color", Type: glhf.Vec4},
	canvasTexture:   {Name: "texture", Type: glhf.Vec2},
	canvasIntensity: {Name: "intensity", Type: glhf.Float},
}

const (
	canvasTransform int = iota
	canvasColorMask
	canvasBounds
	canvasTexBounds
)

var canvasUniformFormat = glhf.AttrFormat{
	canvasTransform: {Name: "transform", Type: glhf.Mat3},
	canvasColorMask: {Name: "colorMask", Type: glhf.Vec4},
	canvasBounds:    {Name: "bounds", Type: glhf.Vec4},
	canvasTexBounds: {Name: "texBounds", Type: glhf.Vec4},
}

var canvasVertexShader = `
#version 330 core

in vec2 position;
in vec4 color;
in vec2 texture;
in float intensity;

out vec4 Color;
out vec2 Texture;
out float Intensity;

uniform mat3 transform;
uniform vec4 bounds;

void main() {
	vec2 transPos = (transform * vec3(position, 1.0)).xy;
	vec2 normPos = (transPos - bounds.xy) / bounds.zw * 2 - vec2(1, 1);
	gl_Position = vec4(normPos, 0.0, 1.0);
	Color = color;
	Texture = texture;
	Intensity = intensity;
}
`

var canvasFragmentShader = `
#version 330 core

in vec4 Color;
in vec2 Texture;
in float Intensity;

out vec4 color;

uniform vec4 colorMask;
uniform vec4 texBounds;
uniform sampler2D tex;

void main() {
	if (Intensity == 0) {
		color = colorMask * Color;
	} else {
		color = vec4(0, 0, 0, 0);
		color += (1 - Intensity) * Color;
		vec2 t = (Texture - texBounds.xy) / texBounds.zw;
		color += Intensity * Color * texture(tex, t);
		color *= colorMask;
	}
}
`
<commit_msg>fix Canvas drawing when bounds don't start at (0, 0)<commit_after>package pixelgl

import (
	"fmt"
	"image/color"

	"github.com/faiface/glhf"
	"github.com/faiface/mainthread"
	"github.com/faiface/pixel"
	"github.com/go-gl/mathgl/mgl32"
	"github.com/pkg/errors"
)

// Canvas is an off-screen rectangular BasicTarget and Picture at the same time, that you can draw
// onto.
//
// It supports TrianglesPosition, TrianglesColor, TrianglesPicture and PictureColor.
type Canvas struct {
	gf     *GLFrame
	shader *glhf.Shader

	cmp    pixel.ComposeMethod
	mat    mgl32.Mat3
	col    mgl32.Vec4
	smooth bool

	sprite *pixel.Sprite
}

var _ pixel.ComposeTarget = (*Canvas)(nil)

// NewCanvas creates a new empty, fully transparent Canvas with given bounds.
func NewCanvas(bounds pixel.Rect) *Canvas {
	c := &Canvas{
		gf:  NewGLFrame(bounds),
		mat: mgl32.Ident3(),
		col: mgl32.Vec4{1, 1, 1, 1},
	}

	c.SetBounds(bounds)

	var shader *glhf.Shader
	mainthread.Call(func() {
		var err error
		shader, err = glhf.NewShader(
			canvasVertexFormat,
			canvasUniformFormat,
			canvasVertexShader,
			canvasFragmentShader,
		)
		if err != nil {
			panic(errors.Wrap(err, "failed to create Canvas, there's a bug in the shader"))
		}
	})
	c.shader = shader

	return c
}

// MakeTriangles creates a specialized copy of the supplied Triangles that draws onto this Canvas.
//
// TrianglesPosition, TrianglesColor and TrianglesPicture are supported.
func (c *Canvas) MakeTriangles(t pixel.Triangles) pixel.TargetTriangles {
	return &canvasTriangles{
		GLTriangles: NewGLTriangles(c.shader, t),
		dst:         c,
	}
}

// MakePicture create a specialized copy of the supplied Picture that draws onto this Canvas.
//
// PictureColor is supported.
func (c *Canvas) MakePicture(p pixel.Picture) pixel.TargetPicture {
	if cp, ok := p.(*canvasPicture); ok {
		return &canvasPicture{
			GLPicture: cp.GLPicture,
			dst:       c,
		}
	}
	if gp, ok := p.(GLPicture); ok {
		return &canvasPicture{
			GLPicture: gp,
			dst:       c,
		}
	}
	return &canvasPicture{
		GLPicture: NewGLPicture(p),
		dst:       c,
	}
}

// SetMatrix sets a Matrix that every point will be projected by.
func (c *Canvas) SetMatrix(m pixel.Matrix) {
	for i := range m {
		c.mat[i] = float32(m[i])
	}
}

// SetColorMask sets a color that every color in triangles or a picture will be multiplied by.
func (c *Canvas) SetColorMask(col color.Color) {
	rgba := pixel.Alpha(1)
	if col != nil {
		rgba = pixel.ToRGBA(col)
	}
	c.col = mgl32.Vec4{
		float32(rgba.R),
		float32(rgba.G),
		float32(rgba.B),
		float32(rgba.A),
	}
}

// SetComposeMethod sets a Porter-Duff composition method to be used in the following draws onto
// this Canvas.
func (c *Canvas) SetComposeMethod(cmp pixel.ComposeMethod) {
	c.cmp = cmp
}

// SetBounds resizes the Canvas to the new bounds. Old content will be preserved.
func (c *Canvas) SetBounds(bounds pixel.Rect) {
	c.gf.SetBounds(bounds)
	if c.sprite == nil {
		c.sprite = pixel.NewSprite(nil, pixel.Rect{})
	}
	c.sprite.Set(c, c.Bounds())
	c.sprite.SetMatrix(pixel.IM.Moved(c.Bounds().Center()))
}

// Bounds returns the rectangular bounds of the Canvas.
func (c *Canvas) Bounds() pixel.Rect {
	return c.gf.Bounds()
}

// SetSmooth sets whether stretched Pictures drawn onto this Canvas should be drawn smooth or
// pixely.
func (c *Canvas) SetSmooth(smooth bool) {
	c.smooth = smooth
}

// Smooth returns whether stretched Pictures drawn onto this Canvas are set to be drawn smooth or
// pixely.
func (c *Canvas) Smooth() bool {
	return c.smooth
}

// must be manually called inside mainthread
func (c *Canvas) setGlhfBounds() {
	_, _, bw, bh := intBounds(c.gf.Bounds())
	glhf.Bounds(0, 0, bw, bh)
}

// must be manually called inside mainthread
func setBlendFunc(cmp pixel.ComposeMethod) {
	switch cmp {
	case pixel.ComposeOver:
		glhf.BlendFunc(glhf.One, glhf.OneMinusSrcAlpha)
	case pixel.ComposeIn:
		glhf.BlendFunc(glhf.DstAlpha, glhf.Zero)
	case pixel.ComposeOut:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.Zero)
	case pixel.ComposeAtop:
		glhf.BlendFunc(glhf.DstAlpha, glhf.OneMinusSrcAlpha)
	case pixel.ComposeRover:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.One)
	case pixel.ComposeRin:
		glhf.BlendFunc(glhf.Zero, glhf.SrcAlpha)
	case pixel.ComposeRout:
		glhf.BlendFunc(glhf.Zero, glhf.OneMinusSrcAlpha)
	case pixel.ComposeRatop:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.SrcAlpha)
	case pixel.ComposeXor:
		glhf.BlendFunc(glhf.OneMinusDstAlpha, glhf.OneMinusSrcAlpha)
	case pixel.ComposePlus:
		glhf.BlendFunc(glhf.One, glhf.One)
	case pixel.ComposeCopy:
		glhf.BlendFunc(glhf.One, glhf.Zero)
	default:
		panic(errors.New("Canvas: invalid compose method"))
	}
}

// Clear fills the whole Canvas with a single color.
func (c *Canvas) Clear(color color.Color) {
	c.gf.Dirty()

	rgba := pixel.ToRGBA(color)

	// color masking
	rgba = rgba.Mul(pixel.RGBA{
		R: float64(c.col[0]),
		G: float64(c.col[1]),
		B: float64(c.col[2]),
		A: float64(c.col[3]),
	})

	mainthread.CallNonBlock(func() {
		c.setGlhfBounds()
		c.gf.Frame().Begin()
		glhf.Clear(
			float32(rgba.R),
			float32(rgba.G),
			float32(rgba.B),
			float32(rgba.A),
		)
		c.gf.Frame().End()
	})
}

// Color returns the color of the pixel over the given position inside the Canvas.
func (c *Canvas) Color(at pixel.Vec) pixel.RGBA {
	return c.gf.Color(at)
}

// Texture returns the underlying OpenGL Texture of this Canvas.
//
// Implements GLPicture interface.
func (c *Canvas) Texture() *glhf.Texture {
	return c.gf.Texture()
}

// Draw draws a rectangle equal to Canvas's Bounds containing the Canvas's content to another
// Target.
//
// Note, that the matrix and the color mask of this Canvas have no effect here.
func (c *Canvas) Draw(t pixel.Target) {
	c.sprite.Draw(t)
}

type canvasTriangles struct {
	*GLTriangles
	dst *Canvas
}

func (ct *canvasTriangles) draw(tex *glhf.Texture, bounds pixel.Rect) {
	ct.dst.gf.Dirty()

	// save the current state vars to avoid race condition
	cmp := ct.dst.cmp
	mat := ct.dst.mat
	col := ct.dst.col
	smt := ct.dst.smooth

	mainthread.CallNonBlock(func() {
		ct.dst.setGlhfBounds()
		setBlendFunc(cmp)

		frame := ct.dst.gf.Frame()
		shader := ct.dst.shader

		frame.Begin()
		shader.Begin()

		dstBounds := ct.dst.Bounds()
		shader.SetUniformAttr(canvasBounds, mgl32.Vec4{
			float32(dstBounds.Min.X()),
			float32(dstBounds.Min.Y()),
			float32(dstBounds.W()),
			float32(dstBounds.H()),
		})
		shader.SetUniformAttr(canvasTransform, mat)
		shader.SetUniformAttr(canvasColorMask, col)

		if tex == nil {
			ct.vs.Begin()
			ct.vs.Draw()
			ct.vs.End()
		} else {
			tex.Begin()

			bx, by, bw, bh := intBounds(bounds)
			shader.SetUniformAttr(canvasTexBounds, mgl32.Vec4{
				float32(bx),
				float32(by),
				float32(bw),
				float32(bh),
			})

			if tex.Smooth() != smt {
				tex.SetSmooth(smt)
			}

			ct.vs.Begin()
			ct.vs.Draw()
			fmt.Println(ct.vs.VertexData())
			ct.vs.End()

			tex.End()
		}

		shader.End()
		frame.End()
	})
}

func (ct *canvasTriangles) Draw() {
	ct.draw(nil, pixel.Rect{})
}

type canvasPicture struct {
	GLPicture
	dst *Canvas
}

func (cp *canvasPicture) Draw(t pixel.TargetTriangles) {
	ct := t.(*canvasTriangles)
	if cp.dst != ct.dst {
		panic(fmt.Errorf("(%T).Draw: TargetTriangles generated by different Canvas", cp))
	}
	ct.draw(cp.GLPicture.Texture(), cp.GLPicture.Bounds())
}

const (
	canvasPosition int = iota
	canvasColor
	canvasTexture
	canvasIntensity
)

var canvasVertexFormat = glhf.AttrFormat{
	canvasPosition:  {Name: "position", Type: glhf.Vec2},
	canvasColor:     {Name: "color", Type: glhf.Vec4},
	canvasTexture:   {Name: "texture", Type: glhf.Vec2},
	canvasIntensity: {Name: "intensity", Type: glhf.Float},
}

const (
	canvasTransform int = iota
	canvasColorMask
	canvasBounds
	canvasTexBounds
)

var canvasUniformFormat = glhf.AttrFormat{
	canvasTransform: {Name: "transform", Type: glhf.Mat3},
	canvasColorMask: {Name: "colorMask", Type: glhf.Vec4},
	canvasBounds:    {Name: "bounds", Type: glhf.Vec4},
	canvasTexBounds: {Name: "texBounds", Type: glhf.Vec4},
}

var canvasVertexShader = `
#version 330 core

in vec2 position;
in vec4 color;
in vec2 texture;
in float intensity;

out vec4 Color;
out vec2 Texture;
out float Intensity;

uniform mat3 transform;
uniform vec4 bounds;

void main() {
	vec2 transPos = (transform * vec3(position, 1.0)).xy;
	vec2 normPos = (transPos - bounds.xy) / bounds.zw * 2 - vec2(1, 1);
	gl_Position = vec4(normPos, 0.0, 1.0);
	Color = color;
	Texture = texture;
	Intensity = intensity;
}
`

var canvasFragmentShader = `
#version 330 core

in vec4 Color;
in vec2 Texture;
in float Intensity;

out vec4 color;

uniform vec4 colorMask;
uniform vec4 texBounds;
uniform sampler2D tex;

void main() {
	if (Intensity == 0) {
		color = colorMask * Color;
	} else {
		color = vec4(0, 0, 0, 0);
		color += (1 - Intensity) * Color;
		vec2 t = (Texture - texBounds.xy) / texBounds.zw;
		color += Intensity * Color * texture(tex, t);
		color *= colorMask;
	}
}
`
<|endoftext|>


<commit_before>package pixy

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/mailgun/log"
	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/mailgun/sarama"
	. "github.com/mailgun/kafka-pixy/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	// Use Shopify/sarama Vagrant box (copied over from https://github.com/Shopify/sarama/blob/master/functional_test.go#L18)
	VagrantKafkaPeers     = "192.168.100.67:9091,192.168.100.67:9092,192.168.100.67:9093,192.168.100.67:9094,192.168.100.67:9095"
	VagrantZookeeperPeers = "192.168.100.67:2181,192.168.100.67:2182,192.168.100.67:2183,192.168.100.67:2184,192.168.100.67:2185"
)

var (
	testKafkaPeers     []string
	testZookeeperPeers []string

	initLogOnce = sync.Once{}
)

func init() {
	kafkaPeersStr := os.Getenv("KAFKA_PEERS")
	if kafkaPeersStr == "" {
		kafkaPeersStr = VagrantKafkaPeers
	}
	testKafkaPeers = strings.Split(kafkaPeersStr, ",")
}

func Test(t *testing.T) {
	TestingT(t)
}

func InitTestLog() {
	initLogOnce.Do(func() {
		consoleLogger, _ := log.NewConsoleLogger(log.Config{Severity: "info"})
		log.Init(consoleLogger)
		InitLibraryLoggers()
	})
}

// NewUDSHTTPClient creates an HTTP client that always connects to the
// specified unix domain socket ignoring the host part of requested HTTP URLs.
func NewUDSHTTPClient(unixSockAddr string) *http.Client {
	dial := func(proto, addr string) (net.Conn, error) {
		return net.Dial("unix", unixSockAddr)
	}
	tr := &http.Transport{Dial: dial}
	return &http.Client{Transport: tr}
}

type TestKafkaClient struct {
	client   sarama.Client
	consumer sarama.Consumer
}

func NewTestKafkaClient(brokers []string) *TestKafkaClient {
	tkc := &TestKafkaClient{}
	clientCfg := sarama.NewConfig()
	clientCfg.ClientID = "unittest-runner"
	err := error(nil)
	if tkc.client, err = sarama.NewClient(brokers, clientCfg); err != nil {
		panic(err)
	}
	if tkc.consumer, err = sarama.NewConsumerFromClient(tkc.client); err != nil {
		panic(err)
	}
	return tkc
}

func (tkc *TestKafkaClient) Close() {
	tkc.consumer.Close()
	tkc.client.Close()
}

func (tkc *TestKafkaClient) getOffsets(topic string) []int64 {
	offsets := []int64{}
	partitions, err := tkc.client.Partitions(topic)
	if err != nil {
		panic(err)
	}
	for _, p := range partitions {
		offset, err := tkc.client.GetOffset(topic, p, sarama.OffsetNewest)
		if err != nil {
			panic(err)
		}
		offsets = append(offsets, offset)
	}
	return offsets
}

func (tkc *TestKafkaClient) getMessages(topic string, begin, end []int64) [][]string {
	writtenMsgs := make([][]string, len(begin))
	for i := range begin {
		p, err := tkc.consumer.ConsumePartition(topic, int32(i), begin[i])
		if err != nil {
			panic(err)
		}
		writtenMsgCount := int(end[i] - begin[i])
		for j := 0; j < writtenMsgCount; j++ {
			connMsg := <-p.Messages()
			writtenMsgs[i] = append(writtenMsgs[i], string(connMsg.Value))
		}
		p.Close()
	}
	return writtenMsgs
}

// GenMessage generates an ASCII message of the specified size.
func GenMessage(size int) string {
	b := bytes.NewBuffer(nil)
	for b.Len() < size {
		b.WriteString(strconv.Itoa(b.Len()))
		b.WriteString("-")
	}
	return string(b.Bytes()[:size])
}

// ChunkReader allows reading its underlying buffer in chunks making the
// specified pauses between the chunks. After each pause `Read()` returns
// `0, nil`. This kind of reader is useful to simulate HTTP requests that
// require several read operations on the request body to get all of it.
type ChunkReader struct {
	chunks   []string
	chunk    string
	pause    time.Duration
	chunkDue time.Time
}

func NewChunkReader(s string, count int, pause time.Duration) *ChunkReader {
	chunkSize := len(s) / count
	chunks := make([]string, count, count+1)
	for i := 0; i < count; i++ {
		begin := chunkSize * i
		end := begin + chunkSize
		chunks[i] = s[begin:end]
	}
	if count*chunkSize != len(s) {
		chunks = append(chunks, s[chunkSize*count:])
	}
	return &ChunkReader{
		chunks:   chunks,
		pause:    pause,
		chunkDue: time.Now().Add(pause),
	}
}

func (cr *ChunkReader) Read(b []byte) (n int, err error) {
	if len(cr.chunk) == 0 {
		if len(cr.chunks) == 0 {
			return 0, io.EOF
		}
		cr.chunk = cr.chunks[0]
		cr.chunks = cr.chunks[1:]
		cr.chunkDue = time.Now().Add(cr.pause)
	}

	waitFor := time.Now().Sub(cr.chunkDue)
	if waitFor > 0 {
		time.Sleep(waitFor)
		return 0, nil
	}

	copied := copy(b, cr.chunk)
	cr.chunk = cr.chunk[copied:]
	return copied, nil
}

func PostChunked(clt *http.Client, url, msg string) *http.Response {
	req, err := http.NewRequest("POST", url, NewChunkReader(msg, 1, 10*time.Millisecond))
	if err != nil {
		panic(fmt.Sprintf("Failed to make a request object: cause=(%v)", err))
	}
	req.Header.Add("Content-Type", "text/plain")
	req.ContentLength = int64(len(msg))
	resp, err := clt.Do(req)
	if err != nil {
		panic(fmt.Sprintf("Failed to do a request: cause=(%v)", err))
	}
	return resp
}

func AssertHTTPResp(c *C, res *http.Response, expectedStatusCode int, expectedBody string) {
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		c.Error(err)
	}
	res.Body.Close()
	c.Assert(res.StatusCode, Equals, expectedStatusCode)
	c.Assert(string(body), Matches, expectedBody)
}

func ProdMsgMetadataSize(key []byte) int {
	size := 26 // the metadata overhead of CRC, flags, etc.
	if key != nil {
		size += sarama.ByteEncoder(key).Length()
	}
	return size
}
<commit_msg>Fix test chunked reader<commit_after>package pixy

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/mailgun/log"
	"github.com/mailgun/kafka-pixy/Godeps/_workspace/src/github.com/mailgun/sarama"
	. "github.com/mailgun/kafka-pixy/Godeps/_workspace/src/gopkg.in/check.v1"
)

const (
	// Use Shopify/sarama Vagrant box (copied over from https://github.com/Shopify/sarama/blob/master/functional_test.go#L18)
	VagrantKafkaPeers     = "192.168.100.67:9091,192.168.100.67:9092,192.168.100.67:9093,192.168.100.67:9094,192.168.100.67:9095"
	VagrantZookeeperPeers = "192.168.100.67:2181,192.168.100.67:2182,192.168.100.67:2183,192.168.100.67:2184,192.168.100.67:2185"
)

var (
	testKafkaPeers     []string
	testZookeeperPeers []string

	initLogOnce = sync.Once{}
)

func init() {
	kafkaPeersStr := os.Getenv("KAFKA_PEERS")
	if kafkaPeersStr == "" {
		kafkaPeersStr = VagrantKafkaPeers
	}
	testKafkaPeers = strings.Split(kafkaPeersStr, ",")
}

func Test(t *testing.T) {
	TestingT(t)
}

func InitTestLog() {
	initLogOnce.Do(func() {
		consoleLogger, _ := log.NewConsoleLogger(log.Config{Severity: "info"})
		log.Init(consoleLogger)
		InitLibraryLoggers()
	})
}

// NewUDSHTTPClient creates an HTTP client that always connects to the
// specified unix domain socket ignoring the host part of requested HTTP URLs.
func NewUDSHTTPClient(unixSockAddr string) *http.Client {
	dial := func(proto, addr string) (net.Conn, error) {
		return net.Dial("unix", unixSockAddr)
	}
	tr := &http.Transport{Dial: dial}
	return &http.Client{Transport: tr}
}

type TestKafkaClient struct {
	client   sarama.Client
	consumer sarama.Consumer
}

func NewTestKafkaClient(brokers []string) *TestKafkaClient {
	tkc := &TestKafkaClient{}
	clientCfg := sarama.NewConfig()
	clientCfg.ClientID = "unittest-runner"
	err := error(nil)
	if tkc.client, err = sarama.NewClient(brokers, clientCfg); err != nil {
		panic(err)
	}
	if tkc.consumer, err = sarama.NewConsumerFromClient(tkc.client); err != nil {
		panic(err)
	}
	return tkc
}

func (tkc *TestKafkaClient) Close() {
	tkc.consumer.Close()
	tkc.client.Close()
}

func (tkc *TestKafkaClient) getOffsets(topic string) []int64 {
	offsets := []int64{}
	partitions, err := tkc.client.Partitions(topic)
	if err != nil {
		panic(err)
	}
	for _, p := range partitions {
		offset, err := tkc.client.GetOffset(topic, p, sarama.OffsetNewest)
		if err != nil {
			panic(err)
		}
		offsets = append(offsets, offset)
	}
	return offsets
}

func (tkc *TestKafkaClient) getMessages(topic string, begin, end []int64) [][]string {
	writtenMsgs := make([][]string, len(begin))
	for i := range begin {
		p, err := tkc.consumer.ConsumePartition(topic, int32(i), begin[i])
		if err != nil {
			panic(err)
		}
		writtenMsgCount := int(end[i] - begin[i])
		for j := 0; j < writtenMsgCount; j++ {
			connMsg := <-p.Messages()
			writtenMsgs[i] = append(writtenMsgs[i], string(connMsg.Value))
		}
		p.Close()
	}
	return writtenMsgs
}

// GenMessage generates an ASCII message of the specified size.
func GenMessage(size int) string {
	b := bytes.NewBuffer(nil)
	for b.Len() < size {
		b.WriteString(strconv.Itoa(b.Len()))
		b.WriteString("-")
	}
	return string(b.Bytes()[:size])
}

// ChunkReader allows reading its underlying buffer in chunks making the
// specified pauses between the chunks. After each pause `Read()` returns
// `0, nil`. This kind of reader is useful to simulate HTTP requests that
// require several read operations on the request body to get all of it.
type ChunkReader struct {
	buf         []byte
	begin       int
	chunkSize   int
	pause       time.Duration
	shouldPause bool
}

func NewChunkReader(s string, count int, pause time.Duration) *ChunkReader {
	return &ChunkReader{
		buf:       []byte(s),
		chunkSize: len(s) / count,
		pause:     pause,
	}
}

func (cr *ChunkReader) Read(b []byte) (n int, err error) {
	if cr.begin == len(cr.buf) {
		return 0, io.EOF
	}
	if cr.shouldPause = !cr.shouldPause; cr.shouldPause {
		time.Sleep(cr.pause)
		return 0, nil
	}
	chunkSize := cr.chunkSize
	if len(b) < chunkSize {
		chunkSize = len(b)
	}
	if len(cr.buf)-cr.begin < chunkSize {
		chunkSize = len(cr.buf) - cr.begin
	}
	end := cr.begin + chunkSize
	copied := copy(b, cr.buf[cr.begin:end])
	cr.begin = end
	return copied, nil
}

func PostChunked(clt *http.Client, url, msg string) *http.Response {
	req, err := http.NewRequest("POST", url, NewChunkReader(msg, 1, 10*time.Millisecond))
	if err != nil {
		panic(fmt.Sprintf("Failed to make a request object: cause=(%v)", err))
	}
	req.Header.Add("Content-Type", "text/plain")
	req.ContentLength = int64(len(msg))
	resp, err := clt.Do(req)
	if err != nil {
		panic(fmt.Sprintf("Failed to do a request: cause=(%v)", err))
	}
	return resp
}

func AssertHTTPResp(c *C, res *http.Response, expectedStatusCode int, expectedBody string) {
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		c.Error(err)
	}
	res.Body.Close()
	c.Assert(res.StatusCode, Equals, expectedStatusCode)
	c.Assert(string(body), Matches, expectedBody)
}

func ProdMsgMetadataSize(key []byte) int {
	size := 26 // the metadata overhead of CRC, flags, etc.
	if key != nil {
		size += sarama.ByteEncoder(key).Length()
	}
	return size
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 acme

import (
	"context"
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	acmecl "golang.org/x/crypto/acme"

	acme "github.com/jetstack/cert-manager/pkg/acme/client"
	acmemw "github.com/jetstack/cert-manager/pkg/acme/client/middleware"
	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1alpha2"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	"github.com/jetstack/cert-manager/pkg/util"
)

// This file implements a basic cache for ACME clients that can be used to
// obtain a reference to an ACME client.
// This can be accessed via the 'helper' defined in helper.go, or directly with
// the ClientWithKey function below.

// ClientWithKey will construct a new ACME client for the provided Issuer, using
// the given RSA private key.
func ClientWithKey(iss cmapi.GenericIssuer, pk *rsa.PrivateKey) (acme.Interface, error) {
	acmeSpec := iss.GetSpec().ACME
	if acmeSpec == nil {
		return nil, fmt.Errorf("issuer %q is not an ACME issuer. Ensure the 'acme' stanza is correctly specified on your Issuer resource", iss.GetObjectMeta().Name)
	}
	acmeCl := lookupClient(acmeSpec, pk)

	return acmemw.NewLogger(acmeCl), nil
}

// clientRepo is a collection of acme clients indexed
// by the options used to create them. This is used so
// that the cert-manager controllers can concurrently access
// the anti-replay nonces and directory information.
var (
	clientRepo   map[repoKey]*acmecl.Client
	clientRepoMu sync.Mutex
)

type repoKey struct {
	skiptls   bool
	server    string
	publickey string
	exponent  int
}

func lookupClient(spec *cmacme.ACMEIssuer, pk *rsa.PrivateKey) *acmecl.Client {
	clientRepoMu.Lock()
	defer clientRepoMu.Unlock()
	if clientRepo == nil {
		clientRepo = make(map[repoKey]*acmecl.Client)
	}
	repokey := repoKey{
		skiptls: spec.SkipTLSVerify,
		server:  spec.Server,
	}
	// Encoding a big.Int cannot fail
	pkbytes, _ := pk.PublicKey.N.GobEncode()
	repokey.publickey = string(pkbytes)
	repokey.exponent = pk.PublicKey.E

	client := clientRepo[repokey]
	if client != nil {
		return client
	}
	acmeCl := &acmecl.Client{
		HTTPClient:   buildHTTPClient(spec.SkipTLSVerify),
		Key:          pk,
		DirectoryURL: spec.Server,
		UserAgent:    util.CertManagerUserAgent,
	}
	clientRepo[repokey] = acmeCl
	return acmeCl
}

func ClearClientCache() {
	clientRepoMu.Lock()
	defer clientRepoMu.Unlock()
	clientRepo = nil
}

// buildHTTPClient returns an HTTP client to be used by the ACME client.
// For the time being, we construct a new HTTP client on each invocation.
// This is because we need to set the 'skipTLSVerify' flag on the HTTP client
// itself.
// In future, we may change to having two global HTTP clients - one that ignores
// TLS connection errors, and the other that does not.
func buildHTTPClient(skipTLSVerify bool) *http.Client {
	return acme.NewInstrumentedClient(&http.Client{
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           dialTimeout,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: skipTLSVerify},
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: time.Second * 30,
	})
}

var timeout = 5 * time.Second

func dialTimeout(ctx context.Context, network, addr string) (net.Conn, error) {
	d := net.Dialer{Timeout: timeout}
	return d.DialContext(ctx, network, addr)
}
<commit_msg>Re-use net.Dialer on repeat calls to DialContext<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 acme

import (
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	acmecl "golang.org/x/crypto/acme"

	acme "github.com/jetstack/cert-manager/pkg/acme/client"
	acmemw "github.com/jetstack/cert-manager/pkg/acme/client/middleware"
	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1alpha2"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	"github.com/jetstack/cert-manager/pkg/util"
)

// This file implements a basic cache for ACME clients that can be used to
// obtain a reference to an ACME client.
// This can be accessed via the 'helper' defined in helper.go, or directly with
// the ClientWithKey function below.

// ClientWithKey will construct a new ACME client for the provided Issuer, using
// the given RSA private key.
func ClientWithKey(iss cmapi.GenericIssuer, pk *rsa.PrivateKey) (acme.Interface, error) {
	acmeSpec := iss.GetSpec().ACME
	if acmeSpec == nil {
		return nil, fmt.Errorf("issuer %q is not an ACME issuer. Ensure the 'acme' stanza is correctly specified on your Issuer resource", iss.GetObjectMeta().Name)
	}
	acmeCl := lookupClient(acmeSpec, pk)

	return acmemw.NewLogger(acmeCl), nil
}

// clientRepo is a collection of acme clients indexed
// by the options used to create them. This is used so
// that the cert-manager controllers can concurrently access
// the anti-replay nonces and directory information.
var (
	clientRepo   map[repoKey]*acmecl.Client
	clientRepoMu sync.Mutex
)

type repoKey struct {
	skiptls   bool
	server    string
	publickey string
	exponent  int
}

func lookupClient(spec *cmacme.ACMEIssuer, pk *rsa.PrivateKey) *acmecl.Client {
	clientRepoMu.Lock()
	defer clientRepoMu.Unlock()
	if clientRepo == nil {
		clientRepo = make(map[repoKey]*acmecl.Client)
	}
	repokey := repoKey{
		skiptls: spec.SkipTLSVerify,
		server:  spec.Server,
	}
	// Encoding a big.Int cannot fail
	pkbytes, _ := pk.PublicKey.N.GobEncode()
	repokey.publickey = string(pkbytes)
	repokey.exponent = pk.PublicKey.E

	client := clientRepo[repokey]
	if client != nil {
		return client
	}
	acmeCl := &acmecl.Client{
		HTTPClient:   buildHTTPClient(spec.SkipTLSVerify),
		Key:          pk,
		DirectoryURL: spec.Server,
		UserAgent:    util.CertManagerUserAgent,
	}
	clientRepo[repokey] = acmeCl
	return acmeCl
}

func ClearClientCache() {
	clientRepoMu.Lock()
	defer clientRepoMu.Unlock()
	clientRepo = nil
}

// buildHTTPClient returns an HTTP client to be used by the ACME client.
// For the time being, we construct a new HTTP client on each invocation.
// This is because we need to set the 'skipTLSVerify' flag on the HTTP client
// itself.
// In future, we may change to having two global HTTP clients - one that ignores
// TLS connection errors, and the other that does not.
func buildHTTPClient(skipTLSVerify bool) *http.Client {
	return acme.NewInstrumentedClient(&http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).DialContext,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: skipTLSVerify},
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: time.Second * 30,
	})
}
<|endoftext|>