<commit_before>package amqp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		b.GetConfig().DefaultQueue,      // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		nil,                             // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, err
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occured on queue: %s. Reconnecting", queueName)
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	connection, err := b.GetOrOpenConnection(signature.RoutingKey,
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil,                           // exchange declare args
		nil,                           // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	errorsChan := make(chan error)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		delivery.Nack(multiple, requeue)
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	connection, err := b.GetOrOpenConnection(queueName,
		queueName,        // queue binding key
		nil,              // exchange declare args
		declareQueueArgs, // queue declare arghs
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct" {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<commit_msg>AMQP - use Custom queue, fallback to default if not exists<commit_after>package amqp

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil, // exchange declare args
		nil, // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, err
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occured on queue: %s. Reconnecting", queueName)
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	connection, err := b.GetOrOpenConnection(signature.RoutingKey,
		b.GetConfig().AMQP.BindingKey, // queue binding key
		nil, // exchange declare args
		nil, // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	errorsChan := make(chan error)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		delivery.Nack(multiple, requeue)
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	connection, err := b.GetOrOpenConnection(queueName,
		queueName,                                       // queue binding key
		nil,                                             // exchange declare args
		declareQueueArgs,                                // queue declare arghs
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	channel := connection.channel

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct" {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"flag"
	"net/http"
	"strings"
	"regexp"
	"github.com/palmergs/tokensearch"
)

var root = tokensearch.NewTokenNode()

var validPath = regexp.MustCompile("^/tokens/([a-zA-Z0-9_-]+)$")

func searchHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	query := r.FormValue("q")
	matches, err := root.Find(query)
	if err != nil {
		http.Error(w, "Not found", 404)
	} else {
		for match := range matches {
			// TODO
			fmt.Printf("Match = %v\n", match)
		}
	}
}

func tokenHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	switch strings.ToUpper(r.Method) {
	case "POST", "PUT":
		insertTokenHandler(w, r)
	case "DELETE":
		deleteTokenHandler(w, r)
	case "GET", "":
		getTokenHandler(w, r)
	}
}

func insertTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := r.FormValue("ident")
	display := r.FormValue("display")
	category := r.FormValue("category")
	token := tokensearch.NewToken(ident, display, category)
	_, err := root.Insert(token)
	if err != nil {
		http.Error(w, err.Error(), 401)
	} else {
		// TODO
	}
}

func deleteTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := r.FormValue("ident")
	display := r.FormValue("display")
	token := tokensearch.NewToken(ident, display, "")
	_, err := root.Remove(token)
	if err != nil {
		http.Error(w, err.Error(), 401)
	} else {
		// TODO
	}
}

func getTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := validPath.FindStringSubmatch(r.URL.Path)
	if ident == nil {
		http.NotFound(w, r)
	} else {
		// TODO: scan through tree looking for matching strings
		// root.FindTokens(ident)
	}
}

func main() {

	serverPort := flag.Int("p", 6060, "server port")
	flag.Parse()

	fmt.Printf("Starting server on port %v...", *serverPort)
	http.HandleFunc("/search/", searchHandler)
	http.HandleFunc("/tokens/", tokenHandler)
	http.ListenAndServe(fmt.Sprintf(":%v", *serverPort), nil)

	fmt.Printf("done\n")
}<commit_msg>pull server address from flags<commit_after>package main

import (
	"fmt"
	"flag"
	"net/http"
	"strings"
	"regexp"
	"github.com/palmergs/tokensearch"
)

var root = tokensearch.NewTokenNode()

var validPath = regexp.MustCompile("^/tokens/([a-zA-Z0-9_-]+)$")

func searchHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	query := r.FormValue("q")
	matches, err := root.Find(query)
	if err != nil {
		http.Error(w, "Not found", 404)
	} else {
		for match := range matches {
			// TODO
			fmt.Printf("Match = %v\n", match)
		}
	}
}

func tokenHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	switch strings.ToUpper(r.Method) {
	case "POST", "PUT":
		insertTokenHandler(w, r)
	case "DELETE":
		deleteTokenHandler(w, r)
	case "GET", "":
		getTokenHandler(w, r)
	}
}

func insertTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := r.FormValue("ident")
	display := r.FormValue("display")
	category := r.FormValue("category")
	token := tokensearch.NewToken(ident, display, category)
	_, err := root.Insert(token)
	if err != nil {
		http.Error(w, err.Error(), 401)
	} else {
		// TODO
	}
}

func deleteTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := r.FormValue("ident")
	display := r.FormValue("display")
	token := tokensearch.NewToken(ident, display, "")
	_, err := root.Remove(token)
	if err != nil {
		http.Error(w, err.Error(), 401)
	} else {
		// TODO
	}
}

func getTokenHandler(w http.ResponseWriter, r *http.Request) {
	ident := validPath.FindStringSubmatch(r.URL.Path)
	if ident == nil {
		http.NotFound(w, r)
	} else {
		// TODO: scan through tree looking for matching strings
		// root.FindTokens(ident)
	}
}

func main() {

	serverPort := flag.Int("p", 6060, "server port")
	flag.Parse()

	serverAddr := fmt.Sprintf(":%v", *serverPort)
	fmt.Printf("Starting server on %v...", serverAddr)

	http.HandleFunc("/search/", searchHandler)
	http.HandleFunc("/tokens/", tokenHandler)
	http.ListenAndServe(serverAddr, nil)

	fmt.Printf("done\n")
}
<|endoftext|>


<commit_before>package sensu

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"strings"

	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/check"
	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/client"
)

const defaultRabbitMQURI string = "amqp://guest:guest@localhost:5672/%2f"

type configFlagSet struct {
	configFile string
	verbose    bool
}

type Config struct {
	flagSet *configFlagSet
	config  *configPayload
}

type configPayload struct {
	Client      *client.Client `json:"client,omitempty"`
	Checks      []*check.Check `json:"checks,omitempty"`
	RabbitMQURI *string        `json:"rabbitmq_uri,omitempty"`
}

func NewConfigFromFlagSet(flagset *configFlagSet) (*Config, error) {
	var cfg = Config{flagset, &configPayload{}}

	if flagset != nil && flagset.configFile != "" {
		buf, err := ioutil.ReadFile(flagset.configFile)

		if err != nil {
			return nil, err
		}

		if err := json.Unmarshal(buf, &cfg.config); err != nil {
			return nil, err
		}
	}

	return &cfg, nil
}

func (c *Config) RabbitMQURI() string {
	if cfg := c.config; cfg != nil && cfg.RabbitMQURI != nil {
		return *cfg.RabbitMQURI
	} else if uri := os.Getenv("RABBITMQ_URI"); uri != "" {
		return uri
	}

	return defaultRabbitMQURI
}

func (c *Config) Client() *client.Client {
	if cfg := c.config; cfg != nil && cfg.Client != nil {
		return cfg.Client
	}

	return &client.Client{
		Name:          os.Getenv("SENSU_CLIENT_NAME"),
		Address:       os.Getenv("SENSU_ADDRESS"),
		Subscriptions: strings.Split(os.Getenv("SENSU_CLIENT_SUBSCRIPTIONS"), ","),
	}
}

func (c *Config) Checks() []*check.Check {
	if cfg := c.config; cfg != nil {
		return cfg.Checks
	}

	return []*check.Check{}
}
<commit_msg>Use SENSU_CLIENT_ADDRESS env var instead of SENSU_ADDRESS<commit_after>package sensu

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"strings"

	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/check"
	"github.com/upfluence/sensu-client-go/Godeps/_workspace/src/github.com/upfluence/sensu-go/sensu/client"
)

const defaultRabbitMQURI string = "amqp://guest:guest@localhost:5672/%2f"

type configFlagSet struct {
	configFile string
	verbose    bool
}

type Config struct {
	flagSet *configFlagSet
	config  *configPayload
}

type configPayload struct {
	Client      *client.Client `json:"client,omitempty"`
	Checks      []*check.Check `json:"checks,omitempty"`
	RabbitMQURI *string        `json:"rabbitmq_uri,omitempty"`
}

func NewConfigFromFlagSet(flagset *configFlagSet) (*Config, error) {
	var cfg = Config{flagset, &configPayload{}}

	if flagset != nil && flagset.configFile != "" {
		buf, err := ioutil.ReadFile(flagset.configFile)

		if err != nil {
			return nil, err
		}

		if err := json.Unmarshal(buf, &cfg.config); err != nil {
			return nil, err
		}
	}

	return &cfg, nil
}

func (c *Config) RabbitMQURI() string {
	if cfg := c.config; cfg != nil && cfg.RabbitMQURI != nil {
		return *cfg.RabbitMQURI
	} else if uri := os.Getenv("RABBITMQ_URI"); uri != "" {
		return uri
	}

	return defaultRabbitMQURI
}

func (c *Config) Client() *client.Client {
	if cfg := c.config; cfg != nil && cfg.Client != nil {
		return cfg.Client
	}

	return &client.Client{
		Name:          os.Getenv("SENSU_CLIENT_NAME"),
		Address:       os.Getenv("SENSU_CLIENT_ADDRESS"),
		Subscriptions: strings.Split(os.Getenv("SENSU_CLIENT_SUBSCRIPTIONS"), ","),
	}
}

func (c *Config) Checks() []*check.Check {
	if cfg := c.config; cfg != nil {
		return cfg.Checks
	}

	return []*check.Check{}
}
<|endoftext|>


<commit_before>package tracker

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"
)

var Now = time.Now()

type Tracker struct {
	Name string
	Days map[string]Day `json:"days"`
}

func New(filename string) (*Tracker, error) {
	t := &Tracker{Name: filename, Days: make(map[string]Day)}
	flags := os.O_CREATE | os.O_RDONLY
	f, err := os.OpenFile(t.Name, flags, 0644)
	defer f.Close()
	if err != nil {
		return t, err
	}
	decoder := json.NewDecoder(f)
	err = decoder.Decode(t)
	if err != nil && err != io.EOF {
		return t, err
	}
	return t, nil
}

func (t *Tracker) Save() error {
	flags := os.O_TRUNC | os.O_WRONLY
	f, err := os.OpenFile(t.Name, flags, 0644)
	defer f.Close()
	if err != nil {
		return err
	}
	encoder := json.NewEncoder(f)
	err = encoder.Encode(t)
	return err
}

func (t *Tracker) NewDay() (Day, error) {
	key := date(Now)
	d, ok := t.Days[key]
	if ok {
		return d, fmt.Errorf("Day %s already exists", key)
	}
	return Day{Start: Now, Tasks: []Task{}, Pauses: []Pause{}}, nil
}

func (t *Tracker) Today() (Day, error) {
	key := date(Now)
	d, ok := t.Days[key]
	if !ok {
		return d, fmt.Errorf("Day %s hasn't been started yet", key)
	}
	return d, nil
}

func (t *Tracker) SaveDay(d Day) {
	key := date(d.Start)
	t.Days[key] = d
}

func date(t time.Time) string {
	return t.Format("2006-01-02")
}
<commit_msg>Prevent file from being create on read<commit_after>package tracker

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"
)

var Now = time.Now()

type Tracker struct {
	Name string
	Days map[string]Day `json:"days"`
}

func New(filename string) (*Tracker, error) {
	t := &Tracker{Name: filename, Days: make(map[string]Day)}
	flags := os.O_RDONLY
	f, err := os.OpenFile(t.Name, flags, 0644)
	if os.IsNotExist(err) {
		return t, nil
	}
	defer f.Close()
	if err != nil {
		return t, err
	}
	decoder := json.NewDecoder(f)
	err = decoder.Decode(t)
	if err != nil && err != io.EOF {
		return t, err
	}
	return t, nil
}

func (t *Tracker) Save() error {
	flags := os.O_CREATE | os.O_TRUNC | os.O_WRONLY
	f, err := os.OpenFile(t.Name, flags, 0644)
	defer f.Close()
	if err != nil {
		return err
	}
	encoder := json.NewEncoder(f)
	err = encoder.Encode(t)
	return err
}

func (t *Tracker) NewDay() (Day, error) {
	key := date(Now)
	d, ok := t.Days[key]
	if ok {
		return d, fmt.Errorf("Day %s already exists", key)
	}
	return Day{Start: Now, Tasks: []Task{}, Pauses: []Pause{}}, nil
}

func (t *Tracker) Today() (Day, error) {
	key := date(Now)
	d, ok := t.Days[key]
	if !ok {
		return d, fmt.Errorf("Day %s hasn't been started yet", key)
	}
	return d, nil
}

func (t *Tracker) SaveDay(d Day) {
	key := date(d.Start)
	t.Days[key] = d
}

func date(t time.Time) string {
	return t.Format("2006-01-02")
}
<|endoftext|>


<commit_before>package test

import (
	"bytes"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/coredns/coredns/plugin"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/prometheus/util/promlint"
)

func TestMetricNaming(t *testing.T) {

	walker := validMetricWalker{}
	err := filepath.Walk("..", walker.walk)

	if err != nil {
		t.Fatal(err)
	}

	if len(walker.Metrics) > 0 {

		buf := &bytes.Buffer{}
		encoder := expfmt.NewEncoder(buf, expfmt.FmtText)
		for _, mf := range walker.Metrics {
			if err := encoder.Encode(mf); err != nil {
				t.Fatalf("Encoding and sending metric family: %s", err)
			}
		}

		l := promlint.New(buf)
		problems, err := l.Lint()
		if err != nil {
			t.Fatalf("Link found error: %s", err)
		}

		if len(problems) > 0 {
			t.Fatalf("A slice of Problems indicating any issues found in the metrics stream: %s", problems)
		}
	}

}

type validMetricWalker struct {
	Metrics []*dto.MetricFamily
}

func (w *validMetricWalker) walk(path string, info os.FileInfo, _ error) error {
	// only for regular files, not starting with a . and those that are go files.
	if !info.Mode().IsRegular() {
		return nil
	}
	// Is it appropriate to compare the file name equals metrics.go directly？
	if strings.HasPrefix(path, "../.") {
		return nil
	}
	if strings.HasSuffix(path, "_test.go") {
		return nil
	}
	if !strings.HasSuffix(path, ".go") {
		return nil
	}

	fs := token.NewFileSet()
	f, err := parser.ParseFile(fs, path, nil, parser.AllErrors)
	if err != nil {
		return err
	}
	l := &metric{}
	ast.Walk(l, f)
	if l.Metric != nil {
		w.Metrics = append(w.Metrics, l.Metric)
	}
	return nil
}

type metric struct {
	Metric *dto.MetricFamily
}

func (l metric) Visit(n ast.Node) ast.Visitor {
	if n == nil {
		return nil
	}
	ce, ok := n.(*ast.CallExpr)
	if !ok {
		return l
	}
	se, ok := ce.Fun.(*ast.SelectorExpr)
	if !ok {
		return l
	}
	id, ok := se.X.(*ast.Ident)
	if !ok {
		return l
	}
	if id.Name != "prometheus" { //prometheus
		return l
	}
	var metricsType dto.MetricType
	switch se.Sel.Name {
	case "NewCounterVec", "NewCounter":
		metricsType = dto.MetricType_COUNTER
	case "NewGaugeVec", "NewGauge":
		metricsType = dto.MetricType_GAUGE
	case "NewHistogramVec", "NewHistogram":
		metricsType = dto.MetricType_HISTOGRAM
	case "NewSummaryVec", "NewSummary":
		metricsType = dto.MetricType_SUMMARY
	default:
		return l
	}
	// Check first arg, that should have basic lit with capital
	if len(ce.Args) < 1 {
		return l
	}
	bl, ok := ce.Args[0].(*ast.CompositeLit)
	if !ok {
		return l
	}

	// parse Namespace Subsystem Name Help
	var subsystem, name, help string
	for _, elt := range bl.Elts {
		expr, ok := elt.(*ast.KeyValueExpr)
		if !ok {
			continue
		}
		object, ok := expr.Key.(*ast.Ident)
		if !ok {
			continue
		}
		value, ok := expr.Value.(*ast.BasicLit)
		if !ok {
			continue
		}
		switch object.Name {
		case "Subsystem":
			subsystem = value.Value
		case "Name":
			name = value.Value
		case "Help":
			help = value.Value
		}
	}

	// validate metrics field
	if len(name) == 0 || len(help) == 0 {
		return l
	}

	var metricName string
	if len(subsystem) > 0 {
		metricName = strings.Join([]string{plugin.Namespace, subsystem, name}, "_")
	} else {
		metricName = strings.Join([]string{plugin.Namespace, name}, "_")
	}
	l.Metric = &dto.MetricFamily{
		Name: &metricName,
		Help: &help,
		Type: &metricsType,
	}
	return l
}
<commit_msg>fix metric naming test (#4017)<commit_after>package test

import (
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/coredns/coredns/plugin"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil/promlint"
	dto "github.com/prometheus/client_model/go"
)

func TestMetricNaming(t *testing.T) {

	walker := validMetricWalker{}
	err := filepath.Walk("..", walker.walk)

	if err != nil {
		t.Fatal(err)
	}

	if len(walker.Metrics) > 0 {
		l := promlint.NewWithMetricFamilies(walker.Metrics)
		problems, err := l.Lint()
		if err != nil {
			t.Fatalf("Link found error: %s", err)
		}

		if len(problems) > 0 {
			t.Fatalf("A slice of Problems indicating any issues found in the metrics stream: %s", problems)
		}
	}

}

type validMetricWalker struct {
	Metrics []*dto.MetricFamily
}

func (w *validMetricWalker) walk(path string, info os.FileInfo, _ error) error {
	// only for regular files, not starting with a . and those that are go files.
	if !info.Mode().IsRegular() {
		return nil
	}
	// Is it appropriate to compare the file name equals metrics.go directly？
	if strings.HasPrefix(path, "../.") {
		return nil
	}
	if strings.HasSuffix(path, "_test.go") {
		return nil
	}
	if !strings.HasSuffix(path, ".go") {
		return nil
	}

	fs := token.NewFileSet()
	f, err := parser.ParseFile(fs, path, nil, parser.AllErrors)
	if err != nil {
		return err
	}
	l := &metric{}
	ast.Walk(l, f)
	if l.Metric != nil {
		w.Metrics = append(w.Metrics, l.Metric)
	}
	return nil
}

type metric struct {
	Metric *dto.MetricFamily
}

func (l *metric) Visit(n ast.Node) ast.Visitor {
	if n == nil {
		return nil
	}
	ce, ok := n.(*ast.CallExpr)
	if !ok {
		return l
	}
	se, ok := ce.Fun.(*ast.SelectorExpr)
	if !ok {
		return l
	}
	id, ok := se.X.(*ast.Ident)
	if !ok {
		return l
	}
	if id.Name != "prometheus" { //prometheus
		return l
	}
	var metricsType dto.MetricType
	switch se.Sel.Name {
	case "NewCounterVec", "NewCounter":
		metricsType = dto.MetricType_COUNTER
	case "NewGaugeVec", "NewGauge":
		metricsType = dto.MetricType_GAUGE
	case "NewHistogramVec", "NewHistogram":
		metricsType = dto.MetricType_HISTOGRAM
	case "NewSummaryVec", "NewSummary":
		metricsType = dto.MetricType_SUMMARY
	default:
		return l
	}
	// Check first arg, that should have basic lit with capital
	if len(ce.Args) < 1 {
		return l
	}
	bl, ok := ce.Args[0].(*ast.CompositeLit)
	if !ok {
		return l
	}

	// parse Namespace Subsystem Name Help
	var subsystem, name, help string
	for _, elt := range bl.Elts {
		expr, ok := elt.(*ast.KeyValueExpr)
		if !ok {
			continue
		}
		object, ok := expr.Key.(*ast.Ident)
		if !ok {
			continue
		}
		value, ok := expr.Value.(*ast.BasicLit)
		if !ok {
			continue
		}

		// remove quotes
		stringLiteral, err := strconv.Unquote(value.Value)
		if err != nil {
			return l
		}

		switch object.Name {
		case "Subsystem":
			subsystem = stringLiteral
		case "Name":
			name = stringLiteral
		case "Help":
			help = stringLiteral
		}
	}

	// validate metrics field
	if len(name) == 0 || len(help) == 0 {
		return l
	}

	metricName := prometheus.BuildFQName(plugin.Namespace, subsystem, name)
	l.Metric = &dto.MetricFamily{
		Name: &metricName,
		Help: &help,
		Type: &metricsType,
	}
	return l
}
<|endoftext|>


<commit_before>package main

import (
	"container/ring"
	"fmt"
)

type Color string

var colors = [...]Color{"white", "blue", "red", "yellow", "orange", "green"}
var edgesForFace = map[Color][]Color{
	"white":  {"red", "green", "orange", "blue"},
	"red":    {"blue", "yellow", "green", "white"},
	"blue":   {"white", "orange", "yellow", "red"},
	"yellow": {"green", "red", "blue", "orange"},
	"orange": {"yellow", "blue", "white", "green"},
	"green":  {"orange", "white", "red", "yellow"},
}

type Face [8]Color

type Edge [12]*Color

type Cube struct {
	faceMap map[Color]*Face
	edgeMap map[Color]Edge
}

type ThreeDTransformer struct {
	faceRing ring.Ring
	edgeRing ring.Ring
}

func main() {
	cube1 := new(Cube)
	face1 := &Face{"red", "red", "red", "red", "red", "red", "red", "red"}
	faceMap1 := make(map[Color]*Face)
	faceMap1["red"] = face1
	cube1.faceMap = faceMap1
	edge1 := Edge{&face1[0], &face1[1], &face1[2], &face1[3], &face1[4], &face1[5],
		&face1[6], &face1[7], &face1[0], &face1[1], &face1[2], &face1[3]}
	edgeMap1 := make(map[Color]Edge)
	edgeMap1["red"] = edge1
	cube1.edgeMap = edgeMap1
	*cube1.edgeMap["red"][0] = "blue"
	*cube1.edgeMap["red"][1] = "green"
	fmt.Println(cube1.faceMap["red"][0])
	fmt.Println(cube1.faceMap["red"][1])
	fmt.Println(cube1.faceMap["red"][2])
	fmt.Println(*cube1.edgeMap["red"][2])
}

<commit_msg>Added a method to the madness<commit_after>package main

import (
	"container/ring"
	"fmt"
)

type Color string

var colors = [...]Color{"white", "blue", "red", "yellow", "orange", "green"}
var edgesForFace = map[Color][]Color{
	"white":  {"red", "green", "orange", "blue"},
	"red":    {"blue", "yellow", "green", "white"},
	"blue":   {"white", "orange", "yellow", "red"},
	"yellow": {"green", "red", "blue", "orange"},
	"orange": {"yellow", "blue", "white", "green"},
	"green":  {"orange", "white", "red", "yellow"},
}

var edgePos = [...]int{0, 1, 2, 4, 3, 2, 4, 5, 6, 6, 7, 0}

type Face [8]Color

type Edge [12]*Color

type Cube struct {
	faceMap map[Color]*Face
	edgeMap map[Color]Edge
}

func New() (*Cube, error) {
	newFaceMap := make(map[Color]*Face)
	newEdgeMap := make(map[Color]Edge)
	for _, color := range colors {
		newFaceMap[color] = &Face{color, color, color, color, color, color, color, color}
	}
	i := 0
	for _, faceColor := range colors {
		var newEdge Edge
		for _, edgeColor := range edgesForFace[faceColor] {
			newEdge[i] = &newFaceMap[edgeColor][edgePos[i]]
			newEdge[i+1] = &newFaceMap[edgeColor][edgePos[i+1]]
			newEdge[i+2] = &newFaceMap[edgeColor][edgePos[i+2]]
			i += 3
		}
		newEdgeMap[faceColor] = newEdge
	}
	return &Cube{newFaceMap, newEdgeMap}, nil
}

type ThreeDTransformer struct {
	faceRing ring.Ring
	edgeRing ring.Ring
}

func main() {
	cube1 := new(Cube)
	face1 := &Face{"red", "red", "red", "red", "red", "red", "red", "red"}
	faceMap1 := make(map[Color]*Face)
	faceMap1["red"] = face1
	cube1.faceMap = faceMap1
	edge1 := Edge{&face1[0], &face1[1], &face1[2], &face1[3], &face1[4], &face1[5],
		&face1[6], &face1[7], &face1[0], &face1[1], &face1[2], &face1[3]}
	edgeMap1 := make(map[Color]Edge)
	edgeMap1["red"] = edge1
	cube1.edgeMap = edgeMap1
	*cube1.edgeMap["red"][0] = "blue"
	*cube1.edgeMap["red"][1] = "green"
	fmt.Println(cube1.faceMap["red"][0])
	fmt.Println(cube1.faceMap["red"][1])
	fmt.Println(cube1.faceMap["red"][2])
	fmt.Println(*cube1.edgeMap["red"][2])
}

<|endoftext|>


<commit_before>package main

import (
	"container/ring"
)

type Color string

var colors = [...]Color {"white","blue","red","yellow","orange","green"}

type Face [8]Color

type Edge [12]*Color

type Cube struct {
    faceMap map[Color]Face
    edgeMap map[Color]Edge
}

type ThreeDTransformer struct {
    faceRing ring.Ring
    edgeRing ring.Ring
}

func main() {
    cube1 := new(Cube)
    face1 := Face {"red","red","red","red","red","red","red","red"}
    faceMap1 := make(map[Color]Face)
    faceMap1["red"] = face1
    cube1.faceMap = faceMap1
}
<commit_msg>Extend painful exercise to test syntax for edges.<commit_after>package main

import (
	"container/ring"
)

type Color string

var colors = [...]Color {"white","blue","red","yellow","orange","green"}

type Face [8]Color

type Edge [12]*Color

type Cube struct {
    faceMap map[Color]Face
    edgeMap map[Color]Edge
}

type ThreeDTransformer struct {
    faceRing ring.Ring
    edgeRing ring.Ring
}

func main() {
    cube1 := new(Cube)
    face1 := Face {"red","red","red","red","red","red","red","red"}
    faceMap1 := make(map[Color]Face)
    faceMap1["red"] = face1
    cube1.faceMap = faceMap1
    edge1 := Edge {&face1[0], &face1[0], &face1[0], &face1[0], &face1[0], &face1[0],
                   &face1[0], &face1[0], &face1[0], &face1[0], &face1[0], &face1[0]}
    edgeMap1 := make(map[Color]Edge)
    edgeMap1["red"] = edge1
    cube1.edgeMap = edgeMap1
}
<|endoftext|>


<commit_before>package mail

import (
	"testing"
)

type getHeadersTest struct {

}

func TestGetHeaders(t *testing.T) {

}

type splitHeadersTest struct {

}

func TestSplitHeader(t *testing.T) {

}

type hdr struct {
	key, val string
}

type parseTest struct {
	orig string
	hdrs []string
	cont string
}

var parseTests = []parseTest{

}

func TestParse(t *testing.T) {
	for _, pt := range parseTests {
		m := Parse(pt.orig)
		m = m
	}
}
<commit_msg>Full splitHeadersTests<commit_after>package mail

import (
	"strings"
	"testing"
)

// Converts all newlines to CRLFs.
func crlf(s string) string {
	return strings.Replace(s, "\n", "\r\n", -1)
}

type getHeadersTest struct {

}

var getHeadersTests = []getHeadersTest{

}

func TestGetHeaders(t *testing.T) {

}

type splitHeadersTest struct {
	orig, key, val string
}

var splitHeadersTests = []splitHeadersTest{
	{
		`a: b`,
		`a`, `b`,
	},
	{
		`A1: cD`,
		`A1`, `cD`,
	},
	{
		crlf(`ab: cd
 ef`),
		`ab`, `cd ef`,
	},
	{
		crlf(`ab: cd
	ef
	dh`),
		`ab`, `cd	ef	dh`,
	},
}

func TestSplitHeader(t *testing.T) {
	for i, ht := range splitHeadersTests {
		k, v := splitHeader(ht.orig)
		if k != ht.key || v != ht.val {
			t.Errorf(`%d. splitHeader gave ("%s", "%s"), wanted ("%s", "%s")`,
				i, k, v, ht.key, ht.val)
		}
	}
}

type hdr struct {
	key, val string
}

type parseTest struct {
	orig string
	hdrs []string
	cont string
}

var parseTests = []parseTest{

}

func TestParse(t *testing.T) {
	for _, pt := range parseTests {
		m := Parse(pt.orig)
		m = m
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"

	I "github.com/brentp/irelate"
)

func init() {
	// so that output stops when piping to head.
	done := make(chan os.Signal, 1)

	signal.Notify(done, os.Interrupt, syscall.SIGIO, syscall.SIGPIPE)
	go func() {
		//for range done {
		// for travis
		for _ = range done {
			os.Exit(0)
		}
	}()
}

func main() {

	cpuProfile := flag.Bool("cpuProfile", false, "perform CPU profiling")
	flag.Parse()
	files := flag.Args()

	streams := make([]I.RelatableChannel, 0)
	for _, f := range files {
		// Streamer automatically returns a Relatalbe Channel for bam/gff/bed(.gz)
		streams = append(streams, I.Streamer(f))
	}

	if *cpuProfile {
		f, err := os.Create("irelate.cpu.pprof")
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	buf := bufio.NewWriter(os.Stdout)

	//for interval := range I.IRelate(merged, I.CheckRelatedByOverlap) {
	for interval := range I.IRelate(I.CheckRelatedByOverlap, false, 0, streams...) {
		// for bam output:
		// bam := *(interval).(*I.Bam)
		fmt.Fprintf(buf, "%s\t%d\t%d\t%d\n", interval.Chrom(), interval.Start(), interval.End(), len(interval.Related()))

	}
	buf.Flush()
}
<commit_msg>update main to new sig<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"

	I "github.com/brentp/irelate"
)

func init() {
	// so that output stops when piping to head.
	done := make(chan os.Signal, 1)

	signal.Notify(done, os.Interrupt, syscall.SIGIO, syscall.SIGPIPE)
	go func() {
		//for range done {
		// for travis
		for _ = range done {
			os.Exit(0)
		}
	}()
}

func main() {

	cpuProfile := flag.Bool("cpuProfile", false, "perform CPU profiling")
	flag.Parse()
	files := flag.Args()

	streams := make([]I.RelatableChannel, 0)
	for _, f := range files {
		// Streamer automatically returns a Relatalbe Channel for bam/gff/bed(.gz)
		streams = append(streams, I.Streamer(f))
	}

	if *cpuProfile {
		f, err := os.Create("irelate.cpu.pprof")
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	buf := bufio.NewWriter(os.Stdout)

	//for interval := range I.IRelate(merged, I.CheckRelatedByOverlap) {
	for interval := range I.IRelate(I.CheckRelatedByOverlap, 0, streams...) {
		// for bam output:
		// bam := *(interval).(*I.Bam)
		fmt.Fprintf(buf, "%s\t%d\t%d\t%d\n", interval.Chrom(), interval.Start(), interval.End(), len(interval.Related()))

	}
	buf.Flush()
}
<|endoftext|>


<commit_before>package morningStar

import (
	"../jsonHttp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"sync"
	"time"
)

const PERFORMANCE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=1&id=`
const VOLATILITE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=2&id=`
const REFRESH_DELAY = 18

var LIST_REQUEST = regexp.MustCompile(`^/list$`)
var PERF_REQUEST = regexp.MustCompile(`^/(.+?)$`)

var ISIN = regexp.MustCompile(`ISIN.:(\S+)`)
var LABEL = regexp.MustCompile(`<h1[^>]*?>((?:.|\n)*?)</h1>`)
var RATING = regexp.MustCompile(`<span\sclass=".*?stars([0-9]).*?">`)
var CATEGORY = regexp.MustCompile(`<span[^>]*?>Catégorie</span>.*?<span[^>]*?>(.*?)</span>`)
var PERF_ONE_MONTH = regexp.MustCompile(`<td[^>]*?>1 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_THREE_MONTH = regexp.MustCompile(`<td[^>]*?>3 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_SIX_MONTH = regexp.MustCompile(`<td[^>]*?>6 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_ONE_YEAR = regexp.MustCompile(`<td[^>]*?>1 an</td><td[^>]*?>(.*?)</td>`)
var VOL_3_YEAR = regexp.MustCompile(`<td[^>]*?>Ecart-type 3 ans.?</td><td[^>]*?>(.*?)</td>`)

type SyncedMap struct {
	sync.RWMutex
	m map[string]Performance
}

func (o SyncedMap) get(key string) (Performance, bool) {
	o.RLock()
	defer o.RUnlock()
	return o.m[key]
}

func (o SyncedMap) push(key string, performance Performance) {
	o.Lock()
	defer o.Unlock()
	o.m[key] = performance
}

var PERFORMANCE_CACHE = SyncedMap{m: make(map[string]Performance)}

type Performance struct {
	Id            string    `json:"id"`
	Isin          string    `json:"isin"`
	Label         string    `json:"label"`
	Category      string    `json:"category"`
	Rating        string    `json:"rating"`
	OneMonth      float64   `json:"1m"`
	ThreeMonth    float64   `json:"3m"`
	SixMonth      float64   `json:"6m"`
	OneYear       float64   `json:"1y"`
	VolThreeYears float64   `json:"v3y"`
	Score         float64   `json:"score"`
	Update        time.Time `json:"ts"`
}

type PerformanceAsync struct {
	performance *Performance
	err         error
}

type Search struct {
	Id    string `json:"i"`
	Label string `json:"n"`
}

type Results struct {
	Results interface{} `json:"results"`
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func getBody(url string) ([]byte, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, errors.New(`Error while retrieving data from ` + url)
	}

	if response.StatusCode >= 400 {
		return nil, errors.New(`Got error ` + strconv.Itoa(response.StatusCode) + ` while getting ` + url)
	}

	body, err := readBody(response.Body)
	if err != nil {
		return nil, errors.New(`Error while reading body of ` + url)
	}

	return body, nil
}

func getLabel(extract *regexp.Regexp, body []byte) []byte {
	match := extract.FindSubmatch(body)
	if match == nil {
		return nil
	}

	return bytes.Replace(match[1], []byte(`&`), []byte(`&`), -1)
}

func getPerformance(extract *regexp.Regexp, body []byte) float64 {
	dotResult := bytes.Replace(getLabel(extract, body), []byte(`,`), []byte(`.`), -1)
	percentageResult := bytes.Replace(dotResult, []byte(`%`), []byte(``), -1)
	trimResult := bytes.TrimSpace(percentageResult)

	result, err := strconv.ParseFloat(string(trimResult), 64)
	if err != nil {
		return 0.0
	}
	return result
}

func SinglePerformance(morningStarId []byte) (*Performance, error) {
	cleanId := string(bytes.ToLower(morningStarId))

	performance, ok := PERFORMANCE_CACHE.get(cleanId)

	if ok && time.Now().Add(time.Hour*-REFRESH_DELAY).Before(performance.Update) {
		return &performance, nil
	}

	performanceBody, err := getBody(PERFORMANCE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	volatiliteBody, err := getBody(VOLATILITE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	isin := string(getLabel(ISIN, performanceBody))
	label := string(getLabel(LABEL, performanceBody))
	rating := string(getLabel(RATING, performanceBody))
	category := string(getLabel(CATEGORY, performanceBody))
	oneMonth := getPerformance(PERF_ONE_MONTH, performanceBody)
	threeMonths := getPerformance(PERF_THREE_MONTH, performanceBody)
	sixMonths := getPerformance(PERF_SIX_MONTH, performanceBody)
	oneYear := getPerformance(PERF_ONE_YEAR, performanceBody)
	volThreeYears := getPerformance(VOL_3_YEAR, volatiliteBody)

	score := (0.25 * oneMonth) + (0.3 * threeMonths) + (0.25 * sixMonths) + (0.2 * oneYear) - (0.1 * volThreeYears)
	scoreTruncated := float64(int(score*100)) / 100

	performance = Performance{cleanId, isin, label, category, rating, oneMonth, threeMonths, sixMonths, oneYear, volThreeYears, scoreTruncated, time.Now()}

	PERFORMANCE_CACHE.push(cleanId, performance)

	return &performance, nil
}

func singlePerformanceAsync(morningStarId []byte, ch chan<- PerformanceAsync) {
	performance, err := SinglePerformance(morningStarId)
	ch <- PerformanceAsync{performance, err}
}

func singlePerformanceHandler(w http.ResponseWriter, morningStarId []byte) {
	performance, err := SinglePerformance(morningStarId)

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		jsonHttp.ResponseJson(w, *performance)
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
	listBody, err := readBody(r.Body)
	if err != nil {
		http.Error(w, `Error while reading body for list`, 500)
		return
	}

	if len(bytes.TrimSpace(listBody)) == 0 {
		jsonHttp.ResponseJson(w, Results{[0]Performance{}})
		return
	}

	ids := bytes.Split(listBody, []byte(`,`))
	size := len(ids)

	ch := make(chan PerformanceAsync, size)
	for _, id := range ids {
		go singlePerformanceAsync(id, ch)
	}

	results := make([]Performance, 0, size)
	for range ids {
		if performanceAsync := <-ch; performanceAsync.err == nil {
			results = append(results, *performanceAsync.performance)
		}
	}

	jsonHttp.ResponseJson(w, Results{results})
}

type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	urlPath := []byte(r.URL.Path)

	if LIST_REQUEST.Match(urlPath) {
		listHandler(w, r)
	} else if PERF_REQUEST.Match(urlPath) {
		singlePerformanceHandler(w, PERF_REQUEST.FindSubmatch(urlPath)[1])
	}
}
<commit_msg>Update morningStar.go<commit_after>package morningStar

import (
	"../jsonHttp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"sync"
	"time"
)

const PERFORMANCE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=1&id=`
const VOLATILITE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=2&id=`
const REFRESH_DELAY = 18

var LIST_REQUEST = regexp.MustCompile(`^/list$`)
var PERF_REQUEST = regexp.MustCompile(`^/(.+?)$`)

var ISIN = regexp.MustCompile(`ISIN.:(\S+)`)
var LABEL = regexp.MustCompile(`<h1[^>]*?>((?:.|\n)*?)</h1>`)
var RATING = regexp.MustCompile(`<span\sclass=".*?stars([0-9]).*?">`)
var CATEGORY = regexp.MustCompile(`<span[^>]*?>Catégorie</span>.*?<span[^>]*?>(.*?)</span>`)
var PERF_ONE_MONTH = regexp.MustCompile(`<td[^>]*?>1 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_THREE_MONTH = regexp.MustCompile(`<td[^>]*?>3 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_SIX_MONTH = regexp.MustCompile(`<td[^>]*?>6 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_ONE_YEAR = regexp.MustCompile(`<td[^>]*?>1 an</td><td[^>]*?>(.*?)</td>`)
var VOL_3_YEAR = regexp.MustCompile(`<td[^>]*?>Ecart-type 3 ans.?</td><td[^>]*?>(.*?)</td>`)

type SyncedMap struct {
	sync.RWMutex
	m map[string]Performance
}

func (o SyncedMap) get(key string) (Performance, bool) {
	o.RLock()
	defer o.RUnlock()

	performance, ok := o.m[key]
	return performance, ok
}

func (o SyncedMap) push(key string, performance Performance) {
	o.Lock()
	defer o.Unlock()
	o.m[key] = performance
}

var PERFORMANCE_CACHE = SyncedMap{m: make(map[string]Performance)}

type Performance struct {
	Id            string    `json:"id"`
	Isin          string    `json:"isin"`
	Label         string    `json:"label"`
	Category      string    `json:"category"`
	Rating        string    `json:"rating"`
	OneMonth      float64   `json:"1m"`
	ThreeMonth    float64   `json:"3m"`
	SixMonth      float64   `json:"6m"`
	OneYear       float64   `json:"1y"`
	VolThreeYears float64   `json:"v3y"`
	Score         float64   `json:"score"`
	Update        time.Time `json:"ts"`
}

type PerformanceAsync struct {
	performance *Performance
	err         error
}

type Search struct {
	Id    string `json:"i"`
	Label string `json:"n"`
}

type Results struct {
	Results interface{} `json:"results"`
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func getBody(url string) ([]byte, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, errors.New(`Error while retrieving data from ` + url)
	}

	if response.StatusCode >= 400 {
		return nil, errors.New(`Got error ` + strconv.Itoa(response.StatusCode) + ` while getting ` + url)
	}

	body, err := readBody(response.Body)
	if err != nil {
		return nil, errors.New(`Error while reading body of ` + url)
	}

	return body, nil
}

func getLabel(extract *regexp.Regexp, body []byte) []byte {
	match := extract.FindSubmatch(body)
	if match == nil {
		return nil
	}

	return bytes.Replace(match[1], []byte(`&`), []byte(`&`), -1)
}

func getPerformance(extract *regexp.Regexp, body []byte) float64 {
	dotResult := bytes.Replace(getLabel(extract, body), []byte(`,`), []byte(`.`), -1)
	percentageResult := bytes.Replace(dotResult, []byte(`%`), []byte(``), -1)
	trimResult := bytes.TrimSpace(percentageResult)

	result, err := strconv.ParseFloat(string(trimResult), 64)
	if err != nil {
		return 0.0
	}
	return result
}

func SinglePerformance(morningStarId []byte) (*Performance, error) {
	cleanId := string(bytes.ToLower(morningStarId))

	performance, ok := PERFORMANCE_CACHE.get(cleanId)

	if ok && time.Now().Add(time.Hour*-REFRESH_DELAY).Before(performance.Update) {
		return &performance, nil
	}

	performanceBody, err := getBody(PERFORMANCE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	volatiliteBody, err := getBody(VOLATILITE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	isin := string(getLabel(ISIN, performanceBody))
	label := string(getLabel(LABEL, performanceBody))
	rating := string(getLabel(RATING, performanceBody))
	category := string(getLabel(CATEGORY, performanceBody))
	oneMonth := getPerformance(PERF_ONE_MONTH, performanceBody)
	threeMonths := getPerformance(PERF_THREE_MONTH, performanceBody)
	sixMonths := getPerformance(PERF_SIX_MONTH, performanceBody)
	oneYear := getPerformance(PERF_ONE_YEAR, performanceBody)
	volThreeYears := getPerformance(VOL_3_YEAR, volatiliteBody)

	score := (0.25 * oneMonth) + (0.3 * threeMonths) + (0.25 * sixMonths) + (0.2 * oneYear) - (0.1 * volThreeYears)
	scoreTruncated := float64(int(score*100)) / 100

	performance = Performance{cleanId, isin, label, category, rating, oneMonth, threeMonths, sixMonths, oneYear, volThreeYears, scoreTruncated, time.Now()}

	PERFORMANCE_CACHE.push(cleanId, performance)

	return &performance, nil
}

func singlePerformanceAsync(morningStarId []byte, ch chan<- PerformanceAsync) {
	performance, err := SinglePerformance(morningStarId)
	ch <- PerformanceAsync{performance, err}
}

func singlePerformanceHandler(w http.ResponseWriter, morningStarId []byte) {
	performance, err := SinglePerformance(morningStarId)

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		jsonHttp.ResponseJson(w, *performance)
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
	listBody, err := readBody(r.Body)
	if err != nil {
		http.Error(w, `Error while reading body for list`, 500)
		return
	}

	if len(bytes.TrimSpace(listBody)) == 0 {
		jsonHttp.ResponseJson(w, Results{[0]Performance{}})
		return
	}

	ids := bytes.Split(listBody, []byte(`,`))
	size := len(ids)

	ch := make(chan PerformanceAsync, size)
	for _, id := range ids {
		go singlePerformanceAsync(id, ch)
	}

	results := make([]Performance, 0, size)
	for range ids {
		if performanceAsync := <-ch; performanceAsync.err == nil {
			results = append(results, *performanceAsync.performance)
		}
	}

	jsonHttp.ResponseJson(w, Results{results})
}

type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	urlPath := []byte(r.URL.Path)

	if LIST_REQUEST.Match(urlPath) {
		listHandler(w, r)
	} else if PERF_REQUEST.Match(urlPath) {
		singlePerformanceHandler(w, PERF_REQUEST.FindSubmatch(urlPath)[1])
	}
}
<|endoftext|>


<commit_before>package runner

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

// https://wfuzz.googlecode.com/svn/trunk/wordlist/Injections/SQL.txt
var fuzzList = `
'
"
#
-
--
'%20--
--';
'%20;
=%20'
=%20;
=%20--
\x23
\x27
\x3D%20\x3B'
\x3D%20\x27
\x27\x4F\x52 SELECT *
\x27\x6F\x72 SELECT *
'or%20select *
admin'--
<>"'%;)(&+
'%20or%20''='
'%20or%20'x'='x
"%20or%20"x"="x
')%20or%20('x'='x
0 or 1=1
' or 0=0 --
" or 0=0 --
or 0=0 --
' or 0=0 #
" or 0=0 #
or 0=0 #
' or 1=1--
" or 1=1--
' or '1'='1'--
"' or 1 --'"
or 1=1--
or%201=1
or%201=1 --
' or 1=1 or ''='
" or 1=1 or ""="
' or a=a--
" or "a"="a
') or ('a'='a
") or ("a"="a
hi" or "a"="a
hi" or 1=1 --
hi' or 1=1 --
hi' or 'a'='a
hi') or ('a'='a
hi") or ("a"="a
'hi' or 'x'='x';
@variable
,@variable
PRINT
PRINT @@variable
select
insert
as
or
procedure
limit
order by
asc
desc
delete
update
distinct
having
truncate
replace
like
handler
bfilename
' or username like '%
' or uname like '%
' or userid like '%
' or uid like '%
' or user like '%
exec xp
exec sp
'; exec master..xp_cmdshell
'; exec xp_regread
t'exec master..xp_cmdshell 'nslookup www.google.com'--
--sp_password
\x27UNION SELECT
' UNION SELECT
' UNION ALL SELECT
' or (EXISTS)
' (select top 1
'||UTL_HTTP.REQUEST
1;SELECT%20*
to_timestamp_tz
tz_offset
<>"'%;)(&+
'%20or%201=1
%27%20or%201=1
%20$(sleep%2050)
%20'sleep%2050'
char%4039%41%2b%40SELECT
'%20OR
'sqlattempt1
(sqlattempt2)
|
%7C
*|
%2A%7C
*(|(mail=*))
%2A%28%7C%28mail%3D%2A%29%29
*(|(objectclass=*))
%2A%28%7C%28objectclass%3D%2A%29%29
(
%28
)
%29
&
%26
!
%21
' or 1=1 or ''='
' or ''='
x' or 1=1 or 'x'='y
/
//
//*
*/*
`

func init() {
	fuzzList += "\b"
	fuzzList += "\n"
	fuzzList += "\n"
	fuzzList += "\r"
	fuzzList += "\t"
	fuzzList += "Hello\tworld"
}

func TestSQLInjectionBuilder(t *testing.T) {
	for i, fuzz := range strings.Split(fuzzList, "\n") {
		if i == 3 {
			return
		}
		if fuzz == "" {
			continue
		}
		fuzz = strings.Trim(fuzz, " \t")

		var id int64
		var comment string
		err := conn.
			InsertInto("comments").
			Columns("comment").
			Values(fuzz).
			SetIsInterpolated(true).
			Returning("id", "comment").
			QueryScalar(&id, &comment)

		assert.True(t, id > 0)
		assert.Equal(t, fuzz, comment)

		var result int
		err = conn.SQL(`
			SELECT 42
			FROM comments
			WHERE id = $1 AND comment = $2
		`, id, comment).QueryScalar(&result)

		assert.NoError(t, err)
		assert.Equal(t, 42, result)
	}
}

func TestSQLInjectionSQL(t *testing.T) {
	for _, fuzz := range strings.Split(fuzzList, "\n") {
		if fuzz == "" {
			continue
		}
		fuzz = strings.Trim(fuzz, " \t")

		var id int64
		var comment string
		err := conn.
			SQL(`
				INSERT INTO comments (comment)
				VALUES ($1)
				RETURNING id, comment
			`, fuzz).
			QueryScalar(&id, &comment)

		assert.True(t, id > 0)
		assert.Equal(t, fuzz, comment)

		var result int
		err = conn.SQL(`
			SELECT 42
			FROM comments
			WHERE id = $1 AND comment = $2
		`, id, comment).QueryScalar(&result)

		assert.NoError(t, err)
		assert.Equal(t, 42, result)
	}
}
<commit_msg>set is interpolated<commit_after>package runner

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

// https://wfuzz.googlecode.com/svn/trunk/wordlist/Injections/SQL.txt
var fuzzList = `
'
"
#
-
--
'%20--
--';
'%20;
=%20'
=%20;
=%20--
\x23
\x27
\x3D%20\x3B'
\x3D%20\x27
\x27\x4F\x52 SELECT *
\x27\x6F\x72 SELECT *
'or%20select *
admin'--
<>"'%;)(&+
'%20or%20''='
'%20or%20'x'='x
"%20or%20"x"="x
')%20or%20('x'='x
0 or 1=1
' or 0=0 --
" or 0=0 --
or 0=0 --
' or 0=0 #
" or 0=0 #
or 0=0 #
' or 1=1--
" or 1=1--
' or '1'='1'--
"' or 1 --'"
or 1=1--
or%201=1
or%201=1 --
' or 1=1 or ''='
" or 1=1 or ""="
' or a=a--
" or "a"="a
') or ('a'='a
") or ("a"="a
hi" or "a"="a
hi" or 1=1 --
hi' or 1=1 --
hi' or 'a'='a
hi') or ('a'='a
hi") or ("a"="a
'hi' or 'x'='x';
@variable
,@variable
PRINT
PRINT @@variable
select
insert
as
or
procedure
limit
order by
asc
desc
delete
update
distinct
having
truncate
replace
like
handler
bfilename
' or username like '%
' or uname like '%
' or userid like '%
' or uid like '%
' or user like '%
exec xp
exec sp
'; exec master..xp_cmdshell
'; exec xp_regread
t'exec master..xp_cmdshell 'nslookup www.google.com'--
--sp_password
\x27UNION SELECT
' UNION SELECT
' UNION ALL SELECT
' or (EXISTS)
' (select top 1
'||UTL_HTTP.REQUEST
1;SELECT%20*
to_timestamp_tz
tz_offset
<>"'%;)(&+
'%20or%201=1
%27%20or%201=1
%20$(sleep%2050)
%20'sleep%2050'
char%4039%41%2b%40SELECT
'%20OR
'sqlattempt1
(sqlattempt2)
|
%7C
*|
%2A%7C
*(|(mail=*))
%2A%28%7C%28mail%3D%2A%29%29
*(|(objectclass=*))
%2A%28%7C%28objectclass%3D%2A%29%29
(
%28
)
%29
&
%26
!
%21
' or 1=1 or ''='
' or ''='
x' or 1=1 or 'x'='y
/
//
//*
*/*
`

func init() {
	fuzzList += "\b"
	fuzzList += "\n"
	fuzzList += "\n"
	fuzzList += "\r"
	fuzzList += "\t"
	fuzzList += "Hello\tworld"
}

func TestSQLInjectionBuilder(t *testing.T) {
	for i, fuzz := range strings.Split(fuzzList, "\n") {
		if i == 3 {
			return
		}
		if fuzz == "" {
			continue
		}
		fuzz = strings.Trim(fuzz, " \t")

		var id int64
		var comment string
		err := conn.
			InsertInto("comments").
			Columns("comment").
			Values(fuzz).
			SetIsInterpolated(true).
			Returning("id", "comment").
			QueryScalar(&id, &comment)

		assert.True(t, id > 0)
		assert.Equal(t, fuzz, comment)

		var result int
		err = conn.SQL(`
			SELECT 42
			FROM comments
			WHERE id = $1 AND comment = $2
		`, id, comment).QueryScalar(&result)

		assert.NoError(t, err)
		assert.Equal(t, 42, result)
	}
}

func TestSQLInjectionSQL(t *testing.T) {
	for _, fuzz := range strings.Split(fuzzList, "\n") {
		if fuzz == "" {
			continue
		}
		fuzz = strings.Trim(fuzz, " \t")

		var id int64
		var comment string
		err := conn.
			SQL(`
				INSERT INTO comments (comment)
				VALUES ($1)
				RETURNING id, comment
			`, fuzz).
			SetIsInterpolated(true).
			QueryScalar(&id, &comment)

		assert.True(t, id > 0)
		assert.Equal(t, fuzz, comment)

		var result int
		err = conn.SQL(`
			SELECT 42
			FROM comments
			WHERE id = $1 AND comment = $2
		`, id, comment).QueryScalar(&result)

		assert.NoError(t, err)
		assert.Equal(t, 42, result)
	}
}
<|endoftext|>


<commit_before>package vantage

import (
	"fmt"
	"io"
	"log"
	"sort"
	"time"
)

const PAGE_COUNT = 513
const RECORDS_PER_PAGE = 5
const DATA_RECORD_LENGTH = 52

// Rev B
type ArchiveRecord struct {
	ArchiveTime     time.Time
	OutsideTemp     float32
	HighOutsideTemp float32
	LowOutsideTemp  float32
	Rainfall        int
	HighRainRate    int
	Barometer       float32
	SolarRad        int
	WindSamples     int
	InsideTemp      float32
	InsideHumidity  int
	OutsideHumidity int
	WindAvg         int
	WindMax         int
	WindMaxDir      int
	WindDir         int
	UVIndexAvg      float32
	ET              float32
	HighSolarRad    int
	UVIndexMax      int
	ForecastRule    int
	LeafTemp        []int //2
	LeafWetness     []int //2
	SoilTemp        []int //4
	RecordType      int
	ExtraHumidities []int //2
	ExtraTemps      []int //3
	SoilMoistures   []int //4
}

type sortedArchive []*ArchiveRecord

func (sa sortedArchive) Len() int           { return len(sa) }
func (sa sortedArchive) Swap(i, j int)      { sa[i], sa[j] = sa[j], sa[i] }
func (sa sortedArchive) Less(i, j int) bool { return sa[i].ArchiveTime.Before(sa[j].ArchiveTime) }

func (vc *Conn) GetArchiveRecords() ([]*ArchiveRecord, error) {
	ars := make(sortedArchive, 0, PAGE_COUNT*RECORDS_PER_PAGE)
	archiveChan := make(chan *ArchiveRecord, 10)
	errChan := make(chan error, 1)

	err := vc.GetArchiveStream(archiveChan, errChan)
	if err != nil {
		return nil, err
	}
	for {
		select {
		case ar := <-archiveChan:
			if ar == nil {
				// Channel closed
				sort.Sort(ars)
				return ars, nil
			}
			ars = append(ars, ar)
		case err = <-errChan:
			return nil, err
		}
	}
}

func (vc *Conn) GetArchiveStream(archiveChan chan *ArchiveRecord, errChan chan error) error {
	err := vc.sendAckCommand("DMP\n")
	if err != nil {
		return fmt.Errorf("DMP command failed: %v", err)
	}
	go vc.dmpArchive(archiveChan, errChan)
	return nil
}

func (vc *Conn) dmpArchive(archiveChan chan *ArchiveRecord, errChan chan error) {
	pkt := make([]byte, 267)
	for i := 0; i < PAGE_COUNT; i++ {
		vc.conn.SetReadDeadline(time.Now().Add(10 * time.Second))
		c, err := io.ReadFull(vc.buf, pkt)
		if err != nil {
			if c > 0 {
				log.Printf("Got bytes: %v", pkt[:c])
			}
			errChan <- fmt.Errorf("Error during DMP read: %v\n", err)
			return
		}
		ars, err := parseArchive(pkt)
		if err != nil {
			//TODO
		}
		for _, ar := range ars {
			archiveChan <- ar
		}
	}
	close(archiveChan)
}

func parseArchive(pkt []byte) ([]*ArchiveRecord, error) {
	ret := make([]*ArchiveRecord, 0, 5)
	for i := 0; i < 5; i++ {
		dr := pkt[i*DATA_RECORD_LENGTH+1 : (i+1)*DATA_RECORD_LENGTH]
		tm := parseArchiveTime(toInt(dr[0], dr[1]), toInt(dr[2], dr[3]))
		if tm == (time.Time{}) {
			continue
		}
		// TODO CRC
		ar := &ArchiveRecord{
			ArchiveTime:     tm,
			OutsideTemp:     float32(toInt(dr[4], dr[5])) / 10,
			HighOutsideTemp: float32(toInt(dr[6], dr[7])) / 10,
			LowOutsideTemp:  float32(toInt(dr[8], dr[9])) / 10,
			Rainfall:        toInt(dr[10], dr[11]),
			HighRainRate:    toInt(dr[12], dr[13]),
			Barometer:       float32(toInt(dr[14], dr[15])) / 1000,
			SolarRad:        toInt(dr[16], dr[17]),
			WindSamples:     toInt(dr[18], dr[19]),
			InsideTemp:      float32(toInt(dr[20], dr[21])) / 10,
			InsideHumidity:  int(dr[22]),
			OutsideHumidity: int(dr[23]),
			WindAvg:         int(dr[24]),
			WindMax:         int(dr[25]),
			WindMaxDir:      archiveDirectionLookup[int(26)],
			WindDir:         archiveDirectionLookup[int(27)],
			UVIndexAvg:      float32(int(dr[28])) / 10,
			ET:              float32(int(dr[29])) / 1000,
			HighSolarRad:    toInt(dr[30], dr[31]),
			UVIndexMax:      int(dr[32]),
			ForecastRule:    int(dr[33]),
			LeafTemp:        nil,
			LeafWetness:     nil,
			SoilTemp:        nil,
			RecordType:      int(dr[42]),
			ExtraHumidities: nil,
			ExtraTemps:      nil,
			SoilMoistures:   nil,
		}
		ret = append(ret, ar)
	}
	return ret, nil
}

var archiveDirectionLookup map[int]int = map[int]int{
	0:   0,   // N
	1:   22,  // NNE
	2:   45,  // NE
	3:   67,  // ENE
	4:   90,  // E
	5:   112, // ESE
	6:   135, // SE
	7:   157, // SSE
	8:   180, // S
	9:   202, // SSW
	10:  225, // SW
	11:  247, // WSW
	12:  270, // W
	13:  292, // WNW
	14:  315, // NW
	15:  337, // NNW
	255: 0,
}

func parseArchiveTime(dt, tm int) time.Time {
	if dt == 0 {
		return time.Time{}
	}
	day := dt & 0x1f                     // lower 5 bits
	month := time.Month((dt >> 5) & 0xF) // 4 bits
	year := (dt >> 9) + 2000             // 7 bits
	hour := tm / 100
	min := tm - hour

	return time.Date(year, month, day, hour, min, 0, 0, time.Local)
}
<commit_msg>adding more timeout<commit_after>package vantage

import (
	"fmt"
	"io"
	"log"
	"sort"
	"time"
)

const PAGE_COUNT = 513
const RECORDS_PER_PAGE = 5
const DATA_RECORD_LENGTH = 52

// Rev B
type ArchiveRecord struct {
	ArchiveTime     time.Time
	OutsideTemp     float32
	HighOutsideTemp float32
	LowOutsideTemp  float32
	Rainfall        int
	HighRainRate    int
	Barometer       float32
	SolarRad        int
	WindSamples     int
	InsideTemp      float32
	InsideHumidity  int
	OutsideHumidity int
	WindAvg         int
	WindMax         int
	WindMaxDir      int
	WindDir         int
	UVIndexAvg      float32
	ET              float32
	HighSolarRad    int
	UVIndexMax      int
	ForecastRule    int
	LeafTemp        []int //2
	LeafWetness     []int //2
	SoilTemp        []int //4
	RecordType      int
	ExtraHumidities []int //2
	ExtraTemps      []int //3
	SoilMoistures   []int //4
}

type sortedArchive []*ArchiveRecord

func (sa sortedArchive) Len() int           { return len(sa) }
func (sa sortedArchive) Swap(i, j int)      { sa[i], sa[j] = sa[j], sa[i] }
func (sa sortedArchive) Less(i, j int) bool { return sa[i].ArchiveTime.Before(sa[j].ArchiveTime) }

func (vc *Conn) GetArchiveRecords() ([]*ArchiveRecord, error) {
	ars := make(sortedArchive, 0, PAGE_COUNT*RECORDS_PER_PAGE)
	archiveChan := make(chan *ArchiveRecord, 10)
	errChan := make(chan error, 1)

	err := vc.GetArchiveStream(archiveChan, errChan)
	if err != nil {
		return nil, err
	}
	for {
		select {
		case ar := <-archiveChan:
			if ar == nil {
				// Channel closed
				sort.Sort(ars)
				return ars, nil
			}
			ars = append(ars, ar)
		case err = <-errChan:
			return nil, err
		}
	}
}

func (vc *Conn) GetArchiveStream(archiveChan chan *ArchiveRecord, errChan chan error) error {
	err := vc.sendAckCommand("DMP\n")
	if err != nil {
		return fmt.Errorf("DMP command failed: %v", err)
	}
	go vc.dmpArchive(archiveChan, errChan)
	return nil
}

func (vc *Conn) dmpArchive(archiveChan chan *ArchiveRecord, errChan chan error) {
	pkt := make([]byte, 267)
	for i := 0; i < PAGE_COUNT; i++ {
		vc.conn.SetReadDeadline(time.Now().Add(30 * time.Second))
		c, err := io.ReadFull(vc.buf, pkt)
		if err != nil {
			if c > 0 {
				log.Printf("Got bytes: %v", pkt[:c])
			}
			errChan <- fmt.Errorf("Error during DMP read: %v\n", err)
			return
		}
		ars, err := parseArchive(pkt)
		if err != nil {
			//TODO
		}
		for _, ar := range ars {
			archiveChan <- ar
		}
	}
	close(archiveChan)
}

func parseArchive(pkt []byte) ([]*ArchiveRecord, error) {
	ret := make([]*ArchiveRecord, 0, 5)
	for i := 0; i < 5; i++ {
		dr := pkt[i*DATA_RECORD_LENGTH+1 : (i+1)*DATA_RECORD_LENGTH]
		tm := parseArchiveTime(toInt(dr[0], dr[1]), toInt(dr[2], dr[3]))
		if tm == (time.Time{}) {
			continue
		}
		// TODO CRC
		ar := &ArchiveRecord{
			ArchiveTime:     tm,
			OutsideTemp:     float32(toInt(dr[4], dr[5])) / 10,
			HighOutsideTemp: float32(toInt(dr[6], dr[7])) / 10,
			LowOutsideTemp:  float32(toInt(dr[8], dr[9])) / 10,
			Rainfall:        toInt(dr[10], dr[11]),
			HighRainRate:    toInt(dr[12], dr[13]),
			Barometer:       float32(toInt(dr[14], dr[15])) / 1000,
			SolarRad:        toInt(dr[16], dr[17]),
			WindSamples:     toInt(dr[18], dr[19]),
			InsideTemp:      float32(toInt(dr[20], dr[21])) / 10,
			InsideHumidity:  int(dr[22]),
			OutsideHumidity: int(dr[23]),
			WindAvg:         int(dr[24]),
			WindMax:         int(dr[25]),
			WindMaxDir:      archiveDirectionLookup[int(26)],
			WindDir:         archiveDirectionLookup[int(27)],
			UVIndexAvg:      float32(int(dr[28])) / 10,
			ET:              float32(int(dr[29])) / 1000,
			HighSolarRad:    toInt(dr[30], dr[31]),
			UVIndexMax:      int(dr[32]),
			ForecastRule:    int(dr[33]),
			LeafTemp:        nil,
			LeafWetness:     nil,
			SoilTemp:        nil,
			RecordType:      int(dr[42]),
			ExtraHumidities: nil,
			ExtraTemps:      nil,
			SoilMoistures:   nil,
		}
		ret = append(ret, ar)
	}
	return ret, nil
}

var archiveDirectionLookup map[int]int = map[int]int{
	0:   0,   // N
	1:   22,  // NNE
	2:   45,  // NE
	3:   67,  // ENE
	4:   90,  // E
	5:   112, // ESE
	6:   135, // SE
	7:   157, // SSE
	8:   180, // S
	9:   202, // SSW
	10:  225, // SW
	11:  247, // WSW
	12:  270, // W
	13:  292, // WNW
	14:  315, // NW
	15:  337, // NNW
	255: 0,
}

func parseArchiveTime(dt, tm int) time.Time {
	if dt == 0 {
		return time.Time{}
	}
	day := dt & 0x1f                     // lower 5 bits
	month := time.Month((dt >> 5) & 0xF) // 4 bits
	year := (dt >> 9) + 2000             // 7 bits
	hour := tm / 100
	min := tm - hour

	return time.Date(year, month, day, hour, min, 0, 0, time.Local)
}
<|endoftext|>


<commit_before>package main

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

var (
	SUPPORTED_HTTP_METHODS         = []string{"GET", "POST", "PUT", "DELETE"}
	HTTP_METHODS_WITH_REQUEST_BODY = []string{"POST", "PUT", "DELETE"}
	server                         *RequestRecordingServer
	TEST_PORT                      = 8000
)

func sutUrl(path string) string{
	return fmt.Sprintf("http://localhost:%d%s",TEST_PORT,path)
}

var _ = BeforeSuite(func() {
	configureLogging()
})

var _ = Describe("Main", func() {

	var (
		exePath string
		err     error
	)

	BeforeEach(func() {
		exePath, err = filepath.Abs("./code-named-something")
		if err != nil {
			panic(err)
		}
		server = CreateRequestRecordingServer(TEST_PORT)
		server.Start()
	})

	AfterEach(func() {
		server.Clear()
		server.Stop()
	})

	It("Generate statistics of data from the execution", func() {
		list := []string{
			fmt.Sprintf(`%s -X POST -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X PUT -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X DELETE -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X GET`,sutUrl("/A")),
		}

		file := CreateFileFromLines(list)
		defer os.Remove(file.Name())
		cmd := exec.Command(exePath, "-f", file.Name())
		output, err := cmd.CombinedOutput()
		fmt.Println(string(output))
		Expect(err).To(BeNil())

		Expect(PathExists("./output.yml")).To(Equal(true))

		var executionOutput ExecutionOutput

		UnmarshalYamlFromFile("./output.yml", &executionOutput)

		Expect(executionOutput.Summary.TotalBytesSent).To(Equal(uint64(1)))
	})

	Describe("Support sending data with http request", func() {
		for _, method := range HTTP_METHODS_WITH_REQUEST_BODY {
			It(fmt.Sprintf("in the body for verb %s", method), func() {
				data := "a=1&b=2&c=3"
				list := []string{fmt.Sprintf(`%s -X %s -d %s`,sutUrl("/A"), method, data)}
				file := CreateFileFromLines(list)
				defer os.Remove(file.Name())
				cmd := exec.Command(exePath, "-f", file.Name())
				output, err := cmd.CombinedOutput()
				fmt.Println(string(output))
				Expect(err).To(BeNil())

				predicates := []HttpRequestPredicate{}
				predicates = append(predicates, RequestWithPath("/A"))
				predicates = append(predicates, RequestWithMethod(method))
				predicates = append(predicates, RequestWithBody(data))
				Expect(server.Find(predicates...)).To(Equal(true))
			})
		}

		It("in the querystring for verb GET", func() {
			method := "GET"
			data := "a=1&b=2&c=3"
			list := []string{fmt.Sprintf(`%s -X %s -d %s"`,sutUrl("/A"), method, data)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())

			predicates := []HttpRequestPredicate{}
			predicates = append(predicates, RequestWithPath("/A"))
			predicates = append(predicates, RequestWithMethod(method))
			predicates = append(predicates, RequestWithQuerystring(data))
			Expect(server.Find(predicates...)).To(Equal(true))
		})
	})

	for _, method := range SUPPORTED_HTTP_METHODS {
		It(fmt.Sprintf("Makes a http %s request with http headers", method), func() {
			applicationJson := "Content-Type:application/json"
			applicationSoapXml := "Accept:application/soap+xml"
			list := []string{fmt.Sprintf(`%s -X %s -H "%s" -H "%s"`,sutUrl("/A"), method, applicationJson, applicationSoapXml)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())

			predicates := []HttpRequestPredicate{}
			predicates = append(predicates, RequestWithPath("/A"))
			predicates = append(predicates, RequestWithMethod(method))
			predicates = append(predicates, RequestWithHeader("Content-Type", "application/json"))
			predicates = append(predicates, RequestWithHeader("Accept", "application/soap+xml"))
			Expect(server.Find(predicates...)).To(Equal(true))
		})
	}

	for _, method := range SUPPORTED_HTTP_METHODS {
		It(fmt.Sprintf("Makes a http %s request", method), func() {
			list := []string{fmt.Sprintf(`%s -X %s`,sutUrl("/A"), method)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())
			Expect(server.Find(RequestWithPath("/A"), RequestWithMethod(method))).To(Equal(true))
		})
	}

	It("Makes a http get request to each url in a file", func() {
		list := []string{
			sutUrl("/A"),
			sutUrl("/B"),
			sutUrl("/C"),
		}
		file := CreateFileFromLines(list)
		defer os.Remove(file.Name())

		cmd := exec.Command(exePath, "-f", file.Name())
		output, err := cmd.CombinedOutput()
		fmt.Println(string(output))

		Expect(err).To(BeNil())
		Expect(server.Find(RequestWithPath("/A"))).To(Equal(true))
		Expect(server.Find(RequestWithPath("/B"))).To(Equal(true))
		Expect(server.Find(RequestWithPath("/C"))).To(Equal(true))
	})
})
<commit_msg>Moved the start and stop of the server into a before and after suite function<commit_after>package main

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

var (
	SUPPORTED_HTTP_METHODS         = []string{"GET", "POST", "PUT", "DELETE"}
	HTTP_METHODS_WITH_REQUEST_BODY = []string{"POST", "PUT", "DELETE"}
	server                         *RequestRecordingServer
	TEST_PORT                      = 8000
)

func sutUrl(path string) string{
	return fmt.Sprintf("http://localhost:%d%s",TEST_PORT,path)
}

var _ = BeforeSuite(func() {
	configureLogging()
	server = CreateRequestRecordingServer(TEST_PORT)
	server.Start()
})

var _ = AfterSuite(func(){
	server.Stop()
})

var _ = Describe("Main", func() {

	var (
		exePath string
		err     error
	)

	BeforeEach(func() {
		exePath, err = filepath.Abs("./code-named-something")
		if err != nil {
			panic(err)
		}
	})

	AfterEach(func() {
		server.Clear()
	})

	It("Generate statistics of data from the execution", func() {
		list := []string{
			fmt.Sprintf(`%s -X POST -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X PUT -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X DELETE -H "Content-type:application/json" -d '{"name":"talula"}'`,sutUrl("/A")),
			fmt.Sprintf(`%s -X GET`,sutUrl("/A")),
		}

		file := CreateFileFromLines(list)
		defer os.Remove(file.Name())
		cmd := exec.Command(exePath, "-f", file.Name())
		output, err := cmd.CombinedOutput()
		fmt.Println(string(output))
		Expect(err).To(BeNil())

		Expect(PathExists("./output.yml")).To(Equal(true))

		var executionOutput ExecutionOutput

		UnmarshalYamlFromFile("./output.yml", &executionOutput)

		Expect(executionOutput.Summary.TotalBytesSent).To(Equal(uint64(1)))
	})

	Describe("Support sending data with http request", func() {
		for _, method := range HTTP_METHODS_WITH_REQUEST_BODY {
			It(fmt.Sprintf("in the body for verb %s", method), func() {
				data := "a=1&b=2&c=3"
				list := []string{fmt.Sprintf(`%s -X %s -d %s`,sutUrl("/A"), method, data)}
				file := CreateFileFromLines(list)
				defer os.Remove(file.Name())
				cmd := exec.Command(exePath, "-f", file.Name())
				output, err := cmd.CombinedOutput()
				fmt.Println(string(output))
				Expect(err).To(BeNil())

				predicates := []HttpRequestPredicate{}
				predicates = append(predicates, RequestWithPath("/A"))
				predicates = append(predicates, RequestWithMethod(method))
				predicates = append(predicates, RequestWithBody(data))
				Expect(server.Find(predicates...)).To(Equal(true))
			})
		}

		It("in the querystring for verb GET", func() {
			method := "GET"
			data := "a=1&b=2&c=3"
			list := []string{fmt.Sprintf(`%s -X %s -d %s"`,sutUrl("/A"), method, data)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())

			predicates := []HttpRequestPredicate{}
			predicates = append(predicates, RequestWithPath("/A"))
			predicates = append(predicates, RequestWithMethod(method))
			predicates = append(predicates, RequestWithQuerystring(data))
			Expect(server.Find(predicates...)).To(Equal(true))
		})
	})

	for _, method := range SUPPORTED_HTTP_METHODS {
		It(fmt.Sprintf("Makes a http %s request with http headers", method), func() {
			applicationJson := "Content-Type:application/json"
			applicationSoapXml := "Accept:application/soap+xml"
			list := []string{fmt.Sprintf(`%s -X %s -H "%s" -H "%s"`,sutUrl("/A"), method, applicationJson, applicationSoapXml)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())

			predicates := []HttpRequestPredicate{}
			predicates = append(predicates, RequestWithPath("/A"))
			predicates = append(predicates, RequestWithMethod(method))
			predicates = append(predicates, RequestWithHeader("Content-Type", "application/json"))
			predicates = append(predicates, RequestWithHeader("Accept", "application/soap+xml"))
			Expect(server.Find(predicates...)).To(Equal(true))
		})
	}

	for _, method := range SUPPORTED_HTTP_METHODS {
		It(fmt.Sprintf("Makes a http %s request", method), func() {
			list := []string{fmt.Sprintf(`%s -X %s`,sutUrl("/A"), method)}
			file := CreateFileFromLines(list)
			defer os.Remove(file.Name())
			cmd := exec.Command(exePath, "-f", file.Name())
			output, err := cmd.CombinedOutput()
			fmt.Println(string(output))
			Expect(err).To(BeNil())
			Expect(server.Find(RequestWithPath("/A"), RequestWithMethod(method))).To(Equal(true))
		})
	}

	It("Makes a http get request to each url in a file", func() {
		list := []string{
			sutUrl("/A"),
			sutUrl("/B"),
			sutUrl("/C"),
		}
		file := CreateFileFromLines(list)
		defer os.Remove(file.Name())

		cmd := exec.Command(exePath, "-f", file.Name())
		output, err := cmd.CombinedOutput()
		fmt.Println(string(output))

		Expect(err).To(BeNil())
		Expect(server.Find(RequestWithPath("/A"))).To(Equal(true))
		Expect(server.Find(RequestWithPath("/B"))).To(Equal(true))
		Expect(server.Find(RequestWithPath("/C"))).To(Equal(true))
	})
})
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestCheck(t *testing.T) {
	tests := []struct {
		startDate               time.Time // of cert
		daysRemaining           int       // of cert
		checkDaysExpiringWithin int
		err                     error
	}{
		{
			time.Now(),
			31,
			30,
			nil,
		},
		{
			time.Now(),
			29,
			30,
			errExpiringSoon,
		},
		{
			time.Now().Add(time.Hour * 24 * -7),
			-1,
			30,
			errExpired,
		},
	}

	for i, test := range tests {
		cert, key, err := genCertAndKey(test.startDate, time.Hour*24*time.Duration(test.daysRemaining))
		if err != nil {
			t.Fatalf("error generating test cert and key: %v", err)
		}

		tlsCert, err := tls.X509KeyPair(cert, key)
		if err != nil {
			t.Fatalf("error parsing test certificate: %v", err)
		}

		s := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
		s.TLS = &tls.Config{Certificates: []tls.Certificate{tlsCert}}
		s.StartTLS()

		_, port, err := net.SplitHostPort(s.Listener.Addr().String())
		if err != nil {
			t.Fatalf("error getting port of test TLS server: %v", err)
		}

		if err := check("127.0.0.1", port, test.checkDaysExpiringWithin, false); err != test.err {
			t.Errorf("%d: want %v, got %v", i, test.err, err)
		}
		s.Close()
	}
}

// Generate a self-signed X.509 cert and private key for testing a TLS server
func genCertAndKey(startDate time.Time, duration time.Duration) ([]byte, []byte, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return nil, nil, err
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, nil, err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Example Corp"},
		},
		NotBefore:   startDate,
		NotAfter:    startDate.Add(duration),
		KeyUsage:    x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}

	hosts := []string{"127.0.0.1", "::1", "example.com"}
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return nil, nil, err
	}

	cert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	key := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	return cert, key, nil
}
<commit_msg>Simplify test cert gen<commit_after>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestCheck(t *testing.T) {
	tests := []struct {
		startDate               time.Time // of cert
		daysRemaining           int       // of cert
		checkDaysExpiringWithin int
		err                     error
	}{
		{
			time.Now(),
			31,
			30,
			nil,
		},
		{
			time.Now(),
			29,
			30,
			errExpiringSoon,
		},
		{
			time.Now().Add(time.Hour * 24 * -7),
			-1,
			30,
			errExpired,
		},
	}

	for i, test := range tests {
		cert, err := genTestCert(test.startDate, time.Hour*24*time.Duration(test.daysRemaining))
		if err != nil {
			t.Fatalf("error generating test cert: %v", err)
		}

		s := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
		s.TLS = &tls.Config{Certificates: []tls.Certificate{cert}}
		s.StartTLS()

		_, port, err := net.SplitHostPort(s.Listener.Addr().String())
		if err != nil {
			t.Fatalf("error getting port of test TLS server: %v", err)
		}

		if err := check("127.0.0.1", port, test.checkDaysExpiringWithin, false); err != test.err {
			t.Errorf("%d: want %v, got %v", i, test.err, err)
		}
		s.Close()
	}
}

// Generate a self-signed cert for testing a TLS server. Start date and duration
// of the cert are configurable.
func genTestCert(startDate time.Time, duration time.Duration) (tls.Certificate, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return tls.Certificate{}, err
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return tls.Certificate{}, err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Example Corp"},
		},
		NotBefore:   startDate,
		NotAfter:    startDate.Add(duration),
		KeyUsage:    x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}

	hosts := []string{"127.0.0.1", "::1", "example.com"}
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return tls.Certificate{}, err
	}

	cert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	key := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return tls.Certificate{}, err
	}

	return tlsCert, nil
}
<|endoftext|>