<commit_before>package gokiq

import (
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"os/signal"
	"reflect"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/garyburd/redigo/redis"
)

type Job struct {
	Type  string        `json:"class"`
	Args  []interface{} `json:"args"`
	Queue string        `json:"queue,omitempty"`
	ID    string        `json:"jid"`

	Retry interface{} `json:"retry"` // can be int (number of retries) or bool (true means default)

	MaxRetries   int    `json:"-"`
	RetryCount   int    `json:"retry_count"`
	ErrorMessage string `json:"error_message,omitempty"`
	ErrorType    string `json:"error_class,omitempty"`
	RetriedAt    string `json:"retried_at,omitempty"`
	FailedAt     string `json:"failed_at,omitempty"`

	StartTime time.Time `json:"-"`
}

func (job *Job) FromJSON(data []byte) error {
	err := json.Unmarshal(data, job)
	if err != nil {
		return err
	}
	if max, ok := job.Retry.(float64); ok {
		job.MaxRetries = int(max)
	} else if r, ok := job.Retry.(bool); ok && !r {
	} else {
		job.MaxRetries = defaultMaxRetries
	}
	return nil
}

func (job *Job) JSON() []byte {
	res, _ := json.Marshal(job)
	return res
}

type message struct {
	job *Job
	die bool
}

const (
	TimestampFormat     = "2006-01-02 15:04:05 MST"
	dateFormat          = "2006-01-02"
	redisTimeout        = 1
	defaultMaxRetries   = 25
	defaultPollInterval = 5
	defaultWorkerCount  = 25
	defaultRedisServer  = "127.0.0.1:6379"
	keyExpiry           = 86400 // one day
)

type QueueConfig map[string]int

func (q QueueConfig) String() string {
	str := ""
	for queue, priority := range q {
		str += fmt.Sprintf("%s=%d,", queue, priority)
	}
	return str[:len(str)-1]
}

type Worker interface {
	Perform([]interface{}) error
}

var Workers = NewWorkerConfig()

type WorkerConfig struct {
	RedisServer    string // TODO: allow specifying redis db
	RedisNamespace string
	Queues         QueueConfig
	WorkerCount    int
	PollInterval   int
	ReportError    func(error, *Job) // TODO: pass in a stack trace for context

	workerMapping map[string]reflect.Type
	randomQueues  []string
	redisPool     *redis.Pool
	workQueue     chan message
	done          sync.WaitGroup
	sync.RWMutex  // R is locked by Run() and scheduler(), W is locked by quitHandler() when it receives a signal
}

func NewWorkerConfig() *WorkerConfig {
	return &WorkerConfig{
		RedisServer:   defaultRedisServer,
		PollInterval:  defaultPollInterval,
		WorkerCount:   defaultWorkerCount,
		Queues:        QueueConfig{"default": 1},
		ReportError:   func(error, *Job) {},
		workerMapping: make(map[string]reflect.Type),
		workQueue:     make(chan message),
	}
}

func (w *WorkerConfig) Register(name string, worker Worker) {
	w.workerMapping[name] = workerType(worker)
}

func (w *WorkerConfig) Run() {
	log.Printf(`state=starting worker_count=%d redis=%s/0/%s queues="%s" pid=%d`, w.WorkerCount, w.RedisServer, w.RedisNamespace, w.Queues, pid)
	w.denormalizeQueues()
	w.connectRedis()

	for i := 0; i < w.WorkerCount; i++ {
		go w.worker(workerID(i))
	}

	go w.scheduler()
	go w.quitHandler()

	log.Printf(`state=started pid=%d`, pid)
	for {
		w.run()
	}
}

func (w *WorkerConfig) run() {
	w.RLock() // don't let quitHandler() stop us in the middle of a job
	defer w.RUnlock()

	msg, err := redis.Values(w.redisQuery("BLPOP", append(w.queueList(), redisTimeout)...))
	if err == redis.ErrNil {
		return
	}
	if err != nil {
		w.handleError(err)
		time.Sleep(redisTimeout * time.Second) // likely a transient redis error, sleep before retrying
		return
	}

	job := &Job{}
	err = job.FromJSON(msg[1].([]byte))
	if err != nil {
		w.handleError(err)
		return
	}
	job.Queue = string(msg[0].([]byte)[len(w.nsKey("queue:")):])
	w.workQueue <- message{job: job}
}

// create a slice of queues with duplicates using the assigned frequencies
func (w *WorkerConfig) denormalizeQueues() {
	for queue, x := range w.Queues {
		for i := 0; i < x; i++ {
			w.randomQueues = append(w.randomQueues, w.nsKey("queue:"+queue))
		}
	}
}

// get a random slice of unique queues from the slice of denormalized queues
func (w *WorkerConfig) queueList() []interface{} {
	size := len(w.Queues)
	res := make([]interface{}, 0, size)
	queues := make(map[string]struct{}, size)

	indices := rand.Perm(len(w.randomQueues))[:size]
	for _, i := range indices {
		queue := w.randomQueues[i]
		if _, ok := queues[queue]; !ok {
			queues[queue] = struct{}{}
			res = append(res, queue)
		}
	}

	return res
}

func (w *WorkerConfig) handleError(err error) {
	log.Printf(`event=error error_type=%T error_message="%s" pid=%d`, err, err, pid)
	w.ReportError(err, nil)
}

// checks the sorted set of scheduled jobs and retries and queues them when it's time
// TODO: move this to a Lua script
func (w *WorkerConfig) scheduler() {
	pollSets := []string{w.nsKey("retry"), w.nsKey("schedule")}

	for _ = range time.Tick(time.Duration(w.PollInterval) * time.Second) {
		w.RLock() // don't let quitHandler() stop us in the middle of a run
		conn := w.redisPool.Get()
		now := fmt.Sprintf("%f", currentTimeFloat())
		for _, set := range pollSets {
			conn.Send("MULTI")
			conn.Send("ZRANGEBYSCORE", set, "-inf", now)
			conn.Send("ZREMRANGEBYSCORE", set, "-inf", now)
			res, err := redis.Values(conn.Do("EXEC"))
			if err != nil {
				w.handleError(err)
				continue
			}

			for _, msg := range res[0].([]interface{}) {
				parsedMsg := &struct {
					Queue string `json:"queue"`
				}{}
				msgBytes := msg.([]byte)
				err := json.Unmarshal(msgBytes, parsedMsg)
				if err != nil {
					w.handleError(err)
					continue
				}
				conn.Send("MULTI")
				conn.Send("SADD", w.nsKey("queues"), parsedMsg.Queue)
				conn.Send("RPUSH", w.nsKey("queue:"+parsedMsg.Queue), msgBytes)
				_, err = conn.Do("EXEC")
				if err != nil {
					w.handleError(err)
				}
			}
		}
		conn.Close()
		w.RUnlock()
	}
}

// listens for SIGINT, SIGTERM, and SIGQUIT to perform a clean shutdown
func (w *WorkerConfig) quitHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, syscall.SIGQUIT)

	for sig := range c {
		log.Printf("state=stopping signal=%s pid=%d", sig, pid)
		w.Lock()           // wait for the current run loop and scheduler iterations to finish
		close(w.workQueue) // tell worker goroutines to stop after they finish their current job
		for i := 0; i < w.WorkerCount; i++ {
			w.done.Wait()
		}
		log.Printf("state=stopped pid=%d", pid)
		os.Exit(0)
	}
}

func (w *WorkerConfig) connectRedis() {
	// TODO: add a mutex for the redis pool
	if w.redisPool != nil {
		w.redisPool.Close()
	}
	w.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", w.RedisServer)
	}, w.WorkerCount+1)
}

func (w *WorkerConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := w.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (w *WorkerConfig) worker(id string) {
	w.done.Add(1)
	for msg := range w.workQueue {
		if msg.die {
			break
		}

		job := msg.job
		typ, ok := w.workerMapping[msg.job.Type]
		if !ok {
			err := UnknownWorkerError{job.Type}
			w.scheduleRetry(job, err)
			continue
		}

		w.logJobStart(job, id)

		// wrap Perform() in a function so that we can recover from panics
		var err error
		func() {
			defer func() {
				if r := recover(); r != nil {
					// TODO: log stack trace
					err = panicToError(r)
				}
			}()
			err = reflect.New(typ).Interface().(Worker).Perform(msg.job.Args)
		}()
		if err != nil {
			w.scheduleRetry(job, err)
		}
		w.logJobFinish(job, id, err == nil)
	}
	w.done.Done()
}

func (w *WorkerConfig) scheduleRetry(job *Job, err error) {
	w.ReportError(err, job)

	now := time.Now().UTC().Format(TimestampFormat)
	if job.FailedAt == "" {
		job.FailedAt = now
	} else {
		job.RetryCount += 1
	}
	if job.RetryCount > 0 {
		job.RetriedAt = now
	}

	log.Printf(`event=job_error job_id=%s job_type=%s queue=%s retries=%d max_retries=%d error_type=%T error_message="%s" pid=%d`, job.ID, job.Type, job.Queue, job.RetryCount, job.MaxRetries, err, err, pid)

	if job.RetryCount < job.MaxRetries {
		job.ErrorType = fmt.Sprintf("%T", err)
		job.ErrorMessage = err.Error()

		nextRetry := currentTimeFloat() + retryDelay(job.RetryCount)

		w.redisQuery("ZADD", w.nsKey("retry"), strconv.FormatFloat(nextRetry, 'f', -1, 64), job.JSON())
	}
}

type runningJob struct {
	Queue     string `json:"queue"`
	Job       *Job   `json:"payload"`
	Timestamp int64  `json:"run_at"`
}

// TODO: make a lua script for this
func (w *WorkerConfig) logJobStart(job *Job, workerID string) {
	conn := w.redisPool.Get()
	defer conn.Close()

	conn.Send("MULTI")
	conn.Send("SADD", w.nsKey("workers"), workerID)
	conn.Send("SETEX", w.nsKey("worker:"+workerID+":started"), keyExpiry, time.Now().UTC().String())
	payload := &runningJob{job.Queue, job, time.Now().Unix()}
	json, _ := json.Marshal(payload)
	conn.Send("SETEX", w.nsKey("worker:"+workerID), keyExpiry, json)
	_, err := conn.Do("EXEC")
	if err != nil {
		w.handleError(err)
	}

	job.StartTime = time.Now()
	log.Printf("event=job_start job_id=%s job_type=%s queue=%s worker_id=%s pid=%d", job.ID, job.Type, job.Queue, workerID, pid)
}

// TODO: make a lua script for this
func (w *WorkerConfig) logJobFinish(job *Job, workerID string, success bool) {
	log.Printf("event=job_finish job_id=%s job_type=%s queue=%s duration=%v success=%t worker_id=%s pid=%d", job.ID, job.Type, job.Queue, time.Since(job.StartTime), success, workerID, pid)

	conn := w.redisPool.Get()
	defer conn.Close()

	date := time.Now().Format(dateFormat)
	conn.Send("MULTI")
	conn.Send("SREM", w.nsKey("workers"), workerID)
	conn.Send("DEL", w.nsKey("worker:"+workerID+":started"))
	conn.Send("DEL", w.nsKey("worker:"+workerID))
	conn.Send("INCR", w.nsKey("stat:processed"))
	conn.Send("INCR", w.nsKey("stat:processed:"+date))
	if !success {
		conn.Send("INCR", w.nsKey("stat:failed"))
		conn.Send("INCR", w.nsKey("stat:failed:"+date))
	}
	_, err := conn.Do("EXEC")
	if err != nil {
		w.handleError(err)
	}
}

func (w *WorkerConfig) nsKey(key string) string {
	if w.RedisNamespace != "" {
		return w.RedisNamespace + ":" + key
	}
	return key
}

// formula from Sidekiq (originally from delayed_job)
func retryDelay(count int) float64 {
	return math.Pow(float64(count), 4) + 15 + (float64(rand.Intn(30)) * (float64(count) + 1))
}

func currentTimeFloat() float64 {
	return float64(time.Now().UnixNano()) / float64(time.Second)
}

func panicToError(err interface{}) error {
	if str, ok := err.(string); ok {
		return fmt.Errorf(str)
	}
	return err.(error)
}

var (
	pid         = os.Getpid()
	hostname, _ = os.Hostname()
)

func workerID(i int) string {
	return fmt.Sprintf("%s:%d-%d", hostname, pid, i)
}

func workerType(worker Worker) reflect.Type {
	return reflect.Indirect(reflect.ValueOf(worker)).Type()
}

type UnknownWorkerError struct{ Type string }

func (e UnknownWorkerError) Error() string {
	return "gokiq: Unknown worker type: " + e.Type
}
<commit_msg>Make PollInterval a time.Duration<commit_after>package gokiq

import (
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"os/signal"
	"reflect"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/garyburd/redigo/redis"
)

type Job struct {
	Type  string        `json:"class"`
	Args  []interface{} `json:"args"`
	Queue string        `json:"queue,omitempty"`
	ID    string        `json:"jid"`

	Retry interface{} `json:"retry"` // can be int (number of retries) or bool (true means default)

	MaxRetries   int    `json:"-"`
	RetryCount   int    `json:"retry_count"`
	ErrorMessage string `json:"error_message,omitempty"`
	ErrorType    string `json:"error_class,omitempty"`
	RetriedAt    string `json:"retried_at,omitempty"`
	FailedAt     string `json:"failed_at,omitempty"`

	StartTime time.Time `json:"-"`
}

func (job *Job) FromJSON(data []byte) error {
	err := json.Unmarshal(data, job)
	if err != nil {
		return err
	}
	if max, ok := job.Retry.(float64); ok {
		job.MaxRetries = int(max)
	} else if r, ok := job.Retry.(bool); ok && !r {
	} else {
		job.MaxRetries = defaultMaxRetries
	}
	return nil
}

func (job *Job) JSON() []byte {
	res, _ := json.Marshal(job)
	return res
}

type message struct {
	job *Job
	die bool
}

const (
	TimestampFormat     = "2006-01-02 15:04:05 MST"
	dateFormat          = "2006-01-02"
	redisTimeout        = 1
	defaultMaxRetries   = 25
	defaultPollInterval = 5 * time.Second
	defaultWorkerCount  = 25
	defaultRedisServer  = "127.0.0.1:6379"
	keyExpiry           = 86400 // one day
)

type QueueConfig map[string]int

func (q QueueConfig) String() string {
	str := ""
	for queue, priority := range q {
		str += fmt.Sprintf("%s=%d,", queue, priority)
	}
	return str[:len(str)-1]
}

type Worker interface {
	Perform([]interface{}) error
}

var Workers = NewWorkerConfig()

type WorkerConfig struct {
	RedisServer    string // TODO: allow specifying redis db
	RedisNamespace string
	Queues         QueueConfig
	WorkerCount    int
	PollInterval   time.Duration
	ReportError    func(error, *Job) // TODO: pass in a stack trace for context

	workerMapping map[string]reflect.Type
	randomQueues  []string
	redisPool     *redis.Pool
	workQueue     chan message
	done          sync.WaitGroup
	sync.RWMutex  // R is locked by Run() and scheduler(), W is locked by quitHandler() when it receives a signal
}

func NewWorkerConfig() *WorkerConfig {
	return &WorkerConfig{
		RedisServer:   defaultRedisServer,
		PollInterval:  defaultPollInterval,
		WorkerCount:   defaultWorkerCount,
		Queues:        QueueConfig{"default": 1},
		ReportError:   func(error, *Job) {},
		workerMapping: make(map[string]reflect.Type),
		workQueue:     make(chan message),
	}
}

func (w *WorkerConfig) Register(name string, worker Worker) {
	w.workerMapping[name] = workerType(worker)
}

func (w *WorkerConfig) Run() {
	log.Printf(`state=starting worker_count=%d redis=%s/0/%s queues="%s" pid=%d`, w.WorkerCount, w.RedisServer, w.RedisNamespace, w.Queues, pid)
	w.denormalizeQueues()
	w.connectRedis()

	for i := 0; i < w.WorkerCount; i++ {
		go w.worker(workerID(i))
	}

	go w.scheduler()
	go w.quitHandler()

	log.Printf(`state=started pid=%d`, pid)
	for {
		w.run()
	}
}

func (w *WorkerConfig) run() {
	w.RLock() // don't let quitHandler() stop us in the middle of a job
	defer w.RUnlock()

	msg, err := redis.Values(w.redisQuery("BLPOP", append(w.queueList(), redisTimeout)...))
	if err == redis.ErrNil {
		return
	}
	if err != nil {
		w.handleError(err)
		time.Sleep(redisTimeout * time.Second) // likely a transient redis error, sleep before retrying
		return
	}

	job := &Job{}
	err = job.FromJSON(msg[1].([]byte))
	if err != nil {
		w.handleError(err)
		return
	}
	job.Queue = string(msg[0].([]byte)[len(w.nsKey("queue:")):])
	w.workQueue <- message{job: job}
}

// create a slice of queues with duplicates using the assigned frequencies
func (w *WorkerConfig) denormalizeQueues() {
	for queue, x := range w.Queues {
		for i := 0; i < x; i++ {
			w.randomQueues = append(w.randomQueues, w.nsKey("queue:"+queue))
		}
	}
}

// get a random slice of unique queues from the slice of denormalized queues
func (w *WorkerConfig) queueList() []interface{} {
	size := len(w.Queues)
	res := make([]interface{}, 0, size)
	queues := make(map[string]struct{}, size)

	indices := rand.Perm(len(w.randomQueues))[:size]
	for _, i := range indices {
		queue := w.randomQueues[i]
		if _, ok := queues[queue]; !ok {
			queues[queue] = struct{}{}
			res = append(res, queue)
		}
	}

	return res
}

func (w *WorkerConfig) handleError(err error) {
	log.Printf(`event=error error_type=%T error_message="%s" pid=%d`, err, err, pid)
	w.ReportError(err, nil)
}

// checks the sorted set of scheduled jobs and retries and queues them when it's time
// TODO: move this to a Lua script
func (w *WorkerConfig) scheduler() {
	pollSets := []string{w.nsKey("retry"), w.nsKey("schedule")}

	for _ = range time.Tick(w.PollInterval) {
		w.RLock() // don't let quitHandler() stop us in the middle of a run
		conn := w.redisPool.Get()
		now := fmt.Sprintf("%f", currentTimeFloat())
		for _, set := range pollSets {
			conn.Send("MULTI")
			conn.Send("ZRANGEBYSCORE", set, "-inf", now)
			conn.Send("ZREMRANGEBYSCORE", set, "-inf", now)
			res, err := redis.Values(conn.Do("EXEC"))
			if err != nil {
				w.handleError(err)
				continue
			}

			for _, msg := range res[0].([]interface{}) {
				parsedMsg := &struct {
					Queue string `json:"queue"`
				}{}
				msgBytes := msg.([]byte)
				err := json.Unmarshal(msgBytes, parsedMsg)
				if err != nil {
					w.handleError(err)
					continue
				}
				conn.Send("MULTI")
				conn.Send("SADD", w.nsKey("queues"), parsedMsg.Queue)
				conn.Send("RPUSH", w.nsKey("queue:"+parsedMsg.Queue), msgBytes)
				_, err = conn.Do("EXEC")
				if err != nil {
					w.handleError(err)
				}
			}
		}
		conn.Close()
		w.RUnlock()
	}
}

// listens for SIGINT, SIGTERM, and SIGQUIT to perform a clean shutdown
func (w *WorkerConfig) quitHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, syscall.SIGQUIT)

	for sig := range c {
		log.Printf("state=stopping signal=%s pid=%d", sig, pid)
		w.Lock()           // wait for the current run loop and scheduler iterations to finish
		close(w.workQueue) // tell worker goroutines to stop after they finish their current job
		for i := 0; i < w.WorkerCount; i++ {
			w.done.Wait()
		}
		log.Printf("state=stopped pid=%d", pid)
		os.Exit(0)
	}
}

func (w *WorkerConfig) connectRedis() {
	// TODO: add a mutex for the redis pool
	if w.redisPool != nil {
		w.redisPool.Close()
	}
	w.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", w.RedisServer)
	}, w.WorkerCount+1)
}

func (w *WorkerConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := w.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (w *WorkerConfig) worker(id string) {
	w.done.Add(1)
	for msg := range w.workQueue {
		if msg.die {
			break
		}

		job := msg.job
		typ, ok := w.workerMapping[msg.job.Type]
		if !ok {
			err := UnknownWorkerError{job.Type}
			w.scheduleRetry(job, err)
			continue
		}

		w.logJobStart(job, id)

		// wrap Perform() in a function so that we can recover from panics
		var err error
		func() {
			defer func() {
				if r := recover(); r != nil {
					// TODO: log stack trace
					err = panicToError(r)
				}
			}()
			err = reflect.New(typ).Interface().(Worker).Perform(msg.job.Args)
		}()
		if err != nil {
			w.scheduleRetry(job, err)
		}
		w.logJobFinish(job, id, err == nil)
	}
	w.done.Done()
}

func (w *WorkerConfig) scheduleRetry(job *Job, err error) {
	w.ReportError(err, job)

	now := time.Now().UTC().Format(TimestampFormat)
	if job.FailedAt == "" {
		job.FailedAt = now
	} else {
		job.RetryCount += 1
	}
	if job.RetryCount > 0 {
		job.RetriedAt = now
	}

	log.Printf(`event=job_error job_id=%s job_type=%s queue=%s retries=%d max_retries=%d error_type=%T error_message="%s" pid=%d`, job.ID, job.Type, job.Queue, job.RetryCount, job.MaxRetries, err, err, pid)

	if job.RetryCount < job.MaxRetries {
		job.ErrorType = fmt.Sprintf("%T", err)
		job.ErrorMessage = err.Error()

		nextRetry := currentTimeFloat() + retryDelay(job.RetryCount)

		w.redisQuery("ZADD", w.nsKey("retry"), strconv.FormatFloat(nextRetry, 'f', -1, 64), job.JSON())
	}
}

type runningJob struct {
	Queue     string `json:"queue"`
	Job       *Job   `json:"payload"`
	Timestamp int64  `json:"run_at"`
}

// TODO: make a lua script for this
func (w *WorkerConfig) logJobStart(job *Job, workerID string) {
	conn := w.redisPool.Get()
	defer conn.Close()

	conn.Send("MULTI")
	conn.Send("SADD", w.nsKey("workers"), workerID)
	conn.Send("SETEX", w.nsKey("worker:"+workerID+":started"), keyExpiry, time.Now().UTC().String())
	payload := &runningJob{job.Queue, job, time.Now().Unix()}
	json, _ := json.Marshal(payload)
	conn.Send("SETEX", w.nsKey("worker:"+workerID), keyExpiry, json)
	_, err := conn.Do("EXEC")
	if err != nil {
		w.handleError(err)
	}

	job.StartTime = time.Now()
	log.Printf("event=job_start job_id=%s job_type=%s queue=%s worker_id=%s pid=%d", job.ID, job.Type, job.Queue, workerID, pid)
}

// TODO: make a lua script for this
func (w *WorkerConfig) logJobFinish(job *Job, workerID string, success bool) {
	log.Printf("event=job_finish job_id=%s job_type=%s queue=%s duration=%v success=%t worker_id=%s pid=%d", job.ID, job.Type, job.Queue, time.Since(job.StartTime), success, workerID, pid)

	conn := w.redisPool.Get()
	defer conn.Close()

	date := time.Now().Format(dateFormat)
	conn.Send("MULTI")
	conn.Send("SREM", w.nsKey("workers"), workerID)
	conn.Send("DEL", w.nsKey("worker:"+workerID+":started"))
	conn.Send("DEL", w.nsKey("worker:"+workerID))
	conn.Send("INCR", w.nsKey("stat:processed"))
	conn.Send("INCR", w.nsKey("stat:processed:"+date))
	if !success {
		conn.Send("INCR", w.nsKey("stat:failed"))
		conn.Send("INCR", w.nsKey("stat:failed:"+date))
	}
	_, err := conn.Do("EXEC")
	if err != nil {
		w.handleError(err)
	}
}

func (w *WorkerConfig) nsKey(key string) string {
	if w.RedisNamespace != "" {
		return w.RedisNamespace + ":" + key
	}
	return key
}

// formula from Sidekiq (originally from delayed_job)
func retryDelay(count int) float64 {
	return math.Pow(float64(count), 4) + 15 + (float64(rand.Intn(30)) * (float64(count) + 1))
}

func currentTimeFloat() float64 {
	return float64(time.Now().UnixNano()) / float64(time.Second)
}

func panicToError(err interface{}) error {
	if str, ok := err.(string); ok {
		return fmt.Errorf(str)
	}
	return err.(error)
}

var (
	pid         = os.Getpid()
	hostname, _ = os.Hostname()
)

func workerID(i int) string {
	return fmt.Sprintf("%s:%d-%d", hostname, pid, i)
}

func workerType(worker Worker) reflect.Type {
	return reflect.Indirect(reflect.ValueOf(worker)).Type()
}

type UnknownWorkerError struct{ Type string }

func (e UnknownWorkerError) Error() string {
	return "gokiq: Unknown worker type: " + e.Type
}
<|endoftext|>


<commit_before>package main

import (
	zmq "github.com/pebbe/zmq4"
	"flag"
	"os"
	"fmt"
)

func Exitf(code int, format string, v... interface{}) {
	fmt.Fprintf(os.Stderr, "Fatal: ")
	fmt.Fprintf(os.Stderr, format, v...)
	os.Exit(code)
}

func main() {
	block := flag.Bool("block", true, "Use a blocking zmq_recv.")
	text := flag.Bool("text", false, "Read messages as strings.")
	multipart := flag.Bool("multipart", false, "Read multipart messages and print the parts individually.")
	flag.Usage = func() {
		Exitf(2, "Usage of %s:\n", os.Args[0])
		Exitf(2, "%s <zmq_endpoint> [opts]\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() < 1 {
		flag.Usage()
	}
	uri := flag.Arg(0)
	sock, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		Exitf(2, "Could not create ZMQ socket: %v", err)
	}
	err = sock.Connect(uri)
	if err != nil {
		Exitf(2, "Could not connect to endpoint %s: %v", uri, err)
	}
	err = sock.SetSubscribe("")
	if err != nil {
		Exitf(2, "Could not subscribe: %v", err)
	}
	var recvFlag zmq.Flag
	if *block {
		recvFlag = 0
	} else {
		recvFlag = zmq.DONTWAIT
	}
	var strMsg string
	var strMsgMulti []string
	var rawMsg []byte
	var rawMsgMulti [][]byte
	for {
		switch {
		case *multipart && *text:
			strMsgMulti, err = sock.RecvMessage(recvFlag)
		case *multipart:
			rawMsgMulti, err = sock.RecvMessageBytes(recvFlag)
		case *text:
			strMsg, err = sock.Recv(recvFlag)
		default:
			rawMsg, err = sock.RecvBytes(recvFlag)
		}
		if err != nil {
			Exitf(2, "Could not receive bytes: %v", err)
		}
		switch {
		case *multipart && *text:
			for _, part := range strMsgMulti {
				fmt.Println(part)
			}
		case *multipart:
			for _, part := range rawMsgMulti {
				os.Stdout.Write(part)
			}
		case *text:
			fmt.Println(strMsg)
		default:
			os.Stdout.Write(rawMsg)
		}
	}
}
<commit_msg>Don't die if we get EAGAIN or EINTR<commit_after>package main

import (
	"flag"
	"fmt"
	"syscall"
	zmq "github.com/pebbe/zmq4"
	"os"
)

func Exitf(code int, format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, "Fatal: ")
	fmt.Fprintf(os.Stderr, format, v...)
	os.Exit(code)
}

func main() {
	block := flag.Bool("block", true, "Use a blocking zmq_recv.")
	text := flag.Bool("text", false, "Read messages as strings.")
	multipart := flag.Bool("multipart", false, "Read multipart messages and print the parts individually.")
	flag.Usage = func() {
		Exitf(2, "Usage of %s:\n", os.Args[0])
		Exitf(2, "%s <zmq_endpoint> [opts]\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() < 1 {
		flag.Usage()
	}
	uri := flag.Arg(0)
	sock, err := zmq.NewSocket(zmq.SUB)
	if err != nil {
		Exitf(2, "Could not create ZMQ socket: %v", err)
	}
	err = sock.Connect(uri)
	if err != nil {
		Exitf(2, "Could not connect to endpoint %s: %v", uri, err)
	}
	err = sock.SetSubscribe("")
	if err != nil {
		Exitf(2, "Could not subscribe: %v", err)
	}
	var recvFlag zmq.Flag
	if *block {
		recvFlag = 0
	} else {
		recvFlag = zmq.DONTWAIT
	}
	var strMsg string
	var strMsgMulti []string
	var rawMsg []byte
	var rawMsgMulti [][]byte
	for {
		switch {
		case *multipart && *text:
			strMsgMulti, err = sock.RecvMessage(recvFlag)
		case *multipart:
			rawMsgMulti, err = sock.RecvMessageBytes(recvFlag)
		case *text:
			strMsg, err = sock.Recv(recvFlag)
		default:
			rawMsg, err = sock.RecvBytes(recvFlag)
		}
		if err != nil && err != syscall.EAGAIN && err != syscall.EINTR {
			Exitf(2, "Could not receive bytes: %v", err)
		}
		switch {
		case *multipart && *text:
			for _, part := range strMsgMulti {
				fmt.Println(part)
			}
		case *multipart:
			for _, part := range rawMsgMulti {
				os.Stdout.Write(part)
			}
		case *text:
			fmt.Println(strMsg)
		default:
			os.Stdout.Write(rawMsg)
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>Allow static credentials in API<commit_after><|endoftext|>


<commit_before><commit_msg>change port<commit_after><|endoftext|>


<commit_before>package cache

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Cache interface {
	Get(key string) (string, error)
	Put(key, directory string) error
}

const (
	DefaultInflatedCacheDir = "/tmp/risu/cache"
)

func NewCache(backend string) Cache {
	switch backend {
	case "local":
		return NewLocalFsCache()
	case "s3":
		return NewS3Cache()
	default:
		return NewLocalFsCache()
	}
}

func DeflateTarGz(tarGzPath, deflateDir string) error {
	tarFile, err := os.Create(tarGzPath)

	if err != nil {
		return err
	}
	defer tarFile.Close()

	gzipWriter := gzip.NewWriter(tarFile)
	defer gzipWriter.Close()

	tarGzWriter := tar.NewWriter(gzipWriter)
	defer tarGzWriter.Close()

	deflateDir, err = filepath.Abs(deflateDir)

	if err != nil {
		return err
	}

	walkDir(deflateDir, deflateDir, tarGzWriter)

	return nil
}

func InflateTarGz(tarGzPath, inflateDir string) error {
	file, err := os.Open(tarGzPath)

	if err != nil {
		return err
	}
	defer file.Close()

	gzfile, err := gzip.NewReader(file)

	if err != nil {
		return err
	}

	inflateDir, err = filepath.Abs(inflateDir)

	if err != nil {
		return err
	}

	if _, err = os.Stat(inflateDir); err != nil {
		os.MkdirAll(inflateDir, 0755)
	}

	reader := tar.NewReader(gzfile)

	for {
		header, err := reader.Next()

		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		buffer := new(bytes.Buffer)
		outPath := inflateDir + string(filepath.Separator) + header.Name

		switch header.Typeflag {
		case tar.TypeDir:
			if _, err = os.Stat(outPath); err != nil {
				os.MkdirAll(outPath, 0755)
			}

		case tar.TypeReg, tar.TypeRegA:
			if _, err = io.Copy(buffer, reader); err != nil {
				return err
			}

			if err = ioutil.WriteFile(outPath, buffer.Bytes(), 0644); err != nil {
				return err
			}
		}
	}

	return nil
}

func walkDir(walkingDir, baseDir string, tarGzWriter *tar.Writer) error {
	dir, err := os.Open(walkingDir)

	if err != nil {
		return err
	}
	defer dir.Close()

	files, err := dir.Readdir(0)

	if err != nil {
		return err
	}

	for _, fileInfo := range files {
		filePath := dir.Name() + string(filepath.Separator) + fileInfo.Name()

		if fileInfo.IsDir() {
			walkDir(filePath, baseDir, tarGzWriter)
		} else {
			writeTarGz(filePath, baseDir, tarGzWriter, fileInfo)
		}
	}

	return nil
}

func writeTarGz(filePath, baseDir string, tarGzWriter *tar.Writer, fileInfo os.FileInfo) error {
	file, err := os.Open(filePath)

	if err != nil {
		return err
	}
	defer file.Close()

	relativePath, err := filepath.Rel(baseDir, filePath)

	if err != nil {
		return err
	}

	header := new(tar.Header)
	header.Name = relativePath
	header.Size = fileInfo.Size()
	header.Mode = int64(fileInfo.Mode())
	header.ModTime = fileInfo.ModTime()

	err = tarGzWriter.WriteHeader(header)

	if err != nil {
		return err
	}

	_, err = io.Copy(tarGzWriter, file)

	if err != nil {
		return err
	}

	return nil
}

func getArchivedCacheFilePath(cacheDir, key string) string {
	return cacheDir + string(filepath.Separator) + key + ".tar.gz"
}

func getInflateDirPath(cacheDir, key string) string {
	return cacheDir + string(filepath.Separator) + key
}
<commit_msg>Remove path separator from generator<commit_after>package cache

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Cache interface {
	Get(key string) (string, error)
	Put(key, directory string) error
}

const (
	DefaultInflatedCacheDir = "/tmp/risu/cache"
)

func NewCache(backend string) Cache {
	switch backend {
	case "local":
		return NewLocalFsCache()
	case "s3":
		return NewS3Cache()
	default:
		return NewLocalFsCache()
	}
}

func DeflateTarGz(tarGzPath, deflateDir string) error {
	tarFile, err := os.Create(tarGzPath)

	if err != nil {
		return err
	}
	defer tarFile.Close()

	gzipWriter := gzip.NewWriter(tarFile)
	defer gzipWriter.Close()

	tarGzWriter := tar.NewWriter(gzipWriter)
	defer tarGzWriter.Close()

	deflateDir, err = filepath.Abs(deflateDir)

	if err != nil {
		return err
	}

	walkDir(deflateDir, deflateDir, tarGzWriter)

	return nil
}

func InflateTarGz(tarGzPath, inflateDir string) error {
	file, err := os.Open(tarGzPath)

	if err != nil {
		return err
	}
	defer file.Close()

	gzfile, err := gzip.NewReader(file)

	if err != nil {
		return err
	}

	inflateDir, err = filepath.Abs(inflateDir)

	if err != nil {
		return err
	}

	if _, err = os.Stat(inflateDir); err != nil {
		os.MkdirAll(inflateDir, 0755)
	}

	reader := tar.NewReader(gzfile)

	for {
		header, err := reader.Next()

		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		buffer := new(bytes.Buffer)
		outPath := inflateDir + string(filepath.Separator) + header.Name

		switch header.Typeflag {
		case tar.TypeDir:
			if _, err = os.Stat(outPath); err != nil {
				os.MkdirAll(outPath, 0755)
			}

		case tar.TypeReg, tar.TypeRegA:
			if _, err = io.Copy(buffer, reader); err != nil {
				return err
			}

			if err = ioutil.WriteFile(outPath, buffer.Bytes(), 0644); err != nil {
				return err
			}
		}
	}

	return nil
}

func walkDir(walkingDir, baseDir string, tarGzWriter *tar.Writer) error {
	dir, err := os.Open(walkingDir)

	if err != nil {
		return err
	}
	defer dir.Close()

	files, err := dir.Readdir(0)

	if err != nil {
		return err
	}

	for _, fileInfo := range files {
		filePath := dir.Name() + string(filepath.Separator) + fileInfo.Name()

		if fileInfo.IsDir() {
			walkDir(filePath, baseDir, tarGzWriter)
		} else {
			writeTarGz(filePath, baseDir, tarGzWriter, fileInfo)
		}
	}

	return nil
}

func writeTarGz(filePath, baseDir string, tarGzWriter *tar.Writer, fileInfo os.FileInfo) error {
	file, err := os.Open(filePath)

	if err != nil {
		return err
	}
	defer file.Close()

	relativePath, err := filepath.Rel(baseDir, filePath)

	if err != nil {
		return err
	}

	header := new(tar.Header)
	header.Name = relativePath
	header.Size = fileInfo.Size()
	header.Mode = int64(fileInfo.Mode())
	header.ModTime = fileInfo.ModTime()

	err = tarGzWriter.WriteHeader(header)

	if err != nil {
		return err
	}

	_, err = io.Copy(tarGzWriter, file)

	if err != nil {
		return err
	}

	return nil
}

func getArchivedCacheFilePath(cacheDir, key string) string {
	return cacheDir + key + ".tar.gz"
}

func getInflateDirPath(cacheDir, key string) string {
	return cacheDir + key
}
<|endoftext|>


<commit_before>package specs

import "os"

// LinuxStateDirectory holds the container's state information
const LinuxStateDirectory = "/run/opencontainer/containers"

// LinuxRuntimeSpec is the full specification for linux containers.
type LinuxRuntimeSpec struct {
	RuntimeSpec
	// LinuxRuntime is platform specific configuration for linux based containers.
	Linux LinuxRuntime `json:"linux"`
}

// LinuxRuntime hosts the Linux-only runtime information
type LinuxRuntime struct {
	// UIDMapping specifies user mappings for supporting user namespaces on linux.
	UIDMappings []IDMapping `json:"uidMappings"`
	// GIDMapping specifies group mappings for supporting user namespaces on linux.
	GIDMappings []IDMapping `json:"gidMappings"`
	// Rlimits specifies rlimit options to apply to the container's process.
	Rlimits []Rlimit `json:"rlimits"`
	// Sysctl are a set of key value pairs that are set for the container on start
	Sysctl map[string]string `json:"sysctl"`
	// Resources contain cgroup information for handling resource constraints
	// for the container
	Resources *Resources `json:"resources"`
	// CgroupsPath specifies the path to cgroups that are created and/or joined by the container.
	// The path is expected to be relative to the cgroups mountpoint.
	// If resources are specified, the cgroups at CgroupsPath will be updated based on resources.
	CgroupsPath string `json:"cgroupsPath"`
	// Namespaces contains the namespaces that are created and/or joined by the container
	Namespaces []Namespace `json:"namespaces"`
	// Devices are a list of device nodes that are created and enabled for the container
	Devices []Device `json:"devices"`
	// ApparmorProfile specified the apparmor profile for the container.
	ApparmorProfile string `json:"apparmorProfile"`
	// SelinuxProcessLabel specifies the selinux context that the container process is run as.
	SelinuxProcessLabel string `json:"selinuxProcessLabel"`
	// Seccomp specifies the seccomp security settings for the container.
	Seccomp Seccomp `json:"seccomp"`
	// RootfsPropagation is the rootfs mount propagation mode for the container
	RootfsPropagation string `json:"rootfsPropagation"`
}

// Namespace is the configuration for a linux namespace
type Namespace struct {
	// Type is the type of Linux namespace
	Type NamespaceType `json:"type"`
	// Path is a path to an existing namespace persisted on disk that can be joined
	// and is of the same type
	Path string `json:"path"`
}

// NamespaceType is one of the linux namespaces
type NamespaceType string

const (
	// PIDNamespace for isolating process IDs
	PIDNamespace NamespaceType = "pid"
	// NetworkNamespace for isolating network devices, stacks, ports, etc
	NetworkNamespace = "network"
	// MountNamespace for isolating mount points
	MountNamespace = "mount"
	// IPCNamespace for isolating System V IPC, POSIX message queues
	IPCNamespace = "ipc"
	// UTSNamespace for isolating hostname and NIS domain name
	UTSNamespace = "uts"
	// UserNamespace for isolating user and group IDs
	UserNamespace = "user"
)

// IDMapping specifies UID/GID mappings
type IDMapping struct {
	// HostID is the UID/GID of the host user or group
	HostID uint32 `json:"hostID"`
	// ContainerID is the UID/GID of the container's user or group
	ContainerID uint32 `json:"containerID"`
	// Size is the length of the range of IDs mapped between the two namespaces
	Size uint32 `json:"size"`
}

// Rlimit type and restrictions
type Rlimit struct {
	// Type of the rlimit to set
	Type string `json:"type"`
	// Hard is the hard limit for the specified type
	Hard uint64 `json:"hard"`
	// Soft is the soft limit for the specified type
	Soft uint64 `json:"soft"`
}

// HugepageLimit structure corresponds to limiting kernel hugepages
type HugepageLimit struct {
	// Pagesize is the hugepage size
	Pagesize string `json:"pageSize"`
	// Limit is the limit of "hugepagesize" hugetlb usage
	Limit uint64 `json:"limit"`
}

// InterfacePriority for network interfaces
type InterfacePriority struct {
	// Name is the name of the network interface
	Name string `json:"name"`
	// Priority for the interface
	Priority uint32 `json:"priority"`
}

// blockIODevice holds major:minor format supported in blkio cgroup
type blockIODevice struct {
	// Major is the device's major number.
	Major int64 `json:"major"`
	// Minor is the device's minor number.
	Minor int64 `json:"minor"`
}

// WeightDevice struct holds a `major:minor weight` pair for blkioWeightDevice
type WeightDevice struct {
	blockIODevice
	// Weight is the bandwidth rate for the device, range is from 10 to 1000
	Weight uint16 `json:"weight"`
	// LeafWeight is the bandwidth rate for the device while competing with the cgroup's child cgroups, range is from 10 to 1000, CFQ scheduler only
	LeafWeight uint16 `json:"leafWeight"`
}

// ThrottleDevice struct holds a `major:minor rate_per_second` pair
type ThrottleDevice struct {
	blockIODevice
	// Rate is the IO rate limit per cgroup per device
	Rate uint64 `json:"rate"`
}

// BlockIO for Linux cgroup 'blkio' resource management
type BlockIO struct {
	// Specifies per cgroup weight, range is from 10 to 1000
	Weight uint16 `json:"blkioWeight"`
	// Specifies tasks' weight in the given cgroup while competing with the cgroup's child cgroups, range is from 10 to 1000, CFQ scheduler only
	LeafWeight uint16 `json:"blkioLeafWeight"`
	// Weight per cgroup per device, can override BlkioWeight
	WeightDevice []*WeightDevice `json:"blkioWeightDevice"`
	// IO read rate limit per cgroup per device, bytes per second
	ThrottleReadBpsDevice []*ThrottleDevice `json:"blkioThrottleReadBpsDevice"`
	// IO write rate limit per cgroup per device, bytes per second
	ThrottleWriteBpsDevice []*ThrottleDevice `json:"blkioThrottleWriteBpsDevice"`
	// IO read rate limit per cgroup per device, IO per second
	ThrottleReadIOPSDevice []*ThrottleDevice `json:"blkioThrottleReadIOPSDevice"`
	// IO write rate limit per cgroup per device, IO per second
	ThrottleWriteIOPSDevice []*ThrottleDevice `json:"blkioThrottleWriteIOPSDevice"`
}

// Memory for Linux cgroup 'memory' resource management
type Memory struct {
	// Memory limit (in bytes)
	Limit uint64 `json:"limit"`
	// Memory reservation or soft_limit (in bytes)
	Reservation uint64 `json:"reservation"`
	// Total memory usage (memory + swap); set `-1' to disable swap
	Swap uint64 `json:"swap"`
	// Kernel memory limit (in bytes)
	Kernel uint64 `json:"kernel"`
	// How aggressive the kernel will swap memory pages. Range from 0 to 100. Set -1 to use system default
	Swappiness uint64 `json:"swappiness"`
}

// CPU for Linux cgroup 'cpu' resource management
type CPU struct {
	// CPU shares (relative weight vs. other cgroups with cpu shares)
	Shares uint64 `json:"shares"`
	// CPU hardcap limit (in usecs). Allowed cpu time in a given period
	Quota uint64 `json:"quota"`
	// CPU period to be used for hardcapping (in usecs). 0 to use system default
	Period uint64 `json:"period"`
	// How many time CPU will use in realtime scheduling (in usecs)
	RealtimeRuntime uint64 `json:"realtimeRuntime"`
	// CPU period to be used for realtime scheduling (in usecs)
	RealtimePeriod uint64 `json:"realtimePeriod"`
	// CPU to use within the cpuset
	Cpus string `json:"cpus"`
	// MEM to use within the cpuset
	Mems string `json:"mems"`
}

// Pids for Linux cgroup 'pids' resource management (Linux 4.3)
type Pids struct {
	// Maximum number of PIDs. A value <= 0 indicates "no limit".
	Limit int64 `json:"limit"`
}

// Network identification and priority configuration
type Network struct {
	// Set class identifier for container's network packets
	// this is actually a string instead of a uint64 to overcome the json
	// limitation of specifying hex numbers
	ClassID string `json:"classID"`
	// Set priority of network traffic for container
	Priorities []InterfacePriority `json:"priorities"`
}

// Resources has container runtime resource constraints
type Resources struct {
	// DisableOOMKiller disables the OOM killer for out of memory conditions
	DisableOOMKiller bool `json:"disableOOMKiller"`
	// Specify an oom_score_adj for the container. Optional.
	OOMScoreAdj int `json:"oomScoreAdj"`
	// Memory restriction configuration
	Memory Memory `json:"memory"`
	// CPU resource restriction configuration
	CPU CPU `json:"cpu"`
	// Task resource restriction configuration.
	Pids Pids `json:"pids"`
	// BlockIO restriction configuration
	BlockIO BlockIO `json:"blockIO"`
	// Hugetlb limit (in bytes)
	HugepageLimits []HugepageLimit `json:"hugepageLimits"`
	// Network restriction configuration
	Network Network `json:"network"`
}

// Device represents the information on a Linux special device file
type Device struct {
	// Path to the device.
	Path string `json:"path"`
	// Device type, block, char, etc.
	Type rune `json:"type"`
	// Major is the device's major number.
	Major int64 `json:"major"`
	// Minor is the device's minor number.
	Minor int64 `json:"minor"`
	// Cgroup permissions format, rwm.
	Permissions string `json:"permissions"`
	// FileMode permission bits for the device.
	FileMode os.FileMode `json:"fileMode"`
	// UID of the device.
	UID uint32 `json:"uid"`
	// Gid of the device.
	GID uint32 `json:"gid"`
}

// Seccomp represents syscall restrictions
type Seccomp struct {
	DefaultAction Action     `json:"defaultAction"`
	Architectures []Arch     `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// Additional architectures permitted to be used for system calls
// By default only the native architecture of the kernel is permitted
type Arch string

const (
	ArchX86         Arch = "SCMP_ARCH_X86"
	ArchX86_64      Arch = "SCMP_ARCH_X86_64"
	ArchX32         Arch = "SCMP_ARCH_X32"
	ArchARM         Arch = "SCMP_ARCH_ARM"
	ArchAARCH64     Arch = "SCMP_ARCH_AARCH64"
	ArchMIPS        Arch = "SCMP_ARCH_MIPS"
	ArchMIPS64      Arch = "SCMP_ARCH_MIPS64"
	ArchMIPS64N32   Arch = "SCMP_ARCH_MIPS64N32"
	ArchMIPSEL      Arch = "SCMP_ARCH_MIPSEL"
	ArchMIPSEL64    Arch = "SCMP_ARCH_MIPSEL64"
	ArchMIPSEL64N32 Arch = "SCMP_ARCH_MIPSEL64N32"
)

// Action taken upon Seccomp rule match
type Action string

const (
	ActKill  Action = "SCMP_ACT_KILL"
	ActTrap  Action = "SCMP_ACT_TRAP"
	ActErrno Action = "SCMP_ACT_ERRNO"
	ActTrace Action = "SCMP_ACT_TRACE"
	ActAllow Action = "SCMP_ACT_ALLOW"
)

// Operator used to match syscall arguments in Seccomp
type Operator string

const (
	OpNotEqual     Operator = "SCMP_CMP_NE"
	OpLessThan     Operator = "SCMP_CMP_LT"
	OpLessEqual    Operator = "SCMP_CMP_LE"
	OpEqualTo      Operator = "SCMP_CMP_EQ"
	OpGreaterEqual Operator = "SCMP_CMP_GE"
	OpGreaterThan  Operator = "SCMP_CMP_GT"
	OpMaskedEqual  Operator = "SCMP_CMP_MASKED_EQ"
)

// Arg used for matching specific syscall arguments in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"valueTwo"`
	Op       Operator `json:"op"`
}

// Syscall is used to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}
<commit_msg>Fix golint warnings<commit_after>package specs

import "os"

// LinuxStateDirectory holds the container's state information
const LinuxStateDirectory = "/run/opencontainer/containers"

// LinuxRuntimeSpec is the full specification for linux containers.
type LinuxRuntimeSpec struct {
	RuntimeSpec
	// LinuxRuntime is platform specific configuration for linux based containers.
	Linux LinuxRuntime `json:"linux"`
}

// LinuxRuntime hosts the Linux-only runtime information
type LinuxRuntime struct {
	// UIDMapping specifies user mappings for supporting user namespaces on linux.
	UIDMappings []IDMapping `json:"uidMappings"`
	// GIDMapping specifies group mappings for supporting user namespaces on linux.
	GIDMappings []IDMapping `json:"gidMappings"`
	// Rlimits specifies rlimit options to apply to the container's process.
	Rlimits []Rlimit `json:"rlimits"`
	// Sysctl are a set of key value pairs that are set for the container on start
	Sysctl map[string]string `json:"sysctl"`
	// Resources contain cgroup information for handling resource constraints
	// for the container
	Resources *Resources `json:"resources"`
	// CgroupsPath specifies the path to cgroups that are created and/or joined by the container.
	// The path is expected to be relative to the cgroups mountpoint.
	// If resources are specified, the cgroups at CgroupsPath will be updated based on resources.
	CgroupsPath string `json:"cgroupsPath"`
	// Namespaces contains the namespaces that are created and/or joined by the container
	Namespaces []Namespace `json:"namespaces"`
	// Devices are a list of device nodes that are created and enabled for the container
	Devices []Device `json:"devices"`
	// ApparmorProfile specified the apparmor profile for the container.
	ApparmorProfile string `json:"apparmorProfile"`
	// SelinuxProcessLabel specifies the selinux context that the container process is run as.
	SelinuxProcessLabel string `json:"selinuxProcessLabel"`
	// Seccomp specifies the seccomp security settings for the container.
	Seccomp Seccomp `json:"seccomp"`
	// RootfsPropagation is the rootfs mount propagation mode for the container
	RootfsPropagation string `json:"rootfsPropagation"`
}

// Namespace is the configuration for a linux namespace
type Namespace struct {
	// Type is the type of Linux namespace
	Type NamespaceType `json:"type"`
	// Path is a path to an existing namespace persisted on disk that can be joined
	// and is of the same type
	Path string `json:"path"`
}

// NamespaceType is one of the linux namespaces
type NamespaceType string

const (
	// PIDNamespace for isolating process IDs
	PIDNamespace NamespaceType = "pid"
	// NetworkNamespace for isolating network devices, stacks, ports, etc
	NetworkNamespace = "network"
	// MountNamespace for isolating mount points
	MountNamespace = "mount"
	// IPCNamespace for isolating System V IPC, POSIX message queues
	IPCNamespace = "ipc"
	// UTSNamespace for isolating hostname and NIS domain name
	UTSNamespace = "uts"
	// UserNamespace for isolating user and group IDs
	UserNamespace = "user"
)

// IDMapping specifies UID/GID mappings
type IDMapping struct {
	// HostID is the UID/GID of the host user or group
	HostID uint32 `json:"hostID"`
	// ContainerID is the UID/GID of the container's user or group
	ContainerID uint32 `json:"containerID"`
	// Size is the length of the range of IDs mapped between the two namespaces
	Size uint32 `json:"size"`
}

// Rlimit type and restrictions
type Rlimit struct {
	// Type of the rlimit to set
	Type string `json:"type"`
	// Hard is the hard limit for the specified type
	Hard uint64 `json:"hard"`
	// Soft is the soft limit for the specified type
	Soft uint64 `json:"soft"`
}

// HugepageLimit structure corresponds to limiting kernel hugepages
type HugepageLimit struct {
	// Pagesize is the hugepage size
	Pagesize string `json:"pageSize"`
	// Limit is the limit of "hugepagesize" hugetlb usage
	Limit uint64 `json:"limit"`
}

// InterfacePriority for network interfaces
type InterfacePriority struct {
	// Name is the name of the network interface
	Name string `json:"name"`
	// Priority for the interface
	Priority uint32 `json:"priority"`
}

// blockIODevice holds major:minor format supported in blkio cgroup
type blockIODevice struct {
	// Major is the device's major number.
	Major int64 `json:"major"`
	// Minor is the device's minor number.
	Minor int64 `json:"minor"`
}

// WeightDevice struct holds a `major:minor weight` pair for blkioWeightDevice
type WeightDevice struct {
	blockIODevice
	// Weight is the bandwidth rate for the device, range is from 10 to 1000
	Weight uint16 `json:"weight"`
	// LeafWeight is the bandwidth rate for the device while competing with the cgroup's child cgroups, range is from 10 to 1000, CFQ scheduler only
	LeafWeight uint16 `json:"leafWeight"`
}

// ThrottleDevice struct holds a `major:minor rate_per_second` pair
type ThrottleDevice struct {
	blockIODevice
	// Rate is the IO rate limit per cgroup per device
	Rate uint64 `json:"rate"`
}

// BlockIO for Linux cgroup 'blkio' resource management
type BlockIO struct {
	// Specifies per cgroup weight, range is from 10 to 1000
	Weight uint16 `json:"blkioWeight"`
	// Specifies tasks' weight in the given cgroup while competing with the cgroup's child cgroups, range is from 10 to 1000, CFQ scheduler only
	LeafWeight uint16 `json:"blkioLeafWeight"`
	// Weight per cgroup per device, can override BlkioWeight
	WeightDevice []*WeightDevice `json:"blkioWeightDevice"`
	// IO read rate limit per cgroup per device, bytes per second
	ThrottleReadBpsDevice []*ThrottleDevice `json:"blkioThrottleReadBpsDevice"`
	// IO write rate limit per cgroup per device, bytes per second
	ThrottleWriteBpsDevice []*ThrottleDevice `json:"blkioThrottleWriteBpsDevice"`
	// IO read rate limit per cgroup per device, IO per second
	ThrottleReadIOPSDevice []*ThrottleDevice `json:"blkioThrottleReadIOPSDevice"`
	// IO write rate limit per cgroup per device, IO per second
	ThrottleWriteIOPSDevice []*ThrottleDevice `json:"blkioThrottleWriteIOPSDevice"`
}

// Memory for Linux cgroup 'memory' resource management
type Memory struct {
	// Memory limit (in bytes)
	Limit uint64 `json:"limit"`
	// Memory reservation or soft_limit (in bytes)
	Reservation uint64 `json:"reservation"`
	// Total memory usage (memory + swap); set `-1' to disable swap
	Swap uint64 `json:"swap"`
	// Kernel memory limit (in bytes)
	Kernel uint64 `json:"kernel"`
	// How aggressive the kernel will swap memory pages. Range from 0 to 100. Set -1 to use system default
	Swappiness uint64 `json:"swappiness"`
}

// CPU for Linux cgroup 'cpu' resource management
type CPU struct {
	// CPU shares (relative weight vs. other cgroups with cpu shares)
	Shares uint64 `json:"shares"`
	// CPU hardcap limit (in usecs). Allowed cpu time in a given period
	Quota uint64 `json:"quota"`
	// CPU period to be used for hardcapping (in usecs). 0 to use system default
	Period uint64 `json:"period"`
	// How many time CPU will use in realtime scheduling (in usecs)
	RealtimeRuntime uint64 `json:"realtimeRuntime"`
	// CPU period to be used for realtime scheduling (in usecs)
	RealtimePeriod uint64 `json:"realtimePeriod"`
	// CPU to use within the cpuset
	Cpus string `json:"cpus"`
	// MEM to use within the cpuset
	Mems string `json:"mems"`
}

// Pids for Linux cgroup 'pids' resource management (Linux 4.3)
type Pids struct {
	// Maximum number of PIDs. A value <= 0 indicates "no limit".
	Limit int64 `json:"limit"`
}

// Network identification and priority configuration
type Network struct {
	// Set class identifier for container's network packets
	// this is actually a string instead of a uint64 to overcome the json
	// limitation of specifying hex numbers
	ClassID string `json:"classID"`
	// Set priority of network traffic for container
	Priorities []InterfacePriority `json:"priorities"`
}

// Resources has container runtime resource constraints
type Resources struct {
	// DisableOOMKiller disables the OOM killer for out of memory conditions
	DisableOOMKiller bool `json:"disableOOMKiller"`
	// Specify an oom_score_adj for the container. Optional.
	OOMScoreAdj int `json:"oomScoreAdj"`
	// Memory restriction configuration
	Memory Memory `json:"memory"`
	// CPU resource restriction configuration
	CPU CPU `json:"cpu"`
	// Task resource restriction configuration.
	Pids Pids `json:"pids"`
	// BlockIO restriction configuration
	BlockIO BlockIO `json:"blockIO"`
	// Hugetlb limit (in bytes)
	HugepageLimits []HugepageLimit `json:"hugepageLimits"`
	// Network restriction configuration
	Network Network `json:"network"`
}

// Device represents the information on a Linux special device file
type Device struct {
	// Path to the device.
	Path string `json:"path"`
	// Device type, block, char, etc.
	Type rune `json:"type"`
	// Major is the device's major number.
	Major int64 `json:"major"`
	// Minor is the device's minor number.
	Minor int64 `json:"minor"`
	// Cgroup permissions format, rwm.
	Permissions string `json:"permissions"`
	// FileMode permission bits for the device.
	FileMode os.FileMode `json:"fileMode"`
	// UID of the device.
	UID uint32 `json:"uid"`
	// Gid of the device.
	GID uint32 `json:"gid"`
}

// Seccomp represents syscall restrictions
type Seccomp struct {
	DefaultAction Action     `json:"defaultAction"`
	Architectures []Arch     `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// Arch used for additional architectures
type Arch string

// Additional architectures permitted to be used for system calls
// By default only the native architecture of the kernel is permitted
const (
	ArchX86         Arch = "SCMP_ARCH_X86"
	ArchX86_64      Arch = "SCMP_ARCH_X86_64"
	ArchX32         Arch = "SCMP_ARCH_X32"
	ArchARM         Arch = "SCMP_ARCH_ARM"
	ArchAARCH64     Arch = "SCMP_ARCH_AARCH64"
	ArchMIPS        Arch = "SCMP_ARCH_MIPS"
	ArchMIPS64      Arch = "SCMP_ARCH_MIPS64"
	ArchMIPS64N32   Arch = "SCMP_ARCH_MIPS64N32"
	ArchMIPSEL      Arch = "SCMP_ARCH_MIPSEL"
	ArchMIPSEL64    Arch = "SCMP_ARCH_MIPSEL64"
	ArchMIPSEL64N32 Arch = "SCMP_ARCH_MIPSEL64N32"
)

// Action taken upon Seccomp rule match
type Action string

// Define actions for Seccomp rules
const (
	ActKill  Action = "SCMP_ACT_KILL"
	ActTrap  Action = "SCMP_ACT_TRAP"
	ActErrno Action = "SCMP_ACT_ERRNO"
	ActTrace Action = "SCMP_ACT_TRACE"
	ActAllow Action = "SCMP_ACT_ALLOW"
)

// Operator used to match syscall arguments in Seccomp
type Operator string

// Define operators for syscall arguments in Seccomp
const (
	OpNotEqual     Operator = "SCMP_CMP_NE"
	OpLessThan     Operator = "SCMP_CMP_LT"
	OpLessEqual    Operator = "SCMP_CMP_LE"
	OpEqualTo      Operator = "SCMP_CMP_EQ"
	OpGreaterEqual Operator = "SCMP_CMP_GE"
	OpGreaterThan  Operator = "SCMP_CMP_GT"
	OpMaskedEqual  Operator = "SCMP_CMP_MASKED_EQ"
)

// Arg used for matching specific syscall arguments in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"valueTwo"`
	Op       Operator `json:"op"`
}

// Syscall is used to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}
<|endoftext|>


<commit_before>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package sa

import (
	"crypto/sha256"
	"crypto/x509"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
)

type SQLStorageAuthority struct {
	db     *sql.DB
	bucket map[string]interface{} // XXX included only for backward compat
	log    *blog.AuditLogger
}

func digest256(data []byte) []byte {
	d := sha256.New()
	_, _ = d.Write(data) // Never returns an error
	return d.Sum(nil)
}

func NewSQLStorageAuthority(logger *blog.AuditLogger, driver string, name string) (ssa *SQLStorageAuthority, err error) {
	logger.Notice("Storage Authority Starting")

	db, err := sql.Open(driver, name)
	if err != nil {
		return
	}
	if err = db.Ping(); err != nil {
		return
	}

	ssa = &SQLStorageAuthority{
		db:     db,
		log:    logger,
		bucket: make(map[string]interface{}),
	}

	err = ssa.InitTables()
	if err != nil {
		return
	}

	return
}

func (ssa *SQLStorageAuthority) InitTables() (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Create registrations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS registrations (id TEXT, thumbprint TEXT, value TEXT);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create pending authorizations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS pending_authz (id TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create finalized authorizations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS authz (sequence INTEGER, id TEXT, digest TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create certificates table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS certificates (serial STRING, digest TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) dumpTables(tx *sql.Tx) {
	fmt.Printf("===== TABLE DUMP =====\n")
	fmt.Printf("\n----- registrations -----\n")
	rows, err := tx.Query("SELECT id, thumbprint, value FROM registrations")
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
	} else {
		defer rows.Close()
		for rows.Next() {
			var id, key, value []byte
			if err := rows.Scan(&id, &key, &value); err == nil {
				fmt.Printf("%s | %s | %s\n", string(id), string(key), hex.EncodeToString(value))
			} else {
				fmt.Printf("ERROR: %v\n", err)
			}
		}
	}

	fmt.Printf("\n----- pending_authz -----\n") // TODO
	fmt.Printf("\n----- authz -----\n")         // TODO
	fmt.Printf("\n----- certificates -----\n")  // TODO
}

func statusIsPending(status core.AcmeStatus) bool {
	return status == core.StatusPending || status == core.StatusProcessing || status == core.StatusUnknown
}

func existingPending(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM pending_authz WHERE id = ?;", id).Scan(&count)
	return
}

func existingFinal(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM authz WHERE id = ?;", id).Scan(&count)
	return
}

func existingRegistration(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM registrations WHERE id = ?;", id).Scan(&count)
	return
}

func (ssa *SQLStorageAuthority) GetRegistration(id string) (reg core.Registration, err error) {
	var jsonReg []byte
	err = ssa.db.QueryRow("SELECT value FROM registrations WHERE id = ?;", id).Scan(&jsonReg)
	if err != nil {
		return
	}

	err = json.Unmarshal(jsonReg, &reg)
	return
}

func (ssa *SQLStorageAuthority) GetAuthorization(id string) (authz core.Authorization, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	var jsonAuthz []byte
	err = tx.QueryRow("SELECT value FROM pending_authz WHERE id = ?;", id).Scan(&jsonAuthz)
	switch {
	case err == sql.ErrNoRows:
		err = tx.QueryRow("SELECT value FROM authz WHERE id = ?;", id).Scan(&jsonAuthz)
		if err != nil {
			tx.Rollback()
			return
		}
	case err != nil:
		tx.Rollback()
		return
	}
	tx.Commit()

	err = json.Unmarshal(jsonAuthz, &authz)
	return
}

// GetCertificate takes an id consisting of the first, sequential half of a
// serial number and returns the first certificate whose full serial number is
// lexically greater than that id. This allows clients to query on the known
// sequential half of our serial numbers to enumerate all certificates.
// TODO: Add index on certificates table
// TODO: Implement error when there are multiple certificates with the same
// sequential half.
func (ssa *SQLStorageAuthority) GetCertificate(id string) (cert []byte, err error) {
	if len(id) != 16 {
		err = errors.New("Invalid certificate serial " + id)
	}
	err = ssa.db.QueryRow(
		"SELECT value FROM certificates WHERE serial LIKE ? LIMIT 1;",
		id + "%").Scan(&cert)
	return
}

func (ssa *SQLStorageAuthority) NewRegistration() (id string, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that it doesn't exist already
	candidate := core.NewToken()
	for existingRegistration(tx, candidate) > 0 {
		candidate = core.NewToken()
	}

	// Insert a stub row in pending
	_, err = tx.Exec("INSERT INTO registrations (id) VALUES (?);", candidate)
	if err != nil {
		tx.Rollback()
		return
	}

	if err = tx.Commit(); err != nil {
		return
	}

	id = candidate
	return
}

func (ssa *SQLStorageAuthority) UpdateRegistration(reg core.Registration) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	if existingRegistration(tx, reg.ID) != 1 {
		err = errors.New("Requested registration not found " + reg.ID)
		tx.Rollback()
		return
	}

	jsonReg, err := json.Marshal(reg)
	if err != nil {
		tx.Rollback()
		return
	}

	_, err = tx.Exec("UPDATE registrations SET thumbprint=?, value=? WHERE id = ?;", reg.Key.Thumbprint, string(jsonReg), reg.ID)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) NewPendingAuthorization() (id string, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that it doesn't exist already
	candidate := core.NewToken()
	for existingPending(tx, candidate) > 0 || existingFinal(tx, candidate) > 0 {
		candidate = core.NewToken()
	}

	// Insert a stub row in pending
	_, err = tx.Exec("INSERT INTO pending_authz (id) VALUES (?);", candidate)
	if err != nil {
		tx.Rollback()
		return
	}

	if err = tx.Commit(); err != nil {
		return
	}

	id = candidate
	return
}

func (ssa *SQLStorageAuthority) UpdatePendingAuthorization(authz core.Authorization) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	if !statusIsPending(authz.Status) {
		err = errors.New("Use Finalize() to update to a final status")
		tx.Rollback()
		return
	}

	if existingFinal(tx, authz.ID) > 0 {
		err = errors.New("Cannot update a final authorization")
		tx.Rollback()
		return
	}

	if existingPending(tx, authz.ID) != 1 {
		err = errors.New("Requested authorization not found " + authz.ID)
		tx.Rollback()
		return
	}

	jsonAuthz, err := json.Marshal(authz)
	if err != nil {
		tx.Rollback()
		return
	}

	_, err = tx.Exec("UPDATE pending_authz SET value = ? WHERE id = ?;", jsonAuthz, authz.ID)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) FinalizeAuthorization(authz core.Authorization) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that a pending authz exists
	if existingPending(tx, authz.ID) != 1 {
		err = errors.New("Cannot finalize a authorization that is not pending")
		tx.Rollback()
		return
	}
	if statusIsPending(authz.Status) {
		err = errors.New("Cannot finalize to a non-final status")
		tx.Rollback()
		return
	}

	// Manually set the index, to avoid AUTOINCREMENT issues
	var sequence int64
	var scanTarget sql.NullInt64
	err = tx.QueryRow("SELECT max(sequence) FROM authz").Scan(&scanTarget)
	switch {
	case !scanTarget.Valid:
		sequence = 0
	case err != nil:
		tx.Rollback()
		return
	default:
		sequence += scanTarget.Int64 + 1
	}

	jsonAuthz, err := json.Marshal(authz)
	if err != nil {
		tx.Rollback()
		return
	}
	digest := core.Fingerprint256(jsonAuthz)

	// Add to final table and delete from pending
	if err != nil {
		tx.Rollback()
		return
	}
	_, err = tx.Exec("INSERT INTO authz (sequence, id, digest, value) VALUES (?, ?, ?, ?);", sequence, authz.ID, digest, jsonAuthz)
	if err != nil {
		tx.Rollback()
		return
	}
	_, err = tx.Exec("DELETE FROM pending_authz WHERE id = ?;", authz.ID)
	if err != nil {
		tx.Rollback()
		return
	}
	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) AddCertificate(certDER []byte) (digest string, err error) {
	var parsedCertificate *x509.Certificate
	parsedCertificate, err = x509.ParseCertificate(certDER)
	if err != nil {
		return
	}
	serial := fmt.Sprintf("%016x", parsedCertificate.SerialNumber)

	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	digest = core.Fingerprint256(certDER)
	_, err = tx.Exec("INSERT INTO certificates (serial, digest, value) VALUES (?,?,?);",
		serial, digest, certDER)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}
<commit_msg>fix error i introduced somewhere<commit_after>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package sa

import (
	"crypto/sha256"
	"crypto/x509"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
)

type SQLStorageAuthority struct {
	db     *sql.DB
	bucket map[string]interface{} // XXX included only for backward compat
	log    *blog.AuditLogger
}

func digest256(data []byte) []byte {
	d := sha256.New()
	_, _ = d.Write(data) // Never returns an error
	return d.Sum(nil)
}

func NewSQLStorageAuthority(logger *blog.AuditLogger, driver string, name string) (ssa *SQLStorageAuthority, err error) {
	logger.Notice("Storage Authority Starting")

	db, err := sql.Open(driver, name)
	if err != nil {
		return
	}
	if err = db.Ping(); err != nil {
		return
	}

	ssa = &SQLStorageAuthority{
		db:     db,
		log:    logger,
		bucket: make(map[string]interface{}),
	}

	err = ssa.InitTables()
	if err != nil {
		return
	}

	return
}

func (ssa *SQLStorageAuthority) InitTables() (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Create registrations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS registrations (id TEXT, thumbprint TEXT, value TEXT);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create pending authorizations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS pending_authz (id TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create finalized authorizations table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS authz (sequence INTEGER, id TEXT, digest TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	// Create certificates table
	_, err = tx.Exec("CREATE TABLE IF NOT EXISTS certificates (serial TEXT, digest TEXT, value BLOB);")
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) dumpTables(tx *sql.Tx) {
	fmt.Printf("===== TABLE DUMP =====\n")
	fmt.Printf("\n----- registrations -----\n")
	rows, err := tx.Query("SELECT id, thumbprint, value FROM registrations")
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
	} else {
		defer rows.Close()
		for rows.Next() {
			var id, key, value []byte
			if err := rows.Scan(&id, &key, &value); err == nil {
				fmt.Printf("%s | %s | %s\n", string(id), string(key), hex.EncodeToString(value))
			} else {
				fmt.Printf("ERROR: %v\n", err)
			}
		}
	}

	fmt.Printf("\n----- pending_authz -----\n") // TODO
	fmt.Printf("\n----- authz -----\n")         // TODO
	fmt.Printf("\n----- certificates -----\n")  // TODO
}

func statusIsPending(status core.AcmeStatus) bool {
	return status == core.StatusPending || status == core.StatusProcessing || status == core.StatusUnknown
}

func existingPending(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM pending_authz WHERE id = ?;", id).Scan(&count)
	return
}

func existingFinal(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM authz WHERE id = ?;", id).Scan(&count)
	return
}

func existingRegistration(tx *sql.Tx, id string) (count int64) {
	tx.QueryRow("SELECT count(*) FROM registrations WHERE id = ?;", id).Scan(&count)
	return
}

func (ssa *SQLStorageAuthority) GetRegistration(id string) (reg core.Registration, err error) {
	var jsonReg []byte
	err = ssa.db.QueryRow("SELECT value FROM registrations WHERE id = ?;", id).Scan(&jsonReg)
	if err != nil {
		return
	}

	err = json.Unmarshal(jsonReg, &reg)
	return
}

func (ssa *SQLStorageAuthority) GetAuthorization(id string) (authz core.Authorization, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	var jsonAuthz []byte
	err = tx.QueryRow("SELECT value FROM pending_authz WHERE id = ?;", id).Scan(&jsonAuthz)
	switch {
	case err == sql.ErrNoRows:
		err = tx.QueryRow("SELECT value FROM authz WHERE id = ?;", id).Scan(&jsonAuthz)
		if err != nil {
			tx.Rollback()
			return
		}
	case err != nil:
		tx.Rollback()
		return
	}
	tx.Commit()

	err = json.Unmarshal(jsonAuthz, &authz)
	return
}

// GetCertificate takes an id consisting of the first, sequential half of a
// serial number and returns the first certificate whose full serial number is
// lexically greater than that id. This allows clients to query on the known
// sequential half of our serial numbers to enumerate all certificates.
// TODO: Add index on certificates table
// TODO: Implement error when there are multiple certificates with the same
// sequential half.
func (ssa *SQLStorageAuthority) GetCertificate(id string) (cert []byte, err error) {
	if len(id) != 16 {
		err = errors.New("Invalid certificate serial " + id)
	}
	err = ssa.db.QueryRow(
		"SELECT value FROM certificates WHERE serial LIKE ? LIMIT 1;",
		id + "%").Scan(&cert)
	return
}

func (ssa *SQLStorageAuthority) NewRegistration() (id string, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that it doesn't exist already
	candidate := core.NewToken()
	for existingRegistration(tx, candidate) > 0 {
		candidate = core.NewToken()
	}

	// Insert a stub row in pending
	_, err = tx.Exec("INSERT INTO registrations (id) VALUES (?);", candidate)
	if err != nil {
		tx.Rollback()
		return
	}

	if err = tx.Commit(); err != nil {
		return
	}

	id = candidate
	return
}

func (ssa *SQLStorageAuthority) UpdateRegistration(reg core.Registration) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	if existingRegistration(tx, reg.ID) != 1 {
		err = errors.New("Requested registration not found " + reg.ID)
		tx.Rollback()
		return
	}

	jsonReg, err := json.Marshal(reg)
	if err != nil {
		tx.Rollback()
		return
	}

	_, err = tx.Exec("UPDATE registrations SET thumbprint=?, value=? WHERE id = ?;", reg.Key.Thumbprint, string(jsonReg), reg.ID)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) NewPendingAuthorization() (id string, err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that it doesn't exist already
	candidate := core.NewToken()
	for existingPending(tx, candidate) > 0 || existingFinal(tx, candidate) > 0 {
		candidate = core.NewToken()
	}

	// Insert a stub row in pending
	_, err = tx.Exec("INSERT INTO pending_authz (id) VALUES (?);", candidate)
	if err != nil {
		tx.Rollback()
		return
	}

	if err = tx.Commit(); err != nil {
		return
	}

	id = candidate
	return
}

func (ssa *SQLStorageAuthority) UpdatePendingAuthorization(authz core.Authorization) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	if !statusIsPending(authz.Status) {
		err = errors.New("Use Finalize() to update to a final status")
		tx.Rollback()
		return
	}

	if existingFinal(tx, authz.ID) > 0 {
		err = errors.New("Cannot update a final authorization")
		tx.Rollback()
		return
	}

	if existingPending(tx, authz.ID) != 1 {
		err = errors.New("Requested authorization not found " + authz.ID)
		tx.Rollback()
		return
	}

	jsonAuthz, err := json.Marshal(authz)
	if err != nil {
		tx.Rollback()
		return
	}

	_, err = tx.Exec("UPDATE pending_authz SET value = ? WHERE id = ?;", jsonAuthz, authz.ID)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) FinalizeAuthorization(authz core.Authorization) (err error) {
	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	// Check that a pending authz exists
	if existingPending(tx, authz.ID) != 1 {
		err = errors.New("Cannot finalize a authorization that is not pending")
		tx.Rollback()
		return
	}
	if statusIsPending(authz.Status) {
		err = errors.New("Cannot finalize to a non-final status")
		tx.Rollback()
		return
	}

	// Manually set the index, to avoid AUTOINCREMENT issues
	var sequence int64
	var scanTarget sql.NullInt64
	err = tx.QueryRow("SELECT max(sequence) FROM authz").Scan(&scanTarget)
	switch {
	case !scanTarget.Valid:
		sequence = 0
	case err != nil:
		tx.Rollback()
		return
	default:
		sequence += scanTarget.Int64 + 1
	}

	jsonAuthz, err := json.Marshal(authz)
	if err != nil {
		tx.Rollback()
		return
	}
	digest := core.Fingerprint256(jsonAuthz)

	// Add to final table and delete from pending
	if err != nil {
		tx.Rollback()
		return
	}
	_, err = tx.Exec("INSERT INTO authz (sequence, id, digest, value) VALUES (?, ?, ?, ?);", sequence, authz.ID, digest, jsonAuthz)
	if err != nil {
		tx.Rollback()
		return
	}
	_, err = tx.Exec("DELETE FROM pending_authz WHERE id = ?;", authz.ID)
	if err != nil {
		tx.Rollback()
		return
	}
	err = tx.Commit()
	return
}

func (ssa *SQLStorageAuthority) AddCertificate(certDER []byte) (digest string, err error) {
	var parsedCertificate *x509.Certificate
	parsedCertificate, err = x509.ParseCertificate(certDER)
	if err != nil {
		return
	}
	serial := fmt.Sprintf("%016x", parsedCertificate.SerialNumber)

	tx, err := ssa.db.Begin()
	if err != nil {
		return
	}

	digest = core.Fingerprint256(certDER)
	_, err = tx.Exec("INSERT INTO certificates (serial, digest, value) VALUES (?,?,?);",
		serial, digest, certDER)
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	return
}
<|endoftext|>


<commit_before>package mackerel

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

func TestCreateGraphAnnotation(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations" {
			t.Error("request URL should be /api/v0/graph-annotations but :", req.URL.Path)
		}

		if req.Method != "POST" {
			t.Error("request method should be GET but :", req.Method)
		}
		body, _ := ioutil.ReadAll(req.Body)

		var data struct {
			Service     string   `json:"service,omitempty"`
			Roles       []string `json:"roles,omitempty"`
			From        int64    `json:"from,omitempty"`
			To          int64    `json:"to,omitempty"`
			Title       string   `json:"title,omitempty"`
			Description string   `json:"description,omitempty"`
		}

		err := json.Unmarshal(body, &data)
		if err != nil {
			t.Fatal("request body should be decoded as json", string(body))
		}
		if data.Service != "My-Blog" {
			t.Errorf("request sends json including Service but: %s", data.Service)
		}
		if !reflect.DeepEqual(data.Roles, []string{"Role1", "Role2"}) {
			t.Error("request sends json including Roles but: ", data.Roles)
		}
		if data.From != 1485675275 {
			t.Errorf("request sends json including From but: %d", data.From)
		}
		if data.To != 1485675299 {
			t.Errorf("request sends json including To but: %d", data.To)
		}
		if data.Title != "Deployed" {
			t.Errorf("request sends json including Title but: %s", data.Title)
		}
		if data.Description != "Deployed my blog" {
			t.Errorf("request sends json including Description but: %s", data.Description)
		}
		respJSON, _ := json.Marshal(map[string]string{
			"result": "OK",
		})
		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	err := client.CreateGraphAnnotation(&GraphAnnotation{
		Service:     "My-Blog",
		Roles:       []string{"Role1", "Role2"},
		From:        1485675275,
		To:          1485675299,
		Title:       "Deployed",
		Description: "Deployed my blog",
	})

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}
}

func TestFindGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations" {
			t.Error("request URL should be /api/v0/graph-annotations but :", req.URL.Path)
		}
		if req.Method != "GET" {
			t.Error("request method should be GET but :", req.Method)
		}

		query := req.URL.Query()
		if query.Get("service") != "My-Blog" {
			t.Error("request query 'service' param should be My-Blog but :", query.Get("service"))
		}
		if query.Get("from") != "1485675275" {
			t.Error("request query 'from' param should be 1485675275 but :", query.Get("from"))
		}
		if query.Get("to") != "1485675299" {
			t.Error("request query 'from' param should be 1485675299 but :", query.Get("from"))
		}

		respJSON, _ := json.Marshal(map[string][]map[string]interface{}{
			"graphAnnotations": {
				{
					"service":     "My-Blog",
					"roles":       []string{"Role1", "Role2"},
					"from":        1485675275,
					"to":          1485675299,
					"title":       "Deployed",
					"description": "Deployed my blog",
				},
			},
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotations, err := client.FindGraphAnnotations("My-Blog", 1485675275, 1485675299)

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}

	if annotations[0].Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotations[0].Service)
	}

	if annotations[0].Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotations[0].Title)
	}
}

func TestUpdateGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations/123456789" {
			t.Error("request URL should be /api/v0/graph-annotations/123456789 but :", req.URL.Path)
		}
		if req.Method != "PUT" {
			t.Error("request method should be PUT but :", req.Method)
		}

		respJSON, _ := json.Marshal(map[string]interface{}{
			"service":     "My-Blog",
			"roles":       []string{"Role1", "Role2"},
			"from":        1485675275,
			"to":          1485675299,
			"title":       "Deployed",
			"description": "Deployed my blog",
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotation, err := client.UpdateGraphAnnotation("123456789",
		&GraphAnnotation{
			Service:     "My-Blog",
			Roles:       []string{"Role1", "Role2"},
			From:        1485675275,
			To:          1485675299,
			Title:       "Deployed",
			Description: "Deployed my blog",
		})

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}

	if annotation.Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotation.Service)
	}

	if annotation.Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotation.Title)
	}
}

func TestDeleteGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations/123456789" {
			t.Error("request URL should be /api/v0/graph-annotations/123456789 but :", req.URL.Path)
		}
		if req.Method != "DELETE" {
			t.Error("request method should be PUT but :", req.Method)
		}

		respJSON, _ := json.Marshal(map[string]interface{}{
			"service":     "My-Blog",
			"roles":       []string{"Role1", "Role2"},
			"from":        1485675275,
			"to":          1485675299,
			"title":       "Deployed",
			"description": "Deployed my blog",
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotation, err := client.DeleteGraphAnnotation("123456789")

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}
	if annotation.Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotation.Service)
	}
	if annotation.Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotation.Title)
	}
}
<commit_msg>add tests for several fields<commit_after>package mackerel

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

func TestCreateGraphAnnotation(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations" {
			t.Error("request URL should be /api/v0/graph-annotations but :", req.URL.Path)
		}

		if req.Method != "POST" {
			t.Error("request method should be GET but :", req.Method)
		}
		body, _ := ioutil.ReadAll(req.Body)

		var data struct {
			Service     string   `json:"service,omitempty"`
			Roles       []string `json:"roles,omitempty"`
			From        int64    `json:"from,omitempty"`
			To          int64    `json:"to,omitempty"`
			Title       string   `json:"title,omitempty"`
			Description string   `json:"description,omitempty"`
		}

		err := json.Unmarshal(body, &data)
		if err != nil {
			t.Fatal("request body should be decoded as json", string(body))
		}
		if data.Service != "My-Blog" {
			t.Errorf("request sends json including Service but: %s", data.Service)
		}
		if !reflect.DeepEqual(data.Roles, []string{"Role1", "Role2"}) {
			t.Error("request sends json including Roles but: ", data.Roles)
		}
		if data.From != 1485675275 {
			t.Errorf("request sends json including From but: %d", data.From)
		}
		if data.To != 1485675299 {
			t.Errorf("request sends json including To but: %d", data.To)
		}
		if data.Title != "Deployed" {
			t.Errorf("request sends json including Title but: %s", data.Title)
		}
		if data.Description != "Deployed my blog" {
			t.Errorf("request sends json including Description but: %s", data.Description)
		}
		respJSON, _ := json.Marshal(map[string]string{
			"result": "OK",
		})
		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	err := client.CreateGraphAnnotation(&GraphAnnotation{
		Service:     "My-Blog",
		Roles:       []string{"Role1", "Role2"},
		From:        1485675275,
		To:          1485675299,
		Title:       "Deployed",
		Description: "Deployed my blog",
	})

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}
}

func TestFindGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations" {
			t.Error("request URL should be /api/v0/graph-annotations but :", req.URL.Path)
		}
		if req.Method != "GET" {
			t.Error("request method should be GET but :", req.Method)
		}

		query := req.URL.Query()
		if query.Get("service") != "My-Blog" {
			t.Error("request query 'service' param should be My-Blog but :", query.Get("service"))
		}
		if query.Get("from") != "1485675275" {
			t.Error("request query 'from' param should be 1485675275 but :", query.Get("from"))
		}
		if query.Get("to") != "1485675299" {
			t.Error("request query 'from' param should be 1485675299 but :", query.Get("from"))
		}

		respJSON, _ := json.Marshal(map[string][]map[string]interface{}{
			"graphAnnotations": {
				{
					"service":     "My-Blog",
					"roles":       []string{"Role1", "Role2"},
					"from":        1485675275,
					"to":          1485675299,
					"title":       "Deployed",
					"description": "Deployed my blog",
				},
			},
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotations, err := client.FindGraphAnnotations("My-Blog", 1485675275, 1485675299)

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}

	if annotations[0].Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotations[0].Service)
	}

	if !reflect.DeepEqual(annotations[0].Roles, []string{"Role1", "Role2"}) {
		t.Error("request sends json including Roles but: ", annotations[0].Roles)
	}

	if annotations[0].From != 1485675275 {
		t.Error("request sends json including From but: ", annotations[0].From)
	}

	if annotations[0].To != 1485675299 {
		t.Error("request sends json including To but: ", annotations[0].To)
	}

	if annotations[0].Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotations[0].Title)
	}

	if annotations[0].Description != "Deployed my blog" {
		t.Error("request sends json including Description but: ", annotations[0].Description)
	}
}

func TestUpdateGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations/123456789" {
			t.Error("request URL should be /api/v0/graph-annotations/123456789 but :", req.URL.Path)
		}
		if req.Method != "PUT" {
			t.Error("request method should be PUT but :", req.Method)
		}

		respJSON, _ := json.Marshal(map[string]interface{}{
			"service":     "My-Blog",
			"roles":       []string{"Role1", "Role2"},
			"from":        1485675275,
			"to":          1485675299,
			"title":       "Deployed",
			"description": "Deployed my blog",
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotation, err := client.UpdateGraphAnnotation("123456789",
		&GraphAnnotation{
			Service:     "My-Blog",
			Roles:       []string{"Role1", "Role2"},
			From:        1485675275,
			To:          1485675299,
			Title:       "Deployed",
			Description: "Deployed my blog",
		})

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}

	if annotation.Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotation.Service)
	}

	if !reflect.DeepEqual(annotation.Roles, []string{"Role1", "Role2"}) {
		t.Error("request sends json including Roles but: ", annotation.Roles)
	}

	if annotation.From != 1485675275 {
		t.Error("request sends json including From but: ", annotation.From)
	}

	if annotation.To != 1485675299 {
		t.Error("request sends json including To but: ", annotation.To)
	}

	if annotation.Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotation.Title)
	}

	if annotation.Description != "Deployed my blog" {
		t.Error("request sends json including Description but: ", annotation.Description)
	}
}

func TestDeleteGraphAnnotations(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/api/v0/graph-annotations/123456789" {
			t.Error("request URL should be /api/v0/graph-annotations/123456789 but :", req.URL.Path)
		}
		if req.Method != "DELETE" {
			t.Error("request method should be PUT but :", req.Method)
		}

		respJSON, _ := json.Marshal(map[string]interface{}{
			"service":     "My-Blog",
			"roles":       []string{"Role1", "Role2"},
			"from":        1485675275,
			"to":          1485675299,
			"title":       "Deployed",
			"description": "Deployed my blog",
		})

		res.Header()["Content-Type"] = []string{"application/json"}
		fmt.Fprint(res, string(respJSON))
	}))
	defer ts.Close()

	client, _ := NewClientWithOptions("dummy-key", ts.URL, false)
	annotation, err := client.DeleteGraphAnnotation("123456789")

	if err != nil {
		t.Error("err shoud be nil but: ", err)
	}

	if annotation.Service != "My-Blog" {
		t.Error("request sends json including Service but: ", annotation.Service)
	}

	if !reflect.DeepEqual(annotation.Roles, []string{"Role1", "Role2"}) {
		t.Error("request sends json including Roles but: ", annotation.Roles)
	}

	if annotation.From != 1485675275 {
		t.Error("request sends json including From but: ", annotation.From)
	}

	if annotation.To != 1485675299 {
		t.Error("request sends json including To but: ", annotation.To)
	}

	if annotation.Title != "Deployed" {
		t.Error("request sends json including Title but: ", annotation.Title)
	}

	if annotation.Description != "Deployed my blog" {
		t.Error("request sends json including Description but: ", annotation.Description)
	}
}
<|endoftext|>


<commit_before>package fwd

import "io"

const (
	// DefaultWriterSize is the
	// default write buffer size.
	DefaultWriterSize = 2048

	minWriterSize = minReaderSize
)

// Writer is a buffered writer
type Writer struct {
	w   io.Writer // writer
	buf []byte    // 0:len(buf) is bufered data
}

// NewWriter returns a new writer
// that writes to 'w' and has a buffer
// that is `DefaultWriterSize` bytes.
func NewWriter(w io.Writer) *Writer {
	if wr, ok := w.(*Writer); ok {
		return wr
	}
	return &Writer{
		w:   w,
		buf: make([]byte, 0, DefaultWriterSize),
	}
}

// NewWriterSize returns a new writer
// that writes to 'w' and has a buffer
// that is 'size' bytes.
func NewWriterSize(w io.Writer, size int) *Writer {
	if wr, ok := w.(*Writer); ok && cap(wr.buf) >= size {
		return wr
	}
	return &Writer{
		w:   w,
		buf: make([]byte, 0, max(size, minWriterSize)),
	}
}

// Buffered returns the number of buffered bytes
// in the reader.
func (w *Writer) Buffered() int { return len(w.buf) }

// BufferSize returns the maximum size of the buffer.
func (w *Writer) BufferSize() int { return cap(w.buf) }

// Flush flushes any buffered bytes
// to the underlying writer.
func (w *Writer) Flush() error {
	l := len(w.buf)
	if l > 0 {
		n, err := w.w.Write(w.buf)

		// if we didn't write the whole
		// thing, copy the unwritten
		// bytes to the beginnning of the
		// buffer.
		if n < l && n > 0 {
			w.pushback(n)
			if err == nil {
				err = io.ErrShortWrite
			}
		}
		if err != nil {
			return err
		}
		w.buf = w.buf[:0]
		return nil
	}
	return nil
}

// Write implements `io.Writer`
func (w *Writer) Write(p []byte) (int, error) {
	c, l, ln := cap(w.buf), len(w.buf), len(p)
	avail := c - l

	// requires flush
	if avail < ln {
		if err := w.Flush(); err != nil {
			return 0, err
		}
		l = len(w.buf)
	}
	// too big to fit in buffer;
	// write directly to w.w
	if c < ln {
		return w.w.Write(p)
	}

	// grow buf slice; copy; return
	w.buf = w.buf[:l+ln]
	return copy(w.buf[l:], p), nil
}

// WriteString is analagous to Write, but it takes a string.
func (w *Writer) WriteString(s string) (int, error) {
	c, l, ln := cap(w.buf), len(w.buf), len(s)
	avail := c - l

	// requires flush
	if avail < ln {
		if err := w.Flush(); err != nil {
			return 0, err
		}
		l = len(w.buf)
	}
	// too big to fit in buffer;
	// write directly to w.w
	//
	// yes, this is unsafe. *but*
	// io.Writer is not allowed
	// to mutate its input or
	// maintain a reference to it,
	// per the spec in package io.
	//
	// plus, if the string is really
	// too big to fit in the buffer, then
	// creating a copy to write it is
	// expensive (and, strictly speaking,
	// unnecessary)
	if c < ln {
		return w.w.Write(unsafestr(s))
	}

	// grow buf slice; copy; return
	w.buf = w.buf[:l+ln]
	return copy(w.buf[l:], s), nil
}

// WriteByte implements `io.ByteWriter`
func (w *Writer) WriteByte(b byte) error {
	if len(w.buf) == cap(w.buf) {
		if err := w.Flush(); err != nil {
			return err
		}
	}
	w.buf = append(w.buf, b)
	return nil
}

// Next returns the next 'n' free bytes
// in the write buffer, flushing the writer
// as necessary. Next will return `io.ErrShortBuffer`
// if 'n' is greater than the size of the write buffer.
// Calls to 'next' increment the write position by
// the size of the returned buffer.
func (w *Writer) Next(n int) ([]byte, error) {
	c, l := cap(w.buf), len(w.buf)
	if n > c {
		return nil, io.ErrShortBuffer
	}
	avail := c - l
	if avail < n {
		if err := w.Flush(); err != nil {
			return nil, err
		}
		l = len(w.buf)
	}
	w.buf = w.buf[:l+n]
	return w.buf[l:], nil
}

// take the bytes from w.buf[n:len(w.buf)]
// and put them at the beginning of w.buf,
// and resize to the length of the copied segment.
func (w *Writer) pushback(n int) {
	w.buf = w.buf[:copy(w.buf, w.buf[n:])]
}

// ReadFrom implements `io.ReaderFrom`
func (w *Writer) ReadFrom(r io.Reader) (int64, error) {
	// anticipatory flush
	if err := w.Flush(); err != nil {
		return 0, err
	}

	w.buf = w.buf[0:cap(w.buf)] // expand buffer

	var nn int64  // written
	var err error // error
	var x int     // read

	// 1:1 reads and writes
	for err == nil {
		x, err = r.Read(w.buf)
		if x > 0 {
			n, werr := w.w.Write(w.buf[:x])
			nn += int64(n)

			if err != nil {
				if n < x && n > 0 {
					w.pushback(n - x)
				}
				return nn, werr
			}
			if n < x {
				w.pushback(n - x)
				return nn, io.ErrShortWrite
			}
		} else if err == nil {
			err = io.ErrNoProgress
			break
		}
	}
	if err != io.EOF {
		return nn, err
	}

	// we only clear here
	// because we are sure
	// the writes have
	// suceeded. otherwise,
	// we retain the data in case
	// future writes succeed.
	w.buf = w.buf[0:0]

	return nn, nil
}
<commit_msg>fix typos<commit_after>package fwd

import "io"

const (
	// DefaultWriterSize is the
	// default write buffer size.
	DefaultWriterSize = 2048

	minWriterSize = minReaderSize
)

// Writer is a buffered writer
type Writer struct {
	w   io.Writer // writer
	buf []byte    // 0:len(buf) is bufered data
}

// NewWriter returns a new writer
// that writes to 'w' and has a buffer
// that is `DefaultWriterSize` bytes.
func NewWriter(w io.Writer) *Writer {
	if wr, ok := w.(*Writer); ok {
		return wr
	}
	return &Writer{
		w:   w,
		buf: make([]byte, 0, DefaultWriterSize),
	}
}

// NewWriterSize returns a new writer
// that writes to 'w' and has a buffer
// that is 'size' bytes.
func NewWriterSize(w io.Writer, size int) *Writer {
	if wr, ok := w.(*Writer); ok && cap(wr.buf) >= size {
		return wr
	}
	return &Writer{
		w:   w,
		buf: make([]byte, 0, max(size, minWriterSize)),
	}
}

// Buffered returns the number of buffered bytes
// in the reader.
func (w *Writer) Buffered() int { return len(w.buf) }

// BufferSize returns the maximum size of the buffer.
func (w *Writer) BufferSize() int { return cap(w.buf) }

// Flush flushes any buffered bytes
// to the underlying writer.
func (w *Writer) Flush() error {
	l := len(w.buf)
	if l > 0 {
		n, err := w.w.Write(w.buf)

		// if we didn't write the whole
		// thing, copy the unwritten
		// bytes to the beginnning of the
		// buffer.
		if n < l && n > 0 {
			w.pushback(n)
			if err == nil {
				err = io.ErrShortWrite
			}
		}
		if err != nil {
			return err
		}
		w.buf = w.buf[:0]
		return nil
	}
	return nil
}

// Write implements `io.Writer`
func (w *Writer) Write(p []byte) (int, error) {
	c, l, ln := cap(w.buf), len(w.buf), len(p)
	avail := c - l

	// requires flush
	if avail < ln {
		if err := w.Flush(); err != nil {
			return 0, err
		}
		l = len(w.buf)
	}
	// too big to fit in buffer;
	// write directly to w.w
	if c < ln {
		return w.w.Write(p)
	}

	// grow buf slice; copy; return
	w.buf = w.buf[:l+ln]
	return copy(w.buf[l:], p), nil
}

// WriteString is analogous to Write, but it takes a string.
func (w *Writer) WriteString(s string) (int, error) {
	c, l, ln := cap(w.buf), len(w.buf), len(s)
	avail := c - l

	// requires flush
	if avail < ln {
		if err := w.Flush(); err != nil {
			return 0, err
		}
		l = len(w.buf)
	}
	// too big to fit in buffer;
	// write directly to w.w
	//
	// yes, this is unsafe. *but*
	// io.Writer is not allowed
	// to mutate its input or
	// maintain a reference to it,
	// per the spec in package io.
	//
	// plus, if the string is really
	// too big to fit in the buffer, then
	// creating a copy to write it is
	// expensive (and, strictly speaking,
	// unnecessary)
	if c < ln {
		return w.w.Write(unsafestr(s))
	}

	// grow buf slice; copy; return
	w.buf = w.buf[:l+ln]
	return copy(w.buf[l:], s), nil
}

// WriteByte implements `io.ByteWriter`
func (w *Writer) WriteByte(b byte) error {
	if len(w.buf) == cap(w.buf) {
		if err := w.Flush(); err != nil {
			return err
		}
	}
	w.buf = append(w.buf, b)
	return nil
}

// Next returns the next 'n' free bytes
// in the write buffer, flushing the writer
// as necessary. Next will return `io.ErrShortBuffer`
// if 'n' is greater than the size of the write buffer.
// Calls to 'next' increment the write position by
// the size of the returned buffer.
func (w *Writer) Next(n int) ([]byte, error) {
	c, l := cap(w.buf), len(w.buf)
	if n > c {
		return nil, io.ErrShortBuffer
	}
	avail := c - l
	if avail < n {
		if err := w.Flush(); err != nil {
			return nil, err
		}
		l = len(w.buf)
	}
	w.buf = w.buf[:l+n]
	return w.buf[l:], nil
}

// take the bytes from w.buf[n:len(w.buf)]
// and put them at the beginning of w.buf,
// and resize to the length of the copied segment.
func (w *Writer) pushback(n int) {
	w.buf = w.buf[:copy(w.buf, w.buf[n:])]
}

// ReadFrom implements `io.ReaderFrom`
func (w *Writer) ReadFrom(r io.Reader) (int64, error) {
	// anticipatory flush
	if err := w.Flush(); err != nil {
		return 0, err
	}

	w.buf = w.buf[0:cap(w.buf)] // expand buffer

	var nn int64  // written
	var err error // error
	var x int     // read

	// 1:1 reads and writes
	for err == nil {
		x, err = r.Read(w.buf)
		if x > 0 {
			n, werr := w.w.Write(w.buf[:x])
			nn += int64(n)

			if err != nil {
				if n < x && n > 0 {
					w.pushback(n - x)
				}
				return nn, werr
			}
			if n < x {
				w.pushback(n - x)
				return nn, io.ErrShortWrite
			}
		} else if err == nil {
			err = io.ErrNoProgress
			break
		}
	}
	if err != io.EOF {
		return nn, err
	}

	// we only clear here
	// because we are sure
	// the writes have
	// succeeded. otherwise,
	// we retain the data in case
	// future writes succeed.
	w.buf = w.buf[0:0]

	return nn, nil
}
<|endoftext|>