<commit_before>package nats_output

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/nats-io/nats.go"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	natsConnectWait         = 2 * time.Second
	natsReconnectBufferSize = 100 * 1024 * 1024
	defaultSubjectName      = "gnmic-telemetry"
	defaultFormat           = "json"
	defaultNumWorkers       = 1
	defaultWriteTimeout     = 10 * time.Second
)

func init() {
	outputs.Register("nats", func() outputs.Output {
		return &NatsOutput{
			Cfg: &Config{},
			wg:  new(sync.WaitGroup),
		}
	})
}

type protoMsg struct {
	m    proto.Message
	meta outputs.Meta
}

// NatsOutput //
type NatsOutput struct {
	Cfg      *Config
	ctx      context.Context
	cancelFn context.CancelFunc
	msgChan  chan *protoMsg
	wg       *sync.WaitGroup
	logger   *log.Logger
	mo       *formatters.MarshalOptions
	evps     []formatters.EventProcessor
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	SubjectPrefix   string        `mapstructure:"subject-prefix,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
	Format          string        `mapstructure:"format,omitempty"`
	NumWorkers      int           `mapstructure:"num-workers,omitempty"`
	WriteTimeout    time.Duration `mapstructure:"write-timeout,omitempty"`
	Debug           bool          `mapstructure:"debug,omitempty"`
	EventProcessors []string      `mapstructure:"event-processors,omitempty"`
}

func (n *NatsOutput) String() string {
	b, err := json.Marshal(n)
	if err != nil {
		return ""
	}
	return string(b)
}
func (n *NatsOutput) SetLogger(logger *log.Logger) {
	if logger != nil {
		n.logger = log.New(logger.Writer(), "nats_output ", logger.Flags())
		return
	}
	n.logger = log.New(os.Stderr, "nats_output ", log.LstdFlags|log.Lmicroseconds)
}

func (n *NatsOutput) SetEventProcessors(ps map[string]map[string]interface{}, log *log.Logger) {
	for _, epName := range n.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], log)
				if err != nil {
					n.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				n.evps = append(n.evps, ep)
				n.logger.Printf("added event processor '%s' of type=%s to nats output", epName, epType)
			}
		}
	}
}

// Init //
func (n *NatsOutput) Init(ctx context.Context, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, n.Cfg)
	if err != nil {
		return err
	}
	err = n.setDefaults()
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(n)
	}
	n.msgChan = make(chan *protoMsg)
	initMetrics()
	n.mo = &formatters.MarshalOptions{Format: n.Cfg.Format}
	n.ctx, n.cancelFn = context.WithCancel(ctx)
	n.wg.Add(n.Cfg.NumWorkers)
	for i := 0; i < n.Cfg.NumWorkers; i++ {
		cfg := *n.Cfg
		cfg.Name = fmt.Sprintf("%s-%d", cfg.Name, i)
		go n.worker(ctx, i, &cfg)
	}

	go func() {
		<-ctx.Done()
		n.Close()
	}()
	return nil
}

func (n *NatsOutput) setDefaults() error {
	if n.Cfg.ConnectTimeWait <= 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Subject == "" && n.Cfg.SubjectPrefix == "" {
		n.Cfg.Subject = defaultSubjectName
	}
	if n.Cfg.Format == "" {
		n.Cfg.Format = defaultFormat
	}
	if !(n.Cfg.Format == "event" || n.Cfg.Format == "protojson" || n.Cfg.Format == "proto" || n.Cfg.Format == "json") {
		return fmt.Errorf("unsupported output format '%s' for output type NATS", n.Cfg.Format)
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if n.Cfg.NumWorkers <= 0 {
		n.Cfg.NumWorkers = defaultNumWorkers
	}
	if n.Cfg.WriteTimeout <= 0 {
		n.Cfg.WriteTimeout = defaultWriteTimeout
	}
	return nil
}

// Write //
func (n *NatsOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil || n.mo == nil {
		return
	}
	select {
	case <-ctx.Done():
		return
	case n.msgChan <- &protoMsg{m: rsp, meta: meta}:
	case <-time.After(n.Cfg.WriteTimeout):
		if n.Cfg.Debug {
			n.logger.Printf("writing expired after %s, NATS output might not be initialized", n.Cfg.WriteTimeout)
		}
		NatsNumberOfFailSendMsgs.WithLabelValues(n.Cfg.Name, "timeout").Inc()
		return
	}
}

// Close //
func (n *NatsOutput) Close() error {
	//	n.conn.Close()
	n.cancelFn()
	n.wg.Wait()
	return nil
}

// Metrics //
func (n *NatsOutput) Metrics() []prometheus.Collector {
	return []prometheus.Collector{
		NatsNumberOfSentMsgs,
		NatsNumberOfSentBytes,
		NatsNumberOfFailSendMsgs,
		NatsSendDuration,
	}
}

func (n *NatsOutput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(n.Cfg.ConnectTimeWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsOutput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Printf("attempting to connect to %s", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}

func (n *NatsOutput) worker(ctx context.Context, i int, cfg *Config) {
	defer n.wg.Done()
	var natsConn *nats.Conn
	var err error
	workerLogPrefix := fmt.Sprintf("worker-%d", i)
	n.logger.Printf("%s starting", workerLogPrefix)
CRCONN:
	natsConn, err = n.createNATSConn(cfg)
	if err != nil {
		n.logger.Printf("%s failed to create connection: %v", workerLogPrefix, err)
		time.Sleep(10 * time.Second)
		goto CRCONN
	}
	defer natsConn.Close()
	n.logger.Printf("%s initialized nats producer: %+v", workerLogPrefix, cfg)
	for {
		select {
		case <-ctx.Done():
			n.logger.Printf("%s flushing", workerLogPrefix)
			natsConn.FlushTimeout(time.Second)
			n.logger.Printf("%s shutting down", workerLogPrefix)
			return
		case m := <-n.msgChan:
			b, err := n.mo.Marshal(m.m, m.meta, n.evps...)
			if err != nil {
				if n.Cfg.Debug {
					n.logger.Printf("%s failed marshaling proto msg: %v", workerLogPrefix, err)
				}
				NatsNumberOfFailSendMsgs.WithLabelValues(cfg.Name, "marshal_error").Inc()
				continue
			}
			subject := n.subjectName(cfg, m.meta)
			start := time.Now()
			err = natsConn.Publish(subject, b)
			if err != nil {
				if n.Cfg.Debug {
					n.logger.Printf("%s failed to write to nats subject '%s': %v", workerLogPrefix, subject, err)
				}
				NatsNumberOfFailSendMsgs.WithLabelValues(cfg.Name, "publish_error").Inc()
				natsConn.Close()
				time.Sleep(10 * time.Second)
				goto CRCONN
			}
			NatsSendDuration.WithLabelValues(cfg.Name).Set(float64(time.Since(start).Nanoseconds()))
			NatsNumberOfSentMsgs.WithLabelValues(cfg.Name, subject).Inc()
			NatsNumberOfSentBytes.WithLabelValues(cfg.Name, subject).Add(float64(len(b)))
		}
	}
}

func (n *NatsOutput) subjectName(c *Config, meta outputs.Meta) string {
	ssb := strings.Builder{}
	ssb.WriteString(c.SubjectPrefix)
	if n.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(s, ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(subname)
		}
	} else if n.Cfg.Subject != "" {
		ssb.WriteString(n.Cfg.Subject)
	}
	return strings.ReplaceAll(ssb.String(), " ", "_")
}
<commit_msg>rework subjectname<commit_after>package nats_output

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/nats-io/nats.go"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	natsConnectWait         = 2 * time.Second
	natsReconnectBufferSize = 100 * 1024 * 1024
	defaultSubjectName      = "gnmic-telemetry"
	defaultFormat           = "json"
	defaultNumWorkers       = 1
	defaultWriteTimeout     = 10 * time.Second
)

func init() {
	outputs.Register("nats", func() outputs.Output {
		return &NatsOutput{
			Cfg: &Config{},
			wg:  new(sync.WaitGroup),
		}
	})
}

type protoMsg struct {
	m    proto.Message
	meta outputs.Meta
}

// NatsOutput //
type NatsOutput struct {
	Cfg      *Config
	ctx      context.Context
	cancelFn context.CancelFunc
	msgChan  chan *protoMsg
	wg       *sync.WaitGroup
	logger   *log.Logger
	mo       *formatters.MarshalOptions
	evps     []formatters.EventProcessor
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	SubjectPrefix   string        `mapstructure:"subject-prefix,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
	Format          string        `mapstructure:"format,omitempty"`
	NumWorkers      int           `mapstructure:"num-workers,omitempty"`
	WriteTimeout    time.Duration `mapstructure:"write-timeout,omitempty"`
	Debug           bool          `mapstructure:"debug,omitempty"`
	EventProcessors []string      `mapstructure:"event-processors,omitempty"`
}

func (n *NatsOutput) String() string {
	b, err := json.Marshal(n)
	if err != nil {
		return ""
	}
	return string(b)
}
func (n *NatsOutput) SetLogger(logger *log.Logger) {
	if logger != nil {
		n.logger = log.New(logger.Writer(), "nats_output ", logger.Flags())
		return
	}
	n.logger = log.New(os.Stderr, "nats_output ", log.LstdFlags|log.Lmicroseconds)
}

func (n *NatsOutput) SetEventProcessors(ps map[string]map[string]interface{}, log *log.Logger) {
	for _, epName := range n.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], log)
				if err != nil {
					n.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				n.evps = append(n.evps, ep)
				n.logger.Printf("added event processor '%s' of type=%s to nats output", epName, epType)
			}
		}
	}
}

// Init //
func (n *NatsOutput) Init(ctx context.Context, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, n.Cfg)
	if err != nil {
		return err
	}
	err = n.setDefaults()
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(n)
	}
	n.msgChan = make(chan *protoMsg)
	initMetrics()
	n.mo = &formatters.MarshalOptions{Format: n.Cfg.Format}
	n.ctx, n.cancelFn = context.WithCancel(ctx)
	n.wg.Add(n.Cfg.NumWorkers)
	for i := 0; i < n.Cfg.NumWorkers; i++ {
		cfg := *n.Cfg
		cfg.Name = fmt.Sprintf("%s-%d", cfg.Name, i)
		go n.worker(ctx, i, &cfg)
	}

	go func() {
		<-ctx.Done()
		n.Close()
	}()
	return nil
}

func (n *NatsOutput) setDefaults() error {
	if n.Cfg.ConnectTimeWait <= 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Subject == "" && n.Cfg.SubjectPrefix == "" {
		n.Cfg.Subject = defaultSubjectName
	}
	if n.Cfg.Format == "" {
		n.Cfg.Format = defaultFormat
	}
	if !(n.Cfg.Format == "event" || n.Cfg.Format == "protojson" || n.Cfg.Format == "proto" || n.Cfg.Format == "json") {
		return fmt.Errorf("unsupported output format '%s' for output type NATS", n.Cfg.Format)
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if n.Cfg.NumWorkers <= 0 {
		n.Cfg.NumWorkers = defaultNumWorkers
	}
	if n.Cfg.WriteTimeout <= 0 {
		n.Cfg.WriteTimeout = defaultWriteTimeout
	}
	return nil
}

// Write //
func (n *NatsOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil || n.mo == nil {
		return
	}
	select {
	case <-ctx.Done():
		return
	case n.msgChan <- &protoMsg{m: rsp, meta: meta}:
	case <-time.After(n.Cfg.WriteTimeout):
		if n.Cfg.Debug {
			n.logger.Printf("writing expired after %s, NATS output might not be initialized", n.Cfg.WriteTimeout)
		}
		NatsNumberOfFailSendMsgs.WithLabelValues(n.Cfg.Name, "timeout").Inc()
		return
	}
}

// Close //
func (n *NatsOutput) Close() error {
	//	n.conn.Close()
	n.cancelFn()
	n.wg.Wait()
	return nil
}

// Metrics //
func (n *NatsOutput) Metrics() []prometheus.Collector {
	return []prometheus.Collector{
		NatsNumberOfSentMsgs,
		NatsNumberOfSentBytes,
		NatsNumberOfFailSendMsgs,
		NatsSendDuration,
	}
}

func (n *NatsOutput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(n.Cfg.ConnectTimeWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsOutput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Printf("attempting to connect to %s", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}

func (n *NatsOutput) worker(ctx context.Context, i int, cfg *Config) {
	defer n.wg.Done()
	var natsConn *nats.Conn
	var err error
	workerLogPrefix := fmt.Sprintf("worker-%d", i)
	n.logger.Printf("%s starting", workerLogPrefix)
CRCONN:
	natsConn, err = n.createNATSConn(cfg)
	if err != nil {
		n.logger.Printf("%s failed to create connection: %v", workerLogPrefix, err)
		time.Sleep(10 * time.Second)
		goto CRCONN
	}
	defer natsConn.Close()
	n.logger.Printf("%s initialized nats producer: %+v", workerLogPrefix, cfg)
	for {
		select {
		case <-ctx.Done():
			n.logger.Printf("%s flushing", workerLogPrefix)
			natsConn.FlushTimeout(time.Second)
			n.logger.Printf("%s shutting down", workerLogPrefix)
			return
		case m := <-n.msgChan:
			b, err := n.mo.Marshal(m.m, m.meta, n.evps...)
			if err != nil {
				if n.Cfg.Debug {
					n.logger.Printf("%s failed marshaling proto msg: %v", workerLogPrefix, err)
				}
				NatsNumberOfFailSendMsgs.WithLabelValues(cfg.Name, "marshal_error").Inc()
				continue
			}
			subject := n.subjectName(cfg, m.meta)
			start := time.Now()
			err = natsConn.Publish(subject, b)
			if err != nil {
				if n.Cfg.Debug {
					n.logger.Printf("%s failed to write to nats subject '%s': %v", workerLogPrefix, subject, err)
				}
				NatsNumberOfFailSendMsgs.WithLabelValues(cfg.Name, "publish_error").Inc()
				natsConn.Close()
				time.Sleep(10 * time.Second)
				goto CRCONN
			}
			NatsSendDuration.WithLabelValues(cfg.Name).Set(float64(time.Since(start).Nanoseconds()))
			NatsNumberOfSentMsgs.WithLabelValues(cfg.Name, subject).Inc()
			NatsNumberOfSentBytes.WithLabelValues(cfg.Name, subject).Add(float64(len(b)))
		}
	}
}

func (n *NatsOutput) subjectName(c *Config, meta outputs.Meta) string {
	if c.SubjectPrefix != "" {
		ssb := strings.Builder{}
		ssb.WriteString(n.Cfg.SubjectPrefix)
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(s, ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(subname)
		}
		return strings.ReplaceAll(ssb.String(), " ", "_")
	}
	return strings.ReplaceAll(n.Cfg.Subject, " ", "_")
}
<|endoftext|>


<commit_before>package scanner

import (
	"bufio"
	"bytes"
	"github.com/github/git-media/pointer"
	"github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

var (
	blobSizeCutoff = 125
	stdoutBufSize  = 16384
	chanBufSize    = 100
)

type wrappedPointer struct {
	Sha1 string
	Name string
	*pointer.Pointer
}

// Scan takes a ref and returns a slice of pointer.Pointer objects
// for all git media pointers it finds for that ref.
func Scan(ref string) ([]*wrappedPointer, error) {
	nameMap := make(map[string]string, 0)
	start := time.Now()

	revs, err := revListShas(ref, ref == "", nameMap)
	if err != nil {
		return nil, err
	}

	smallShas, err := catFileBatchCheck(revs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if name, ok := nameMap[p.Sha1]; ok {
			p.Name = name
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan", start)

	return pointers, nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(ref string, all bool, nameMap map[string]string) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, "--no-walk")
		refArgs = append(refArgs, ref)
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			sha1 := line[0:40]
			if len(line) > 40 {
				nameMap[sha1] = line[41:len(line)]
			}
			revs <- sha1
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			// Format is:
			// <sha1> <type> <size>
			// type is at a fixed spot, if we see that it's "blob", we can avoid
			// splitting the line just to get the size.
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a git media pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer, chanBufSize)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil {
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{string(fields[0]), "", p}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {
				break
			}
		}
		close(pointers)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	tracerx.Printf("run_command: %s %s", command, strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<commit_msg>Give just a little more room for the size comparison<commit_after>package scanner

import (
	"bufio"
	"bytes"
	"github.com/github/git-media/pointer"
	"github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

var (
	blobSizeCutoff = 130
	stdoutBufSize  = 16384
	chanBufSize    = 100
)

type wrappedPointer struct {
	Sha1 string
	Name string
	*pointer.Pointer
}

// Scan takes a ref and returns a slice of pointer.Pointer objects
// for all git media pointers it finds for that ref.
func Scan(ref string) ([]*wrappedPointer, error) {
	nameMap := make(map[string]string, 0)
	start := time.Now()

	revs, err := revListShas(ref, ref == "", nameMap)
	if err != nil {
		return nil, err
	}

	smallShas, err := catFileBatchCheck(revs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if name, ok := nameMap[p.Sha1]; ok {
			p.Name = name
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan", start)

	return pointers, nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(ref string, all bool, nameMap map[string]string) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, "--no-walk")
		refArgs = append(refArgs, ref)
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			sha1 := line[0:40]
			if len(line) > 40 {
				nameMap[sha1] = line[41:len(line)]
			}
			revs <- sha1
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			// Format is:
			// <sha1> <type> <size>
			// type is at a fixed spot, if we see that it's "blob", we can avoid
			// splitting the line just to get the size.
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a git media pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer, chanBufSize)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil {
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{string(fields[0]), "", p}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {
				break
			}
		}
		close(pointers)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	tracerx.Printf("run_command: %s %s", command, strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<|endoftext|>


<commit_before>package scanner

import (
	"bufio"
	"bytes"
	"github.com/hawser/git-hawser/pointer"
	"github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	// blobSizeCutoff is used to determine which files to scan for hawser pointers.
	// Any file with a size below this cutoff will be scanned.
	blobSizeCutoff = 130

	// stdoutBufSize is the size of the buffers given to a sub-process stdout
	stdoutBufSize = 16384

	// chanBufSize is the size of the channels used to pass data from one sub-process
	// to another.
	chanBufSize = 100
)

// wrappedPointer wraps a pointer.Pointer and provides the git sha1
// and the file name associated with the object, taken from the
// rev-list output.
type wrappedPointer struct {
	Sha1    string
	Name    string
	SrcName string
	Size    int64
	Status  string
	*pointer.Pointer
}

// indexFile is used when scanning the index. It stores the name of
// the file, the status of the file in the index, and, in the case of
// a moved or copied file, the original name of the file.
type indexFile struct {
	Name    string
	SrcName string
	Status  string
}

var z40 = regexp.MustCompile(`\^?0{40}`)

// Scan takes a ref and returns a slice of wrappedPointer objects
// for all hawser pointers it finds for that ref.
func Scan(refLeft, refRight string) ([]*wrappedPointer, error) {
	nameMap := make(map[string]string, 0)
	start := time.Now()

	revs, err := revListShas(refLeft, refRight, refLeft == "", nameMap)
	if err != nil {
		return nil, err
	}

	smallShas, err := catFileBatchCheck(revs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if name, ok := nameMap[p.Sha1]; ok {
			p.Name = name
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan", start)

	return pointers, nil
}

// ScanIndex returns a slice of wrappedPointer objects for all
// hawser pointers it finds in the index.
func ScanIndex() ([]*wrappedPointer, error) {
	nameMap := make(map[string]*indexFile, 0)
	start := time.Now()

	revs, err := revListIndex(false, nameMap)
	if err != nil {
		return nil, err
	}

	cachedRevs, err := revListIndex(true, nameMap)
	if err != nil {
		return nil, err
	}

	allRevs := make(chan string)
	go func() {
		seenRevs := make(map[string]bool, 0)

		for rev := range revs {
			seenRevs[rev] = true
			allRevs <- rev
		}

		for rev := range cachedRevs {
			if _, ok := seenRevs[rev]; !ok {
				allRevs <- rev
			}
		}
		close(allRevs)
	}()

	smallShas, err := catFileBatchCheck(allRevs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if e, ok := nameMap[p.Sha1]; ok {
			p.Name = e.Name
			p.Status = e.Status
			p.SrcName = e.SrcName
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan-staging", start)

	return pointers, nil

}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(refLeft, refRight string, all bool, nameMap map[string]string) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, "--no-walk")
		refArgs = append(refArgs, refLeft)
		if refRight != "" && !z40.MatchString(refRight) {
			refArgs = append(refArgs, refRight)
		}
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if len(line) < 40 {
				continue
			}

			sha1 := line[0:40]
			if len(line) > 40 {
				nameMap[sha1] = line[41:len(line)]
			}
			revs <- sha1
		}
		close(revs)
	}()

	return revs, nil
}

// revListIndex uses git diff-index to return the list of object sha1s
// for in the indexf. It returns a channel from which sha1 strings can be read.
// The namMap will be filled indexFile pointers mapping sha1s to indexFiles.
func revListIndex(cache bool, nameMap map[string]*indexFile) (chan string, error) {
	cmdArgs := []string{"diff-index", "-M"}
	if cache {
		cmdArgs = append(cmdArgs, "--cached")
	}
	cmdArgs = append(cmdArgs, "HEAD")

	cmd, err := startCommand("git", cmdArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			// Format is:
			// :100644 100644 c5b3d83a7542255ec7856487baa5e83d65b1624c 9e82ac1b514be060945392291b5b3108c22f6fe3 M foo.gif
			// :<old mode> <new mode> <old sha1> <new sha1> <status>\t<file name>[\t<file name>]
			line := scanner.Text()
			parts := strings.Split(line, "\t")
			if len(parts) < 2 {
				continue
			}

			description := strings.Split(parts[0], " ")
			files := parts[1:len(parts)]

			if len(description) >= 5 {
				status := description[4][0:1]
				sha1 := description[3]
				if status == "M" {
					sha1 = description[2] // This one is modified but not added
				}
				nameMap[sha1] = &indexFile{files[len(files)-1], files[0], status}
				revs <- sha1
			}
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			// Format is:
			// <sha1> <type> <size>
			// type is at a fixed spot, if we see that it's "blob", we can avoid
			// splitting the line just to get the size.
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a hawser pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer, chanBufSize)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil {
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{
					Sha1:    string(fields[0]),
					Size:    p.Size,
					Pointer: p,
				}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {
				break
			}
		}
		close(pointers)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	tracerx.Printf("run_command: %s %s", command, strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<commit_msg>アーア アアアア アーアー<commit_after>package scanner

import (
	"bufio"
	"bytes"
	"github.com/hawser/git-hawser/pointer"
	"github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	// blobSizeCutoff is used to determine which files to scan for hawser pointers.
	// Any file with a size below this cutoff will be scanned.
	blobSizeCutoff = 140

	// stdoutBufSize is the size of the buffers given to a sub-process stdout
	stdoutBufSize = 16384

	// chanBufSize is the size of the channels used to pass data from one sub-process
	// to another.
	chanBufSize = 100
)

// wrappedPointer wraps a pointer.Pointer and provides the git sha1
// and the file name associated with the object, taken from the
// rev-list output.
type wrappedPointer struct {
	Sha1    string
	Name    string
	SrcName string
	Size    int64
	Status  string
	*pointer.Pointer
}

// indexFile is used when scanning the index. It stores the name of
// the file, the status of the file in the index, and, in the case of
// a moved or copied file, the original name of the file.
type indexFile struct {
	Name    string
	SrcName string
	Status  string
}

var z40 = regexp.MustCompile(`\^?0{40}`)

// Scan takes a ref and returns a slice of wrappedPointer objects
// for all hawser pointers it finds for that ref.
func Scan(refLeft, refRight string) ([]*wrappedPointer, error) {
	nameMap := make(map[string]string, 0)
	start := time.Now()

	revs, err := revListShas(refLeft, refRight, refLeft == "", nameMap)
	if err != nil {
		return nil, err
	}

	smallShas, err := catFileBatchCheck(revs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if name, ok := nameMap[p.Sha1]; ok {
			p.Name = name
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan", start)

	return pointers, nil
}

// ScanIndex returns a slice of wrappedPointer objects for all
// hawser pointers it finds in the index.
func ScanIndex() ([]*wrappedPointer, error) {
	nameMap := make(map[string]*indexFile, 0)
	start := time.Now()

	revs, err := revListIndex(false, nameMap)
	if err != nil {
		return nil, err
	}

	cachedRevs, err := revListIndex(true, nameMap)
	if err != nil {
		return nil, err
	}

	allRevs := make(chan string)
	go func() {
		seenRevs := make(map[string]bool, 0)

		for rev := range revs {
			seenRevs[rev] = true
			allRevs <- rev
		}

		for rev := range cachedRevs {
			if _, ok := seenRevs[rev]; !ok {
				allRevs <- rev
			}
		}
		close(allRevs)
	}()

	smallShas, err := catFileBatchCheck(allRevs)
	if err != nil {
		return nil, err
	}

	pointerc, err := catFileBatch(smallShas)
	if err != nil {
		return nil, err
	}

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		if e, ok := nameMap[p.Sha1]; ok {
			p.Name = e.Name
			p.Status = e.Status
			p.SrcName = e.SrcName
		}
		pointers = append(pointers, p)
	}

	tracerx.PerformanceSince("scan-staging", start)

	return pointers, nil

}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(refLeft, refRight string, all bool, nameMap map[string]string) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, "--no-walk")
		refArgs = append(refArgs, refLeft)
		if refRight != "" && !z40.MatchString(refRight) {
			refArgs = append(refArgs, refRight)
		}
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if len(line) < 40 {
				continue
			}

			sha1 := line[0:40]
			if len(line) > 40 {
				nameMap[sha1] = line[41:len(line)]
			}
			revs <- sha1
		}
		close(revs)
	}()

	return revs, nil
}

// revListIndex uses git diff-index to return the list of object sha1s
// for in the indexf. It returns a channel from which sha1 strings can be read.
// The namMap will be filled indexFile pointers mapping sha1s to indexFiles.
func revListIndex(cache bool, nameMap map[string]*indexFile) (chan string, error) {
	cmdArgs := []string{"diff-index", "-M"}
	if cache {
		cmdArgs = append(cmdArgs, "--cached")
	}
	cmdArgs = append(cmdArgs, "HEAD")

	cmd, err := startCommand("git", cmdArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			// Format is:
			// :100644 100644 c5b3d83a7542255ec7856487baa5e83d65b1624c 9e82ac1b514be060945392291b5b3108c22f6fe3 M foo.gif
			// :<old mode> <new mode> <old sha1> <new sha1> <status>\t<file name>[\t<file name>]
			line := scanner.Text()
			parts := strings.Split(line, "\t")
			if len(parts) < 2 {
				continue
			}

			description := strings.Split(parts[0], " ")
			files := parts[1:len(parts)]

			if len(description) >= 5 {
				status := description[4][0:1]
				sha1 := description[3]
				if status == "M" {
					sha1 = description[2] // This one is modified but not added
				}
				nameMap[sha1] = &indexFile{files[len(files)-1], files[0], status}
				revs <- sha1
			}
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string, chanBufSize)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			// Format is:
			// <sha1> <type> <size>
			// type is at a fixed spot, if we see that it's "blob", we can avoid
			// splitting the line just to get the size.
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a hawser pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer, chanBufSize)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil {
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{
					Sha1:    string(fields[0]),
					Size:    p.Size,
					Pointer: p,
				}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {
				break
			}
		}
		close(pointers)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	tracerx.Printf("run_command: %s %s", command, strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<|endoftext|>


<commit_before>package controllers

type AuthController struct {
	AB backends.AuthBackend
}

func GetNewAuthenticationController(authBackend backends.AuthBackend) *AuthController {
	return &AuthController{AB: authBackend}
}

func (a *AuthController) Login(w http.ResponseWriter, r *http.Request) {
	requestUser := new(authentication.User)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(&requestUser)

	responseStatus, token := authentication.Login(requestUser, a.AB)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(responseStatus)
	w.Write(token)
}

func (a *AuthController) RefreshToken(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	requestUser := new(authentication.User)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(&requestUser)

	w.Header().Set("Content-Type", "application/json")
	w.Write(authentication.RefreshToken(requestUser, a.AB))
}

<commit_msg>logout controller<commit_after>package controllers

import (
	"encoding/json"
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/backends"
	"net/http"
)

type AuthController struct {
	AB backends.AuthBackend
}

func GetNewAuthenticationController(authBackend backends.AuthBackend) *AuthController {
	return &AuthController{AB: authBackend}
}

func (a *AuthController) Login(w http.ResponseWriter, r *http.Request) {
	requestUser := new(authentication.User)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(&requestUser)

	responseStatus, token := authentication.Login(requestUser, a.AB)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(responseStatus)
	w.Write(token)
}

func (a *AuthController) RefreshToken(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	requestUser := new(authentication.User)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(&requestUser)

	w.Header().Set("Content-Type", "application/json")
	w.Write(authentication.RefreshToken(requestUser, a.AB))
}

func (a *AuthController) Logout(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	err := authentication.Logout(r, a.AB)
	w.Header().Set("Content-Type", "application/json")
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	} else {
		w.WriteHeader(http.StatusOK)
	}
}
<|endoftext|>


<commit_before>package brightbox

import (
	"fmt"
	"log"
	"strings"

	"github.com/brightbox/gobrightbox"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceBrightboxFirewallRule() *schema.Resource {
	return &schema.Resource{
		Create: resourceBrightboxFirewallRuleCreate,
		Read:   resourceBrightboxFirewallRuleRead,
		Update: resourceBrightboxFirewallRuleUpdate,
		Delete: resourceBrightboxFirewallRuleDelete,

		Schema: map[string]*schema.Schema{
			"firewall_policy": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"protocol": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"source": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"source_port": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"destination": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"destination_port": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"icmp_type_name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
		},
	}
}

func resourceBrightboxFirewallRuleCreate(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	log.Printf("[INFO] Creating Firewall Rule")
	firewall_rule_opts := &brightbox.FirewallRuleOptions{
		FirewallPolicy: d.Get("firewall_policy").(string),
	}
	err := addUpdateableFirewallRuleOptions(d, firewall_rule_opts)
	if err != nil {
		return err
	}

	log.Printf("[INFO] Firewall Rule create configuration: %#v", firewall_rule_opts)

	firewall_rule, err := client.CreateFirewallRule(firewall_rule_opts)
	if err != nil {
		return fmt.Errorf("Error creating Firewall Rule: %s", err)
	}

	d.SetId(firewall_rule.Id)

	setFirewallRuleAttributes(d, firewall_rule)

	return nil
}

func setFirewallRuleAttributes(
	d *schema.ResourceData,
	firewall_rule *brightbox.FirewallRule,
) {
	d.Set("firewall_policy", firewall_rule.FirewallPolicy)
	d.Set("protocol", firewall_rule.Protocol)
	d.Set("source", firewall_rule.Source)
	d.Set("source_port", firewall_rule.SourcePort)
	d.Set("destination", firewall_rule.Destination)
	d.Set("destination_port", firewall_rule.DestinationPort)
	d.Set("icmp_type_name", firewall_rule.IcmpTypeName)
	d.Set("description", firewall_rule.Description)
}

func resourceBrightboxFirewallRuleRead(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	firewall_rule, err := client.FirewallRule(d.Id())
	if err != nil {
		if strings.HasPrefix(err.Error(), "missing_resource:") {
			log.Printf("[WARN] Firewall Rule not found, removing from state: %s", d.Id())
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving Firewall Rule details: %s", err)
	}

	setFirewallRuleAttributes(d, firewall_rule)

	return nil
}

func resourceBrightboxFirewallRuleDelete(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	log.Printf("[INFO] Deleting Firewall Rule %s", d.Id())
	err := client.DestroyFirewallRule(d.Id())
	if err != nil {
		return fmt.Errorf("Error deleting Firewall Rule (%s): %s", d.Id(), err)
	}
	return nil
}

func resourceBrightboxFirewallRuleUpdate(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	firewall_rule_opts := &brightbox.FirewallRuleOptions{
		Id: d.Id(),
	}
	err := addUpdateableFirewallRuleOptions(d, firewall_rule_opts)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Firewall Rule update configuration: %#v", firewall_rule_opts)

	firewall_rule, err := client.UpdateFirewallRule(firewall_rule_opts)
	if err != nil {
		return fmt.Errorf("Error updating Firewall Rule (%s): %s", firewall_rule_opts.Id, err)
	}

	setFirewallRuleAttributes(d, firewall_rule)
	return nil
}

func addUpdateableFirewallRuleOptions(
	d *schema.ResourceData,
	opts *brightbox.FirewallRuleOptions,
) error {
	assign_string(d, &opts.Protocol, "protocol")
	assign_string(d, &opts.Source, "source")
	assign_string(d, &opts.SourcePort, "source_port")
	assign_string(d, &opts.Destination, "destination")
	assign_string(d, &opts.DestinationPort, "destination_port")
	assign_string(d, &opts.IcmpTypeName, "icmp_type_name")
	assign_string(d, &opts.Description, "description")
	return nil
}
<commit_msg>Refactor Firewall Rule<commit_after>package brightbox

import (
	"fmt"
	"log"
	"strings"

	"github.com/brightbox/gobrightbox"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceBrightboxFirewallRule() *schema.Resource {
	return &schema.Resource{
		Create: resourceBrightboxFirewallRuleCreate,
		Read:   resourceBrightboxFirewallRuleRead,
		Update: resourceBrightboxFirewallRuleUpdate,
		Delete: resourceBrightboxFirewallRuleDelete,

		Schema: map[string]*schema.Schema{
			"firewall_policy": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"protocol": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"source": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"source_port": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"destination": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"destination_port": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"icmp_type_name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
		},
	}
}

func resourceBrightboxFirewallRuleCreate(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	log.Printf("[INFO] Creating Firewall Rule")
	firewall_rule_opts := &brightbox.FirewallRuleOptions{
		FirewallPolicy: d.Get("firewall_policy").(string),
	}
	err := addUpdateableFirewallRuleOptions(d, firewall_rule_opts)
	if err != nil {
		return err
	}

	log.Printf("[INFO] Firewall Rule create configuration: %#v", firewall_rule_opts)

	firewall_rule, err := client.CreateFirewallRule(firewall_rule_opts)
	if err != nil {
		return fmt.Errorf("Error creating Firewall Rule: %s", err)
	}

	d.SetId(firewall_rule.Id)

	return setFirewallRuleAttributes(d, firewall_rule)
}

func resourceBrightboxFirewallRuleRead(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	firewall_rule, err := client.FirewallRule(d.Id())
	if err != nil {
		if strings.HasPrefix(err.Error(), "missing_resource:") {
			log.Printf("[WARN] Firewall Rule not found, removing from state: %s", d.Id())
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving Firewall Rule details: %s", err)
	}

	return setFirewallRuleAttributes(d, firewall_rule)
}

func resourceBrightboxFirewallRuleDelete(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	log.Printf("[INFO] Deleting Firewall Rule %s", d.Id())
	err := client.DestroyFirewallRule(d.Id())
	if err != nil {
		return fmt.Errorf("Error deleting Firewall Rule (%s): %s", d.Id(), err)
	}
	return nil
}

func resourceBrightboxFirewallRuleUpdate(
	d *schema.ResourceData,
	meta interface{},
) error {
	client := meta.(*CompositeClient).ApiClient

	firewall_rule_opts := &brightbox.FirewallRuleOptions{
		Id: d.Id(),
	}
	err := addUpdateableFirewallRuleOptions(d, firewall_rule_opts)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Firewall Rule update configuration: %#v", firewall_rule_opts)

	firewall_rule, err := client.UpdateFirewallRule(firewall_rule_opts)
	if err != nil {
		return fmt.Errorf("Error updating Firewall Rule (%s): %s", firewall_rule_opts.Id, err)
	}

	return setFirewallRuleAttributes(d, firewall_rule)
}

func addUpdateableFirewallRuleOptions(
	d *schema.ResourceData,
	opts *brightbox.FirewallRuleOptions,
) error {
	assign_string(d, &opts.Protocol, "protocol")
	assign_string(d, &opts.Source, "source")
	assign_string(d, &opts.SourcePort, "source_port")
	assign_string(d, &opts.Destination, "destination")
	assign_string(d, &opts.DestinationPort, "destination_port")
	assign_string(d, &opts.IcmpTypeName, "icmp_type_name")
	assign_string(d, &opts.Description, "description")
	return nil
}

func setFirewallRuleAttributes(
	d *schema.ResourceData,
	firewall_rule *brightbox.FirewallRule,
) error {
	d.Set("firewall_policy", firewall_rule.FirewallPolicy)
	d.Set("protocol", firewall_rule.Protocol)
	d.Set("source", firewall_rule.Source)
	d.Set("source_port", firewall_rule.SourcePort)
	d.Set("destination", firewall_rule.Destination)
	d.Set("destination_port", firewall_rule.DestinationPort)
	d.Set("icmp_type_name", firewall_rule.IcmpTypeName)
	d.Set("description", firewall_rule.Description)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"time"

	"cloud.google.com/go/pubsub"
	cloudtrace "cloud.google.com/go/trace"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager"
	"github.com/google/go-github/github"
	flags "github.com/jessevdk/go-flags"
	"github.com/pivotal-cf/paraphernalia/operate/admin"
	"github.com/pivotal-cf/paraphernalia/secure/tlsconfig"
	"github.com/pivotal-cf/paraphernalia/serve/grpcrunner"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
	"golang.org/x/oauth2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"cred-alert/config"
	"cred-alert/crypto"
	"cred-alert/db"
	"cred-alert/db/migrations"
	"cred-alert/gitclient"
	"cred-alert/metrics"
	"cred-alert/notifications"
	"cred-alert/queue"
	"cred-alert/revok"
	"cred-alert/revok/stats"
	"cred-alert/revokpb"
	"cred-alert/search"
	"cred-alert/sniff"
	"rolodex/rolodexpb"
)

var info = admin.ServiceInfo{
	Name:        "revok",
	Description: "A service which fetches new Git commits and scans them for credentials.",
	Team:        "PCF Security Enablement",
}

func main() {
	var cfg *config.WorkerConfig
	var flagOpts config.WorkerOpts

	logger := lager.NewLogger("revok-worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.DEBUG))

	logger.Info("starting")

	_, err := flags.Parse(&flagOpts)
	if err != nil {
		os.Exit(1)
	}

	bs, err := ioutil.ReadFile(string(flagOpts.ConfigFile))
	if err != nil {
		logger.Error("failed-to-open-config-file", err)
		os.Exit(1)
	}

	cfg, err = config.LoadWorkerConfig(bs)
	if err != nil {
		logger.Error("failed-to-load-config-file", err)
		os.Exit(1)
	}

	errs := cfg.Validate()
	if errs != nil {
		for _, err := range errs {
			fmt.Println(err.Error())
		}
		os.Exit(1)
	}

	if cfg.Metrics.SentryDSN != "" {
		logger.RegisterSink(revok.NewSentrySink(cfg.Metrics.SentryDSN, cfg.Metrics.Environment))
	}

	workdir := cfg.WorkDir
	_, err = os.Lstat(workdir)
	if err != nil {
		log.Fatalf("workdir error: %s", err)
	}

	dbCertificate, dbCaCertPool := loadCerts(
		cfg.MySQL.CertificatePath,
		cfg.MySQL.PrivateKeyPath,
		cfg.MySQL.PrivateKeyPassphrase,
		cfg.MySQL.CACertificatePath,
	)

	dbURI := db.NewDSN(
		cfg.MySQL.Username,
		cfg.MySQL.Password,
		cfg.MySQL.DBName,
		cfg.MySQL.Hostname,
		int(cfg.MySQL.Port),
		cfg.MySQL.ServerName,
		dbCertificate,
		dbCaCertPool,
	)

	database, err := migrations.LockDBAndMigrate(logger, "mysql", dbURI)
	if err != nil {
		log.Fatalf("db error: %s", err)
	}

	database.LogMode(false)

	clk := clock.NewClock()

	cloneMsgCh := make(chan revok.CloneMsg)

	scanRepository := db.NewScanRepository(database, clk)
	repositoryRepository := db.NewRepositoryRepository(database)
	fetchRepository := db.NewFetchRepository(database)
	credentialRepository := db.NewCredentialRepository(database)
	branchRepository := db.NewBranchRepository(database)

	emitter := metrics.BuildEmitter(cfg.Metrics.DatadogAPIKey, cfg.Metrics.Environment)
	gitClient := gitclient.New(cfg.GitHub.PrivateKeyPath, cfg.GitHub.PublicKeyPath)
	repoWhitelist := notifications.BuildWhitelist(cfg.Whitelist...)
	formatter := notifications.NewSlackNotificationFormatter()

	traceClient, err := cloudtrace.NewClient(context.Background(), cfg.Trace.ProjectName)
	if err != nil {
		logger.Error("failed-to-create-trace-client", err)
	}

	slackHTTPClient := &http.Client{
		Timeout: 3 * time.Second,
	}
	notifier := notifications.NewSlackNotifier(clk, slackHTTPClient, formatter)

	certificate, caCertPool := loadCerts(
		cfg.Identity.CertificatePath,
		cfg.Identity.PrivateKeyPath,
		cfg.Identity.PrivateKeyPassphrase,
		cfg.Identity.CACertificatePath,
	)

	rolodexServerAddr := fmt.Sprintf("%s:%d", cfg.Rolodex.ServerAddress, cfg.Rolodex.ServerPort)

	tlsConfig := tlsconfig.Build(
		tlsconfig.WithPivotalDefaults(),
		tlsconfig.WithIdentity(certificate),
	)

	transportCreds := credentials.NewTLS(tlsConfig.Client(tlsconfig.WithAuthority(caCertPool)))

	conn, err := grpc.Dial(
		rolodexServerAddr,
		grpc.WithDialer(keepAliveDial),
		grpc.WithTransportCredentials(transportCreds),
		grpc.WithUnaryInterceptor(cloudtrace.GRPCClientInterceptor()),
	)

	rolodexClient := rolodexpb.NewRolodexClient(conn)

	teamURLs := notifications.NewTeamURLs(
		cfg.Slack.DefaultURL,
		cfg.Slack.DefaultChannel,
		cfg.Slack.TeamURLs,
	)

	addressBook := notifications.NewRolodex(
		rolodexClient,
		teamURLs,
	)

	router := notifications.NewRouter(
		notifier,
		addressBook,
		repoWhitelist,
	)

	sniffer := sniff.NewDefaultSniffer()
	scanner := revok.NewScanner(
		gitClient,
		repositoryRepository,
		scanRepository,
		credentialRepository,
		sniffer,
	)

	notificationComposer := revok.NewNotificationComposer(
		repositoryRepository,
		router,
		scanner,
	)

	changeFetcher := revok.NewChangeFetcher(
		logger,
		gitClient,
		notificationComposer,
		repositoryRepository,
		fetchRepository,
		emitter,
	)

	changeScheduleRunner := revok.NewScheduleRunner(logger)

	changeScheduler := revok.NewChangeScheduler(
		logger,
		repositoryRepository,
		changeScheduleRunner,
		changeFetcher,
	)

	cloner := revok.NewCloner(
		logger,
		workdir,
		cloneMsgCh,
		gitClient,
		repositoryRepository,
		notificationComposer,
		emitter,
		changeScheduler,
	)

	dirscanUpdater := revok.NewRescanner(
		logger,
		scanRepository,
		credentialRepository,
		scanner,
		router,
		emitter,
	)

	statsReporter := stats.NewReporter(
		logger,
		clk,
		60*time.Second,
		db.NewStatsRepository(database),
		emitter,
	)

	headCredentialCounter := revok.NewHeadCredentialCounter(
		logger,
		branchRepository,
		repositoryRepository,
		clk,
		cfg.CredentialCounterInterval,
		gitClient,
		sniffer,
	)

	debug := admin.Runner(
		"6060",
		admin.WithInfo(info),
		admin.WithUptime(),
	)

	members := []grouper.Member{
		{Name: "cloner", Runner: cloner},
		{Name: "dirscan-updater", Runner: dirscanUpdater},
		{Name: "stats-reporter", Runner: statsReporter},
		{Name: "head-credential-counter", Runner: headCredentialCounter},
		{Name: "change-schedule-runner", Runner: changeScheduleRunner},
		{Name: "debug", Runner: debug},
	}

	looper := gitclient.NewLooper()
	searcher := search.NewSearcher(repositoryRepository, looper)
	handler := revok.NewServer(logger, searcher, repositoryRepository, branchRepository)

	serverTls := tlsConfig.Server(tlsconfig.WithClientAuthentication(caCertPool))

	grpcServer := grpcrunner.New(
		logger,
		fmt.Sprintf("%s:%d", cfg.API.BindIP, cfg.API.BindPort),
		func(server *grpc.Server) {
			revokpb.RegisterRevokServer(server, handler)
		},
		grpc.Creds(credentials.NewTLS(serverTls)),
	)

	members = append(members, grouper.Member{
		Name:   "grpc-server",
		Runner: grpcServer,
	})

	pubSubClient, err := pubsub.NewClient(context.Background(), cfg.PubSub.ProjectName)
	if err != nil {
		logger.Fatal("failed", err)
		os.Exit(1)
	}

	subscription := pubSubClient.Subscription(cfg.PubSub.FetchHint.Subscription)

	publicKey, err := crypto.ReadRSAPublicKey(cfg.PubSub.PublicKeyPath)
	if err != nil {
		logger.Fatal("failed", err)
		os.Exit(1)
	}

	pushEventProcessor := queue.NewPushEventProcessor(
		changeFetcher,
		emitter,
		clk,
		traceClient,
	)

	signatureChecker := queue.NewSignatureCheck(crypto.NewRSAVerifier(publicKey), emitter, pushEventProcessor)

	members = append(members, grouper.Member{
		Name:   "github-hint-handler",
		Runner: queue.NewPubSubSubscriber(logger, subscription, signatureChecker, emitter),
	})

	if cfg.GitHub.AccessToken != "" {
		githubHTTPClient := &http.Client{
			Timeout: 30 * time.Second,
			Transport: &oauth2.Transport{
				Source: oauth2.StaticTokenSource(
					&oauth2.Token{AccessToken: cfg.GitHub.AccessToken},
				),
				Base: &http.Transport{
					DisableKeepAlives: true,
				},
			},
		}

		ghClient := revok.NewGitHubClient(github.NewClient(githubHTTPClient))

		repoDiscoverer := revok.NewRepoDiscoverer(
			logger,
			workdir,
			cloneMsgCh,
			ghClient,
			clk,
			cfg.RepositoryDiscovery.Interval,
			cfg.RepositoryDiscovery.Organizations,
			cfg.RepositoryDiscovery.Users,
			repositoryRepository,
		)

		members = append(members, grouper.Member{
			Name:   "repo-discoverer",
			Runner: repoDiscoverer,
		})
	}

	startupTasks := []grouper.Member{
		{
			Name:   "schedule-fetches",
			Runner: changeScheduler,
		},
	}

	system := []grouper.Member{
		{
			Name:   "servers",
			Runner: grouper.NewParallel(os.Interrupt, members),
		},
		{
			Name:   "startup-tasks",
			Runner: grouper.NewParallel(os.Interrupt, startupTasks),
		},
	}

	runner := sigmon.New(grouper.NewOrdered(os.Interrupt, system))

	err = <-ifrit.Invoke(runner).Wait()
	if err != nil {
		log.Fatalf("failed-to-start: %s", err)
	}
}

func loadCerts(certificatePath, privateKeyPath, privateKeyPassphrase, caCertificatePath string) (tls.Certificate, *x509.CertPool) {
	certificate, err := config.LoadCertificate(
		certificatePath,
		privateKeyPath,
		privateKeyPassphrase,
	)
	if err != nil {
		log.Fatalln(err)
	}

	caCertPool, err := config.LoadCertificatePool(caCertificatePath)
	if err != nil {
		log.Fatalln(err)
	}

	return certificate, caCertPool
}

func keepAliveDial(addr string, timeout time.Duration) (net.Conn, error) {
	d := net.Dialer{
		Timeout:   timeout,
		KeepAlive: 60 * time.Second,
	}
	return d.Dial("tcp", addr)
}
<commit_msg>bump slack sending timeout<commit_after>package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"time"

	"cloud.google.com/go/pubsub"
	cloudtrace "cloud.google.com/go/trace"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager"
	"github.com/google/go-github/github"
	flags "github.com/jessevdk/go-flags"
	"github.com/pivotal-cf/paraphernalia/operate/admin"
	"github.com/pivotal-cf/paraphernalia/secure/tlsconfig"
	"github.com/pivotal-cf/paraphernalia/serve/grpcrunner"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
	"golang.org/x/oauth2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"cred-alert/config"
	"cred-alert/crypto"
	"cred-alert/db"
	"cred-alert/db/migrations"
	"cred-alert/gitclient"
	"cred-alert/metrics"
	"cred-alert/notifications"
	"cred-alert/queue"
	"cred-alert/revok"
	"cred-alert/revok/stats"
	"cred-alert/revokpb"
	"cred-alert/search"
	"cred-alert/sniff"
	"rolodex/rolodexpb"
)

var info = admin.ServiceInfo{
	Name:        "revok",
	Description: "A service which fetches new Git commits and scans them for credentials.",
	Team:        "PCF Security Enablement",
}

func main() {
	var cfg *config.WorkerConfig
	var flagOpts config.WorkerOpts

	logger := lager.NewLogger("revok-worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.DEBUG))

	logger.Info("starting")

	_, err := flags.Parse(&flagOpts)
	if err != nil {
		os.Exit(1)
	}

	bs, err := ioutil.ReadFile(string(flagOpts.ConfigFile))
	if err != nil {
		logger.Error("failed-to-open-config-file", err)
		os.Exit(1)
	}

	cfg, err = config.LoadWorkerConfig(bs)
	if err != nil {
		logger.Error("failed-to-load-config-file", err)
		os.Exit(1)
	}

	errs := cfg.Validate()
	if errs != nil {
		for _, err := range errs {
			fmt.Println(err.Error())
		}
		os.Exit(1)
	}

	if cfg.Metrics.SentryDSN != "" {
		logger.RegisterSink(revok.NewSentrySink(cfg.Metrics.SentryDSN, cfg.Metrics.Environment))
	}

	workdir := cfg.WorkDir
	_, err = os.Lstat(workdir)
	if err != nil {
		log.Fatalf("workdir error: %s", err)
	}

	dbCertificate, dbCaCertPool := loadCerts(
		cfg.MySQL.CertificatePath,
		cfg.MySQL.PrivateKeyPath,
		cfg.MySQL.PrivateKeyPassphrase,
		cfg.MySQL.CACertificatePath,
	)

	dbURI := db.NewDSN(
		cfg.MySQL.Username,
		cfg.MySQL.Password,
		cfg.MySQL.DBName,
		cfg.MySQL.Hostname,
		int(cfg.MySQL.Port),
		cfg.MySQL.ServerName,
		dbCertificate,
		dbCaCertPool,
	)

	database, err := migrations.LockDBAndMigrate(logger, "mysql", dbURI)
	if err != nil {
		log.Fatalf("db error: %s", err)
	}

	database.LogMode(false)

	clk := clock.NewClock()

	cloneMsgCh := make(chan revok.CloneMsg)

	scanRepository := db.NewScanRepository(database, clk)
	repositoryRepository := db.NewRepositoryRepository(database)
	fetchRepository := db.NewFetchRepository(database)
	credentialRepository := db.NewCredentialRepository(database)
	branchRepository := db.NewBranchRepository(database)

	emitter := metrics.BuildEmitter(cfg.Metrics.DatadogAPIKey, cfg.Metrics.Environment)
	gitClient := gitclient.New(cfg.GitHub.PrivateKeyPath, cfg.GitHub.PublicKeyPath)
	repoWhitelist := notifications.BuildWhitelist(cfg.Whitelist...)
	formatter := notifications.NewSlackNotificationFormatter()

	traceClient, err := cloudtrace.NewClient(context.Background(), cfg.Trace.ProjectName)
	if err != nil {
		logger.Error("failed-to-create-trace-client", err)
	}

	slackHTTPClient := &http.Client{
		Timeout: 5 * time.Second,
	}
	notifier := notifications.NewSlackNotifier(clk, slackHTTPClient, formatter)

	certificate, caCertPool := loadCerts(
		cfg.Identity.CertificatePath,
		cfg.Identity.PrivateKeyPath,
		cfg.Identity.PrivateKeyPassphrase,
		cfg.Identity.CACertificatePath,
	)

	rolodexServerAddr := fmt.Sprintf("%s:%d", cfg.Rolodex.ServerAddress, cfg.Rolodex.ServerPort)

	tlsConfig := tlsconfig.Build(
		tlsconfig.WithPivotalDefaults(),
		tlsconfig.WithIdentity(certificate),
	)

	transportCreds := credentials.NewTLS(tlsConfig.Client(tlsconfig.WithAuthority(caCertPool)))

	conn, err := grpc.Dial(
		rolodexServerAddr,
		grpc.WithDialer(keepAliveDial),
		grpc.WithTransportCredentials(transportCreds),
		grpc.WithUnaryInterceptor(cloudtrace.GRPCClientInterceptor()),
	)

	rolodexClient := rolodexpb.NewRolodexClient(conn)

	teamURLs := notifications.NewTeamURLs(
		cfg.Slack.DefaultURL,
		cfg.Slack.DefaultChannel,
		cfg.Slack.TeamURLs,
	)

	addressBook := notifications.NewRolodex(
		rolodexClient,
		teamURLs,
	)

	router := notifications.NewRouter(
		notifier,
		addressBook,
		repoWhitelist,
	)

	sniffer := sniff.NewDefaultSniffer()
	scanner := revok.NewScanner(
		gitClient,
		repositoryRepository,
		scanRepository,
		credentialRepository,
		sniffer,
	)

	notificationComposer := revok.NewNotificationComposer(
		repositoryRepository,
		router,
		scanner,
	)

	changeFetcher := revok.NewChangeFetcher(
		logger,
		gitClient,
		notificationComposer,
		repositoryRepository,
		fetchRepository,
		emitter,
	)

	changeScheduleRunner := revok.NewScheduleRunner(logger)

	changeScheduler := revok.NewChangeScheduler(
		logger,
		repositoryRepository,
		changeScheduleRunner,
		changeFetcher,
	)

	cloner := revok.NewCloner(
		logger,
		workdir,
		cloneMsgCh,
		gitClient,
		repositoryRepository,
		notificationComposer,
		emitter,
		changeScheduler,
	)

	dirscanUpdater := revok.NewRescanner(
		logger,
		scanRepository,
		credentialRepository,
		scanner,
		router,
		emitter,
	)

	statsReporter := stats.NewReporter(
		logger,
		clk,
		60*time.Second,
		db.NewStatsRepository(database),
		emitter,
	)

	headCredentialCounter := revok.NewHeadCredentialCounter(
		logger,
		branchRepository,
		repositoryRepository,
		clk,
		cfg.CredentialCounterInterval,
		gitClient,
		sniffer,
	)

	debug := admin.Runner(
		"6060",
		admin.WithInfo(info),
		admin.WithUptime(),
	)

	members := []grouper.Member{
		{Name: "cloner", Runner: cloner},
		{Name: "dirscan-updater", Runner: dirscanUpdater},
		{Name: "stats-reporter", Runner: statsReporter},
		{Name: "head-credential-counter", Runner: headCredentialCounter},
		{Name: "change-schedule-runner", Runner: changeScheduleRunner},
		{Name: "debug", Runner: debug},
	}

	looper := gitclient.NewLooper()
	searcher := search.NewSearcher(repositoryRepository, looper)
	handler := revok.NewServer(logger, searcher, repositoryRepository, branchRepository)

	serverTls := tlsConfig.Server(tlsconfig.WithClientAuthentication(caCertPool))

	grpcServer := grpcrunner.New(
		logger,
		fmt.Sprintf("%s:%d", cfg.API.BindIP, cfg.API.BindPort),
		func(server *grpc.Server) {
			revokpb.RegisterRevokServer(server, handler)
		},
		grpc.Creds(credentials.NewTLS(serverTls)),
	)

	members = append(members, grouper.Member{
		Name:   "grpc-server",
		Runner: grpcServer,
	})

	pubSubClient, err := pubsub.NewClient(context.Background(), cfg.PubSub.ProjectName)
	if err != nil {
		logger.Fatal("failed", err)
		os.Exit(1)
	}

	subscription := pubSubClient.Subscription(cfg.PubSub.FetchHint.Subscription)

	publicKey, err := crypto.ReadRSAPublicKey(cfg.PubSub.PublicKeyPath)
	if err != nil {
		logger.Fatal("failed", err)
		os.Exit(1)
	}

	pushEventProcessor := queue.NewPushEventProcessor(
		changeFetcher,
		emitter,
		clk,
		traceClient,
	)

	signatureChecker := queue.NewSignatureCheck(crypto.NewRSAVerifier(publicKey), emitter, pushEventProcessor)

	members = append(members, grouper.Member{
		Name:   "github-hint-handler",
		Runner: queue.NewPubSubSubscriber(logger, subscription, signatureChecker, emitter),
	})

	if cfg.GitHub.AccessToken != "" {
		githubHTTPClient := &http.Client{
			Timeout: 30 * time.Second,
			Transport: &oauth2.Transport{
				Source: oauth2.StaticTokenSource(
					&oauth2.Token{AccessToken: cfg.GitHub.AccessToken},
				),
				Base: &http.Transport{
					DisableKeepAlives: true,
				},
			},
		}

		ghClient := revok.NewGitHubClient(github.NewClient(githubHTTPClient))

		repoDiscoverer := revok.NewRepoDiscoverer(
			logger,
			workdir,
			cloneMsgCh,
			ghClient,
			clk,
			cfg.RepositoryDiscovery.Interval,
			cfg.RepositoryDiscovery.Organizations,
			cfg.RepositoryDiscovery.Users,
			repositoryRepository,
		)

		members = append(members, grouper.Member{
			Name:   "repo-discoverer",
			Runner: repoDiscoverer,
		})
	}

	startupTasks := []grouper.Member{
		{
			Name:   "schedule-fetches",
			Runner: changeScheduler,
		},
	}

	system := []grouper.Member{
		{
			Name:   "servers",
			Runner: grouper.NewParallel(os.Interrupt, members),
		},
		{
			Name:   "startup-tasks",
			Runner: grouper.NewParallel(os.Interrupt, startupTasks),
		},
	}

	runner := sigmon.New(grouper.NewOrdered(os.Interrupt, system))

	err = <-ifrit.Invoke(runner).Wait()
	if err != nil {
		log.Fatalf("failed-to-start: %s", err)
	}
}

func loadCerts(certificatePath, privateKeyPath, privateKeyPassphrase, caCertificatePath string) (tls.Certificate, *x509.CertPool) {
	certificate, err := config.LoadCertificate(
		certificatePath,
		privateKeyPath,
		privateKeyPassphrase,
	)
	if err != nil {
		log.Fatalln(err)
	}

	caCertPool, err := config.LoadCertificatePool(caCertificatePath)
	if err != nil {
		log.Fatalln(err)
	}

	return certificate, caCertPool
}

func keepAliveDial(addr string, timeout time.Duration) (net.Conn, error) {
	d := net.Dialer{
		Timeout:   timeout,
		KeepAlive: 60 * time.Second,
	}
	return d.Dial("tcp", addr)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Peter Mrekaj. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.txt file.

package arrays

import (
	"math/rand"
	"reflect"
	"testing"
)

func TestNext(t *testing.T) {
	for _, test := range []struct {
		in   []int
		want []int
	}{
		{[]int{1}, []int{}},
		{[]int{1, 2}, []int{2, 1}},
		{[]int{2, 1}, []int{}},
		{[]int{1, 2, 3}, []int{1, 3, 2}},
		{[]int{1, 3, 2}, []int{2, 1, 3}},
		{[]int{2, 1, 3}, []int{2, 3, 1}},
		{[]int{2, 3, 1}, []int{3, 1, 2}},
		{[]int{3, 1, 2}, []int{3, 2, 1}},
		{[]int{3, 2, 1}, []int{}},
	} {
		if got := NextPerm(test.in); !reflect.DeepEqual(got, test.want) {
			t.Errorf("NextPerm(%d) = %d; want %d", test.in, got, test.want)
		}
	}
}

func benchNext(b *testing.B, size int) {
	b.StopTimer()
	for i := 0; i < b.N; i++ {
		data := rand.New(rand.NewSource(int64(i))).Perm(size)
		b.StartTimer()
		NextPerm(data)
		b.StopTimer()
	}
}

func BenchmarkNext1e2(b *testing.B) { benchNext(b, 1e2) }
func BenchmarkNext1e4(b *testing.B) { benchNext(b, 1e4) }
func BenchmarkNext1e6(b *testing.B) { benchNext(b, 1e6) }
<commit_msg>Speedup and fix benchmark names for arrays.NextPerm function<commit_after>// Copyright (c) 2015, Peter Mrekaj. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.txt file.

package arrays

import (
	"math/rand"
	"reflect"
	"testing"
)

func TestNextPerm(t *testing.T) {
	for _, test := range []struct {
		in   []int
		want []int
	}{
		{[]int{1}, []int{}},
		{[]int{1, 2}, []int{2, 1}},
		{[]int{2, 1}, []int{}},
		{[]int{1, 2, 3}, []int{1, 3, 2}},
		{[]int{1, 3, 2}, []int{2, 1, 3}},
		{[]int{2, 1, 3}, []int{2, 3, 1}},
		{[]int{2, 3, 1}, []int{3, 1, 2}},
		{[]int{3, 1, 2}, []int{3, 2, 1}},
		{[]int{3, 2, 1}, []int{}},
	} {
		if got := NextPerm(test.in); !reflect.DeepEqual(got, test.want) {
			t.Errorf("NextPerm(%d) = %d; want %d", test.in, got, test.want)
		}
	}
}

func benchNextPerm(b *testing.B, size int) {
	data := rand.New(rand.NewSource(int64(size))).Perm(size)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NextPerm(data)
	}
}

func BenchmarkNextPerm1e2(b *testing.B) { benchNextPerm(b, 1e2) }
func BenchmarkNextPerm1e4(b *testing.B) { benchNextPerm(b, 1e4) }
func BenchmarkNextPerm1e6(b *testing.B) { benchNextPerm(b, 1e6) }
<|endoftext|>


<commit_before>package asm

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"testing"

	qt "github.com/frankban/quicktest"
)

var test64bitImmProg = []byte{
	// r0 = math.MinInt32 - 1
	0x18, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0x7f,
	0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff,
}

func TestRead64bitImmediate(t *testing.T) {
	var ins Instruction
	n, err := ins.Unmarshal(bytes.NewReader(test64bitImmProg), binary.LittleEndian)
	if err != nil {
		t.Fatal(err)
	}
	if want := uint64(InstructionSize * 2); n != want {
		t.Errorf("Expected %d bytes to be read, got %d", want, n)
	}

	if c := ins.Constant; c != math.MinInt32-1 {
		t.Errorf("Expected immediate to be %v, got %v", int64(math.MinInt32)-1, c)
	}
}

func TestWrite64bitImmediate(t *testing.T) {
	insns := Instructions{
		LoadImm(R0, math.MinInt32-1, DWord),
	}

	var buf bytes.Buffer
	if err := insns.Marshal(&buf, binary.LittleEndian); err != nil {
		t.Fatal(err)
	}

	if prog := buf.Bytes(); !bytes.Equal(prog, test64bitImmProg) {
		t.Errorf("Marshalled program does not match:\n%s", hex.Dump(prog))
	}
}

func TestSignedJump(t *testing.T) {
	insns := Instructions{
		JSGT.Imm(R0, -1, "foo"),
	}

	insns[0].Offset = 1

	err := insns.Marshal(io.Discard, binary.LittleEndian)
	if err != nil {
		t.Error("Can't marshal signed jump:", err)
	}
}

func TestInstructionRewriteMapConstant(t *testing.T) {
	ins := LoadMapValue(R0, 123, 321)

	qt.Assert(t, ins.MapPtr(), qt.Equals, 123)
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(321))

	qt.Assert(t, ins.RewriteMapPtr(-1), qt.IsNil)
	qt.Assert(t, ins.MapPtr(), qt.Equals, -1)

	qt.Assert(t, ins.RewriteMapPtr(1), qt.IsNil)
	qt.Assert(t, ins.MapPtr(), qt.Equals, 1)

	// mapOffset should be unchanged after rewriting the pointer.
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(321))

	qt.Assert(t, ins.RewriteMapOffset(123), qt.IsNil)
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(123))

	// MapPtr should be unchanged.
	qt.Assert(t, ins.MapPtr(), qt.Equals, 1)

	ins = Mov.Imm(R1, 32)
	if err := ins.RewriteMapPtr(1); err == nil {
		t.Error("RewriteMapPtr rewriting bogus instruction")
	}
	if err := ins.RewriteMapOffset(1); err == nil {
		t.Error("RewriteMapOffset rewriting bogus instruction")
	}
}

func TestInstructionLoadMapValue(t *testing.T) {
	ins := LoadMapValue(R0, 1, 123)
	if !ins.IsLoadFromMap() {
		t.Error("isLoadFromMap returns false")
	}
	if fd := ins.MapPtr(); fd != 1 {
		t.Error("Expected map fd to be 1, got", fd)
	}
	if off := ins.mapOffset(); off != 123 {
		t.Fatal("Expected map offset to be 123 after changin the pointer, got", off)
	}
}

func TestInstructionsRewriteMapPtr(t *testing.T) {
	insns := Instructions{
		LoadMapPtr(R1, 0),
		Return(),
	}
	insns[0].Reference = "good"

	if err := insns.RewriteMapPtr("good", 1); err != nil {
		t.Fatal(err)
	}

	if insns[0].Constant != 1 {
		t.Error("Constant should be 1, have", insns[0].Constant)
	}

	if err := insns.RewriteMapPtr("good", 2); err != nil {
		t.Fatal(err)
	}

	if insns[0].Constant != 2 {
		t.Error("Constant should be 2, have", insns[0].Constant)
	}

	if err := insns.RewriteMapPtr("bad", 1); !IsUnreferencedSymbol(err) {
		t.Error("Rewriting unreferenced map doesn't return appropriate error")
	}
}

// You can use format flags to change the way an eBPF
// program is stringified.
func ExampleInstructions_Format() {
	insns := Instructions{
		FnMapLookupElem.Call().Sym("my_func"),
		LoadImm(R0, 42, DWord),
		Return(),
	}

	fmt.Println("Default format:")
	fmt.Printf("%v\n", insns)

	fmt.Println("Don't indent instructions:")
	fmt.Printf("%.0v\n", insns)

	fmt.Println("Indent using spaces:")
	fmt.Printf("% v\n", insns)

	fmt.Println("Control symbol indentation:")
	fmt.Printf("%2v\n", insns)

	// Output: Default format:
	// my_func:
	// 	0: Call FnMapLookupElem
	// 	1: LdImmDW dst: r0 imm: 42
	// 	3: Exit
	//
	// Don't indent instructions:
	// my_func:
	// 0: Call FnMapLookupElem
	// 1: LdImmDW dst: r0 imm: 42
	// 3: Exit
	//
	// Indent using spaces:
	// my_func:
	//  0: Call FnMapLookupElem
	//  1: LdImmDW dst: r0 imm: 42
	//  3: Exit
	//
	// Control symbol indentation:
	// 		my_func:
	// 	0: Call FnMapLookupElem
	// 	1: LdImmDW dst: r0 imm: 42
	// 	3: Exit
}

func TestReadSrcDst(t *testing.T) {
	testSrcDstProg := []byte{
		// on little-endian: r0 = r1
		// on big-endian: be: r1 = r0
		0xbf, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	}

	testcases := []struct {
		bo       binary.ByteOrder
		dst, src Register
	}{
		{binary.BigEndian, R1, R0},
		{binary.LittleEndian, R0, R1},
	}

	for _, tc := range testcases {
		t.Run(tc.bo.String(), func(t *testing.T) {
			var ins Instruction
			_, err := ins.Unmarshal(bytes.NewReader(testSrcDstProg), tc.bo)
			if err != nil {
				t.Fatal(err)
			}
			if ins.Dst != tc.dst {
				t.Errorf("Expected destination to be %v, got %v", tc.dst, ins.Dst)
			}
			if ins.Src != tc.src {
				t.Errorf("Expected source to be %v, got %v", tc.src, ins.Src)
			}
		})
	}
}

func TestInstructionIterator(t *testing.T) {
	insns := Instructions{
		LoadImm(R0, 0, Word),
		LoadImm(R0, 0, DWord),
		Return(),
	}
	offsets := []RawInstructionOffset{0, 1, 3}

	iter := insns.Iterate()
	for i := 0; i < len(insns); i++ {
		if !iter.Next() {
			t.Fatalf("Expected %dth call to Next to return true", i)
		}

		if iter.Ins == nil {
			t.Errorf("Expected iter.Ins to be non-nil")
		}
		if iter.Index != i {
			t.Errorf("Expected iter.Index to be %d, got %d", i, iter.Index)
		}
		if iter.Offset != offsets[i] {
			t.Errorf("Expected iter.Offset to be %d, got %d", offsets[i], iter.Offset)
		}
	}
}
<commit_msg>asm: benchmark Instruction.(Un)Marshal<commit_after>package asm

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"testing"

	qt "github.com/frankban/quicktest"
)

var test64bitImmProg = []byte{
	// r0 = math.MinInt32 - 1
	0x18, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0x7f,
	0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff,
}

func TestRead64bitImmediate(t *testing.T) {
	var ins Instruction
	n, err := ins.Unmarshal(bytes.NewReader(test64bitImmProg), binary.LittleEndian)
	if err != nil {
		t.Fatal(err)
	}
	if want := uint64(InstructionSize * 2); n != want {
		t.Errorf("Expected %d bytes to be read, got %d", want, n)
	}

	if c := ins.Constant; c != math.MinInt32-1 {
		t.Errorf("Expected immediate to be %v, got %v", int64(math.MinInt32)-1, c)
	}
}

func BenchmarkRead64bitImmediate(b *testing.B) {
	r := &bytes.Reader{}
	for i := 0; i < b.N; i++ {
		r.Reset(test64bitImmProg)

		var ins Instruction
		if _, err := ins.Unmarshal(r, binary.LittleEndian); err != nil {
			b.Fatal(err)
		}
	}
}

func TestWrite64bitImmediate(t *testing.T) {
	insns := Instructions{
		LoadImm(R0, math.MinInt32-1, DWord),
	}

	var buf bytes.Buffer
	if err := insns.Marshal(&buf, binary.LittleEndian); err != nil {
		t.Fatal(err)
	}

	if prog := buf.Bytes(); !bytes.Equal(prog, test64bitImmProg) {
		t.Errorf("Marshalled program does not match:\n%s", hex.Dump(prog))
	}
}

func BenchmarkWrite64BitImmediate(b *testing.B) {
	ins := LoadImm(R0, math.MinInt32-1, DWord)

	var buf bytes.Buffer
	for i := 0; i < b.N; i++ {
		buf.Reset()

		if _, err := ins.Marshal(&buf, binary.LittleEndian); err != nil {
			b.Fatal(err)
		}
	}
}

func TestSignedJump(t *testing.T) {
	insns := Instructions{
		JSGT.Imm(R0, -1, "foo"),
	}

	insns[0].Offset = 1

	err := insns.Marshal(io.Discard, binary.LittleEndian)
	if err != nil {
		t.Error("Can't marshal signed jump:", err)
	}
}

func TestInstructionRewriteMapConstant(t *testing.T) {
	ins := LoadMapValue(R0, 123, 321)

	qt.Assert(t, ins.MapPtr(), qt.Equals, 123)
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(321))

	qt.Assert(t, ins.RewriteMapPtr(-1), qt.IsNil)
	qt.Assert(t, ins.MapPtr(), qt.Equals, -1)

	qt.Assert(t, ins.RewriteMapPtr(1), qt.IsNil)
	qt.Assert(t, ins.MapPtr(), qt.Equals, 1)

	// mapOffset should be unchanged after rewriting the pointer.
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(321))

	qt.Assert(t, ins.RewriteMapOffset(123), qt.IsNil)
	qt.Assert(t, ins.mapOffset(), qt.Equals, uint32(123))

	// MapPtr should be unchanged.
	qt.Assert(t, ins.MapPtr(), qt.Equals, 1)

	ins = Mov.Imm(R1, 32)
	if err := ins.RewriteMapPtr(1); err == nil {
		t.Error("RewriteMapPtr rewriting bogus instruction")
	}
	if err := ins.RewriteMapOffset(1); err == nil {
		t.Error("RewriteMapOffset rewriting bogus instruction")
	}
}

func TestInstructionLoadMapValue(t *testing.T) {
	ins := LoadMapValue(R0, 1, 123)
	if !ins.IsLoadFromMap() {
		t.Error("isLoadFromMap returns false")
	}
	if fd := ins.MapPtr(); fd != 1 {
		t.Error("Expected map fd to be 1, got", fd)
	}
	if off := ins.mapOffset(); off != 123 {
		t.Fatal("Expected map offset to be 123 after changin the pointer, got", off)
	}
}

func TestInstructionsRewriteMapPtr(t *testing.T) {
	insns := Instructions{
		LoadMapPtr(R1, 0),
		Return(),
	}
	insns[0].Reference = "good"

	if err := insns.RewriteMapPtr("good", 1); err != nil {
		t.Fatal(err)
	}

	if insns[0].Constant != 1 {
		t.Error("Constant should be 1, have", insns[0].Constant)
	}

	if err := insns.RewriteMapPtr("good", 2); err != nil {
		t.Fatal(err)
	}

	if insns[0].Constant != 2 {
		t.Error("Constant should be 2, have", insns[0].Constant)
	}

	if err := insns.RewriteMapPtr("bad", 1); !IsUnreferencedSymbol(err) {
		t.Error("Rewriting unreferenced map doesn't return appropriate error")
	}
}

// You can use format flags to change the way an eBPF
// program is stringified.
func ExampleInstructions_Format() {
	insns := Instructions{
		FnMapLookupElem.Call().Sym("my_func"),
		LoadImm(R0, 42, DWord),
		Return(),
	}

	fmt.Println("Default format:")
	fmt.Printf("%v\n", insns)

	fmt.Println("Don't indent instructions:")
	fmt.Printf("%.0v\n", insns)

	fmt.Println("Indent using spaces:")
	fmt.Printf("% v\n", insns)

	fmt.Println("Control symbol indentation:")
	fmt.Printf("%2v\n", insns)

	// Output: Default format:
	// my_func:
	// 	0: Call FnMapLookupElem
	// 	1: LdImmDW dst: r0 imm: 42
	// 	3: Exit
	//
	// Don't indent instructions:
	// my_func:
	// 0: Call FnMapLookupElem
	// 1: LdImmDW dst: r0 imm: 42
	// 3: Exit
	//
	// Indent using spaces:
	// my_func:
	//  0: Call FnMapLookupElem
	//  1: LdImmDW dst: r0 imm: 42
	//  3: Exit
	//
	// Control symbol indentation:
	// 		my_func:
	// 	0: Call FnMapLookupElem
	// 	1: LdImmDW dst: r0 imm: 42
	// 	3: Exit
}

func TestReadSrcDst(t *testing.T) {
	testSrcDstProg := []byte{
		// on little-endian: r0 = r1
		// on big-endian: be: r1 = r0
		0xbf, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	}

	testcases := []struct {
		bo       binary.ByteOrder
		dst, src Register
	}{
		{binary.BigEndian, R1, R0},
		{binary.LittleEndian, R0, R1},
	}

	for _, tc := range testcases {
		t.Run(tc.bo.String(), func(t *testing.T) {
			var ins Instruction
			_, err := ins.Unmarshal(bytes.NewReader(testSrcDstProg), tc.bo)
			if err != nil {
				t.Fatal(err)
			}
			if ins.Dst != tc.dst {
				t.Errorf("Expected destination to be %v, got %v", tc.dst, ins.Dst)
			}
			if ins.Src != tc.src {
				t.Errorf("Expected source to be %v, got %v", tc.src, ins.Src)
			}
		})
	}
}

func TestInstructionIterator(t *testing.T) {
	insns := Instructions{
		LoadImm(R0, 0, Word),
		LoadImm(R0, 0, DWord),
		Return(),
	}
	offsets := []RawInstructionOffset{0, 1, 3}

	iter := insns.Iterate()
	for i := 0; i < len(insns); i++ {
		if !iter.Next() {
			t.Fatalf("Expected %dth call to Next to return true", i)
		}

		if iter.Ins == nil {
			t.Errorf("Expected iter.Ins to be non-nil")
		}
		if iter.Index != i {
			t.Errorf("Expected iter.Index to be %d, got %d", i, iter.Index)
		}
		if iter.Offset != offsets[i] {
			t.Errorf("Expected iter.Offset to be %d, got %d", offsets[i], iter.Offset)
		}
	}
}
<|endoftext|>


<commit_before>package ast

type SelectStatement struct {
	Cores  []Node
	Orders []Node
	Limit  Node
	Offset Node
	Lock   Node
	With   Node
}

func (n SelectStatement) Visit(v Visitor) (s string) {
	s = v.GetSelectStatement(n)
	return
}

func NewSelectStatement() (n SelectStatement) {
	n = *new(SelectStatement)
	n.Cores = []Node{SelectCore{}}
	return
}
<commit_msg>removed unnecessary instantiation<commit_after>package ast

type SelectStatement struct {
	Cores  []Node
	Orders []Node
	Limit  Node
	Offset Node
	Lock   Node
	With   Node
}

func (n SelectStatement) Visit(v Visitor) (s string) {
	s = v.GetSelectStatement(n)
	return
}

func NewSelectStatement() (n SelectStatement) {
	n.Cores = []Node{SelectCore{}}
	return
}
<|endoftext|>


<commit_before>package main

/*
Script that watched etcd and rewrites configuration files on change in etcd
*/

// http://blog.gopheracademy.com/advent-2013/day-06-service-discovery-with-etcd/
import (
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"github.com/tcotav/etcdhooks/config"
	"github.com/tcotav/etcdhooks/etcd"
	"github.com/tcotav/etcdhooks/nagios"
	"github.com/tcotav/etcdhooks/web"
	"log"
	"os"
	"strings"
	"time"
)

// think we want to dump a lot of this into a config
// stuff like the etcd info
//
var nagios_host_file = "/tmp/hosts.cfg"
var nagios_group_file = "/tmp/groups.cfg"
var host_list_file = "/tmp/host_list.cfg"

// updateHost wrapper containing async function calls to update the internal map
// as well as the config files
func updateHost(k string, v string) {
	hostMap := etcdWatcher.Map()
	_, containsHost := hostMap[k]
	go etcdWatcher.UpdateMap(k, v)
	// regenerate these files ONLY if it is a new host
	if !containsHost {
		regenHosts()
	}
}

func writeHostMap(hostMap map[string]int) {
	f, err := os.Create(host_list_file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	for host := range hostMap {
		f.WriteString(fmt.Sprintf("%s\n", host))
	}
}

var limiterOn = false
const fileRewriteInterval = 30
var lastFileWrite = time.Now().Add(time.Second * 1000 * -1) // initialize to some point in the past

// regenHostFiles utility function that calls regen methods for files/persistence that contain only
// host data.  We pass along up/down and in/out service info too -- that should be handled with a different
// method.  Currently limited so that we don't write More than fileRewriteInterval seconds.
func regenHosts() {
	if limiterOn { // we're already waiting on a file rewrite
		return
	}

	// do some date math here -- have we waited long enough to write our file?
	if time.Now().Before(lastFileWrite.Add(time.Second * fileRewriteInterval))  {
		log.Println("limiter kicked in")
		limiterOn=true
		// these statements cause us to wait fileRewriteInterval seconds before continuing
		limiter := time.Tick(time.Second * fileRewriteInterval)
    <-limiter
	}

	// flip back our counters
	limiterOn=false
	lastFileWrite = time.Now()

	log.Println("generating files")
	// do the work
	hostMap := etcdWatcher.Map()
	go nagios.GenerateFiles(hostMap, nagios_host_file, nagios_group_file)
	go writeHostMap(hostMap)
}

func removeHost(k string) {
	go etcdWatcher.DeleteFromMap(k)
	// remove from map
	// run the updateNagios command
	regenHosts()
}

func main() {
	config := config.ParseConfig("daemon.cfg")
	nagios_host_file = config["nagios_host_file"]
	nagios_group_file = config["nagios_groups_file"]
	host_list_file = config["host_list_file"]

	// expect this to be csv or single entry
	etcd_server_list := strings.Split(config["etcd_server_list"], ",")
	log.Println("got error list")
	client := etcd.NewClient(etcd_server_list)
	log.Println("got client")
	etcdWatcher.InitDataMap(client)
	log.Println("Dumping map contents for verification")
	etcdWatcher.DumpMap()
	log.Println("Generating initial config files")
	regenHosts()
	//
	// spin up the web server
	//
	go webservice.StartWebService(config["web_listen_port"])
	watchChan := make(chan *etcd.Response)
	go client.Watch(config["base_etcd_url"], 0, true, watchChan, nil)
	log.Println("Waiting for an update...")
	for {
		select {
		case r := <-watchChan:
			// do something with it here
			action := r.Action
			k := r.Node.Key
			v := r.Node.Value
			switch action {
			case "delete":
				log.Printf("delete of key: %s", k)
				go removeHost(k)
			case "set":
				log.Printf("update of key: %s, value: %s", k, v)
				go updateHost(k, v)
			}
		}
	}
	// we don't really care what changed in this case so...
	//DumpServices(client)
}
<commit_msg>f'd up some async<commit_after>package main

/*
Script that watched etcd and rewrites configuration files on change in etcd
*/

// http://blog.gopheracademy.com/advent-2013/day-06-service-discovery-with-etcd/
import (
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"github.com/tcotav/etcdhooks/config"
	"github.com/tcotav/etcdhooks/etcd"
	"github.com/tcotav/etcdhooks/nagios"
	"github.com/tcotav/etcdhooks/web"
	"log"
	"os"
	"strings"
	"time"
)

// think we want to dump a lot of this into a config
// stuff like the etcd info
//
var nagios_host_file = "/tmp/hosts.cfg"
var nagios_group_file = "/tmp/groups.cfg"
var host_list_file = "/tmp/host_list.cfg"

// updateHost wrapper containing async function calls to update the internal map
// as well as the config files
func updateHost(k string, v string) {
	hostMap := etcdWatcher.Map()
	_, containsHost := hostMap[k]
	go etcdWatcher.UpdateMap(k, v)
	// regenerate these files ONLY if it is a new host
	if !containsHost {
		regenHosts()
	}
}

func writeHostMap(hostMap map[string]int) {
	f, err := os.Create(host_list_file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	for host := range hostMap {
		f.WriteString(fmt.Sprintf("%s\n", host))
	}
}

var limiterOn = false

const fileRewriteInterval = 30

var lastFileWrite = time.Now().Add(time.Second * 1000 * -1) // initialize to some point in the past

// regenHostFiles utility function that calls regen methods for files/persistence that contain only
// host data.  We pass along up/down and in/out service info too -- that should be handled with a different
// method.  Currently limited so that we don't write More than fileRewriteInterval seconds.
func regenHosts() {
	if limiterOn { // we're already waiting on a file rewrite
		return
	}

	// do some date math here -- have we waited long enough to write our file?
	if time.Now().Before(lastFileWrite.Add(time.Second * fileRewriteInterval)) {
		log.Println("limiter kicked in")
		limiterOn = true
		// these statements cause us to wait fileRewriteInterval seconds before continuing
		limiter := time.Tick(time.Second * fileRewriteInterval)
		<-limiter
	}

	// flip back our counters
	limiterOn = false
	lastFileWrite = time.Now()

	log.Println("generating files")
	// do the work
	hostMap := etcdWatcher.Map()
	go nagios.GenerateFiles(hostMap, nagios_host_file, nagios_group_file)
	go writeHostMap(hostMap)
}

func removeHost(k string) {
	etcdWatcher.DeleteFromMap(k)
	// remove from map
	// run the updateNagios command
	regenHosts()
}

func main() {
	config := config.ParseConfig("daemon.cfg")
	nagios_host_file = config["nagios_host_file"]
	nagios_group_file = config["nagios_groups_file"]
	host_list_file = config["host_list_file"]

	// expect this to be csv or single entry
	etcd_server_list := strings.Split(config["etcd_server_list"], ",")
	log.Println("got error list")
	client := etcd.NewClient(etcd_server_list)
	log.Println("got client")
	etcdWatcher.InitDataMap(client)
	log.Println("Dumping map contents for verification")
	etcdWatcher.DumpMap()
	log.Println("Generating initial config files")
	regenHosts()
	//
	// spin up the web server
	//
	go webservice.StartWebService(config["web_listen_port"])
	watchChan := make(chan *etcd.Response)
	go client.Watch(config["base_etcd_url"], 0, true, watchChan, nil)
	log.Println("Waiting for an update...")
	for {
		select {
		case r := <-watchChan:
			// do something with it here
			action := r.Action
			k := r.Node.Key
			v := r.Node.Value
			switch action {
			case "delete":
				log.Printf("delete of key: %s", k)
				go removeHost(k)
			case "set":
				log.Printf("update of key: %s, value: %s", k, v)
				go updateHost(k, v)
			}
		}
	}
	// we don't really care what changed in this case so...
	//DumpServices(client)
}
<|endoftext|>