<commit_before>package valuestore

import (
	"encoding/binary"
	"io"
	"time"
)

const _MSG_BULK_SET_ACK = 0x39589f4746844e3b

type bulkSetAckState struct {
	inMsgChan      chan *bulkSetAckMsg
	inFreeMsgChan  chan *bulkSetAckMsg
	outFreeMsgChan chan *bulkSetAckMsg
	inMsgTimeout   time.Duration
}

type bulkSetAckMsg struct {
	vs   *DefaultValueStore
	body []byte
}

func (vs *DefaultValueStore) bulkSetAckInit(cfg *Config) {
	if vs.msgRing != nil {
		vs.msgRing.SetMsgHandler(_MSG_BULK_SET_ACK, vs.newInBulkSetAckMsg)
		vs.bulkSetAckState.inMsgChan = make(chan *bulkSetAckMsg, cfg.InBulkSetAckMsgs)
		vs.bulkSetAckState.inFreeMsgChan = make(chan *bulkSetAckMsg, cfg.InBulkSetAckMsgs)
		for i := 0; i < cap(vs.bulkSetAckState.inFreeMsgChan); i++ {
			vs.bulkSetAckState.inFreeMsgChan <- &bulkSetAckMsg{
				vs:   vs,
				body: make([]byte, cfg.BulkSetAckMsgCap),
			}
		}
		for i := 0; i < cfg.InBulkSetAckWorkers; i++ {
			go vs.inBulkSetAck()
		}
		vs.bulkSetAckState.outFreeMsgChan = make(chan *bulkSetAckMsg, cfg.OutBulkSetAckMsgs)
		for i := 0; i < cap(vs.bulkSetAckState.outFreeMsgChan); i++ {
			vs.bulkSetAckState.outFreeMsgChan <- &bulkSetAckMsg{
				vs:   vs,
				body: make([]byte, cfg.BulkSetAckMsgCap),
			}
		}
		vs.bulkSetAckState.inMsgTimeout = time.Duration(cfg.InBulkSetAckMsgTimeout) * time.Second
	}
}

func (vs *DefaultValueStore) inBulkSetAck() {
	for {
		bsam := <-vs.bulkSetAckState.inMsgChan
		ring := vs.msgRing.Ring()
		version := ring.Version()
		rightwardPartitionShift := 64 - ring.PartitionBitCount()
		b := bsam.body
		l := len(b)
		for o := 0; o < l; o += 24 {
			if version != ring.Version() {
				version = ring.Version()
				rightwardPartitionShift = 64 - ring.PartitionBitCount()
			}
			keyA := binary.BigEndian.Uint64(b[o:])
			if !ring.Responsible(uint32(keyA >> rightwardPartitionShift)) {
				vs.write(keyA, binary.BigEndian.Uint64(b[o+8:]), binary.BigEndian.Uint64(b[o+16:])|_TSB_LOCAL_REMOVAL, nil)
			}
		}
		vs.bulkSetAckState.inFreeMsgChan <- bsam
	}
}

func (vs *DefaultValueStore) newInBulkSetAckMsg(r io.Reader, l uint64) (uint64, error) {
	var bsam *bulkSetAckMsg
	select {
	case bsam = <-vs.bulkSetAckState.inFreeMsgChan:
	case <-time.After(vs.bulkSetAckState.inMsgTimeout):
		left := l
		var sn int
		var err error
		for left > 0 {
			t := toss
			if left < uint64(len(t)) {
				t = t[:left]
			}
			sn, err = r.Read(t)
			left -= uint64(sn)
			if err != nil {
				return l - left, err
			}
		}
		return l, nil
	}
	var n int
	var sn int
	var err error
	if l > uint64(cap(bsam.body)) {
		bsam.body = make([]byte, l)
	}
	bsam.body = bsam.body[:l]
	n = 0
	for n != len(bsam.body) {
		sn, err = r.Read(bsam.body[n:])
		n += sn
		if err != nil {
			return uint64(n), err
		}
	}
	vs.bulkSetAckState.inMsgChan <- bsam
	return l, nil
}

func (vs *DefaultValueStore) newOutBulkSetAckMsg() *bulkSetAckMsg {
	bsam := <-vs.bulkSetAckState.outFreeMsgChan
	bsam.body = bsam.body[:0]
	return bsam
}

func (bsam *bulkSetAckMsg) MsgType() uint64 {
	return _MSG_BULK_SET_ACK
}

func (bsam *bulkSetAckMsg) MsgLength() uint64 {
	return uint64(len(bsam.body))
}

func (bsam *bulkSetAckMsg) WriteContent(w io.Writer) (uint64, error) {
	n, err := w.Write(bsam.body)
	return uint64(n), err
}

func (bsam *bulkSetAckMsg) Done() {
	bsam.vs.bulkSetAckState.outFreeMsgChan <- bsam
}

func (bsam *bulkSetAckMsg) add(keyA uint64, keyB uint64, timestampbits uint64) bool {
	o := len(bsam.body)
	if o+24 >= cap(bsam.body) {
		return false
	}
	bsam.body = bsam.body[:o+24]
	binary.BigEndian.PutUint64(bsam.body[o:], keyA)
	binary.BigEndian.PutUint64(bsam.body[o+8:], keyB)
	binary.BigEndian.PutUint64(bsam.body[o+16:], timestampbits)
	return true
}
<commit_msg>bulksetack.go commentary<commit_after>package valuestore

import (
	"encoding/binary"
	"io"
	"time"
)

// bsam: entries:n
// bsam entry: keyA:8, keyB:8, timestampbits:8
const _BULK_SET_ACK_MSG_TYPE = 0x39589f4746844e3b
const _BULK_SET_ACK_MSG_ENTRY_LENGTH = 24

type bulkSetAckState struct {
	inMsgChan      chan *bulkSetAckMsg
	inFreeMsgChan  chan *bulkSetAckMsg
	outFreeMsgChan chan *bulkSetAckMsg
	inMsgTimeout   time.Duration
}

type bulkSetAckMsg struct {
	vs   *DefaultValueStore
	body []byte
}

func (vs *DefaultValueStore) bulkSetAckInit(cfg *Config) {
	if vs.msgRing != nil {
		vs.msgRing.SetMsgHandler(_BULK_SET_ACK_MSG_TYPE, vs.newInBulkSetAckMsg)
		vs.bulkSetAckState.inMsgChan = make(chan *bulkSetAckMsg, cfg.InBulkSetAckMsgs)
		vs.bulkSetAckState.inFreeMsgChan = make(chan *bulkSetAckMsg, cfg.InBulkSetAckMsgs)
		for i := 0; i < cap(vs.bulkSetAckState.inFreeMsgChan); i++ {
			vs.bulkSetAckState.inFreeMsgChan <- &bulkSetAckMsg{
				vs:   vs,
				body: make([]byte, cfg.BulkSetAckMsgCap),
			}
		}
		for i := 0; i < cfg.InBulkSetAckWorkers; i++ {
			go vs.inBulkSetAck()
		}
		vs.bulkSetAckState.outFreeMsgChan = make(chan *bulkSetAckMsg, cfg.OutBulkSetAckMsgs)
		for i := 0; i < cap(vs.bulkSetAckState.outFreeMsgChan); i++ {
			vs.bulkSetAckState.outFreeMsgChan <- &bulkSetAckMsg{
				vs:   vs,
				body: make([]byte, cfg.BulkSetAckMsgCap),
			}
		}
		vs.bulkSetAckState.inMsgTimeout = time.Duration(cfg.InBulkSetAckMsgTimeout) * time.Second
	}
}

// newInBulkSetAckMsg reads bulk-set-ack messages from the MsgRing and puts
// them on the inMsgChan for the inBulkSetAck workers to work on.
func (vs *DefaultValueStore) newInBulkSetAckMsg(r io.Reader, l uint64) (uint64, error) {
	var bsam *bulkSetAckMsg
	select {
	case bsam = <-vs.bulkSetAckState.inFreeMsgChan:
		// If there isn't a free bulkSetAckMsg after some time, give up and
		// just read and discard the incoming bulk-set-ack message.
	case <-time.After(vs.bulkSetAckState.inMsgTimeout):
		left := l
		var sn int
		var err error
		for left > 0 {
			t := toss
			if left < uint64(len(t)) {
				t = t[:left]
			}
			sn, err = r.Read(t)
			left -= uint64(sn)
			if err != nil {
				return l - left, err
			}
		}
		return l, nil
	}
	var n int
	var sn int
	var err error
	// TODO: I think we should cap the body size to the message cap but that
	// also means that the inBulkSetAck worker will need to handle the likely
	// trailing truncated entry. Once all this is done, the overall cluster
	// should work even if the caps are set differently from node to node
	// (definitely not recommended though), as the bulk-set-ack messages would
	// eventually start falling under the minimum cap as the front-end data is
	// tranferred and acknowledged. Anyway, I think this is needed in case
	// someone accidentally screws up the cap on one node, making it way too
	// big. Rather just have that one node abuse/run-out-of memory instead of
	// it causing every other node it sends bulk-set-ack messages to also have
	// memory issues. This is less likely of an issue than the same issue with
	// bulk-set messaging.
	if l > uint64(cap(bsam.body)) {
		bsam.body = make([]byte, l)
	}
	bsam.body = bsam.body[:l]
	n = 0
	for n != len(bsam.body) {
		sn, err = r.Read(bsam.body[n:])
		n += sn
		if err != nil {
			return uint64(n), err
		}
	}
	vs.bulkSetAckState.inMsgChan <- bsam
	return l, nil
}

// inBulkSetAck actually processes incoming bulk-set-ack messages; there may be
// more than one of these workers.
func (vs *DefaultValueStore) inBulkSetAck() {
	for {
		bsam := <-vs.bulkSetAckState.inMsgChan
		ring := vs.msgRing.Ring()
		rightwardPartitionShift := 64 - ring.PartitionBitCount()
		b := bsam.body
		l := len(b)
		for o := 0; o < l; o += _BULK_SET_ACK_MSG_ENTRY_LENGTH {
			keyA := binary.BigEndian.Uint64(b[o:])
			if !ring.Responsible(uint32(keyA >> rightwardPartitionShift)) {
				vs.write(keyA, binary.BigEndian.Uint64(b[o+8:]), binary.BigEndian.Uint64(b[o+16:])|_TSB_LOCAL_REMOVAL, nil)
			}
		}
		vs.bulkSetAckState.inFreeMsgChan <- bsam
	}
}

// newOutBulkSetAckMsg gives an initialized bulkSetAckMsg for filling out and
// eventually sending using the MsgRing. The MsgRing (or someone else if the
// message doesn't end up with the MsgRing) will call bulkSetAckMsg.Done()
// eventually and the bulkSetAckMsg will be requeued for reuse later. There is
// a fixed number of outgoing bulkSetAckMsg instances that can exist at any
// given time, capping memory usage. Once the limit is reached, this method
// will block until a bulkSetAckMsg is available to return.
func (vs *DefaultValueStore) newOutBulkSetAckMsg() *bulkSetAckMsg {
	bsam := <-vs.bulkSetAckState.outFreeMsgChan
	bsam.body = bsam.body[:0]
	return bsam
}

func (bsam *bulkSetAckMsg) MsgType() uint64 {
	return _BULK_SET_ACK_MSG_TYPE
}

func (bsam *bulkSetAckMsg) MsgLength() uint64 {
	return uint64(len(bsam.body))
}

func (bsam *bulkSetAckMsg) WriteContent(w io.Writer) (uint64, error) {
	n, err := w.Write(bsam.body)
	return uint64(n), err
}

func (bsam *bulkSetAckMsg) Done() {
	bsam.vs.bulkSetAckState.outFreeMsgChan <- bsam
}

func (bsam *bulkSetAckMsg) add(keyA uint64, keyB uint64, timestampbits uint64) bool {
	o := len(bsam.body)
	if o+_BULK_SET_ACK_MSG_ENTRY_LENGTH >= cap(bsam.body) {
		return false
	}
	bsam.body = bsam.body[:o+_BULK_SET_ACK_MSG_ENTRY_LENGTH]
	binary.BigEndian.PutUint64(bsam.body[o:], keyA)
	binary.BigEndian.PutUint64(bsam.body[o+8:], keyB)
	binary.BigEndian.PutUint64(bsam.body[o+16:], timestampbits)
	return true
}
<|endoftext|>


<commit_before>package output

import (
	"encoding/json"
	"fmt"
	"sort"
)

type Config struct {
	Format string
}

type Output struct {
	Config Config
}

func (o *Output) Generate(elements interface{}) (string, error) {
	if elements == nil {
		return "", nil
	}

	switch o.Config.Format {
	case "json":
		return o.JSONOutput(elements)
	case "shell":
		return o.ShellOutput(elements)
	}

	return "", fmt.Errorf("Unrecognized output format: %s", o.Config.Format)
}

func (o *Output) JSONOutput(elements interface{}) (string, error) {
	if _, ok := elements.([]interface{}); ok {
		if j, err := json.MarshalIndent(elements, " ", " "); err != nil {
			return "", err
		} else {
			return string(j), nil
		}
	} else {
		if _, ok := elements.(map[string]interface{}); ok {
			if j, err := json.MarshalIndent(elements, " ", " "); err != nil {
				return "", err
			} else {
				return string(j), nil
			}
		} else {
			return fmt.Sprintf("%s", elements), nil
		}
	}
}

func (o *Output) ShellOutput(elements interface{}) (string, error) {
	var keys []string
	var output string

	parsed := make(map[string]string)
	_, _, parsed = parseElementsForShell("elements", elements, parsed)

	for k := range parsed {
		keys = append(keys, k)
	}

	sort.Strings(keys)

	for _, k := range keys {
		output = fmt.Sprintf("%s%s=\"%s\"\n", output, k, parsed[k])
	}

	return output, nil
}

func parseElementsForShell(label string, elements interface{}, parsed map[string]string) (string, interface{}, map[string]string) {
	var e interface{}
	switch elements.(type) {
	case map[string]interface{}:
		for key, subElements := range elements.(map[string]interface{}) {
			l := fmt.Sprintf("%s_%s", label, key)
			_, e, parsed = parseElementsForShell(l, subElements, parsed)
		}
	case []interface{}:
		for i, subElements := range elements.([]interface{}) {
			l := fmt.Sprintf("%s_%d", label, i)
			_, e, parsed = parseElementsForShell(l, subElements, parsed)
		}
	default:
		parsed[label] = fmt.Sprintf("%v", elements)
		return label, nil, parsed
	}

	return label, e, parsed
}
<commit_msg>Munge non-portable characters for shell output<commit_after>package output

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

type Config struct {
	Format string
}

type Output struct {
	Config Config
}

func (o *Output) Generate(elements interface{}) (string, error) {
	if elements == nil {
		return "", nil
	}

	switch o.Config.Format {
	case "json":
		return o.JSONOutput(elements)
	case "shell":
		return o.ShellOutput(elements)
	}

	return "", fmt.Errorf("Unrecognized output format: %s", o.Config.Format)
}

func (o *Output) JSONOutput(elements interface{}) (string, error) {
	if _, ok := elements.([]interface{}); ok {
		if j, err := json.MarshalIndent(elements, " ", " "); err != nil {
			return "", err
		} else {
			return string(j), nil
		}
	} else {
		if _, ok := elements.(map[string]interface{}); ok {
			if j, err := json.MarshalIndent(elements, " ", " "); err != nil {
				return "", err
			} else {
				return string(j), nil
			}
		} else {
			return fmt.Sprintf("%s", elements), nil
		}
	}
}

func (o *Output) ShellOutput(elements interface{}) (string, error) {
	var keys []string
	var output string

	parsed := make(map[string]string)
	_, _, parsed = parseElementsForShell("elements", elements, parsed)

	for k := range parsed {
		keys = append(keys, k)
	}

	sort.Strings(keys)

	for _, k := range keys {
		output = fmt.Sprintf("%s%s=\"%s\"\n", output, k, parsed[k])
	}

	return output, nil
}

func parseElementsForShell(label string, elements interface{}, parsed map[string]string) (string, interface{}, map[string]string) {
	var e interface{}
	switch elements.(type) {
	case map[string]interface{}:
		for key, subElements := range elements.(map[string]interface{}) {
			l := fmt.Sprintf("%s_%s", label, sanitizeKeyForShell(key))
			_, e, parsed = parseElementsForShell(l, subElements, parsed)
		}
	case []interface{}:
		for i, subElements := range elements.([]interface{}) {
			l := fmt.Sprintf("%s_%d", label, i)
			_, e, parsed = parseElementsForShell(l, subElements, parsed)
		}
	default:
		parsed[label] = fmt.Sprintf("%v", elements)
		return label, nil, parsed
	}

	return label, e, parsed
}

func sanitizeKeyForShell(key string) string {
	// replace non-compliant shell characters with underscores,
	// per IEEE Std 1003.1-2001
	r := strings.NewReplacer(
		"`", "_",
		"~", "_",
		"!", "_",
		"@", "_",
		"#", "_",
		"$", "_",
		"%", "_",
		"^", "_",
		"&", "_",
		"*", "_",
		"(", "_",
		")", "_",
		"-", "_",
		"+", "_",
		"=", "_",
		"{", "_",
		"}", "_",
		"[", "_",
		"]", "_",
		"|", "_",
		":", "_",
		";", "_",
		"'", "_",
		"\"", "_",
		"/", "_",
		"?", "_",
		",", "_",
		".", "_",
		"<", "_",
		">", "_",
	)
	return r.Replace(key)
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package os

import "syscall"

func isSymlink(stat *syscall.Stat_t) bool {
	return stat.Mode&syscall.S_IFMT == syscall.S_IFLNK
}

func fileInfoFromStat(name string, fi *FileInfo, lstat, stat *syscall.Stat_t) *FileInfo {
	fi.Dev = uint64(stat.Dev)
	fi.Ino = uint64(stat.Ino)
	fi.Nlink = uint64(stat.Nlink)
	fi.Mode = uint32(stat.Mode)
	fi.Uid = int(stat.Uid)
	fi.Gid = int(stat.Gid)
	fi.Rdev = uint64(stat.Rdev)
	fi.Size = uint64(stat.Size)
	fi.Blksize = int64(stat.Blksize)
	fi.Blocks = stat.Blocks
	fi.Atime_ns = syscall.TimespecToNsec(stat.Atimespec)
	fi.Mtime_ns = syscall.TimespecToNsec(stat.Mtimespec)
	fi.Ctime_ns = syscall.TimespecToNsec(stat.Ctimespec)
	for i := len(name) - 1; i >= 0; i-- {
		if name[i] == '/' {
			name = name[i+1:]
			break
		}
	}
	fi.Name = name
	if isSymlink(lstat) && !isSymlink(stat) {
		fi.FollowedSymlink = true
	}
	return fi
}
<commit_msg>os/stat_freebsd.go: maybe fix build for freebsd<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package os

import "syscall"

func isSymlink(stat *syscall.Stat_t) bool {
	return stat.Mode&syscall.S_IFMT == syscall.S_IFLNK
}

func fileInfoFromStat(name string, fi *FileInfo, lstat, stat *syscall.Stat_t) *FileInfo {
	fi.Dev = uint64(stat.Dev)
	fi.Ino = uint64(stat.Ino)
	fi.Nlink = uint64(stat.Nlink)
	fi.Mode = uint32(stat.Mode)
	fi.Uid = int(stat.Uid)
	fi.Gid = int(stat.Gid)
	fi.Rdev = uint64(stat.Rdev)
	fi.Size = int64(stat.Size)
	fi.Blksize = int64(stat.Blksize)
	fi.Blocks = stat.Blocks
	fi.Atime_ns = syscall.TimespecToNsec(stat.Atimespec)
	fi.Mtime_ns = syscall.TimespecToNsec(stat.Mtimespec)
	fi.Ctime_ns = syscall.TimespecToNsec(stat.Ctimespec)
	for i := len(name) - 1; i >= 0; i-- {
		if name[i] == '/' {
			name = name[i+1:]
			break
		}
	}
	fi.Name = name
	if isSymlink(lstat) && !isSymlink(stat) {
		fi.FollowedSymlink = true
	}
	return fi
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"testing"
)

func TestSetKey(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")

	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	value, err := dbClient.cache.Get(k)
	expect(t, err, nil)
	expect(t, string(value), string(v))
	dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
}

func TestPayloadSetGet(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	key := []byte("keySetGetCache")
	resp := response{
		Status: 200,
		Body:   "body here",
	}

	payload := Payload{Response: resp}
	bts, err := json.Marshal(payload)
	expect(t, err, nil)

	err = dbClient.cache.Set(key, bts)
	expect(t, err, nil)

	var p Payload
	payloadBts, err := dbClient.cache.Get(key)
	err = json.Unmarshal(payloadBts, &p)
	expect(t, err, nil)
	expect(t, payload.Response.Body, p.Response.Body)

	dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
}

func TestGetNonExistingBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	dbClient.cache.requestsBucket = []byte("some_random_bucket")

	_, err := dbClient.cache.Get([]byte("whatever"))
	expect(t, err.Error(), "Bucket \"some_random_bucket\" not found!")
}

func TestDeleteBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")
	// checking whether bucket is okay
	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	value, err := dbClient.cache.Get(k)
	expect(t, err, nil)
	expect(t, string(value), string(v))

	// deleting bucket
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	expect(t, err, nil)

	// deleting it again
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	refute(t, err, nil)
}

func TestGetAllRequestNoBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	dbClient.cache.requestsBucket = []byte("no_bucket_for_TestGetAllRequestNoBucket")
	_, err := dbClient.cache.GetAllRequests()
	// expecting nil since this would mean that records were wiped
	expect(t, err, nil)
}

func TestCorruptedPayloads(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")

	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	// corrupted payloads should be just skipped
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 0)

}

func TestGetMultipleRecords(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	// getting requests
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)

	for _, payload := range payloads {
		expect(t, payload.Request.Method, "GET")
		expect(t, payload.Response.Status, 201)
	}

	dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
}
<commit_msg>test for error for non existing key<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"testing"
)

func TestSetKey(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")

	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	value, err := dbClient.cache.Get(k)
	expect(t, err, nil)
	expect(t, string(value), string(v))
	dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
}

func TestPayloadSetGet(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	key := []byte("keySetGetCache")
	resp := response{
		Status: 200,
		Body:   "body here",
	}

	payload := Payload{Response: resp}
	bts, err := json.Marshal(payload)
	expect(t, err, nil)

	err = dbClient.cache.Set(key, bts)
	expect(t, err, nil)

	var p Payload
	payloadBts, err := dbClient.cache.Get(key)
	err = json.Unmarshal(payloadBts, &p)
	expect(t, err, nil)
	expect(t, payload.Response.Body, p.Response.Body)

	dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
}

func TestGetNonExistingBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	dbClient.cache.requestsBucket = []byte("some_random_bucket")

	_, err := dbClient.cache.Get([]byte("whatever"))
	expect(t, err.Error(), "Bucket \"some_random_bucket\" not found!")
}

func TestDeleteBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")
	// checking whether bucket is okay
	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	value, err := dbClient.cache.Get(k)
	expect(t, err, nil)
	expect(t, string(value), string(v))

	// deleting bucket
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	expect(t, err, nil)

	// deleting it again
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	refute(t, err, nil)
}

func TestGetAllRequestNoBucket(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	dbClient.cache.requestsBucket = []byte("no_bucket_for_TestGetAllRequestNoBucket")
	_, err := dbClient.cache.GetAllRequests()
	// expecting nil since this would mean that records were wiped
	expect(t, err, nil)
}

func TestCorruptedPayloads(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()

	k := []byte("randomkeyhere")
	v := []byte("value")

	err := dbClient.cache.Set(k, v)
	expect(t, err, nil)

	// corrupted payloads should be just skipped
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 0)

}

func TestGetMultipleRecords(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	// getting requests
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)

	for _, payload := range payloads {
		expect(t, payload.Request.Method, "GET")
		expect(t, payload.Response.Status, 201)
	}
}

func TestGetNonExistingKey(t *testing.T) {
	server, dbClient := testTools(201, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// getting key
	_, err := dbClient.cache.Get([]byte("should not be here"))
	refute(t, err, nil)
}

}
<|endoftext|>


<commit_before>package database

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"sync"
	"time"

	"github.com/asdine/storm"
	"github.com/asdine/storm/index"
	"github.com/boltdb/bolt"
	"github.com/hashicorp/raft"
)

type (
	// Database is the lowest level of the gansoi database, it represent the
	// on-disk database. Database implements raft.FSM.
	Database struct {
		dbMutex       sync.RWMutex
		db            *storm.DB
		listenersLock sync.RWMutex
		listeners     []Listener
	}
)

var (
	// ErrNotFound is returned when the specified record is not saved.
	ErrNotFound = storm.ErrNotFound
)

const (
	// stableBucket is the bucket name used for raft stable storage.
	stableBucket = "raft.StableStore"

	// logBucket is used for raft log storage.
	logBucket = "raft.LogStore"
)

// NewDatabase will instantiate a new database placed in filepath.
func NewDatabase(filepath string) (*Database, error) {
	d := &Database{}

	err := d.open(path.Join(filepath, "gansoi.db"))
	if err != nil {
		return nil, err
	}

	return d, nil
}

// open will open the underlying file storage.
func (d *Database) open(filepath string) error {
	db, err := storm.Open(
		filepath,
		storm.BoltOptions(0600, &bolt.Options{Timeout: 1 * time.Second}),
		storm.AutoIncrement(),
	)
	if err != nil {
		return err
	}

	d.db = db

	d.Storm().Bolt.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte(logBucket))

		return nil
	})

	return nil
}

// Storm will return the underlying Storm database.
func (d *Database) Storm() *storm.DB {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db
}

// ProcessLogEntry will process the log entry and apply whatever needs doing.
func (d *Database) ProcessLogEntry(entry *LogEntry) error {
	var err error

	var v interface{}

	switch entry.Command {
	case CommandSave:
		v, _ = entry.Payload()
		err = d.Save(v)
	case CommandDelete:
		v, _ = entry.Payload()
		err = d.db.DeleteStruct(v)
	default:
		err = fmt.Errorf("not implemented")
	}

	go func(command Command, data interface{}, err error) {
		d.listenersLock.RLock()

		for _, listener := range d.listeners {
			listener.PostLocalApply(command, data, err)
		}

		d.listenersLock.RUnlock()
	}(entry.Command, v, err)

	return err
}

// Apply implements raft.FSM.
func (d *Database) Apply(l *raft.Log) interface{} {
	entry := &LogEntry{}
	err := json.Unmarshal(l.Data, entry)
	if err != nil {
		// This should not happen..?
		fmt.Printf("%s: '%s'\n", err.Error(), string(l.Data))
		return nil
	}

	return d.ProcessLogEntry(entry)
}

// Snapshot implements raft.FSM.
func (d *Database) Snapshot() (raft.FSMSnapshot, error) {
	return &Snapshot{db: d}, nil
}

// Restore implements raft.FSM.
func (d *Database) Restore(source io.ReadCloser) error {
	db := d.Storm().Bolt
	d.dbMutex.Lock()
	defer d.dbMutex.Unlock()
	defer source.Close()

	path := db.Path()
	restorePath := path + ".restoretmp"

	file, err := os.Create(restorePath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, source)
	if err != nil {
		return err
	}

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

	err = os.Remove(path)
	if err != nil {
		return err
	}

	err = os.Rename(restorePath, path)
	if err != nil {
		return err
	}

	err = d.open(path)
	if err != nil {
		return err
	}

	return nil
}

// Save will save an object to the database.
func (d *Database) Save(data interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.Save(data)
}

// One will retrieve one (or zero) record from the database.
func (d *Database) One(fieldName string, value interface{}, to interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.One(fieldName, value, to)
}

// All lists all kinds of a type.
func (d *Database) All(to interface{}, limit int, skip int, reverse bool) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.All(to, func(opts *index.Options) {
		opts.Limit = limit
		opts.Skip = skip
		opts.Reverse = reverse
	})
}

// RegisterListener will register a listener for new changes to the database.
func (d *Database) RegisterListener(listener Listener) {
	d.listenersLock.Lock()
	defer d.listenersLock.Unlock()

	d.listeners = append(d.listeners, listener)
}

// Set implements raft.StableStore.
func (d *Database) Set(k, v []byte) error {
	return d.Storm().Set(stableBucket, k, v)
}

// Get implements raft.StableStore.
func (d *Database) Get(k []byte) ([]byte, error) {
	var value []byte
	err := d.Storm().Get(stableBucket, k, &value)
	if err != nil {
		return nil, err
	}

	return value, nil
}

// SetUint64 is like Set, but handles uint64 values
func (d *Database) SetUint64(key []byte, val uint64) error {
	return d.Set(key, uint64ToBytes(val))
}

// GetUint64 is like Get, but handles uint64 values
func (d *Database) GetUint64(key []byte) (uint64, error) {
	val, err := d.Get(key)
	if err != nil {
		return 0, err
	}

	return bytesToUint64(val), nil
}

// FirstIndex implements raft.LogStore.
func (d *Database) FirstIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	first, _ := curs.First()
	if first == nil {
		return 0, nil
	}

	return bytesToUint64(first), nil
}

// LastIndex implements raft.LogStore.
func (d *Database) LastIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	last, _ := curs.Last()
	if last == nil {
		return 0, nil
	}

	return bytesToUint64(last), nil
}

// GetLog implements raft.LogStore.
func (d *Database) GetLog(idx uint64, log *raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	value := tx.Bucket([]byte(logBucket)).Get(uint64ToBytes(idx))

	return json.Unmarshal(value, log)
}

// StoreLog implements raft.LogStore.
func (d *Database) StoreLog(log *raft.Log) error {
	return d.StoreLogs([]*raft.Log{log})
}

// StoreLogs implements raft.LogStore.
func (d *Database) StoreLogs(logs []*raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	for _, log := range logs {
		key := uint64ToBytes(log.Index)
		value, err := json.Marshal(log)
		if err != nil {
			return err
		}

		bucket := tx.Bucket([]byte(logBucket))
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// DeleteRange implements raft.LogStore.
func (d *Database) DeleteRange(min, max uint64) error {
	minKey := uint64ToBytes(min)

	// We can safely use standar Bolt functions, we're not interested in the
	// encoded values, we're only using delete, so Storm encoding/decoding
	// doesn't matter.
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()
	for k, _ := curs.Seek(minKey); k != nil; k, _ = curs.Next() {
		if bytesToUint64(k) > max {
			break
		}

		if err := curs.Delete(); err != nil {
			return err
		}
	}

	return tx.Commit()
}
<commit_msg>Added database close function.<commit_after>package database

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"sync"
	"time"

	"github.com/asdine/storm"
	"github.com/asdine/storm/index"
	"github.com/boltdb/bolt"
	"github.com/hashicorp/raft"
)

type (
	// Database is the lowest level of the gansoi database, it represent the
	// on-disk database. Database implements raft.FSM.
	Database struct {
		dbMutex       sync.RWMutex
		db            *storm.DB
		listenersLock sync.RWMutex
		listeners     []Listener
	}
)

var (
	// ErrNotFound is returned when the specified record is not saved.
	ErrNotFound = storm.ErrNotFound
)

const (
	// stableBucket is the bucket name used for raft stable storage.
	stableBucket = "raft.StableStore"

	// logBucket is used for raft log storage.
	logBucket = "raft.LogStore"
)

// NewDatabase will instantiate a new database placed in filepath.
func NewDatabase(filepath string) (*Database, error) {
	d := &Database{}

	err := d.open(path.Join(filepath, "gansoi.db"))
	if err != nil {
		return nil, err
	}

	return d, nil
}

// Close will close the database. Accessing the database after this will
// result in a deadlock.
func (d *Database) Close() error {
	d.dbMutex.RLock()
	return d.db.Close()
}

// open will open the underlying file storage.
func (d *Database) open(filepath string) error {
	db, err := storm.Open(
		filepath,
		storm.BoltOptions(0600, &bolt.Options{Timeout: 1 * time.Second}),
		storm.AutoIncrement(),
	)
	if err != nil {
		return err
	}

	d.db = db

	d.Storm().Bolt.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte(logBucket))

		return nil
	})

	return nil
}

// Storm will return the underlying Storm database.
func (d *Database) Storm() *storm.DB {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db
}

// ProcessLogEntry will process the log entry and apply whatever needs doing.
func (d *Database) ProcessLogEntry(entry *LogEntry) error {
	var err error

	var v interface{}

	switch entry.Command {
	case CommandSave:
		v, _ = entry.Payload()
		err = d.Save(v)
	case CommandDelete:
		v, _ = entry.Payload()
		err = d.db.DeleteStruct(v)
	default:
		err = fmt.Errorf("not implemented")
	}

	go func(command Command, data interface{}, err error) {
		d.listenersLock.RLock()

		for _, listener := range d.listeners {
			listener.PostLocalApply(command, data, err)
		}

		d.listenersLock.RUnlock()
	}(entry.Command, v, err)

	return err
}

// Apply implements raft.FSM.
func (d *Database) Apply(l *raft.Log) interface{} {
	entry := &LogEntry{}
	err := json.Unmarshal(l.Data, entry)
	if err != nil {
		// This should not happen..?
		fmt.Printf("%s: '%s'\n", err.Error(), string(l.Data))
		return nil
	}

	return d.ProcessLogEntry(entry)
}

// Snapshot implements raft.FSM.
func (d *Database) Snapshot() (raft.FSMSnapshot, error) {
	return &Snapshot{db: d}, nil
}

// Restore implements raft.FSM.
func (d *Database) Restore(source io.ReadCloser) error {
	db := d.Storm().Bolt
	d.dbMutex.Lock()
	defer d.dbMutex.Unlock()
	defer source.Close()

	path := db.Path()
	restorePath := path + ".restoretmp"

	file, err := os.Create(restorePath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, source)
	if err != nil {
		return err
	}

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

	err = os.Remove(path)
	if err != nil {
		return err
	}

	err = os.Rename(restorePath, path)
	if err != nil {
		return err
	}

	err = d.open(path)
	if err != nil {
		return err
	}

	return nil
}

// Save will save an object to the database.
func (d *Database) Save(data interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.Save(data)
}

// One will retrieve one (or zero) record from the database.
func (d *Database) One(fieldName string, value interface{}, to interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.One(fieldName, value, to)
}

// All lists all kinds of a type.
func (d *Database) All(to interface{}, limit int, skip int, reverse bool) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.All(to, func(opts *index.Options) {
		opts.Limit = limit
		opts.Skip = skip
		opts.Reverse = reverse
	})
}

// RegisterListener will register a listener for new changes to the database.
func (d *Database) RegisterListener(listener Listener) {
	d.listenersLock.Lock()
	defer d.listenersLock.Unlock()

	d.listeners = append(d.listeners, listener)
}

// Set implements raft.StableStore.
func (d *Database) Set(k, v []byte) error {
	return d.Storm().Set(stableBucket, k, v)
}

// Get implements raft.StableStore.
func (d *Database) Get(k []byte) ([]byte, error) {
	var value []byte
	err := d.Storm().Get(stableBucket, k, &value)
	if err != nil {
		return nil, err
	}

	return value, nil
}

// SetUint64 is like Set, but handles uint64 values
func (d *Database) SetUint64(key []byte, val uint64) error {
	return d.Set(key, uint64ToBytes(val))
}

// GetUint64 is like Get, but handles uint64 values
func (d *Database) GetUint64(key []byte) (uint64, error) {
	val, err := d.Get(key)
	if err != nil {
		return 0, err
	}

	return bytesToUint64(val), nil
}

// FirstIndex implements raft.LogStore.
func (d *Database) FirstIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	first, _ := curs.First()
	if first == nil {
		return 0, nil
	}

	return bytesToUint64(first), nil
}

// LastIndex implements raft.LogStore.
func (d *Database) LastIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	last, _ := curs.Last()
	if last == nil {
		return 0, nil
	}

	return bytesToUint64(last), nil
}

// GetLog implements raft.LogStore.
func (d *Database) GetLog(idx uint64, log *raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	value := tx.Bucket([]byte(logBucket)).Get(uint64ToBytes(idx))

	return json.Unmarshal(value, log)
}

// StoreLog implements raft.LogStore.
func (d *Database) StoreLog(log *raft.Log) error {
	return d.StoreLogs([]*raft.Log{log})
}

// StoreLogs implements raft.LogStore.
func (d *Database) StoreLogs(logs []*raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	for _, log := range logs {
		key := uint64ToBytes(log.Index)
		value, err := json.Marshal(log)
		if err != nil {
			return err
		}

		bucket := tx.Bucket([]byte(logBucket))
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// DeleteRange implements raft.LogStore.
func (d *Database) DeleteRange(min, max uint64) error {
	minKey := uint64ToBytes(min)

	// We can safely use standar Bolt functions, we're not interested in the
	// encoded values, we're only using delete, so Storm encoding/decoding
	// doesn't matter.
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()
	for k, _ := curs.Seek(minKey); k != nil; k, _ = curs.Next() {
		if bytesToUint64(k) > max {
			break
		}

		if err := curs.Delete(); err != nil {
			return err
		}
	}

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


<commit_before>package database

import "database/sql"

const (
	sqlCreateTables string = `
CREATE TABLE config (
	identifier TEXT,
	value TEXT
);

CREATE TABLE clients (
	identifier TEXT,
	ip_address TEXT,
	last_ping  INTEGER
);

CREATE TABLE client_url_list (
    client_id INTEGER,
    url_list_id INTEGER
);
CREATE INDEX client_id_index ON client_url_list (client_id);

CREATE TABLE url_lists (
    id INTEGER PRIMARY KEY,
    url_id INTEGER,
    name TEXT
);

CREATE TABLE urls (
	id INTEGER PRIMARY KEY,
	url TEXT
);

`

	sqlInsertUrl string = "INSERT INTO urls(url) VALUES(?);"
	sqlFetchUrls string = "SELECT url FROM urls;"
	sqlDeleteUrl string = "DELETE FROM urls WHERE url = ?;"

	sqlInsertConfig string = "INSERT INTO config(identifier, value) VALUES(?, ?);"
)

type Database struct {
	Conn *sql.DB
}

func (db *Database) Close() (err error) {
	err = db.Conn.Close()
	return
}

func (db *Database) InsertUrl(url string) (err error) {
	_, err = db.Conn.Exec(sqlInsertUrl, url)
	return
}

func (db *Database) DeleteUrl(url string) (err error) {
	_, err = db.Conn.Exec(sqlDeleteUrl, url)
	return
}

func (db *Database) InsertConfig(identifier string, value string) (err error) {
	_, err = db.Conn.Exec(sqlInsertConfig, identifier, value)
	return
}

func (db *Database) FetchUrls() (urls []string, err error) {
	rows, err := db.Conn.Query(sqlFetchUrls)
	defer rows.Close()
	if err != nil {
		return
	}

	for rows.Next() {
		var url string
		err = rows.Scan(&url)
		urls = append(urls, url)
	}
	err = rows.Err()

	return
}

func (db *Database) CreateTables() (err error) {
	// Run SQL to create necessary schema
	_, err = db.Conn.Exec(sqlCreateTables)
	return
}

func Connect(database string) (db *Database, err error) {
	c, err := sql.Open("sqlite3", database)
	if err == nil {
		db = &Database{c}
	}

	return
}
<commit_msg>Add missing schema and DB helper functions<commit_after>package database

import "database/sql"

const (
	sqlCreateTables string = `
CREATE TABLE config (
	identifier TEXT,
	value TEXT
);

CREATE TABLE clients (
	identifier TEXT,
	ip_address TEXT,
	last_ping  INTEGER
);

CREATE TABLE client_url_list (
    client_id INTEGER,
    url_list_id INTEGER
);
CREATE INDEX client_id_index ON client_url_list (client_id);

CREATE TABLE url_lists (
    id INTEGER PRIMARY KEY,
    name TEXT
);

CREATE TABLE urls (
	id INTEGER PRIMARY KEY,
	url TEXT
);

CREATE TABLE url_list_url (
	id INTEGER PRIMARY KEY,
	url_id INTEGER,
	url_list_id INTEGER
);

`

	sqlInsertUrl string = "INSERT INTO urls(url) VALUES(?);"
	sqlFetchUrls string = "SELECT url FROM urls;"
	sqlDeleteUrl string = "DELETE FROM urls WHERE url = ?;"

	sqlFindListId string = "SELECT id FROM url_lists WHERE name = ?;"
	sqlInsertList string = "INSERT INTO url_lists(name) VALUES(?);"
	sqlFetchLists string = "SELECT name FROM url_lists;"
	sqlDeleteList string = "DELETE FROM url_lists WHERE name = ?;"

	sqlInsertListUrl string = "INSERT INTO url_list_url(url_list_id, url_id) VALUES(?, ?);"
	sqlDeleteListUrl string = "DELETE FROM url_list_url WHERE id = ?;"
	sqlFetchListUrls string = "SELECT * FROM url_list_url WHERE url_list_id = ?;"

	sqlInsertConfig string = "INSERT INTO config(identifier, value) VALUES(?, ?);"
)

type Database struct {
	Conn *sql.DB
}

func (db *Database) Close() (err error) {
	err = db.Conn.Close()
	return
}

func (db *Database) InsertUrl(url string) (err error) {
	_, err = db.Conn.Exec(sqlInsertUrl, url)
	return
}

func (db *Database) DeleteUrl(url string) (err error) {
	_, err = db.Conn.Exec(sqlDeleteUrl, url)
	return
}

func (db *Database) InsertList(name string) (err error) {
	_, err = db.Conn.Exec(sqlInsertList, name)
	return
}

func (db *Database) DeleteList(name string) (err error) {
	_, err = db.Conn.Exec(sqlDeleteList, name)
	return
}

func (db *Database) InsertConfig(identifier string, value string) (err error) {
	_, err = db.Conn.Exec(sqlInsertConfig, identifier, value)
	return
}

func (db *Database) FindListId(name string) (id int, err error) {
	row := db.Conn.QueryRow(sqlFindListId, name)

	err = row.Scan(&id)

	return
}

func (db *Database) FetchUrls() (urls []string, err error) {
	rows, err := db.Conn.Query(sqlFetchUrls)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var url string

		err = rows.Scan(&url)
		if err != nil {
			return
		}

		urls = append(urls, url)
	}

	err = rows.Err()

	return
}

func (db *Database) FetchLists() (lists []string, err error) {
	rows, err := db.Conn.Query(sqlFetchLists)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var name string

		err = rows.Scan(&name)
		if err != nil {
			return
		}

		lists = append(lists, name)
	}

	err = rows.Err()

	return
}

func (db *Database) FetchListUrls(name string) (urls []string, err error) {
	list_id, err := db.FindListId(name)
	if err != nil {
		return
	}

	rows, err := db.Conn.Query(sqlFetchListUrls, list_id)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var url string

		err = rows.Scan(&url)
		if err != nil {
			return
		}

		urls = append(urls, url)
	}

	err = rows.Err()

	return
}
func (db *Database) CreateTables() (err error) {
	// Run SQL to create necessary schema
	_, err = db.Conn.Exec(sqlCreateTables)
	return
}

func Connect(database string) (db *Database, err error) {
	c, err := sql.Open("sqlite3", database)
	if err == nil {
		db = &Database{c}
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2013 Am Laher.
// This code is adapted from code within the Go tree.
// See Go's licence information below:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package ar implements access to ar archives.
// At this stage it only implements the 'common' format as used for .deb files.
package ar

import (
	"errors"
	"os"
	"path"
	"time"
)

/*

Sample ar data showing file entries:

!<arch>
debian-binary   1282478016  0     0     100644  4         `
2.0
control.tar.gz  1282478016  0     0     100644  444       `
.....binary-data.....

*/

const (
// the size of a file header
	headerSize   = 60
// the length of an 'ar' archive header
	arHeaderSize = 8
// the length of the filename header field
	fileNameSize = 16
// the length of the modTime header field
	modTimeSize  = 12
// the length of the UID header field
	uidSize      = 6
// the length of the GID header field
	gidSize      = 6
// the length of the Mode header field
	modeSize     = 8
// the length of the Size header field
	sizeSize     = 10
// the length of the 'magic' number
	magicSize    = 2
)

var (
	zeroBlock = make([]byte, headerSize)
// error describing an invalid ar file header
	ErrHeader = errors.New("ar: invalid ar header")
)

type Header struct {
	// Name is the name of the file.
	// It must be a relative path: it must not start with a drive
	// letter (e.g. C:) or leading slash, and only forward slashes
	// are allowed.
	Name    string
	ModTime time.Time
	Uid     int
	Gid     int
	Mode    int64
	Size    int64
}

type slicer []byte

func (sp *slicer) next(n int) (b []byte) {
	s := *sp
	b, *sp = s[0:n], s[n:]
	return
}

// FileInfo returns an os.FileInfo for the Header.
func (h *Header) FileInfo() os.FileInfo {
	return headerFileInfo{h}
}

// headerFileInfo implements os.FileInfo.
type headerFileInfo struct {
	h *Header
}

func (fi headerFileInfo) Size() int64 { return fi.h.Size }
func (fi headerFileInfo) IsDir() bool        { return fi.Mode().IsDir() }
func (fi headerFileInfo) ModTime() time.Time { return fi.h.ModTime }
func (fi headerFileInfo) Sys() interface{}   { return fi.h }

// Name returns the base name of the file.
func (fi headerFileInfo) Name() string {
	if fi.IsDir() {
		return path.Base(path.Clean(fi.h.Name))
	}
	return path.Base(fi.h.Name)
}

// Mode returns the permission and mode bits for the headerFileInfo.
func (fi headerFileInfo) Mode() (mode os.FileMode) {
	// Set file permission bits.
	mode = os.FileMode(fi.h.Mode).Perm()
	return mode
}

// FileInfoHeader creates a partially-populated Header from an
// os.FileInfo.
// Because os.FileInfo's Name method returns only the base name of
// the file it describes, it may be necessary to modify the Name field
// of the returned header to provide the full path name of the file.
func FileInfoHeader(fi os.FileInfo) (*Header, error) {
	size := fi.Size()
	h := &Header{
		Name:               fi.Name(),
		Size: int64(size),
	}
	fm := fi.Mode()
	h.ModTime = fi.ModTime()
	h.Mode = int64(fm.Perm())

	return h, nil
}
<commit_msg>improving docs<commit_after>// Copyright 2013 Am Laher.
// This code is adapted from code within the Go tree.
// See Go's licence information below:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package ar implements access to ar archives.
// At this stage it only implements the 'common' format as used for .deb files.
package ar

import (
	"errors"
	"os"
	"path"
	"time"
)

/*

  Sample ar data showing file entries:

  !<arch>
  debian-binary   1282478016  0     0     100644  4         `
  2.0
  control.tar.gz  1282478016  0     0     100644  444       `
  .....binary-data.....

*/

const (
// the size of a file header
	headerSize   = 60
// the length of an 'ar' archive header
	arHeaderSize = 8
// the length of the filename header field
	fileNameSize = 16
// the length of the modTime header field
	modTimeSize  = 12
// the length of the UID header field
	uidSize      = 6
// the length of the GID header field
	gidSize      = 6
// the length of the Mode header field
	modeSize     = 8
// the length of the Size header field
	sizeSize     = 10
// the length of the 'magic' number
	magicSize    = 2
)

var (
	zeroBlock = make([]byte, headerSize)
// error describing an invalid ar file header
	ErrHeader = errors.New("ar: invalid ar header")
)

// A Header represents a single header in an ar archive.
// Some fields may not be populated.
type Header struct {
	// Name is the name of the file.
	// It must be a relative path: it must not start with a drive
	// letter (e.g. C:) or leading slash, and only forward slashes
	// are allowed.
	Name    string    // name of header file entry
	ModTime time.Time // modified time
	Uid     int       // user id of owner
	Gid     int       // group id of owner
	Mode    int64     // permission and mode bits
	Size    int64     // length in bytes
}

type slicer []byte

func (sp *slicer) next(n int) (b []byte) {
	s := *sp
	b, *sp = s[0:n], s[n:]
	return
}

// FileInfo returns an os.FileInfo for the Header.
func (h *Header) FileInfo() os.FileInfo {
	return headerFileInfo{h}
}

// headerFileInfo implements os.FileInfo.
type headerFileInfo struct {
	h *Header
}

func (fi headerFileInfo) Size() int64 { return fi.h.Size }
func (fi headerFileInfo) IsDir() bool        { return fi.Mode().IsDir() }
func (fi headerFileInfo) ModTime() time.Time { return fi.h.ModTime }
func (fi headerFileInfo) Sys() interface{}   { return fi.h }

// Name returns the base name of the file.
func (fi headerFileInfo) Name() string {
	if fi.IsDir() {
		return path.Base(path.Clean(fi.h.Name))
	}
	return path.Base(fi.h.Name)
}

// Mode returns the permission and mode bits for the headerFileInfo.
func (fi headerFileInfo) Mode() (mode os.FileMode) {
	// Set file permission bits.
	mode = os.FileMode(fi.h.Mode).Perm()
	return mode
}

// FileInfoHeader creates a partially-populated Header from an
// os.FileInfo.
// Because os.FileInfo's Name method returns only the base name of
// the file it describes, it may be necessary to modify the Name field
// of the returned header to provide the full path name of the file.
func FileInfoHeader(fi os.FileInfo) (*Header, error) {
	size := fi.Size()
	h := &Header{
		Name:               fi.Name(),
		Size: int64(size),
	}
	fm := fi.Mode()
	h.ModTime = fi.ModTime()
	h.Mode = int64(fm.Perm())

	return h, nil
}
<|endoftext|>


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

package auth

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"

	"github.com/couchbase/sync_gateway/base"
	ch "github.com/couchbase/sync_gateway/channels"
)

/** A group that users can belong to, with associated channel permisisons. */
type roleImpl struct {
	Name_             string      `json:"name,omitempty"`
	ExplicitChannels_ ch.TimedSet `json:"admin_channels,omitempty"`
	Channels_         ch.TimedSet `json:"all_channels"`
	Sequence_         uint64      `json:"sequence"`
	PreviousChannels_ ch.TimedSet `json:"previous_channels,omitempty"`
	vbNo              *uint16
}

var kValidNameRegexp *regexp.Regexp

func init() {
	var err error
	kValidNameRegexp, err = regexp.Compile(`^[-+.@%\w]*$`)
	if err != nil {
		panic("Bad kValidNameRegexp")
	}
}

func (role *roleImpl) initRole(name string, channels base.Set) error {
	channels = ch.ExpandingStar(channels)
	role.Name_ = name
	role.ExplicitChannels_ = ch.AtSequence(channels, 1)
	return role.validate()
}

// Is this string a valid name for a User/Role? (Valid chars are alphanumeric and any of "_-+.@")
func IsValidPrincipalName(name string) bool {
	return kValidNameRegexp.MatchString(name)
}

// Creates a new Role object.
func (auth *Authenticator) NewRole(name string, channels base.Set) (Role, error) {
	role := &roleImpl{}
	if err := role.initRole(name, channels); err != nil {
		return nil, err
	}
	if err := auth.rebuildChannels(role); err != nil {
		return nil, err
	}
	return role, nil
}

func (auth *Authenticator) UnmarshalRole(data []byte, defaultName string, defaultSeq uint64) (Role, error) {
	role := &roleImpl{}
	if err := json.Unmarshal(data, role); err != nil {
		return nil, err
	}
	if role.Name_ == "" {
		role.Name_ = defaultName
	}

	defaultVbSeq := ch.NewVbSimpleSequence(defaultSeq)
	for channel, seq := range role.ExplicitChannels_ {
		if seq.Sequence == 0 {
			role.ExplicitChannels_[channel] = defaultVbSeq
		}
	}
	if err := role.validate(); err != nil {
		return nil, err
	}

	return role, nil
}

// Key prefix reserved for role documents in the bucket
const RoleKeyPrefix = "_sync:role:"

func docIDForRole(name string) string {
	return RoleKeyPrefix + name
}

func (role *roleImpl) DocID() string {
	return docIDForRole(role.Name_)
}

// Key used in 'access' view (not same meaning as doc ID)
func (role *roleImpl) accessViewKey() string {
	return "role:" + role.Name_
}

//////// ACCESSORS:

func (role *roleImpl) Name() string {
	return role.Name_
}

func (role *roleImpl) Sequence() uint64 {
	return role.Sequence_
}
func (role *roleImpl) SetSequence(sequence uint64) {
	role.Sequence_ = sequence
}

func (role *roleImpl) Channels() ch.TimedSet {
	return role.Channels_
}

func (role *roleImpl) setChannels(channels ch.TimedSet) {
	role.Channels_ = channels
}

func (role *roleImpl) ExplicitChannels() ch.TimedSet {
	return role.ExplicitChannels_
}

func (role *roleImpl) SetExplicitChannels(channels ch.TimedSet) {
	role.ExplicitChannels_ = channels
	role.setChannels(nil)
}

func (role *roleImpl) PreviousChannels() ch.TimedSet {
	return role.PreviousChannels_
}

func (role *roleImpl) SetPreviousChannels(channels ch.TimedSet) {
	role.PreviousChannels_ = channels
}

// Checks whether this role object contains valid data; if not, returns an error.
func (role *roleImpl) validate() error {
	if !IsValidPrincipalName(role.Name_) {
		return base.HTTPErrorf(http.StatusBadRequest, "Invalid name %q", role.Name_)
	}
	return role.ExplicitChannels_.Validate()
}

//////// CHANNEL AUTHORIZATION:

func (role *roleImpl) UnauthError(message string) error {
	if role.Name_ == "" {
		return base.HTTPErrorf(http.StatusUnauthorized, "login required: "+message)
	}
	return base.HTTPErrorf(http.StatusForbidden, message)
}

// Returns true if the Role is allowed to access the channel.
// A nil Role means access control is disabled, so the function will return true.
func (role *roleImpl) CanSeeChannel(channel string) bool {
	return role == nil || role.Channels_.Contains(channel) || role.Channels_.Contains(ch.UserStarChannel)
}

// Returns the sequence number since which the Role has been able to access the channel, else zero.
func (role *roleImpl) CanSeeChannelSince(channel string) uint64 {
	seq := role.Channels_[channel]
	if seq.Sequence == 0 {
		seq = role.Channels_[ch.UserStarChannel]
	}
	return seq.Sequence
}

// Returns the sequence number since which the Role has been able to access the channel, else zero.  Sets the vb
// for an admin channel grant, if needed.
func (role *roleImpl) CanSeeChannelSinceVbSeq(channel string, numVbuckets int) (base.VbSeq, bool) {
	seq, ok := role.Channels_[channel]
	if !ok {
		seq, ok = role.Channels_[ch.UserStarChannel]
		if !ok {
			return base.VbSeq{}, false
		}
	}
	if seq.VbNo == nil {
		roleDocVbNo := role.getVbNo(numVbuckets)
		seq.VbNo = &roleDocVbNo
	}
	return base.VbSeq{*seq.VbNo, seq.Sequence}, true
}

func (role *roleImpl) AuthorizeAllChannels(channels base.Set) error {
	return authorizeAllChannels(role, channels)
}

func (role *roleImpl) AuthorizeAnyChannel(channels base.Set) error {
	return authorizeAnyChannel(role, channels)
}

func (role *roleImpl) ValidateGrant(vbSeq *ch.VbSequence, numVbuckets int) bool {

	// If the sequence is zero, this is an admin grant that hasn't been updated by accel - ignore
	if vbSeq.Sequence == 0 {
		return false
	}

	// If vbSeq is nil, this is an admin grant.  Set the vb to the vb of the user doc
	if vbSeq.VbNo == nil {
		calculatedVbNo := role.getVbNo(numVbuckets)
		vbSeq.VbNo = &calculatedVbNo
	}
	return true
}

func (role *roleImpl) getVbNo(numVbuckets int) uint16 {
	if role.vbNo == nil {
		calculatedVbNo := uint16(base.VBHash(role.DocID(), numVbuckets))
		role.vbNo = &calculatedVbNo
	}
	return *role.vbNo
}

// Returns an HTTP 403 error if the Principal is not allowed to access all the given channels.
// A nil Principal means access control is disabled, so the function will return nil.
func authorizeAllChannels(princ Principal, channels base.Set) error {
	var forbidden []string
	for channel := range channels {
		if !princ.CanSeeChannel(channel) {
			if forbidden == nil {
				forbidden = make([]string, 0, len(channels))
			}
			forbidden = append(forbidden, channel)
		}
	}
	if forbidden != nil {
		return princ.UnauthError(fmt.Sprintf("You are not allowed to see channels %v", forbidden))
	}
	return nil
}

// Returns an HTTP 403 error if the Principal is not allowed to access any of the given channels.
// A nil Role means access control is disabled, so the function will return nil.
func authorizeAnyChannel(princ Principal, channels base.Set) error {
	if len(channels) > 0 {
		for channel := range channels {
			if princ.CanSeeChannel(channel) {
				return nil
			}
		}
	} else if princ.Channels().Contains(ch.UserStarChannel) {
		return nil
	}
	return princ.UnauthError("You are not allowed to see this")
}
<commit_msg>Fixes #2307 Regex lock bottleneck (#2308)<commit_after>//  Copyright (c) 2012 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package auth

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"

	"github.com/couchbase/sync_gateway/base"
	ch "github.com/couchbase/sync_gateway/channels"
)

/** A group that users can belong to, with associated channel permisisons. */
type roleImpl struct {
	Name_             string      `json:"name,omitempty"`
	ExplicitChannels_ ch.TimedSet `json:"admin_channels,omitempty"`
	Channels_         ch.TimedSet `json:"all_channels"`
	Sequence_         uint64      `json:"sequence"`
	PreviousChannels_ ch.TimedSet `json:"previous_channels,omitempty"`
	vbNo              *uint16
}

var kValidNameRegexp *regexp.Regexp

func init() {
	var err error
	kValidNameRegexp, err = regexp.Compile(`^[-+.@%\w]*$`)
	if err != nil {
		panic("Bad kValidNameRegexp")
	}
}

func (role *roleImpl) initRole(name string, channels base.Set) error {
	channels = ch.ExpandingStar(channels)
	role.Name_ = name
	role.ExplicitChannels_ = ch.AtSequence(channels, 1)
	return role.validate()
}

// Is this string a valid name for a User/Role? (Valid chars are alphanumeric and any of "_-+.@")
func IsValidPrincipalName(name string) bool {
	return kValidNameRegexp.Copy().MatchString(name)
}

// Creates a new Role object.
func (auth *Authenticator) NewRole(name string, channels base.Set) (Role, error) {
	role := &roleImpl{}
	if err := role.initRole(name, channels); err != nil {
		return nil, err
	}
	if err := auth.rebuildChannels(role); err != nil {
		return nil, err
	}
	return role, nil
}

func (auth *Authenticator) UnmarshalRole(data []byte, defaultName string, defaultSeq uint64) (Role, error) {
	role := &roleImpl{}
	if err := json.Unmarshal(data, role); err != nil {
		return nil, err
	}
	if role.Name_ == "" {
		role.Name_ = defaultName
	}

	defaultVbSeq := ch.NewVbSimpleSequence(defaultSeq)
	for channel, seq := range role.ExplicitChannels_ {
		if seq.Sequence == 0 {
			role.ExplicitChannels_[channel] = defaultVbSeq
		}
	}
	if err := role.validate(); err != nil {
		return nil, err
	}

	return role, nil
}

// Key prefix reserved for role documents in the bucket
const RoleKeyPrefix = "_sync:role:"

func docIDForRole(name string) string {
	return RoleKeyPrefix + name
}

func (role *roleImpl) DocID() string {
	return docIDForRole(role.Name_)
}

// Key used in 'access' view (not same meaning as doc ID)
func (role *roleImpl) accessViewKey() string {
	return "role:" + role.Name_
}

//////// ACCESSORS:

func (role *roleImpl) Name() string {
	return role.Name_
}

func (role *roleImpl) Sequence() uint64 {
	return role.Sequence_
}
func (role *roleImpl) SetSequence(sequence uint64) {
	role.Sequence_ = sequence
}

func (role *roleImpl) Channels() ch.TimedSet {
	return role.Channels_
}

func (role *roleImpl) setChannels(channels ch.TimedSet) {
	role.Channels_ = channels
}

func (role *roleImpl) ExplicitChannels() ch.TimedSet {
	return role.ExplicitChannels_
}

func (role *roleImpl) SetExplicitChannels(channels ch.TimedSet) {
	role.ExplicitChannels_ = channels
	role.setChannels(nil)
}

func (role *roleImpl) PreviousChannels() ch.TimedSet {
	return role.PreviousChannels_
}

func (role *roleImpl) SetPreviousChannels(channels ch.TimedSet) {
	role.PreviousChannels_ = channels
}

// Checks whether this role object contains valid data; if not, returns an error.
func (role *roleImpl) validate() error {
	if !IsValidPrincipalName(role.Name_) {
		return base.HTTPErrorf(http.StatusBadRequest, "Invalid name %q", role.Name_)
	}
	return role.ExplicitChannels_.Validate()
}

//////// CHANNEL AUTHORIZATION:

func (role *roleImpl) UnauthError(message string) error {
	if role.Name_ == "" {
		return base.HTTPErrorf(http.StatusUnauthorized, "login required: "+message)
	}
	return base.HTTPErrorf(http.StatusForbidden, message)
}

// Returns true if the Role is allowed to access the channel.
// A nil Role means access control is disabled, so the function will return true.
func (role *roleImpl) CanSeeChannel(channel string) bool {
	return role == nil || role.Channels_.Contains(channel) || role.Channels_.Contains(ch.UserStarChannel)
}

// Returns the sequence number since which the Role has been able to access the channel, else zero.
func (role *roleImpl) CanSeeChannelSince(channel string) uint64 {
	seq := role.Channels_[channel]
	if seq.Sequence == 0 {
		seq = role.Channels_[ch.UserStarChannel]
	}
	return seq.Sequence
}

// Returns the sequence number since which the Role has been able to access the channel, else zero.  Sets the vb
// for an admin channel grant, if needed.
func (role *roleImpl) CanSeeChannelSinceVbSeq(channel string, numVbuckets int) (base.VbSeq, bool) {
	seq, ok := role.Channels_[channel]
	if !ok {
		seq, ok = role.Channels_[ch.UserStarChannel]
		if !ok {
			return base.VbSeq{}, false
		}
	}
	if seq.VbNo == nil {
		roleDocVbNo := role.getVbNo(numVbuckets)
		seq.VbNo = &roleDocVbNo
	}
	return base.VbSeq{*seq.VbNo, seq.Sequence}, true
}

func (role *roleImpl) AuthorizeAllChannels(channels base.Set) error {
	return authorizeAllChannels(role, channels)
}

func (role *roleImpl) AuthorizeAnyChannel(channels base.Set) error {
	return authorizeAnyChannel(role, channels)
}

func (role *roleImpl) ValidateGrant(vbSeq *ch.VbSequence, numVbuckets int) bool {

	// If the sequence is zero, this is an admin grant that hasn't been updated by accel - ignore
	if vbSeq.Sequence == 0 {
		return false
	}

	// If vbSeq is nil, this is an admin grant.  Set the vb to the vb of the user doc
	if vbSeq.VbNo == nil {
		calculatedVbNo := role.getVbNo(numVbuckets)
		vbSeq.VbNo = &calculatedVbNo
	}
	return true
}

func (role *roleImpl) getVbNo(numVbuckets int) uint16 {
	if role.vbNo == nil {
		calculatedVbNo := uint16(base.VBHash(role.DocID(), numVbuckets))
		role.vbNo = &calculatedVbNo
	}
	return *role.vbNo
}

// Returns an HTTP 403 error if the Principal is not allowed to access all the given channels.
// A nil Principal means access control is disabled, so the function will return nil.
func authorizeAllChannels(princ Principal, channels base.Set) error {
	var forbidden []string
	for channel := range channels {
		if !princ.CanSeeChannel(channel) {
			if forbidden == nil {
				forbidden = make([]string, 0, len(channels))
			}
			forbidden = append(forbidden, channel)
		}
	}
	if forbidden != nil {
		return princ.UnauthError(fmt.Sprintf("You are not allowed to see channels %v", forbidden))
	}
	return nil
}

// Returns an HTTP 403 error if the Principal is not allowed to access any of the given channels.
// A nil Role means access control is disabled, so the function will return nil.
func authorizeAnyChannel(princ Principal, channels base.Set) error {
	if len(channels) > 0 {
		for channel := range channels {
			if princ.CanSeeChannel(channel) {
				return nil
			}
		}
	} else if princ.Channels().Contains(ch.UserStarChannel) {
		return nil
	}
	return princ.UnauthError("You are not allowed to see this")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"expvar"
	"fmt"
	"github.com/AkihiroSuda/go-netfilter-queue"
	"github.com/google/gopacket/layers"
	"github.com/op/go-logging"
	"github.com/patrickmn/go-cache"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("autoblock")

var blockReason = map[string]string{
	"127.0.0.2":  "amavis[blocklist.de] or bogon",
	"127.0.0.3":  "apacheddos",
	"127.0.0.4":  "asterisk",
	"127.0.0.5":  "badbot",
	"127.0.0.6":  "ftp",
	"127.0.0.7":  "imap",
	"127.0.0.8":  "ircbot",
	"127.0.0.9":  "mail",
	"127.0.0.10": "pop3",
	"127.0.0.11": "regbot",
	"127.0.0.12": "rfi-attack",
	"127.0.0.13": "sasl",
	"127.0.0.14": "ssh",
	"127.0.0.15": "w00tw00t",
	"127.0.0.16": "portflood",
	"127.0.0.17": "sql-injection",
	"127.0.0.18": "webmin",
	"127.0.0.19": "trigger-spam",
	"127.0.0.20": "manual",
	"127.0.0.21": "bruteforcelogin",
	"127.0.0.22": "mysql",
}

var blocklists4 = map[string]string{
	"blocklist.de": ".all.bl.blocklist.de",
	"cymru bogons": ".v4.fullbogons.cymru.com",
}

var (
	countDNSqueries, countDNSqueriesSuccess                   expvar.Int
	count4Packets, count4PacketsBlocked                       expvar.Int
	countCacheHitsBlock, countCacheHitsPass, countCacheMisses expvar.Int
)

func init() {

	m := expvar.NewMap("autoblock")
	m.Set("dnsQueries", &countDNSqueries)
	m.Set("dnsQueriesSuccess", &countDNSqueriesSuccess)
	m.Set("ipv4Packets", &count4Packets)
	m.Set("ipv4PacketsBlocked", &count4PacketsBlocked)
	m.Set("cacheHitsBlock", &countCacheHitsBlock)
	m.Set("cacheHitsPass", &countCacheHitsPass)
	m.Set("cacheHitsMisses", &countCacheMisses)
}

func main() {

	beStdErr := logging.NewLogBackend(os.Stderr, "", 0)
	beSyslog, _ := logging.NewSyslogBackend("autoblock")
	besysl := logging.AddModuleLevel(beSyslog)
	bestdErrl := logging.AddModuleLevel(beStdErr)
	level, err := logging.LogLevel("DEBUG")
	if err != nil {
		log.Fatal("Log level is not valid")
	}
	besysl.SetLevel(level, "")
	bestdErrl.SetLevel(level, "")
	logging.SetBackend(besysl, bestdErrl)

	log.Info("Starting...")

	go http.ListenAndServe(":8080", nil) // expvar serving, not for anything else

	lookupChannel := make(chan net.IP, 1000)

	// Create a cache with a default expiration time of 5 minutes, and which
	// purges expired items every 30 seconds
	// This makes some sense as DNS blocklist entries expire, and IP traffic per
	// entry would normally be low (few retries as they don't get a response)
	c := cache.New(5*time.Minute, 30*time.Second)

	// one worker per queue
	for i := 0; i < 4; i++ {
		go gofilter(i, lookupChannel, c)
	}
	golookup(lookupChannel, c)
}

func golookup(ch chan net.IP, c *cache.Cache) {
	for true {
		select {
		case ip := <-ch:
			_, found := c.Get(ip.String()) // check again, it may have been enqueued multiple times.
			if !found {
				result, _ := checkBlocklists(ip)
				if result != nil {
					c.Set(ip.String(), result, cache.DefaultExpiration)
				} else {
					c.Set(ip.String(), net.IPv4(0, 0, 0, 0), cache.DefaultExpiration) // cache as not blocked
				}
			}
		}
	}
}

func gofilter(i int, ch chan net.IP, c *cache.Cache) {
	nfq, err := netfilter.NewNFQueue(uint16(i), 100, netfilter.NF_DEFAULT_PACKET_SIZE)
	if err != nil {
		log.Critical(err)
		os.Exit(1)
	}
	defer nfq.Close()
	packets := nfq.GetPackets()

	for true {
		select {
		case p := <-packets:

			if ipLayer := p.Packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
				count4Packets.Add(1)
				ipl, _ := ipLayer.(*layers.IPv4)
				reason, found := c.Get(ipl.SrcIP.String())
				if found && !reason.(net.IP).Equal(net.IPv4(0, 0, 0, 0)) {
					countCacheHitsBlock.Add(1)
					p.SetVerdict(netfilter.NF_DROP)
					log.Info(ipl.SrcIP.String() + " to " + ipl.DstIP.String() + " blocked for reason: " + blockReason[reason.(net.IP).String()])
					count4PacketsBlocked.Add(1)
				} else {
					if !found {
						countCacheMisses.Add(1)
						ch <- ipl.SrcIP // enqueue for checks, out of this processing band
						log.Debug(ipl.SrcIP.String(), "passed okay")
					} else {
						countCacheHitsPass.Add(1)
					}
					p.SetVerdict(netfilter.NF_ACCEPT)
				}
			} else {
				log.Error("Not IP packet")
			}
		}
	}
}

func checkBlocklists(src net.IP) (net.IP, error) {
	log.Debug("checkBlocklist:", src)
	src4 := src.To4()
	if src4 == nil {
		return nil, errors.New("Not IPv4")
	}
	//reverse the address
	reversedaddr := fmt.Sprintf("%d.%d.%d.%d", src4[3], src4[2], src4[1], src4[0])

	for name, addr := range blocklists4 {
		ips, err := net.LookupIP(reversedaddr + addr)
		countDNSqueries.Add(1)
		if err != nil {
			log.Debug(src.String(), "not found in blocklist", name, "result:", err)

		} else {
			countDNSqueriesSuccess.Add(1)
			log.Debug(src.String(), "found in blocklist", name, "cause:", blockReason[ips[0].String()])
			return ips[0], err // only the first one required
		}
	}

	return nil, nil
}
<commit_msg>Fixed stats name from cacheHitsMisses to the less poetic cacheMisses<commit_after>package main

import (
	"errors"
	"expvar"
	"fmt"
	"github.com/AkihiroSuda/go-netfilter-queue"
	"github.com/google/gopacket/layers"
	"github.com/op/go-logging"
	"github.com/patrickmn/go-cache"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("autoblock")

var blockReason = map[string]string{
	"127.0.0.2":  "amavis[blocklist.de] or bogon",
	"127.0.0.3":  "apacheddos",
	"127.0.0.4":  "asterisk",
	"127.0.0.5":  "badbot",
	"127.0.0.6":  "ftp",
	"127.0.0.7":  "imap",
	"127.0.0.8":  "ircbot",
	"127.0.0.9":  "mail",
	"127.0.0.10": "pop3",
	"127.0.0.11": "regbot",
	"127.0.0.12": "rfi-attack",
	"127.0.0.13": "sasl",
	"127.0.0.14": "ssh",
	"127.0.0.15": "w00tw00t",
	"127.0.0.16": "portflood",
	"127.0.0.17": "sql-injection",
	"127.0.0.18": "webmin",
	"127.0.0.19": "trigger-spam",
	"127.0.0.20": "manual",
	"127.0.0.21": "bruteforcelogin",
	"127.0.0.22": "mysql",
}

var blocklists4 = map[string]string{
	"blocklist.de": ".all.bl.blocklist.de",
	"cymru bogons": ".v4.fullbogons.cymru.com",
}

var (
	countDNSqueries, countDNSqueriesSuccess                   expvar.Int
	count4Packets, count4PacketsBlocked                       expvar.Int
	countCacheHitsBlock, countCacheHitsPass, countCacheMisses expvar.Int
)

func init() {

	m := expvar.NewMap("autoblock")
	m.Set("dnsQueries", &countDNSqueries)
	m.Set("dnsQueriesSuccess", &countDNSqueriesSuccess)
	m.Set("ipv4Packets", &count4Packets)
	m.Set("ipv4PacketsBlocked", &count4PacketsBlocked)
	m.Set("cacheHitsBlock", &countCacheHitsBlock)
	m.Set("cacheHitsPass", &countCacheHitsPass)
	m.Set("cacheMisses", &countCacheMisses)
}

func main() {

	beStdErr := logging.NewLogBackend(os.Stderr, "", 0)
	beSyslog, _ := logging.NewSyslogBackend("autoblock")
	besysl := logging.AddModuleLevel(beSyslog)
	bestdErrl := logging.AddModuleLevel(beStdErr)
	level, err := logging.LogLevel("DEBUG")
	if err != nil {
		log.Fatal("Log level is not valid")
	}
	besysl.SetLevel(level, "")
	bestdErrl.SetLevel(level, "")
	logging.SetBackend(besysl, bestdErrl)

	log.Info("Starting...")

	go http.ListenAndServe(":8080", nil) // expvar serving, not for anything else

	lookupChannel := make(chan net.IP, 1000)

	// Create a cache with a default expiration time of 5 minutes, and which
	// purges expired items every 30 seconds
	// This makes some sense as DNS blocklist entries expire, and IP traffic per
	// entry would normally be low (few retries as they don't get a response)
	c := cache.New(5*time.Minute, 30*time.Second)

	// one worker per queue
	for i := 0; i < 4; i++ {
		go gofilter(i, lookupChannel, c)
	}
	golookup(lookupChannel, c)
}

func golookup(ch chan net.IP, c *cache.Cache) {
	for true {
		select {
		case ip := <-ch:
			_, found := c.Get(ip.String()) // check again, it may have been enqueued multiple times.
			if !found {
				result, _ := checkBlocklists(ip)
				if result != nil {
					c.Set(ip.String(), result, cache.DefaultExpiration)
				} else {
					c.Set(ip.String(), net.IPv4(0, 0, 0, 0), cache.DefaultExpiration) // cache as not blocked
				}
			}
		}
	}
}

func gofilter(i int, ch chan net.IP, c *cache.Cache) {
	nfq, err := netfilter.NewNFQueue(uint16(i), 100, netfilter.NF_DEFAULT_PACKET_SIZE)
	if err != nil {
		log.Critical(err)
		os.Exit(1)
	}
	defer nfq.Close()
	packets := nfq.GetPackets()

	for true {
		select {
		case p := <-packets:

			if ipLayer := p.Packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
				count4Packets.Add(1)
				ipl, _ := ipLayer.(*layers.IPv4)
				reason, found := c.Get(ipl.SrcIP.String())
				if found && !reason.(net.IP).Equal(net.IPv4(0, 0, 0, 0)) {
					countCacheHitsBlock.Add(1)
					p.SetVerdict(netfilter.NF_DROP)
					log.Info(ipl.SrcIP.String() + " to " + ipl.DstIP.String() + " blocked for reason: " + blockReason[reason.(net.IP).String()])
					count4PacketsBlocked.Add(1)
				} else {
					if !found {
						countCacheMisses.Add(1)
						ch <- ipl.SrcIP // enqueue for checks, out of this processing band
						log.Debug(ipl.SrcIP.String(), "passed okay")
					} else {
						countCacheHitsPass.Add(1)
					}
					p.SetVerdict(netfilter.NF_ACCEPT)
				}
			} else {
				log.Error("Not IP packet")
			}
		}
	}
}

func checkBlocklists(src net.IP) (net.IP, error) {
	log.Debug("checkBlocklist:", src)
	src4 := src.To4()
	if src4 == nil {
		return nil, errors.New("Not IPv4")
	}
	//reverse the address
	reversedaddr := fmt.Sprintf("%d.%d.%d.%d", src4[3], src4[2], src4[1], src4[0])

	for name, addr := range blocklists4 {
		ips, err := net.LookupIP(reversedaddr + addr)
		countDNSqueries.Add(1)
		if err != nil {
			log.Debug(src.String(), "not found in blocklist", name, "result:", err)

		} else {
			countDNSqueriesSuccess.Add(1)
			log.Debug(src.String(), "found in blocklist", name, "cause:", blockReason[ips[0].String()])
			return ips[0], err // only the first one required
		}
	}

	return nil, nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package syscall contains an interface to the low-level operating system
// primitives.  The details vary depending on the underlying system.
// Its primary use is inside other packages that provide a more portable
// interface to the system, such as "os", "time" and "net".  Use those
// packages rather than this one if you can.
// For details of the functions and data types in this package consult
// the manuals for the appropriate operating system.
// These calls return err == nil to indicate success; otherwise
// err is an operating system error describing the failure.
// On most systems, that error has type syscall.Errno.
package syscall

// StringByteSlice returns a NUL-terminated slice of bytes
// containing the text of s.
func StringByteSlice(s string) []byte {
	a := make([]byte, len(s)+1)
	copy(a, s)
	return a
}

// StringBytePtr returns a pointer to a NUL-terminated array of bytes
// containing the text of s.
func StringBytePtr(s string) *byte { return &StringByteSlice(s)[0] }

// Single-word zero for use when we need a valid pointer to 0 bytes.
// See mksyscall.pl.
var _zero uintptr

func (ts *Timespec) Unix() (sec int64, nsec int64) {
	return int64(ts.Sec), int64(ts.Nsec)
}

func (tv *Timeval) Unix() (sec int64, nsec int64) {
	return int64(tv.Sec), int64(tv.Usec) * 1000
}
<commit_msg>syscall: add Timeval.Nano, Timespec.Nano, for conversion to Duration<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package syscall contains an interface to the low-level operating system
// primitives.  The details vary depending on the underlying system.
// Its primary use is inside other packages that provide a more portable
// interface to the system, such as "os", "time" and "net".  Use those
// packages rather than this one if you can.
// For details of the functions and data types in this package consult
// the manuals for the appropriate operating system.
// These calls return err == nil to indicate success; otherwise
// err is an operating system error describing the failure.
// On most systems, that error has type syscall.Errno.
package syscall

// StringByteSlice returns a NUL-terminated slice of bytes
// containing the text of s.
func StringByteSlice(s string) []byte {
	a := make([]byte, len(s)+1)
	copy(a, s)
	return a
}

// StringBytePtr returns a pointer to a NUL-terminated array of bytes
// containing the text of s.
func StringBytePtr(s string) *byte { return &StringByteSlice(s)[0] }

// Single-word zero for use when we need a valid pointer to 0 bytes.
// See mksyscall.pl.
var _zero uintptr

func (ts *Timespec) Unix() (sec int64, nsec int64) {
	return int64(ts.Sec), int64(ts.Nsec)
}

func (tv *Timeval) Unix() (sec int64, nsec int64) {
	return int64(tv.Sec), int64(tv.Usec) * 1000
}

func (ts *Timespec) Nano() int64 {
	return int64(ts.Sec)*1e9 + int64(ts.Nsec)
}

func (tv *Timeval) Nano() int64 {
	return int64(tv.Sec)*1e9 + int64(tv.Usec)*1000
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package testing provides support for automated testing of Go packages.
// It is intended to be used in concert with the ``gotest'' utility, which automates
// execution of any function of the form
//     func TestXxx(*testing.T)
// where Xxx can be any alphanumeric string (but the first letter must not be in
// [a-z]) and serves to identify the test routine.
// These TestXxx routines should be declared within the package they are testing.
//
// Functions of the form
//     func BenchmarkXxx(*testing.B)
// are considered benchmarks, and are executed by gotest when the -test.bench
// flag is provided.
//
// A sample benchmark function looks like this:
//     func BenchmarkHello(b *testing.B) {
//         for i := 0; i < b.N; i++ {
//             fmt.Sprintf("hello")
//         }
//     }
// The benchmark package will vary b.N until the benchmark function lasts
// long enough to be timed reliably.  The output
//     testing.BenchmarkHello    10000000    282 ns/op
// means that the loop ran 10000000 times at a speed of 282 ns per loop.
//
// If a benchmark needs some expensive setup before running, the timer
// may be stopped:
//     func BenchmarkBigLen(b *testing.B) {
//         b.StopTimer()
//         big := NewBig()
//         b.StartTimer()
//         for i := 0; i < b.N; i++ {
//             big.Len()
//         }
//     }
package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"
)

var (
	// The short flag requests that tests run more quickly, but its functionality
	// is provided by test writers themselves.  The testing package is just its
	// home.  The all.bash installation script sets it to make installation more
	// efficient, but by default the flag is off so a plain "gotest" will do a
	// full test of the package.
	short = flag.Bool("test.short", false, "run smaller test suite to save time")

	// Report as tests are run; default is silent for success.
	chatty         = flag.Bool("test.v", false, "verbose: print additional output")
	match          = flag.String("test.run", "", "regular expression to select tests to run")
	memProfile     = flag.String("test.memprofile", "", "write a memory profile to the named file after execution")
	memProfileRate = flag.Int("test.memprofilerate", 0, "if >=0, sets runtime.MemProfileRate")
	cpuProfile     = flag.String("test.cpuprofile", "", "write a cpu profile to the named file during execution")
	timeout        = flag.Int64("test.timeout", 0, "if > 0, sets time limit for tests in seconds")
	cpuListStr     = flag.String("test.cpu", "", "comma-separated list of number of CPUs to use for each test")
	parallel       = flag.Int("test.parallel", runtime.GOMAXPROCS(0), "maximum test parallelism")

	cpuList []int
)

// common holds the elements common between T and B and
// captures common methods such as Errorf.
type common struct {
	output   []byte    // Output generated by test or benchmark.
	failed   bool      // Test or benchmark has failed.
	start    time.Time // Time test or benchmark started
	duration time.Duration
	self     interface{}      // To be sent on signal channel when done.
	signal   chan interface{} // Output for serial tests.
}

// Short reports whether the -test.short flag is set.
func Short() bool {
	return *short
}

// decorate inserts the final newline if needed and indentation tabs for formatting.
// If addFileLine is true, it also prefixes the string with the file and line of the call site.
func decorate(s string, addFileLine bool) string {
	if addFileLine {
		_, file, line, ok := runtime.Caller(3) // decorate + log + public function.
		if ok {
			// Truncate file name at last file name separator.
			if index := strings.LastIndex(file, "/"); index >= 0 {
				file = file[index+1:]
			} else if index = strings.LastIndex(file, "\\"); index >= 0 {
				file = file[index+1:]
			}
		} else {
			file = "???"
			line = 1
		}
		s = fmt.Sprintf("%s:%d: %s", file, line, s)
	}
	s = "\t" + s // Every line is indented at least one tab.
	n := len(s)
	if n > 0 && s[n-1] != '\n' {
		s += "\n"
		n++
	}
	for i := 0; i < n-1; i++ { // -1 to avoid final newline
		if s[i] == '\n' {
			// Second and subsequent lines are indented an extra tab.
			return s[0:i+1] + "\t" + decorate(s[i+1:n], false)
		}
	}
	return s
}

// T is a type passed to Test functions to manage test state and support formatted test logs.
// Logs are accumulated during execution and dumped to standard error when done.
type T struct {
	common
	name          string    // Name of test.
	startParallel chan bool // Parallel tests will wait on this.
}

// Fail marks the function as having failed but continues execution.
func (c *common) Fail() { c.failed = true }

// Failed returns whether the function has failed.
func (c *common) Failed() bool { return c.failed }

// FailNow marks the function as having failed and stops its execution.
// Execution will continue at the next Test.
func (c *common) FailNow() {
	c.duration = time.Now().Sub(c.start)
	c.Fail()
	c.signal <- c.self
	runtime.Goexit()
}

// log generates the output. It's always at the same stack depth.
func (c *common) log(s string) {
	c.output = append(c.output, decorate(s, true)...)
}

// Log formats its arguments using default formatting, analogous to Println(),
// and records the text in the error log.
func (c *common) Log(args ...interface{}) { c.log(fmt.Sprintln(args...)) }

// Logf formats its arguments according to the format, analogous to Printf(),
// and records the text in the error log.
func (c *common) Logf(format string, args ...interface{}) { c.log(fmt.Sprintf(format, args...)) }

// Error is equivalent to Log() followed by Fail().
func (c *common) Error(args ...interface{}) {
	c.log(fmt.Sprintln(args...))
	c.Fail()
}

// Errorf is equivalent to Logf() followed by Fail().
func (c *common) Errorf(format string, args ...interface{}) {
	c.log(fmt.Sprintf(format, args...))
	c.Fail()
}

// Fatal is equivalent to Log() followed by FailNow().
func (c *common) Fatal(args ...interface{}) {
	c.log(fmt.Sprintln(args...))
	c.FailNow()
}

// Fatalf is equivalent to Logf() followed by FailNow().
func (c *common) Fatalf(format string, args ...interface{}) {
	c.log(fmt.Sprintf(format, args...))
	c.FailNow()
}

// Parallel signals that this test is to be run in parallel with (and only with)
// other parallel tests in this CPU group.
func (t *T) Parallel() {
	t.signal <- (*T)(nil) // Release main testing loop
	<-t.startParallel     // Wait for serial tests to finish
}

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalTest struct {
	Name string
	F    func(*T)
}

func tRunner(t *T, test *InternalTest) {
	t.start = time.Now()
	test.F(t)
	t.duration = time.Now().Sub(t.start)
	t.signal <- t
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func Main(matchString func(pat, str string) (bool, error), tests []InternalTest, benchmarks []InternalBenchmark, examples []InternalExample) {
	flag.Parse()
	parseCpuList()

	before()
	startAlarm()
	testOk := RunTests(matchString, tests)
	exampleOk := RunExamples(examples)
	if !testOk || !exampleOk {
		fmt.Println("FAIL")
		os.Exit(1)
	}
	fmt.Println("PASS")
	stopAlarm()
	RunBenchmarks(matchString, benchmarks)
	after()
}

func (t *T) report() {
	tstr := fmt.Sprintf("(%.2f seconds)", t.duration.Seconds())
	format := "--- %s: %s %s\n%s"
	if t.failed {
		fmt.Printf(format, "FAIL", t.name, tstr, t.output)
	} else if *chatty {
		fmt.Printf(format, "PASS", t.name, tstr, t.output)
	}
}

func RunTests(matchString func(pat, str string) (bool, error), tests []InternalTest) (ok bool) {
	ok = true
	if len(tests) == 0 {
		fmt.Fprintln(os.Stderr, "testing: warning: no tests to run")
		return
	}
	for _, procs := range cpuList {
		runtime.GOMAXPROCS(procs)
		// We build a new channel tree for each run of the loop.
		// collector merges in one channel all the upstream signals from parallel tests.
		// If all tests pump to the same channel, a bug can occur where a goroutine
		// kicks off a test, fails, and still delivers a completion signal, which skews the
		// counting.
		var collector = make(chan interface{})

		numParallel := 0
		startParallel := make(chan bool)

		for i := 0; i < len(tests); i++ {
			matched, err := matchString(*match, tests[i].Name)
			if err != nil {
				fmt.Fprintf(os.Stderr, "testing: invalid regexp for -test.run: %s\n", err)
				os.Exit(1)
			}
			if !matched {
				continue
			}
			testName := tests[i].Name
			if procs != 1 {
				testName = fmt.Sprintf("%s-%d", tests[i].Name, procs)
			}
			t := &T{
				common: common{
					signal: make(chan interface{}),
				},
				name:          testName,
				startParallel: startParallel,
			}
			t.self = t
			if *chatty {
				fmt.Printf("=== RUN %s\n", t.name)
			}
			go tRunner(t, &tests[i])
			out := (<-t.signal).(*T)
			if out == nil { // Parallel run.
				go func() {
					collector <- <-t.signal
				}()
				numParallel++
				continue
			}
			t.report()
			ok = ok && !out.failed
		}

		running := 0
		for numParallel+running > 0 {
			if running < *parallel && numParallel > 0 {
				startParallel <- true
				running++
				numParallel--
				continue
			}
			t := (<-collector).(*T)
			t.report()
			ok = ok && !t.failed
			running--
		}
	}
	return
}

// before runs before all testing.
func before() {
	if *memProfileRate > 0 {
		runtime.MemProfileRate = *memProfileRate
	}
	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "testing: %s", err)
			return
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "testing: can't start cpu profile: %s", err)
			f.Close()
			return
		}
		// Could save f so after can call f.Close; not worth the effort.
	}

}

// after runs after all testing.
func after() {
	if *cpuProfile != "" {
		pprof.StopCPUProfile() // flushes profile to disk
	}
	if *memProfile != "" {
		f, err := os.Create(*memProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "testing: %s", err)
			return
		}
		if err = pprof.WriteHeapProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "testing: can't write %s: %s", *memProfile, err)
		}
		f.Close()
	}
}

var timer *time.Timer

// startAlarm starts an alarm if requested.
func startAlarm() {
	if *timeout > 0 {
		timer = time.AfterFunc(time.Duration(*timeout)*time.Second, alarm)
	}
}

// stopAlarm turns off the alarm.
func stopAlarm() {
	if *timeout > 0 {
		timer.Stop()
	}
}

// alarm is called if the timeout expires.
func alarm() {
	panic("test timed out")
}

func parseCpuList() {
	if len(*cpuListStr) == 0 {
		cpuList = append(cpuList, runtime.GOMAXPROCS(-1))
	} else {
		for _, val := range strings.Split(*cpuListStr, ",") {
			cpu, err := strconv.Atoi(val)
			if err != nil || cpu <= 0 {
				fmt.Fprintf(os.Stderr, "testing: invalid value %q for -test.cpu", val)
				os.Exit(1)
			}
			cpuList = append(cpuList, cpu)
		}
	}
}
<commit_msg>testing: fix wording in explanation of potential bug TBR=rsc<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package testing provides support for automated testing of Go packages.
// It is intended to be used in concert with the ``gotest'' utility, which automates
// execution of any function of the form
//     func TestXxx(*testing.T)
// where Xxx can be any alphanumeric string (but the first letter must not be in
// [a-z]) and serves to identify the test routine.
// These TestXxx routines should be declared within the package they are testing.
//
// Functions of the form
//     func BenchmarkXxx(*testing.B)
// are considered benchmarks, and are executed by gotest when the -test.bench
// flag is provided.
//
// A sample benchmark function looks like this:
//     func BenchmarkHello(b *testing.B) {
//         for i := 0; i < b.N; i++ {
//             fmt.Sprintf("hello")
//         }
//     }
// The benchmark package will vary b.N until the benchmark function lasts
// long enough to be timed reliably.  The output
//     testing.BenchmarkHello    10000000    282 ns/op
// means that the loop ran 10000000 times at a speed of 282 ns per loop.
//
// If a benchmark needs some expensive setup before running, the timer
// may be stopped:
//     func BenchmarkBigLen(b *testing.B) {
//         b.StopTimer()
//         big := NewBig()
//         b.StartTimer()
//         for i := 0; i < b.N; i++ {
//             big.Len()
//         }
//     }
package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"
)

var (
	// The short flag requests that tests run more quickly, but its functionality
	// is provided by test writers themselves.  The testing package is just its
	// home.  The all.bash installation script sets it to make installation more
	// efficient, but by default the flag is off so a plain "gotest" will do a
	// full test of the package.
	short = flag.Bool("test.short", false, "run smaller test suite to save time")

	// Report as tests are run; default is silent for success.
	chatty         = flag.Bool("test.v", false, "verbose: print additional output")
	match          = flag.String("test.run", "", "regular expression to select tests to run")
	memProfile     = flag.String("test.memprofile", "", "write a memory profile to the named file after execution")
	memProfileRate = flag.Int("test.memprofilerate", 0, "if >=0, sets runtime.MemProfileRate")
	cpuProfile     = flag.String("test.cpuprofile", "", "write a cpu profile to the named file during execution")
	timeout        = flag.Int64("test.timeout", 0, "if > 0, sets time limit for tests in seconds")
	cpuListStr     = flag.String("test.cpu", "", "comma-separated list of number of CPUs to use for each test")
	parallel       = flag.Int("test.parallel", runtime.GOMAXPROCS(0), "maximum test parallelism")

	cpuList []int
)

// common holds the elements common between T and B and
// captures common methods such as Errorf.
type common struct {
	output   []byte    // Output generated by test or benchmark.
	failed   bool      // Test or benchmark has failed.
	start    time.Time // Time test or benchmark started
	duration time.Duration
	self     interface{}      // To be sent on signal channel when done.
	signal   chan interface{} // Output for serial tests.
}

// Short reports whether the -test.short flag is set.
func Short() bool {
	return *short
}

// decorate inserts the final newline if needed and indentation tabs for formatting.
// If addFileLine is true, it also prefixes the string with the file and line of the call site.
func decorate(s string, addFileLine bool) string {
	if addFileLine {
		_, file, line, ok := runtime.Caller(3) // decorate + log + public function.
		if ok {
			// Truncate file name at last file name separator.
			if index := strings.LastIndex(file, "/"); index >= 0 {
				file = file[index+1:]
			} else if index = strings.LastIndex(file, "\\"); index >= 0 {
				file = file[index+1:]
			}
		} else {
			file = "???"
			line = 1
		}
		s = fmt.Sprintf("%s:%d: %s", file, line, s)
	}
	s = "\t" + s // Every line is indented at least one tab.
	n := len(s)
	if n > 0 && s[n-1] != '\n' {
		s += "\n"
		n++
	}
	for i := 0; i < n-1; i++ { // -1 to avoid final newline
		if s[i] == '\n' {
			// Second and subsequent lines are indented an extra tab.
			return s[0:i+1] + "\t" + decorate(s[i+1:n], false)
		}
	}
	return s
}

// T is a type passed to Test functions to manage test state and support formatted test logs.
// Logs are accumulated during execution and dumped to standard error when done.
type T struct {
	common
	name          string    // Name of test.
	startParallel chan bool // Parallel tests will wait on this.
}

// Fail marks the function as having failed but continues execution.
func (c *common) Fail() { c.failed = true }

// Failed returns whether the function has failed.
func (c *common) Failed() bool { return c.failed }

// FailNow marks the function as having failed and stops its execution.
// Execution will continue at the next Test.
func (c *common) FailNow() {
	c.duration = time.Now().Sub(c.start)
	c.Fail()
	c.signal <- c.self
	runtime.Goexit()
}

// log generates the output. It's always at the same stack depth.
func (c *common) log(s string) {
	c.output = append(c.output, decorate(s, true)...)
}

// Log formats its arguments using default formatting, analogous to Println(),
// and records the text in the error log.
func (c *common) Log(args ...interface{}) { c.log(fmt.Sprintln(args...)) }

// Logf formats its arguments according to the format, analogous to Printf(),
// and records the text in the error log.
func (c *common) Logf(format string, args ...interface{}) { c.log(fmt.Sprintf(format, args...)) }

// Error is equivalent to Log() followed by Fail().
func (c *common) Error(args ...interface{}) {
	c.log(fmt.Sprintln(args...))
	c.Fail()
}

// Errorf is equivalent to Logf() followed by Fail().
func (c *common) Errorf(format string, args ...interface{}) {
	c.log(fmt.Sprintf(format, args...))
	c.Fail()
}

// Fatal is equivalent to Log() followed by FailNow().
func (c *common) Fatal(args ...interface{}) {
	c.log(fmt.Sprintln(args...))
	c.FailNow()
}

// Fatalf is equivalent to Logf() followed by FailNow().
func (c *common) Fatalf(format string, args ...interface{}) {
	c.log(fmt.Sprintf(format, args...))
	c.FailNow()
}

// Parallel signals that this test is to be run in parallel with (and only with)
// other parallel tests in this CPU group.
func (t *T) Parallel() {
	t.signal <- (*T)(nil) // Release main testing loop
	<-t.startParallel     // Wait for serial tests to finish
}

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalTest struct {
	Name string
	F    func(*T)
}

func tRunner(t *T, test *InternalTest) {
	t.start = time.Now()
	test.F(t)
	t.duration = time.Now().Sub(t.start)
	t.signal <- t
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func Main(matchString func(pat, str string) (bool, error), tests []InternalTest, benchmarks []InternalBenchmark, examples []InternalExample) {
	flag.Parse()
	parseCpuList()

	before()
	startAlarm()
	testOk := RunTests(matchString, tests)
	exampleOk := RunExamples(examples)
	if !testOk || !exampleOk {
		fmt.Println("FAIL")
		os.Exit(1)
	}
	fmt.Println("PASS")
	stopAlarm()
	RunBenchmarks(matchString, benchmarks)
	after()
}

func (t *T) report() {
	tstr := fmt.Sprintf("(%.2f seconds)", t.duration.Seconds())
	format := "--- %s: %s %s\n%s"
	if t.failed {
		fmt.Printf(format, "FAIL", t.name, tstr, t.output)
	} else if *chatty {
		fmt.Printf(format, "PASS", t.name, tstr, t.output)
	}
}

func RunTests(matchString func(pat, str string) (bool, error), tests []InternalTest) (ok bool) {
	ok = true
	if len(tests) == 0 {
		fmt.Fprintln(os.Stderr, "testing: warning: no tests to run")
		return
	}
	for _, procs := range cpuList {
		runtime.GOMAXPROCS(procs)
		// We build a new channel tree for each run of the loop.
		// collector merges in one channel all the upstream signals from parallel tests.
		// If all tests pump to the same channel, a bug can occur where a test
		// kicks off a gouroutine that Fails, yet the test still delivers a completion signal,
		// which skews the counting.
		var collector = make(chan interface{})

		numParallel := 0
		startParallel := make(chan bool)

		for i := 0; i < len(tests); i++ {
			matched, err := matchString(*match, tests[i].Name)
			if err != nil {
				fmt.Fprintf(os.Stderr, "testing: invalid regexp for -test.run: %s\n", err)
				os.Exit(1)
			}
			if !matched {
				continue
			}
			testName := tests[i].Name
			if procs != 1 {
				testName = fmt.Sprintf("%s-%d", tests[i].Name, procs)
			}
			t := &T{
				common: common{
					signal: make(chan interface{}),
				},
				name:          testName,
				startParallel: startParallel,
			}
			t.self = t
			if *chatty {
				fmt.Printf("=== RUN %s\n", t.name)
			}
			go tRunner(t, &tests[i])
			out := (<-t.signal).(*T)
			if out == nil { // Parallel run.
				go func() {
					collector <- <-t.signal
				}()
				numParallel++
				continue
			}
			t.report()
			ok = ok && !out.failed
		}

		running := 0
		for numParallel+running > 0 {
			if running < *parallel && numParallel > 0 {
				startParallel <- true
				running++
				numParallel--
				continue
			}
			t := (<-collector).(*T)
			t.report()
			ok = ok && !t.failed
			running--
		}
	}
	return
}

// before runs before all testing.
func before() {
	if *memProfileRate > 0 {
		runtime.MemProfileRate = *memProfileRate
	}
	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "testing: %s", err)
			return
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "testing: can't start cpu profile: %s", err)
			f.Close()
			return
		}
		// Could save f so after can call f.Close; not worth the effort.
	}

}

// after runs after all testing.
func after() {
	if *cpuProfile != "" {
		pprof.StopCPUProfile() // flushes profile to disk
	}
	if *memProfile != "" {
		f, err := os.Create(*memProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "testing: %s", err)
			return
		}
		if err = pprof.WriteHeapProfile(f); err != nil {
			fmt.Fprintf(os.Stderr, "testing: can't write %s: %s", *memProfile, err)
		}
		f.Close()
	}
}

var timer *time.Timer

// startAlarm starts an alarm if requested.
func startAlarm() {
	if *timeout > 0 {
		timer = time.AfterFunc(time.Duration(*timeout)*time.Second, alarm)
	}
}

// stopAlarm turns off the alarm.
func stopAlarm() {
	if *timeout > 0 {
		timer.Stop()
	}
}

// alarm is called if the timeout expires.
func alarm() {
	panic("test timed out")
}

func parseCpuList() {
	if len(*cpuListStr) == 0 {
		cpuList = append(cpuList, runtime.GOMAXPROCS(-1))
	} else {
		for _, val := range strings.Split(*cpuListStr, ",") {
			cpu, err := strconv.Atoi(val)
			if err != nil || cpu <= 0 {
				fmt.Fprintf(os.Stderr, "testing: invalid value %q for -test.cpu", val)
				os.Exit(1)
			}
			cpuList = append(cpuList, cpu)
		}
	}
}
<|endoftext|>