<commit_before>package unionfs

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"io/ioutil"
	"log"
	"os"
	"testing"
)

var _ = fmt.Print
var _ = log.Print

const entryTtl = 0.1

var testAOpts = AutoUnionFsOptions{
	UnionFsOptions: testOpts,
	FileSystemOptions: fuse.FileSystemOptions{
		EntryTimeout:    entryTtl,
		AttrTimeout:     entryTtl,
		NegativeTimeout: 0,
	},
}

func WriteFile(name string, contents string) {
	err := ioutil.WriteFile(name, []byte(contents), 0644)
	CheckSuccess(err)
}

func setup(t *testing.T) (workdir string, cleanup func()) {
	wd, _ := ioutil.TempDir("", "")
	err := os.Mkdir(wd+"/mnt", 0700)
	fuse.CheckSuccess(err)

	err = os.Mkdir(wd+"/store", 0700)
	fuse.CheckSuccess(err)

	os.Mkdir(wd+"/ro", 0700)
	fuse.CheckSuccess(err)
	WriteFile(wd+"/ro/file1", "file1")
	WriteFile(wd+"/ro/file2", "file2")

	fs := NewAutoUnionFs(wd+"/store", testAOpts)
	state, conn, err := fuse.MountPathFileSystem(wd+"/mnt", fs, &testAOpts.FileSystemOptions)
	CheckSuccess(err)
	state.Debug = fuse.VerboseTest()
	conn.Debug = fuse.VerboseTest()
	go state.Loop()

	return wd, func() {
		state.Unmount()
		os.RemoveAll(wd)
	}
}

func TestVersion(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	c, err := ioutil.ReadFile(wd + "/mnt/status/gounionfs_version")
	CheckSuccess(err)
	if len(c) == 0 {
		t.Fatal("No version found.")
	}
	log.Println("Found version:", string(c))
}

func TestAutoFsSymlink(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/backing1", 0755)
	CheckSuccess(err)

	err = os.Symlink(wd+"/ro", wd+"/store/backing1/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/backing1", wd+"/mnt/config/manual1")
	CheckSuccess(err)

	fi, err := os.Lstat(wd + "/mnt/manual1/file1")
	CheckSuccess(err)

	entries, err := ioutil.ReadDir(wd + "/mnt")
	CheckSuccess(err)
	if len(entries) != 3 {
		t.Error("readdir mismatch", entries)
	}

	err = os.Remove(wd + "/mnt/config/manual1")
	CheckSuccess(err)

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	fi, _ = os.Lstat(wd + "/mnt/manual1")
	if fi != nil {
		t.Error("Should not have file:", fi)
	}

	_, err = ioutil.ReadDir(wd + "/mnt/config")
	CheckSuccess(err)

	_, err = os.Lstat(wd + "/mnt/backing1/file1")
	CheckSuccess(err)
}

func TestDetectSymlinkedDirectories(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/backing1", 0755)
	CheckSuccess(err)

	err = os.Symlink(wd+"/ro", wd+"/backing1/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/backing1", wd+"/store/backing1")
	CheckSuccess(err)

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	_, err = os.Lstat(wd + "/mnt/backing1")
	CheckSuccess(err)
}

func TestExplicitScan(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/backing1", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/backing1/READONLY")
	CheckSuccess(err)

	fi, _ := os.Lstat(wd + "/mnt/backing1")
	if fi != nil {
		t.Error("Should not have file:", fi)
	}

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	_, err = os.Lstat(scan)
	if err != nil {
		t.Error(".scan_config missing:", err)
	}

	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	_, err = os.Lstat(wd + "/mnt/backing1")
	if err != nil {
		t.Error("Should have workspace backing1:", err)
	}
}

func TestCreationChecks(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/foo", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/foo/READONLY")
	CheckSuccess(err)

	err = os.Mkdir(wd+"/store/ws2", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/ws2/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/foo", wd+"/mnt/config/bar")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/foo", wd+"/mnt/config/foo")
	code := fuse.ToStatus(err)
	if code != fuse.EBUSY {
		t.Error("Should return EBUSY", err)
	}

	err = os.Symlink(wd+"/store/ws2", wd+"/mnt/config/config")
	code = fuse.ToStatus(err)
	if code != fuse.EINVAL {
		t.Error("Should return EINVAL", err)
	}
}
<commit_msg>Make test compile again.<commit_after>package unionfs

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"io/ioutil"
	"log"
	"os"
	"testing"
)

var _ = fmt.Print
var _ = log.Print

var CheckSuccess = fuse.CheckSuccess
const entryTtl = 0.1
var testAOpts = AutoUnionFsOptions{
	UnionFsOptions: testOpts,
	FileSystemOptions: fuse.FileSystemOptions{
		EntryTimeout:    entryTtl,
		AttrTimeout:     entryTtl,
		NegativeTimeout: 0,
	},
}

func WriteFile(name string, contents string) {
	err := ioutil.WriteFile(name, []byte(contents), 0644)
	CheckSuccess(err)
}

func setup(t *testing.T) (workdir string, cleanup func()) {
	wd, _ := ioutil.TempDir("", "")
	err := os.Mkdir(wd+"/mnt", 0700)
	fuse.CheckSuccess(err)

	err = os.Mkdir(wd+"/store", 0700)
	fuse.CheckSuccess(err)

	os.Mkdir(wd+"/ro", 0700)
	fuse.CheckSuccess(err)
	WriteFile(wd+"/ro/file1", "file1")
	WriteFile(wd+"/ro/file2", "file2")

	fs := NewAutoUnionFs(wd+"/store", testAOpts)
	state, conn, err := fuse.MountPathFileSystem(wd+"/mnt", fs, &testAOpts.FileSystemOptions)
	CheckSuccess(err)
	state.Debug = fuse.VerboseTest()
	conn.Debug = fuse.VerboseTest()
	go state.Loop()

	return wd, func() {
		state.Unmount()
		os.RemoveAll(wd)
	}
}

func TestVersion(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	c, err := ioutil.ReadFile(wd + "/mnt/status/gounionfs_version")
	CheckSuccess(err)
	if len(c) == 0 {
		t.Fatal("No version found.")
	}
	log.Println("Found version:", string(c))
}

func TestAutoFsSymlink(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/backing1", 0755)
	CheckSuccess(err)

	err = os.Symlink(wd+"/ro", wd+"/store/backing1/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/backing1", wd+"/mnt/config/manual1")
	CheckSuccess(err)

	fi, err := os.Lstat(wd + "/mnt/manual1/file1")
	CheckSuccess(err)

	entries, err := ioutil.ReadDir(wd + "/mnt")
	CheckSuccess(err)
	if len(entries) != 3 {
		t.Error("readdir mismatch", entries)
	}

	err = os.Remove(wd + "/mnt/config/manual1")
	CheckSuccess(err)

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	fi, _ = os.Lstat(wd + "/mnt/manual1")
	if fi != nil {
		t.Error("Should not have file:", fi)
	}

	_, err = ioutil.ReadDir(wd + "/mnt/config")
	CheckSuccess(err)

	_, err = os.Lstat(wd + "/mnt/backing1/file1")
	CheckSuccess(err)
}

func TestDetectSymlinkedDirectories(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/backing1", 0755)
	CheckSuccess(err)

	err = os.Symlink(wd+"/ro", wd+"/backing1/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/backing1", wd+"/store/backing1")
	CheckSuccess(err)

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	_, err = os.Lstat(wd + "/mnt/backing1")
	CheckSuccess(err)
}

func TestExplicitScan(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/backing1", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/backing1/READONLY")
	CheckSuccess(err)

	fi, _ := os.Lstat(wd + "/mnt/backing1")
	if fi != nil {
		t.Error("Should not have file:", fi)
	}

	scan := wd + "/mnt/config/" + _SCAN_CONFIG
	_, err = os.Lstat(scan)
	if err != nil {
		t.Error(".scan_config missing:", err)
	}

	err = ioutil.WriteFile(scan, []byte("something"), 0644)
	if err != nil {
		t.Error("error writing:", err)
	}

	_, err = os.Lstat(wd + "/mnt/backing1")
	if err != nil {
		t.Error("Should have workspace backing1:", err)
	}
}

func TestCreationChecks(t *testing.T) {
	wd, clean := setup(t)
	defer clean()

	err := os.Mkdir(wd+"/store/foo", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/foo/READONLY")
	CheckSuccess(err)

	err = os.Mkdir(wd+"/store/ws2", 0755)
	CheckSuccess(err)
	os.Symlink(wd+"/ro", wd+"/store/ws2/READONLY")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/foo", wd+"/mnt/config/bar")
	CheckSuccess(err)

	err = os.Symlink(wd+"/store/foo", wd+"/mnt/config/foo")
	code := fuse.ToStatus(err)
	if code != fuse.EBUSY {
		t.Error("Should return EBUSY", err)
	}

	err = os.Symlink(wd+"/store/ws2", wd+"/mnt/config/config")
	code = fuse.ToStatus(err)
	if code != fuse.EINVAL {
		t.Error("Should return EINVAL", err)
	}
}
<|endoftext|>


<commit_before>package unionfs

import (
	"log"
	"sync"
	"time"
)

var _ = log.Println

type cacheEntry struct {
	data interface{}

	// expiryNs is the absolute timestamp of the expiry.
	expiryNs int64
}

// TimedIntCache caches the result of fetch() for some time.  It is
// thread-safe.
type TimedCache struct {
	fetch func(name string) interface{}

	// ttlNs is a duration of the cache.
	ttlNs int64

	cacheMapMutex sync.RWMutex
	cacheMap      map[string]*cacheEntry

	PurgeTimer *time.Timer
}

const layerCacheTimeoutNs = 1e9

// Creates a new cache with the given TTL.  If TTL <= 0, the caching is
// indefinite.
func NewTimedCache(fetcher func(name string) interface{}, ttlNs int64) *TimedCache {
	l := new(TimedCache)
	l.ttlNs = ttlNs
	l.fetch = fetcher
	l.cacheMap = make(map[string]*cacheEntry)
	return l
}

func (me *TimedCache) Get(name string) interface{} {
	me.cacheMapMutex.RLock()
	info, ok := me.cacheMap[name]
	me.cacheMapMutex.RUnlock()

	valid := ok && (me.ttlNs <= 0 || info.expiryNs > time.Nanoseconds())
	if valid {
		return info.data
	}
	return me.GetFresh(name)
}

func (me *TimedCache) Set(name string, val interface{}) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	me.cacheMap[name] = &cacheEntry{
		data:     val,
		expiryNs: time.Nanoseconds() + me.ttlNs,
	}
}

func (me *TimedCache) DropEntry(name string) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	me.cacheMap[name] = nil, false
}

func (me *TimedCache) GetFresh(name string) interface{} {
	data := me.fetch(name)
	me.Set(name, data)
	return data
}

// Drop all expired entries.
func (me *TimedCache) Purge() {
	keys := make([]string, 0, len(me.cacheMap))
	now := time.Nanoseconds()

	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()
	for k, v := range me.cacheMap {
		if v.expiryNs < now {
			keys = append(keys, k)
		}
	}
	for _, k := range keys {
		me.cacheMap[k] = nil, false
	}
}

func (me *TimedCache) RecurringPurge() {
	if me.ttlNs <= 0 {
		return
	}

	me.Purge()
	me.PurgeTimer = time.AfterFunc(me.ttlNs*5,
		func() { me.RecurringPurge() })
}

func (me *TimedCache) DropAll(names []string) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	if names == nil {
		me.cacheMap = make(map[string]*cacheEntry, len(me.cacheMap))
	} else {
		for _, nm := range names {
			me.cacheMap[nm] = nil, false
		}
	}
}
<commit_msg>Add some more notes about TimedCache.<commit_after>package unionfs

import (
	"log"
	"sync"
	"time"
)

var _ = log.Println

type cacheEntry struct {
	data interface{}

	// expiryNs is the absolute timestamp of the expiry.
	expiryNs int64
}

// TimedIntCache caches the result of fetch() for some time.  It is
// thread-safe.  Calls of fetch() do no happen inside a critical
// section, so when multiple concurrent Get()s happen for the same
// key, multiple fetch() calls may be issued for the same key.
type TimedCache struct {
	fetch func(name string) interface{}

	// ttlNs is a duration of the cache.
	ttlNs int64

	cacheMapMutex sync.RWMutex
	cacheMap      map[string]*cacheEntry

	PurgeTimer *time.Timer
}

const layerCacheTimeoutNs = 1e9

// Creates a new cache with the given TTL.  If TTL <= 0, the caching is
// indefinite.
func NewTimedCache(fetcher func(name string) interface{}, ttlNs int64) *TimedCache {
	l := new(TimedCache)
	l.ttlNs = ttlNs
	l.fetch = fetcher
	l.cacheMap = make(map[string]*cacheEntry)
	return l
}

func (me *TimedCache) Get(name string) interface{} {
	me.cacheMapMutex.RLock()
	info, ok := me.cacheMap[name]
	me.cacheMapMutex.RUnlock()

	valid := ok && (me.ttlNs <= 0 || info.expiryNs > time.Nanoseconds())
	if valid {
		return info.data
	}
	return me.GetFresh(name)
}

func (me *TimedCache) Set(name string, val interface{}) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	me.cacheMap[name] = &cacheEntry{
		data:     val,
		expiryNs: time.Nanoseconds() + me.ttlNs,
	}
}

func (me *TimedCache) DropEntry(name string) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	me.cacheMap[name] = nil, false
}

func (me *TimedCache) GetFresh(name string) interface{} {
	data := me.fetch(name)
	me.Set(name, data)
	return data
}

// Drop all expired entries.
func (me *TimedCache) Purge() {
	keys := make([]string, 0, len(me.cacheMap))
	now := time.Nanoseconds()

	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()
	for k, v := range me.cacheMap {
		if v.expiryNs < now {
			keys = append(keys, k)
		}
	}
	for _, k := range keys {
		me.cacheMap[k] = nil, false
	}
}

func (me *TimedCache) RecurringPurge() {
	if me.ttlNs <= 0 {
		return
	}

	me.Purge()
	me.PurgeTimer = time.AfterFunc(me.ttlNs*5,
		func() { me.RecurringPurge() })
}

func (me *TimedCache) DropAll(names []string) {
	me.cacheMapMutex.Lock()
	defer me.cacheMapMutex.Unlock()

	if names == nil {
		me.cacheMap = make(map[string]*cacheEntry, len(me.cacheMap))
	} else {
		for _, nm := range names {
			me.cacheMap[nm] = nil, false
		}
	}
}
<|endoftext|>


<commit_before>package router

import (
	"fmt"
	"io"
	"net/http"

	"golang.org/x/net/websocket"
)

type Result interface {
	SetRequest(*http.Request)
	Execute(http.ResponseWriter)
	String() string
}

type Rendered struct {
	Content io.Reader
	Status  int
}

func (Rendered) SetRequest(*http.Request) {
}

func (r Rendered) Execute(w http.ResponseWriter) {
	io.Copy(w, r.Content)
}

func (r Rendered) String() string {
	return "Rendered Data"
}

type String struct {
	Content string
	Status  int
}

func (String) SetRequest(*http.Request) {
}
func (r String) Execute(w http.ResponseWriter) {
	io.WriteString(w, r.Content)
}

func (r String) String() string {
	return "String Data"
}

func RedirectTo(location string) Result {
	return &Redirect{URL: location}
}

type Redirect struct {
	Request *http.Request
	URL     string
	Status  int
}

func (r *Redirect) SetRequest(req *http.Request) {
	r.Request = req
}
func (r Redirect) Execute(w http.ResponseWriter) {
	if r.Status == 0 {
		r.Status = 303
	}
	http.Redirect(w, r.Request, r.URL, r.Status)
}

func (r Redirect) String() string {
	return fmt.Sprintf("Redirect To %s", r.URL)
}

func Disallow(location string) Result {
	return &NotAllowed{Fallback: location}
}

type NotAllowed struct {
	Request  *http.Request
	Content  io.Reader
	Fallback string
	Status   int
}

func (na *NotAllowed) SetRequest(req *http.Request) {
	na.Request = req
}
func (na NotAllowed) Execute(w http.ResponseWriter) {
	if na.Fallback != "" {
		if na.Status == 0 {
			na.Status = 302
		}
		http.Redirect(w, na.Request, na.Fallback, na.Status)
	} else {
		if na.Status != 0 {
			w.WriteHeader(na.Status)
		}
		io.Copy(w, na.Content)
	}
}

func (na NotAllowed) String() string {
	if na.Fallback != "" {
		return fmt.Sprintf("Not Allowed, Redirecting To %s", na.Fallback)
	} else {
		return fmt.Sprintf("Not Allowed, Rendered a Response")
	}
}

type InternalError struct {
	Error error
}

func (InternalError) SetRequest(*http.Request) {
}
func (ie InternalError) Execute(w http.ResponseWriter) {
	io.WriteString(w, "
Internal Server Error
")
}

func (ie InternalError) String() string {
	return fmt.Sprintf("Error: %s", ie.Error.Error())
}

type NothingResult struct{}

func (NothingResult) SetRequest(*http.Request) {
}
func (NothingResult) Execute(http.ResponseWriter) {
}

func (NothingResult) String() string {
	return "Being Handled Elsewhere"
}

type NotFound struct{}

func (NotFound) SetRequest(*http.Request) {
}
func (NotFound) Execute(http.ResponseWriter) {
}

func (NotFound) String() string {
	return "NotFound/NotApplicable"
}

type WSResult struct {
	Handler websocket.Handler
	request *http.Request
}

func (ws *WSResult) SetRequest(r *http.Request) {
	ws.request = r
}
func (ws WSResult) Execute(w http.ResponseWriter) {
	ws.Handler.ServeHTTP(w, ws.request)
}

func (WSResult) String() string {
	return "Websocket Connection"
}

type UniqueHandler struct {
	Handler http.Handler
	request *http.Request
}

func (uh *UniqueHandler) SetRequest(r *http.Request) {
	uh.request = r
}

func (uh UniqueHandler) Execute(w http.ResponseWriter) {
	go uh.Handler.ServeHTTP(w, uh.request)
}

func (UniqueHandler) String() string {
	return "One-off Handler"
}
<commit_msg>Add JSON Result<commit_after>package router

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

	"golang.org/x/net/websocket"
)

type Result interface {
	SetRequest(*http.Request)
	Execute(http.ResponseWriter)
	String() string
}

type Rendered struct {
	Content io.Reader
	Status  int
}

func (Rendered) SetRequest(*http.Request) {
}

func (r Rendered) Execute(w http.ResponseWriter) {
	if r.Status != 0 {
		w.WriteHeader(r.Status)
	}
	io.Copy(w, r.Content)
}

func (r Rendered) String() string {
	return "Rendered Data"
}

func JSON(data interface{}) Result {
	return JSONData{Data: data}
}

type JSONData struct {
	Data   interface{}
	Status int
}

func (JSONData) SetRequest(*http.Request) {
}

func (r JSONData) Execute(w http.ResponseWriter) {
	if r.Status != 0 {
		w.WriteHeader(r.Status)
	}
	json.NewEncoder(w).Encode(r.Data)
}

func (r JSONData) String() string {
	return "JSON Data"
}

type String struct {
	Content string
	Status  int
}

func (String) SetRequest(*http.Request) {
}
func (r String) Execute(w http.ResponseWriter) {
	io.WriteString(w, r.Content)
}

func (r String) String() string {
	return "String Data"
}

func RedirectTo(location string) Result {
	return &Redirect{URL: location}
}

type Redirect struct {
	Request *http.Request
	URL     string
	Status  int
}

func (r *Redirect) SetRequest(req *http.Request) {
	r.Request = req
}
func (r Redirect) Execute(w http.ResponseWriter) {
	if r.Status == 0 {
		r.Status = 303
	}
	http.Redirect(w, r.Request, r.URL, r.Status)
}

func (r Redirect) String() string {
	return fmt.Sprintf("Redirect To %s", r.URL)
}

func Disallow(location string) Result {
	return &NotAllowed{Fallback: location}
}

type NotAllowed struct {
	Request  *http.Request
	Content  io.Reader
	Fallback string
	Status   int
}

func (na *NotAllowed) SetRequest(req *http.Request) {
	na.Request = req
}
func (na NotAllowed) Execute(w http.ResponseWriter) {
	if na.Fallback != "" {
		if na.Status == 0 {
			na.Status = 302
		}
		http.Redirect(w, na.Request, na.Fallback, na.Status)
	} else {
		if na.Status != 0 {
			w.WriteHeader(na.Status)
		}
		io.Copy(w, na.Content)
	}
}

func (na NotAllowed) String() string {
	if na.Fallback != "" {
		return fmt.Sprintf("Not Allowed, Redirecting To %s", na.Fallback)
	} else {
		return fmt.Sprintf("Not Allowed, Rendered a Response")
	}
}

type InternalError struct {
	Error error
}

func (InternalError) SetRequest(*http.Request) {
}
func (ie InternalError) Execute(w http.ResponseWriter) {
	io.WriteString(w, "
Internal Server Error
")
}

func (ie InternalError) String() string {
	return fmt.Sprintf("Error: %s", ie.Error.Error())
}

type NothingResult struct{}

func (NothingResult) SetRequest(*http.Request) {
}
func (NothingResult) Execute(http.ResponseWriter) {
}

func (NothingResult) String() string {
	return "Being Handled Elsewhere"
}

type NotFound struct{}

func (NotFound) SetRequest(*http.Request) {
}
func (NotFound) Execute(http.ResponseWriter) {
}

func (NotFound) String() string {
	return "NotFound/NotApplicable"
}

type WSResult struct {
	Handler websocket.Handler
	request *http.Request
}

func (ws *WSResult) SetRequest(r *http.Request) {
	ws.request = r
}
func (ws WSResult) Execute(w http.ResponseWriter) {
	ws.Handler.ServeHTTP(w, ws.request)
}

func (WSResult) String() string {
	return "Websocket Connection"
}

type UniqueHandler struct {
	Handler http.Handler
	request *http.Request
}

func (uh *UniqueHandler) SetRequest(r *http.Request) {
	uh.request = r
}

func (uh UniqueHandler) Execute(w http.ResponseWriter) {
	go uh.Handler.ServeHTTP(w, uh.request)
}

func (UniqueHandler) String() string {
	return "One-off Handler"
}
<|endoftext|>


<commit_before>package updatectl

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"encoding/xml"
	"fmt"
	"github.com/coreos/go-omaha/omaha"
	update "github.com/coreos/updatectl/client/update/v1"
	"github.com/deis/deisctl/client"
	"github.com/deis/deisctl/cmd"
	"github.com/deis/deisctl/utils"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"text/tabwriter"
	"time"
)

const (
	initialInterval = time.Second * 10
	maxInterval     = time.Minute * 7
	downloadDir     = "/home/core/deis/systemd/"
)

var (
	instanceFlags struct {
		groupId       string
		appId         string
		start         int64
		end           int64
		verbose       bool
		clientsPerApp int
		minSleep      int
		maxSleep      int
		errorRate     int
		OEM           string
		pingOnly      int
		version       string
	}

	cmdInstance = &Command{
		Name:    "instance",
		Usage:   "[OPTION]...",
		Summary: "Operations to view instances.",
		Subcommands: []*Command{
			cmdInstanceDeis,
		},
	}

	cmdInstanceDeis = &Command{
		Name:        "instance deis",
		Usage:       "[OPTION]...",
		Description: "Simulate single deis to update instances.",
		Run:         instanceDeis,
	}
)

func init() {
	cmdInstanceDeis.Flags.BoolVar(&instanceFlags.verbose, "verbose", false, "Print out the request bodies")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.clientsPerApp, "clients-per-app", 1, "Number of fake fents per appid.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.minSleep, "min-sleep", 5, "Minimum time between update checks.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.maxSleep, "max-sleep", 10, "Maximum time between update checks.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.errorRate, "errorrate", 1, "Chance of error (0-100)%.")
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.OEM, "oem", "deisclient", "oem to report")
	// simulate reboot lock.
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.pingOnly, "ping-only", 0, "halt update and just send ping requests this many times.")
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.appId, "app-id", os.Getenv("DEISCTL_APP_ID"), "Application ID to update.")
	//instanceFlags.appId.required = true
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.groupId, "group-id", os.Getenv("DEISCTL_GROUP_ID"), "Group ID to update.")
	//instanceFlags.groupId.required = true
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.version, "version", os.Getenv("DEISCTL_APP_VERSION"), "Version to report.")
}

//+ downloadDir + "deis.tar.gz"

type serverConfig struct {
	server string
}

type Client struct {
	Id             string
	SessionId      string
	Version        string
	AppId          string
	Track          string
	config         *serverConfig
	errorRate      int
	pingsRemaining int
}

func (c *Client) Log(format string, v ...interface{}) {
	format = c.Id + ": " + format
	fmt.Printf(format, v...)
}

func (c *Client) getCodebaseUrl(uc *omaha.UpdateCheck) string {
	return uc.Urls.Urls[0].CodeBase
}

func (c *Client) updateservice() {
	fmt.Println("starting systemd units")
	files, _ := utils.ListFiles(downloadDir + "*.service")
	fmt.Println(files)
	deis, _ := client.NewClient()
	localServices := deis.GetLocaljobs()
	Services := utils.GetServices()
	if localServices.Len() == 0 {
		fmt.Println("no local services")
		return
	}
	for _, service := range localServices {
		cmd.Uninstall(deis, []string{strings.Split(strings.Split(service, "-")[1], ".")[0]})
		cmd.Install(deis, []string{strings.Split(strings.Split(service, "-")[1], ".")[0]})
	}
	var count int
	for _, service := range Services {
		count = 0
		for _, lserv := range localServices {
			if strings.Contains(lserv, strings.Split(strings.Split(service, "-")[1], ".")[0]) {
				count = count + 1
			}
		}
		if count == 0 {
			go func() {
				_ = cmd.PullImage(service)
			}()
		}
	}

	// pre-install hook (download all new docker images)
	// use systemd to list local deis-* units, ignore -data units
	// cmd.Unistall([]string{"router.3"})
	// cmd.Install([]string{"router.3"})
	// post-install hook (make sure upgrade was successful)
}

func (c *Client) downloadFromUrl(url, fileName string) (err error) {
	url = url + "deis.tar.gz"
	fmt.Printf("Downloading %s to %s", url, fileName)

	// TODO: check file existence first with io.IsExist
	output, err := os.Create(downloadDir + fileName)
	if err != nil {
		fmt.Println("Error while creating", fileName, "-", err)
		return
	}
	defer output.Close()
	response, err := http.Get(url)
	if err != nil {
		fmt.Println("Error while downloading", url, "-", err)
		return
	}
	defer response.Body.Close()

	n, err := io.Copy(output, response.Body)
	if err != nil {
		fmt.Println("Error while downloading", url, "-", err)
		return
	}

	fmt.Println(n, "bytes downloaded.")
	return
}

func (c *Client) OmahaRequest(otype, result string, updateCheck, isPing bool) *omaha.Request {
	req := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := req.AddApp(c.AppId, c.Version)
	app.MachineID = c.Id
	app.BootId = c.SessionId
	app.Track = c.Track
	app.OEM = instanceFlags.OEM

	if updateCheck {
		app.AddUpdateCheck()
	}

	if isPing {
		app.AddPing()
		app.Ping.LastReportDays = "1"
		app.Ping.Status = "1"
	}

	if otype != "" {
		event := app.AddEvent()
		event.Type = otype
		event.Result = result
		if result == "0" {
			event.ErrorCode = "2000"
		} else {
			event.ErrorCode = ""
		}
	}

	return req
}

func (c *Client) MakeRequest(otype, result string, updateCheck, isPing bool) (*omaha.Response, error) {
	client := &http.Client{}
	req := c.OmahaRequest(otype, result, updateCheck, isPing)
	raw, err := xml.MarshalIndent(req, "", " ")
	if err != nil {
		return nil, err
	}

	resp, err := client.Post(c.config.server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	oresp := new(omaha.Response)
	err = xml.NewDecoder(resp.Body).Decode(oresp)
	if err != nil {
		return nil, err
	}

	if instanceFlags.verbose {
		raw, _ := xml.MarshalIndent(req, "", " ")
		c.Log("request: %s\n", string(raw))
		raw, _ = xml.MarshalIndent(oresp, "", " ")
		c.Log("response: %s\n", string(raw))
	}

	return oresp, nil
}

func (c *Client) SetVersion(resp *omaha.Response) {
	// A field can potentially be nil.
	defer func() {
		if err := recover(); err != nil {
			c.Log("%s: error setting version: %v", c.Id, err)
		}
	}()
	uc := resp.Apps[0].UpdateCheck
	url := c.getCodebaseUrl(uc)
	c.MakeRequest("13", "1", false, false)
	c.downloadFromUrl(url, "deis.tar.gz")
	utils.Extract(downloadDir+"deis.tar.gz", downloadDir)
	c.MakeRequest("14", "1", false, false)
	c.updateservice()
	fmt.Println("updated done")
	c.MakeRequest("3", "1", false, false)
	// installed
	fmt.Println("updated done")
	// simulate reboot lock for a while
	for c.pingsRemaining > 0 {
		c.MakeRequest("", "", false, true)
		c.pingsRemaining--
		time.Sleep(1 * time.Second)
	}

	c.Log("updated from %s to %s\n", c.Version, uc.Manifest.Version)

	c.Version = uc.Manifest.Version

	_, err := c.MakeRequest("3", "2", false, false) // Send complete with new version.
	if err != nil {
		log.Println(err)
	}

	c.SessionId = uuid.New()
}

// Sleep between n and m seconds
func (c *Client) Loop(n, m int) {
	for {
		randSleep(n, m)
		resp, err := c.MakeRequest("3", "2", true, false)
		if err != nil {
			log.Println(err)
			continue
		}
		uc := resp.Apps[0].UpdateCheck
		if uc.Status != "ok" {
			c.Log("update check status: %s\n", uc.Status)
		} else {
			c.SetVersion(resp)
		}
	}
}

// Sleeps randomly between n and m seconds.
func randSleep(n, m int) {
	r := m
	if m-n > 0 {
		r = rand.Intn(m-n) + n
	}
	time.Sleep(time.Duration(r) * time.Second)
}

func instanceDeis(args []string, service *update.Service, out *tabwriter.Writer) int {
	if instanceFlags.appId == "" || instanceFlags.groupId == "" {
		return ERROR_USAGE
	}

	conf := &serverConfig{
		server: globalFlags.Server,
	}

	c := &Client{
		Id:             fmt.Sprintf("{update-client-" + utils.NewUuid()),
		SessionId:      uuid.New(),
		Version:        instanceFlags.version,
		AppId:          instanceFlags.appId,
		Track:          instanceFlags.groupId,
		config:         conf,
		errorRate:      instanceFlags.errorRate,
		pingsRemaining: instanceFlags.pingOnly,
	}
	go c.Loop(instanceFlags.minSleep, instanceFlags.maxSleep)

	// run forever
	wait := make(chan bool)
	<-wait
	return OK
}
<commit_msg>use machineID where possible<commit_after>package updatectl

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"encoding/xml"
	"fmt"
	"github.com/coreos/go-omaha/omaha"
	update "github.com/coreos/updatectl/client/update/v1"
	"github.com/deis/deisctl/client"
	"github.com/deis/deisctl/cmd"
	"github.com/deis/deisctl/utils"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"text/tabwriter"
	"time"
)

const (
	initialInterval = time.Second * 10
	maxInterval     = time.Minute * 7
	downloadDir     = "/home/core/deis/systemd/"
)

var (
	instanceFlags struct {
		groupId       string
		appId         string
		start         int64
		end           int64
		verbose       bool
		clientsPerApp int
		minSleep      int
		maxSleep      int
		errorRate     int
		OEM           string
		pingOnly      int
		version       string
	}

	cmdInstance = &Command{
		Name:    "instance",
		Usage:   "[OPTION]...",
		Summary: "Operations to view instances.",
		Subcommands: []*Command{
			cmdInstanceDeis,
		},
	}

	cmdInstanceDeis = &Command{
		Name:        "instance deis",
		Usage:       "[OPTION]...",
		Description: "Simulate single deis to update instances.",
		Run:         instanceDeis,
	}
)

func init() {
	cmdInstanceDeis.Flags.BoolVar(&instanceFlags.verbose, "verbose", false, "Print out the request bodies")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.clientsPerApp, "clients-per-app", 1, "Number of fake fents per appid.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.minSleep, "min-sleep", 5, "Minimum time between update checks.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.maxSleep, "max-sleep", 10, "Maximum time between update checks.")
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.errorRate, "errorrate", 1, "Chance of error (0-100)%.")
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.OEM, "oem", "deisclient", "oem to report")
	// simulate reboot lock.
	cmdInstanceDeis.Flags.IntVar(&instanceFlags.pingOnly, "ping-only", 0, "halt update and just send ping requests this many times.")
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.appId, "app-id", os.Getenv("DEISCTL_APP_ID"), "Application ID to update.")
	//instanceFlags.appId.required = true
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.groupId, "group-id", os.Getenv("DEISCTL_GROUP_ID"), "Group ID to update.")
	//instanceFlags.groupId.required = true
	cmdInstanceDeis.Flags.StringVar(&instanceFlags.version, "version", os.Getenv("DEISCTL_APP_VERSION"), "Version to report.")
}

//+ downloadDir + "deis.tar.gz"

type serverConfig struct {
	server string
}

type Client struct {
	Id             string
	SessionId      string
	Version        string
	AppId          string
	Track          string
	config         *serverConfig
	errorRate      int
	pingsRemaining int
}

func (c *Client) Log(format string, v ...interface{}) {
	format = c.Id + ": " + format
	fmt.Printf(format, v...)
}

func (c *Client) getCodebaseUrl(uc *omaha.UpdateCheck) string {
	return uc.Urls.Urls[0].CodeBase
}

func (c *Client) updateservice() {
	fmt.Println("starting systemd units")
	files, _ := utils.ListFiles(downloadDir + "*.service")
	fmt.Println(files)
	deis, _ := client.NewClient()
	localServices := deis.GetLocaljobs()
	Services := utils.GetServices()
	if localServices.Len() == 0 {
		fmt.Println("no local services")
		return
	}
	for _, service := range localServices {
		cmd.Uninstall(deis, []string{strings.Split(strings.Split(service, "-")[1], ".")[0]})
		cmd.Install(deis, []string{strings.Split(strings.Split(service, "-")[1], ".")[0]})
	}
	var count int
	for _, service := range Services {
		count = 0
		for _, lserv := range localServices {
			if strings.Contains(lserv, strings.Split(strings.Split(service, "-")[1], ".")[0]) {
				count = count + 1
			}
		}
		if count == 0 {
			go func() {
				_ = cmd.PullImage(service)
			}()
		}
	}

	// pre-install hook (download all new docker images)
	// use systemd to list local deis-* units, ignore -data units
	// cmd.Unistall([]string{"router.3"})
	// cmd.Install([]string{"router.3"})
	// post-install hook (make sure upgrade was successful)
}

func (c *Client) downloadFromUrl(url, fileName string) (err error) {
	url = url + "deis.tar.gz"
	fmt.Printf("Downloading %s to %s", url, fileName)

	// TODO: check file existence first with io.IsExist
	output, err := os.Create(downloadDir + fileName)
	if err != nil {
		fmt.Println("Error while creating", fileName, "-", err)
		return
	}
	defer output.Close()
	response, err := http.Get(url)
	if err != nil {
		fmt.Println("Error while downloading", url, "-", err)
		return
	}
	defer response.Body.Close()

	n, err := io.Copy(output, response.Body)
	if err != nil {
		fmt.Println("Error while downloading", url, "-", err)
		return
	}

	fmt.Println(n, "bytes downloaded.")
	return
}

func (c *Client) OmahaRequest(otype, result string, updateCheck, isPing bool) *omaha.Request {
	req := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := req.AddApp(c.AppId, c.Version)
	app.MachineID = c.Id
	app.BootId = c.SessionId
	app.Track = c.Track
	app.OEM = instanceFlags.OEM

	if updateCheck {
		app.AddUpdateCheck()
	}

	if isPing {
		app.AddPing()
		app.Ping.LastReportDays = "1"
		app.Ping.Status = "1"
	}

	if otype != "" {
		event := app.AddEvent()
		event.Type = otype
		event.Result = result
		if result == "0" {
			event.ErrorCode = "2000"
		} else {
			event.ErrorCode = ""
		}
	}

	return req
}

func (c *Client) MakeRequest(otype, result string, updateCheck, isPing bool) (*omaha.Response, error) {
	client := &http.Client{}
	req := c.OmahaRequest(otype, result, updateCheck, isPing)
	raw, err := xml.MarshalIndent(req, "", " ")
	if err != nil {
		return nil, err
	}

	resp, err := client.Post(c.config.server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	oresp := new(omaha.Response)
	err = xml.NewDecoder(resp.Body).Decode(oresp)
	if err != nil {
		return nil, err
	}

	if instanceFlags.verbose {
		raw, _ := xml.MarshalIndent(req, "", " ")
		c.Log("request: %s\n", string(raw))
		raw, _ = xml.MarshalIndent(oresp, "", " ")
		c.Log("response: %s\n", string(raw))
	}

	return oresp, nil
}

func (c *Client) SetVersion(resp *omaha.Response) {
	// A field can potentially be nil.
	defer func() {
		if err := recover(); err != nil {
			c.Log("%s: error setting version: %v", c.Id, err)
		}
	}()
	uc := resp.Apps[0].UpdateCheck
	url := c.getCodebaseUrl(uc)
	c.MakeRequest("13", "1", false, false)
	c.downloadFromUrl(url, "deis.tar.gz")
	utils.Extract(downloadDir+"deis.tar.gz", downloadDir)
	c.MakeRequest("14", "1", false, false)
	c.updateservice()
	fmt.Println("updated done")
	c.MakeRequest("3", "1", false, false)
	// installed
	fmt.Println("updated done")
	// simulate reboot lock for a while
	for c.pingsRemaining > 0 {
		c.MakeRequest("", "", false, true)
		c.pingsRemaining--
		time.Sleep(1 * time.Second)
	}

	c.Log("updated from %s to %s\n", c.Version, uc.Manifest.Version)

	c.Version = uc.Manifest.Version

	_, err := c.MakeRequest("3", "2", false, false) // Send complete with new version.
	if err != nil {
		log.Println(err)
	}

	c.SessionId = uuid.New()
}

// Sleep between n and m seconds
func (c *Client) Loop(n, m int) {
	for {
		randSleep(n, m)
		resp, err := c.MakeRequest("3", "2", true, false)
		if err != nil {
			log.Println(err)
			continue
		}
		uc := resp.Apps[0].UpdateCheck
		if uc.Status != "ok" {
			c.Log("update check status: %s\n", uc.Status)
		} else {
			c.SetVersion(resp)
		}
	}
}

// Sleeps randomly between n and m seconds.
func randSleep(n, m int) {
	r := m
	if m-n > 0 {
		r = rand.Intn(m-n) + n
	}
	time.Sleep(time.Duration(r) * time.Second)
}

// getClientID returns the CoreOS Machine ID or an unknown UUID string
func getClientID() string {
	machineID := utils.GetMachineID("/")
	if machineID == "" {
		return fmt.Sprintf("{unknown-" + utils.NewUuid() + "}")
	}
	return machineID

}
func instanceDeis(args []string, service *update.Service, out *tabwriter.Writer) int {
	if instanceFlags.appId == "" || instanceFlags.groupId == "" {
		return ERROR_USAGE
	}

	conf := &serverConfig{
		server: globalFlags.Server,
	}

	c := &Client{
		Id:             getClientID(),
		SessionId:      uuid.New(),
		Version:        instanceFlags.version,
		AppId:          instanceFlags.appId,
		Track:          instanceFlags.groupId,
		config:         conf,
		errorRate:      instanceFlags.errorRate,
		pingsRemaining: instanceFlags.pingOnly,
	}
	go c.Loop(instanceFlags.minSleep, instanceFlags.maxSleep)

	// run forever
	wait := make(chan bool)
	<-wait
	return OK
}
<|endoftext|>


<commit_before>package authentication

import (
	"crypto/rsa"
	"crypto/x509"
	"crypto/sha1"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"regexp"
	"time"

	"github.com/satori/go.uuid"

	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"

	"github.com/ONSdigital/go-launch-a-survey/settings"
)

// KeyLoadError describes an error that can occur during key loading
type KeyLoadError struct {
	// Op is the operation which caused the error, such as
	// "read", "parse" or "cast".
	Op string

	// Err is a description of the error that occurred during the operation.
	Err string
}

func (e *KeyLoadError) Error() string {
	if e == nil {
		return "<nil>"
	}
	return e.Op + ": " + e.Err
}

func loadEncryptionKey() (*rsa.PublicKey, string, *KeyLoadError) {
	encryptionKeyPath := settings.Get("JWT_ENCRYPTION_KEY_PATH")

	keyData, err := ioutil.ReadFile(encryptionKeyPath)
	if err != nil {
		return nil, "", &KeyLoadError{Op: "read", Err: "Failed to read encryption key from file: " + encryptionKeyPath}
	}

	block, _ := pem.Decode(keyData)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, "", &KeyLoadError{Op: "parse", Err: "Failed to parse encryption key PEM"}
	}

	kid := fmt.Sprintf("%x", sha1.Sum(keyData))

	publicKey, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, "", &KeyLoadError{Op: "cast", Err: "Failed to cast key to rsa.PublicKey"}
	}

	return publicKey, kid, nil
}

func loadSigningKey() (*rsa.PrivateKey, string, *KeyLoadError) {
	signingKeyPath := settings.Get("JWT_SIGNING_KEY_PATH")
	keyData, err := ioutil.ReadFile(signingKeyPath)
	if err != nil {
		return nil, "", &KeyLoadError{Op: "read", Err: "Failed to read signing key from file: " + signingKeyPath}
	}

	block, _ := pem.Decode(keyData)
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, "", &KeyLoadError{Op: "parse", Err: "Failed to parse signing key from PEM"}
	}

	PublicKey, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, "", &KeyLoadError{Op: "marshal", Err: "Failed to marshal public key"}
	}

	pubBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: PublicKey,
	})

	kid := fmt.Sprintf("%x", sha1.Sum(pubBytes))

	return privateKey, kid, nil
}

type eqClaims struct {
	jwt.Claims
	UserID                string `json:"user_id"`
	EqID                  string `json:"eq_id"`
	PeriodID              string `json:"period_id"`
	PeriodStr             string `json:"period_str"`
	CollectionExerciseSid string `json:"collection_exercise_sid"`
	RuRef                 string `json:"ru_ref"`
	RuName                string `json:"ru_name"`
	RefPStartDate         string `json:"ref_p_start_date"` // iso_8601_date
	RefPEndDate           string `json:"ref_p_end_date"`   // iso_8601_date
	FormType              string `json:"form_type"`
	ReturnBy              string `json:"return_by"`
	TradAs                string `json:"trad_as"`
	EmploymentDate        string `json:"employment_date"` // iso_8601_date
	RegionCode            string `json:"region_code"`
	LanguageCode          string `json:"language_code"`
	VariantFlags          variantFlags `json:"variant_flags"`
	Roles                 string `json:"roles"`
	TxID                  string `json:"tx_id"`
}

type variantFlags struct {
	SexualIdentity		string `json:"sexual_identity"`
}

var eqIDFormTypeRegex = regexp.MustCompile(`^(?P<eq_id>[a-z0-9]+)_(?P<form_type>\w+)\.json`)

func extractEqIDFormType(schema string) (EqID, formType string) {
	match := eqIDFormTypeRegex.FindStringSubmatch(schema)
	if match != nil {
		EqID = match[1]
		formType = match[2]
	}
	return
}

func generateClaims(postValues url.Values) (claims eqClaims) {
	issued := time.Now()
	expires := issued.Add(time.Minute * 10) // TODO: Support custom exp: r.PostForm.Get("exp")

	schema := postValues.Get("schema")
	EqID, formType := extractEqIDFormType(schema)

	return eqClaims{
		Claims: jwt.Claims{
			IssuedAt: jwt.NewNumericDate(issued),
			Expiry:   jwt.NewNumericDate(expires),
			ID:       uuid.NewV4().String(),
		},
		EqID:                  EqID,
		FormType:              formType,
		UserID:                postValues.Get("user_id"),
		PeriodID:              postValues.Get("period_id"),
		PeriodStr:             postValues.Get("period_str"),
		CollectionExerciseSid: postValues.Get("collection_exercise_sid"),
		RuRef:          postValues.Get("ru_ref"),
		RuName:         postValues.Get("ru_name"),
		RefPStartDate:  postValues.Get("ref_p_start_date"),
		RefPEndDate:    postValues.Get("ref_p_end_date"),
		ReturnBy:       postValues.Get("return_by"),
		TradAs:         postValues.Get("trad_as"),
		EmploymentDate: postValues.Get("employment_date"),
		RegionCode:     postValues.Get("region_code"),
		LanguageCode:   postValues.Get("language_code"),
		TxID:           uuid.NewV4().String(),
		Roles:		postValues.Get("roles"),
		VariantFlags:	variantFlags{
			SexualIdentity:	postValues.Get("sexual_identity"),
		},
	}
}

// TokenError describes an error that can occur during JWT generation
type TokenError struct {
	// Err is a description of the error that occurred.
	Desc string

	// From is optionally the original error from which this one was caused.
	From error
}

func (e *TokenError) Error() string {
	if e == nil {
		return "<nil>"
	}
	err := e.Desc
	if e.From != nil {
		err += " (" + e.From.Error() + ")"
	}
	return err
}

// ConvertPostToToken coverts a set of POST values into a JWT
func ConvertPostToToken(postValues url.Values) (string, *TokenError) {
	log.Println("POST received...", postValues)

	cl := generateClaims(postValues)

	signingKey, signingKeyKid, keyErr := loadSigningKey()
	if keyErr != nil {
		return "", &TokenError{Desc: "Error loading signing key", From: keyErr}
	}

	encryptionKey, encryptionKeyKid, keyErr := loadEncryptionKey()
	if keyErr != nil {
		return "", &TokenError{Desc: "Error loading encryption key", From: keyErr}
	}

	opts := jose.SignerOptions{}
	opts.WithType("JWT")
	opts.WithHeader("kid", signingKeyKid)

	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: signingKey}, &opts)
	if err != nil {
		return "", &TokenError{Desc: "Error creating JWT signer", From: err}
	}

	encryptor, err := jose.NewEncrypter(
		jose.A256GCM,
		jose.Recipient{Algorithm: jose.RSA_OAEP, Key: encryptionKey, KeyID: encryptionKeyKid},
		(&jose.EncrypterOptions{}).WithType("JWT").WithContentType("JWT"))

	if err != nil {
		return "", &TokenError{Desc: "Error creating JWT signer", From: err}
	}

	token, err := jwt.SignedAndEncrypted(signer, encryptor).Claims(cl).CompactSerialize()

	if err != nil {
		return "", &TokenError{Desc: "Error signing and encrypting JWT", From: err}
	}

	fmt.Printf("Created signed/encrypted JWT: %v", token)
	return token, nil
}
<commit_msg>Return keyResult type<commit_after>package authentication

import (
	"crypto/rsa"
	"crypto/x509"
	"crypto/sha1"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"regexp"
	"time"

	"github.com/satori/go.uuid"

	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"

	"github.com/ONSdigital/go-launch-a-survey/settings"
)

// KeyLoadError describes an error that can occur during key loading
type KeyLoadError struct {
	// Op is the operation which caused the error, such as
	// "read", "parse" or "cast".
	Op string

	// Err is a description of the error that occurred during the operation.
	Err string
}

func (e *KeyLoadError) Error() string {
	if e == nil {
		return "<nil>"
	}
	return e.Op + ": " + e.Err
}

type PublicKeyResult struct {
	key rsa.PublicKey
	kid string
}

type PrivateKeyResult struct {
	key rsa.PrivateKey
	kid string
}

func loadEncryptionKey() (*PublicKeyResult, *KeyLoadError) {
	encryptionKeyPath := settings.Get("JWT_ENCRYPTION_KEY_PATH")

	keyData, err := ioutil.ReadFile(encryptionKeyPath)
	if err != nil {
		return nil, &KeyLoadError{Op: "read", Err: "Failed to read encryption key from file: " + encryptionKeyPath}
	}

	block, _ := pem.Decode(keyData)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, &KeyLoadError{Op: "parse", Err: "Failed to parse encryption key PEM"}
	}

	kid := fmt.Sprintf("%x", sha1.Sum(keyData))

	publicKey, ok := pub.(rsa.PublicKey)
	if !ok {
		return nil, &KeyLoadError{Op: "cast", Err: "Failed to cast key to rsa.PublicKey"}
	}

	return &PublicKeyResult{publicKey, kid}, nil
}

func loadSigningKey() (*PrivateKeyResult, *KeyLoadError) {
	signingKeyPath := settings.Get("JWT_SIGNING_KEY_PATH")
	keyData, err := ioutil.ReadFile(signingKeyPath)
	if err != nil {
		return nil, &KeyLoadError{Op: "read", Err: "Failed to read signing key from file: " + signingKeyPath}
	}

	block, _ := pem.Decode(keyData)
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, &KeyLoadError{Op: "parse", Err: "Failed to parse signing key from PEM"}
	}

	PublicKey, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, &KeyLoadError{Op: "marshal", Err: "Failed to marshal public key"}
	}

	pubBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: PublicKey,
	})

	kid := fmt.Sprintf("%x", sha1.Sum(pubBytes))

	return &PrivateKeyResult{*privateKey, kid}, nil
}

type eqClaims struct {
	jwt.Claims
	UserID                string `json:"user_id"`
	EqID                  string `json:"eq_id"`
	PeriodID              string `json:"period_id"`
	PeriodStr             string `json:"period_str"`
	CollectionExerciseSid string `json:"collection_exercise_sid"`
	RuRef                 string `json:"ru_ref"`
	RuName                string `json:"ru_name"`
	RefPStartDate         string `json:"ref_p_start_date"` // iso_8601_date
	RefPEndDate           string `json:"ref_p_end_date"`   // iso_8601_date
	FormType              string `json:"form_type"`
	ReturnBy              string `json:"return_by"`
	TradAs                string `json:"trad_as"`
	EmploymentDate        string `json:"employment_date"` // iso_8601_date
	RegionCode            string `json:"region_code"`
	LanguageCode          string `json:"language_code"`
	VariantFlags          variantFlags `json:"variant_flags"`
	Roles                 string `json:"roles"`
	TxID                  string `json:"tx_id"`
}

type variantFlags struct {
	SexualIdentity		string `json:"sexual_identity"`
}

var eqIDFormTypeRegex = regexp.MustCompile(`^(?P<eq_id>[a-z0-9]+)_(?P<form_type>\w+)\.json`)

func extractEqIDFormType(schema string) (EqID, formType string) {
	match := eqIDFormTypeRegex.FindStringSubmatch(schema)
	if match != nil {
		EqID = match[1]
		formType = match[2]
	}
	return
}

func generateClaims(postValues url.Values) (claims eqClaims) {
	issued := time.Now()
	expires := issued.Add(time.Minute * 10) // TODO: Support custom exp: r.PostForm.Get("exp")

	schema := postValues.Get("schema")
	EqID, formType := extractEqIDFormType(schema)

	return eqClaims{
		Claims: jwt.Claims{
			IssuedAt: jwt.NewNumericDate(issued),
			Expiry:   jwt.NewNumericDate(expires),
			ID:       uuid.NewV4().String(),
		},
		EqID:                  EqID,
		FormType:              formType,
		UserID:                postValues.Get("user_id"),
		PeriodID:              postValues.Get("period_id"),
		PeriodStr:             postValues.Get("period_str"),
		CollectionExerciseSid: postValues.Get("collection_exercise_sid"),
		RuRef:          postValues.Get("ru_ref"),
		RuName:         postValues.Get("ru_name"),
		RefPStartDate:  postValues.Get("ref_p_start_date"),
		RefPEndDate:    postValues.Get("ref_p_end_date"),
		ReturnBy:       postValues.Get("return_by"),
		TradAs:         postValues.Get("trad_as"),
		EmploymentDate: postValues.Get("employment_date"),
		RegionCode:     postValues.Get("region_code"),
		LanguageCode:   postValues.Get("language_code"),
		TxID:           uuid.NewV4().String(),
		Roles:		postValues.Get("roles"),
		VariantFlags:	variantFlags{
			SexualIdentity:	postValues.Get("sexual_identity"),
		},
	}
}

// TokenError describes an error that can occur during JWT generation
type TokenError struct {
	// Err is a description of the error that occurred.
	Desc string

	// From is optionally the original error from which this one was caused.
	From error
}

func (e *TokenError) Error() string {
	if e == nil {
		return "<nil>"
	}
	err := e.Desc
	if e.From != nil {
		err += " (" + e.From.Error() + ")"
	}
	return err
}

// ConvertPostToToken coverts a set of POST values into a JWT
func ConvertPostToToken(postValues url.Values) (string, *TokenError) {
	log.Println("POST received...", postValues)

	cl := generateClaims(postValues)

	privateKeyResult, keyErr := loadSigningKey()
	if keyErr != nil {
		return "", &TokenError{Desc: "Error loading signing key", From: keyErr}
	}

	publicKeyResult, keyErr := loadEncryptionKey()
	if keyErr != nil {
		return "", &TokenError{Desc: "Error loading encryption key", From: keyErr}
	}

	opts := jose.SignerOptions{}
	opts.WithType("JWT")
	opts.WithHeader("kid", privateKeyResult.kid)

	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: privateKeyResult.key}, &opts)
	if err != nil {
		return "", &TokenError{Desc: "Error creating JWT signer", From: err}
	}

	encryptor, err := jose.NewEncrypter(
		jose.A256GCM,
		jose.Recipient{Algorithm: jose.RSA_OAEP, Key: publicKeyResult.key, KeyID: publicKeyResult.kid},
		(&jose.EncrypterOptions{}).WithType("JWT").WithContentType("JWT"))

	if err != nil {
		return "", &TokenError{Desc: "Error creating JWT signer", From: err}
	}

	token, err := jwt.SignedAndEncrypted(signer, encryptor).Claims(cl).CompactSerialize()

	if err != nil {
		return "", &TokenError{Desc: "Error signing and encrypting JWT", From: err}
	}

	fmt.Printf("Created signed/encrypted JWT: %v", token)
	return token, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Brighcove 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 rq provides a simple queue abstraction that is backed by Redis.
package rq

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

type Queue struct {
	pooledConnection *redis.Pool
	key              string
}

// Connect to the Redis server at the specified address and create a queue
// corresponding to the given key
func QueueConnect(pooledConnection *redis.Pool, key string) *Queue {
	return &Queue{pooledConnection: pooledConnection, key: key}
}

// Push will perform a right-push onto a Redis list/queue with the supplied
// key and value.  An error will be returned if the operation failed.
func (queue *Queue) Push(value string) error {
	c := queue.pooledConnection.Get()
	defer c.Close()

	err := c.Send("LPUSH", queue.key, value)
	if err == nil {
		return c.Flush()
	} else {
		return err
	}
}

// Pop will perform a blocking left-pop from a Redis list/queue with the supplied
// key.  An error will be returned if the operation failed.
func (queue *Queue) Pop(timeout int) (string, error) {
	c := queue.pooledConnection.Get()
	defer c.Close()

	rep, err := redis.Strings(c.Do("BRPOP", queue.key, timeout))
	if err == nil {
		return rep[1], nil
	} else {
		return "", err
	}
}

// Length will return the number of items in the specified list/queue
func (queue *Queue) Length() (int, error) {
	c := queue.pooledConnection.Get()
	defer c.Close()

	rep, err := redis.Int(c.Do("LLEN", queue.key))
	if err == nil {
		return rep, nil
	} else {
		return 0, err
	}
}
<commit_msg>Updated comments<commit_after>// Copyright 2014 Brighcove 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 rq provides a simple queue abstraction that is backed by Redis.
package rq

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

type Queue struct {
	pooledConnection *redis.Pool
	key              string
}

// Connect to the Redis server at the specified address and create a queue
// corresponding to the given key
func QueueConnect(pooledConnection *redis.Pool, key string) *Queue {
	return &Queue{pooledConnection: pooledConnection, key: key}
}

// Push will perform a left-push onto a Redis list/queue with the supplied
// key and value.  An error will be returned if the operation failed.
func (queue *Queue) Push(value string) error {
	c := queue.pooledConnection.Get()
	defer c.Close()

	err := c.Send("LPUSH", queue.key, value)
	if err == nil {
		return c.Flush()
	} else {
		return err
	}
}

// Pop will perform a blocking right-pop from a Redis list/queue with the
// supplied key.  An error will be returned if the operation failed.
func (queue *Queue) Pop(timeout int) (string, error) {
	c := queue.pooledConnection.Get()
	defer c.Close()

	rep, err := redis.Strings(c.Do("BRPOP", queue.key, timeout))
	if err == nil {
		return rep[1], nil
	} else {
		return "", err
	}
}

// Length will return the number of items in the specified list/queue
func (queue *Queue) Length() (int, error) {
	c := queue.pooledConnection.Get()
	defer c.Close()

	rep, err := redis.Int(c.Do("LLEN", queue.key))
	if err == nil {
		return rep, nil
	} else {
		return 0, err
	}
}
<|endoftext|>


<commit_before>package async

import (
	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	"github.com/jbenet/go-ipfs/blocks"
)

// ForwardN forwards up to |num| blocks to the returned channel.
func ForwardN(ctx context.Context, in <-chan *blocks.Block, num int) <-chan *blocks.Block {
	out := make(chan *blocks.Block)
	go func() {
		defer close(out)
		for i := 0; i < num; i++ {
			select {
			case block, ok := <-in:
				if !ok {
					return
				}
				select {
				case out <- block:
				case <-ctx.Done():
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()
	return out
}
<commit_msg>doc<commit_after>package async

import (
	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	"github.com/jbenet/go-ipfs/blocks"
)

// ForwardN forwards up to |num| blocks to the returned channel.
func ForwardN(ctx context.Context, in <-chan *blocks.Block, num int) <-chan *blocks.Block {
	out := make(chan *blocks.Block)
	go func() {
		defer close(out)
		for i := 0; i < num; i++ {
			select {
			case block, ok := <-in:
				if !ok {
					return // otherwise nil value is forwarded to output
				}
				select {
				case out <- block:
				case <-ctx.Done():
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()
	return out
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"strconv"

	"github.com/codegangsta/cli"
)

func (u *SomaUtil) GetCliArgumentCount(c *cli.Context) int {
	a := c.Args()
	if !a.Present() {
		return 0
	}
	return len(a.Tail()) + 1
}

func (u *SomaUtil) ValidateCliArgument(c *cli.Context, pos uint8, s string) {
	a := c.Args()
	if a.Get(int(pos)-1) != s {
		u.Abort(fmt.Sprintf("Syntax error, missing keyword: ", s))
	}
}

func (u *SomaUtil) ValidateCliMinArgumentCount(c *cli.Context, i uint8) {
	ct := u.GetCliArgumentCount(c)
	if ct < int(i) {
		u.Abort("Syntax error, incorrect argument count")
	}
}

func (u *SomaUtil) ValidateCliArgumentCount(c *cli.Context, i uint8) {
	a := c.Args()
	if i == 0 {
		if a.Present() {
			u.Abort("Syntax error, command takes no arguments")
		}
	} else {
		if !a.Present() || len(a.Tail()) != (int(i)-1) {
			u.Abort("Syntax error, incorrect argument count")
		}
	}
}

func (u *SomaUtil) GetFullArgumentSlice(c *cli.Context) []string {
	sl := []string{c.Args().First()}
	sl = append(sl, c.Args().Tail()...)
	return sl
}

func (u *SomaUtil) ParseVariableArguments(keys []string, rKeys []string, args []string) (map[string]string, []string) {
	// return map of the parse result
	result := make(map[string]string)
	// map to test which required keys were found
	argumentCheck := make(map[string]bool)
	// return slice which optional keys were found
	optionalKeys := make([]string, 0)
	// no required keys is valid
	if len(rKeys) > 0 {
		for _, key := range rKeys {
			argumentCheck[key] = false
		}
	}
	skipNext := false

	for pos, val := range args {
		// skip current argument if last argument was a keyword
		if skipNext {
			skipNext = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check back-to-back keywords
			u.CheckStringNotAKeyword(args[pos+1], keys)
			result[val] = args[pos+1]
			argumentCheck[val] = true
			skipNext = true
			if !u.SliceContainsString(val, rKeys) {
				optionalKeys = append(optionalKeys, val)
			}
			continue
		}
		// keywords trigger continue, arguments are skipped over.
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check we managed to collect all required keywords
	for _, v := range argumentCheck {
		if !v {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", v))
		}
	}

	return result, optionalKeys
}

/*
 * This function parses whitespace separated argument lists of
 * keyword/value pairs were keywords can be specified multiple
 * times, some keywords are required and some only allowed once.
 * Sequence of multiple keywords are detected and lead to abort
 *
 * multKeys => [ "port", "transport" ]
 * uniqKeys => [ "team" ]
 * reqKeys  => [ "team" ]
 * args     => [ "port", "53", "transport", "tcp", "transport",
 *               "udp", "team", "ITOMI" ]
 *
 * result => result["team"] = [ "ITOMI" ]
 *           result["port"] = [ "53" ]
 *           result["transport"] = [ "tcp", "udp" ]
 */
func (u *SomaUtil) ParseVariadicArguments(
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) map[string][]string {
	// returns a map of slices of string
	result := make(map[string][]string)

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		// ok is false if slice is nil
		if _, ok := result[key]; !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result
}

func (u *SomaUtil) ParseVariadicCheckArguments(args []string) (
	map[string][]string,
	[]somaproto.CheckConfigurationConstraint,
	[]somaproto.CheckConfigurationThreshold,
) {
	// create return objects
	result := make(map[string][]string)
	constraints := []somaproto.CheckConfigurationConstraint{}
	thresholds := []somaproto.CheckConfigurationThreshold{}
	var err error

	multiple := []string{"threshold", "constraint"}
	unique := []string{"in", "on", "with", "interval", "inheritance", "childrenonly", "extern"}
	required := []string{"in", "on", "with", "interval"}

	constraintTypes := []string{"service", "attribute", "system", "custom", "oncall", "native"}
	// merge key slices
	keys := append(multiple, unique...)

	// iteration helper
	skip := false
	skipcount := 0

argloop:
	for pos, val := range args {
		// skip current arg if it was already consumed
		if skip {
			skipcount--
			if skipcount == 0 {
				skip = false
			}
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			switch val {
			case "threshold":
				if len(args[pos+1:]) < 6 {
					u.Abort("Syntax error, incomplete threshold specification")
				}
				t := u.ParseVariadicArguments(
					[]string{},
					[]string{"predicate", "level", "value"},
					[]string{"predicate", "level", "value"},
					args[pos+1:pos+7])
				thr := somaproto.CheckConfigurationThreshold{}
				thr.Predicate.Predicate = t["predicate"][0]
				thr.Level.Name = t["level"][0]
				if thr.Value, err = strconv.ParseInt(t["value"][0], 10, 64); err != nil {
					u.Abort(fmt.Sprintf("Syntax error, value argument not numeric: %s",
						t["value"][0]))
				}
				thresholds = append(thresholds, thr)
				skip = true
				skipcount = 6
				continue argloop
			case "constraint":
				// argument is the start of a constraint specification.
				// check we have enough arguments left
				if len(args[pos+1:]) < 3 {
					u.Abort("Syntax error, incomplete constraint specification")
				}
				// check constraint type specification
				if !u.SliceContainsString(args[pos+2], constraintTypes) {
					u.Abort(fmt.Sprintf("Syntax error, unknown contraint type: %s",
						args[pos+1]))
				}
				constr := somaproto.CheckConfigurationConstraint{}
				constr.ConstraintType = args[pos+1]
				switch constr.ConstraintType {
				case "service":
					constr.Service = &somaproto.TreePropertyService{}
					switch args[pos+2] {
					case "name":
						constr.Service.Name = args[pos+3]
					default:
						u.Abort(fmt.Sprintf("Syntax error, can not constraint service to %s",
							args[pos+2]))
					}
				case "attribute":
					constr.Attribute = &somaproto.TreeServiceAttribute{
						Attribute: args[pos+2],
						Value:     args[pos+3],
					}
				case "system":
					constr.System = &somaproto.TreePropertySystem{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				case "custom":
					constr.Custom = &somaproto.TreePropertyCustom{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				case "oncall":
					constr.Oncall = &somaproto.TreePropertyOncall{}
					switch args[pos+2] {
					case "id":
						constr.Oncall.OncallId = args[pos+3]
					case "name":
						constr.Oncall.Name = args[pos+3]
					default:
						u.Abort(fmt.Sprintf("Syntax error, can not constraint oncall to %s",
							args[pos+2]))
					}
				case "native":
					constr.Native = &somaproto.TreePropertyNative{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				}
				constraints = append(constraints, constr)
				skip = true
				skipcount = 3
				continue argloop
			case "on":
				result["on/type"] = append(result["on/type"], args[pos+1])
				result["on/object"] = append(result["on/object"], args[pos+2])
				result[val] = append(result[val], fmt.Sprintf("%s::%s", args[pos+1], args[pos+2]))
				skip = true
				skipcount = 2
				continue argloop
			default:
				// regular key/value keyword
				result[val] = append(result[val], args[pos+1])
				skip = true
				skipcount = 1
				continue argloop
			}
		}
		// error is reached if argument was not skipped and not a
		// recognized keyword
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if all required keywords were collected
	for _, key := range required {
		if _, ok := result[key]; !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specuified once
	for _, key := range unique {
		// check ok since unique may still be optional
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result, constraints, thresholds
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Improve cli argument count errors<commit_after>package util

import (
	"fmt"
	"strconv"

	"github.com/codegangsta/cli"
)

func (u *SomaUtil) GetCliArgumentCount(c *cli.Context) int {
	a := c.Args()
	if !a.Present() {
		return 0
	}
	return len(a.Tail()) + 1
}

func (u *SomaUtil) ValidateCliArgument(c *cli.Context, pos uint8, s string) {
	a := c.Args()
	if a.Get(int(pos)-1) != s {
		u.Abort(fmt.Sprintf("Syntax error, missing keyword: ", s))
	}
}

func (u *SomaUtil) ValidateCliMinArgumentCount(c *cli.Context, i uint8) {
	ct := u.GetCliArgumentCount(c)
	if ct < int(i) {
		u.Abort(fmt.Sprintf(
			"Syntax error, incorrect argument count (%d < %d+ expected)",
			ct,
			i,
		))
	}
}

func (u *SomaUtil) ValidateCliArgumentCount(c *cli.Context, i uint8) {
	a := c.Args()
	if i == 0 {
		if a.Present() {
			u.Abort("Syntax error, command takes no arguments")
		}
	} else {
		if !a.Present() || len(a.Tail()) != (int(i)-1) {
			u.Abort(fmt.Sprintf(
				"Syntax error, incorrect argument count (expected: %d, received %d)",
				i,
				len(a.Tail()),
			))
		}
	}
}

func (u *SomaUtil) GetFullArgumentSlice(c *cli.Context) []string {
	sl := []string{c.Args().First()}
	sl = append(sl, c.Args().Tail()...)
	return sl
}

func (u *SomaUtil) ParseVariableArguments(keys []string, rKeys []string, args []string) (map[string]string, []string) {
	// return map of the parse result
	result := make(map[string]string)
	// map to test which required keys were found
	argumentCheck := make(map[string]bool)
	// return slice which optional keys were found
	optionalKeys := make([]string, 0)
	// no required keys is valid
	if len(rKeys) > 0 {
		for _, key := range rKeys {
			argumentCheck[key] = false
		}
	}
	skipNext := false

	for pos, val := range args {
		// skip current argument if last argument was a keyword
		if skipNext {
			skipNext = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check back-to-back keywords
			u.CheckStringNotAKeyword(args[pos+1], keys)
			result[val] = args[pos+1]
			argumentCheck[val] = true
			skipNext = true
			if !u.SliceContainsString(val, rKeys) {
				optionalKeys = append(optionalKeys, val)
			}
			continue
		}
		// keywords trigger continue, arguments are skipped over.
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check we managed to collect all required keywords
	for _, v := range argumentCheck {
		if !v {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", v))
		}
	}

	return result, optionalKeys
}

/*
 * This function parses whitespace separated argument lists of
 * keyword/value pairs were keywords can be specified multiple
 * times, some keywords are required and some only allowed once.
 * Sequence of multiple keywords are detected and lead to abort
 *
 * multKeys => [ "port", "transport" ]
 * uniqKeys => [ "team" ]
 * reqKeys  => [ "team" ]
 * args     => [ "port", "53", "transport", "tcp", "transport",
 *               "udp", "team", "ITOMI" ]
 *
 * result => result["team"] = [ "ITOMI" ]
 *           result["port"] = [ "53" ]
 *           result["transport"] = [ "tcp", "udp" ]
 */
func (u *SomaUtil) ParseVariadicArguments(
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) map[string][]string {
	// returns a map of slices of string
	result := make(map[string][]string)

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		// ok is false if slice is nil
		if _, ok := result[key]; !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result
}

func (u *SomaUtil) ParseVariadicCheckArguments(args []string) (
	map[string][]string,
	[]somaproto.CheckConfigurationConstraint,
	[]somaproto.CheckConfigurationThreshold,
) {
	// create return objects
	result := make(map[string][]string)
	constraints := []somaproto.CheckConfigurationConstraint{}
	thresholds := []somaproto.CheckConfigurationThreshold{}
	var err error

	multiple := []string{"threshold", "constraint"}
	unique := []string{"in", "on", "with", "interval", "inheritance", "childrenonly", "extern"}
	required := []string{"in", "on", "with", "interval"}

	constraintTypes := []string{"service", "attribute", "system", "custom", "oncall", "native"}
	// merge key slices
	keys := append(multiple, unique...)

	// iteration helper
	skip := false
	skipcount := 0

argloop:
	for pos, val := range args {
		// skip current arg if it was already consumed
		if skip {
			skipcount--
			if skipcount == 0 {
				skip = false
			}
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			switch val {
			case "threshold":
				if len(args[pos+1:]) < 6 {
					u.Abort("Syntax error, incomplete threshold specification")
				}
				t := u.ParseVariadicArguments(
					[]string{},
					[]string{"predicate", "level", "value"},
					[]string{"predicate", "level", "value"},
					args[pos+1:pos+7])
				thr := somaproto.CheckConfigurationThreshold{}
				thr.Predicate.Predicate = t["predicate"][0]
				thr.Level.Name = t["level"][0]
				if thr.Value, err = strconv.ParseInt(t["value"][0], 10, 64); err != nil {
					u.Abort(fmt.Sprintf("Syntax error, value argument not numeric: %s",
						t["value"][0]))
				}
				thresholds = append(thresholds, thr)
				skip = true
				skipcount = 6
				continue argloop
			case "constraint":
				// argument is the start of a constraint specification.
				// check we have enough arguments left
				if len(args[pos+1:]) < 3 {
					u.Abort("Syntax error, incomplete constraint specification")
				}
				// check constraint type specification
				if !u.SliceContainsString(args[pos+2], constraintTypes) {
					u.Abort(fmt.Sprintf("Syntax error, unknown contraint type: %s",
						args[pos+1]))
				}
				constr := somaproto.CheckConfigurationConstraint{}
				constr.ConstraintType = args[pos+1]
				switch constr.ConstraintType {
				case "service":
					constr.Service = &somaproto.TreePropertyService{}
					switch args[pos+2] {
					case "name":
						constr.Service.Name = args[pos+3]
					default:
						u.Abort(fmt.Sprintf("Syntax error, can not constraint service to %s",
							args[pos+2]))
					}
				case "attribute":
					constr.Attribute = &somaproto.TreeServiceAttribute{
						Attribute: args[pos+2],
						Value:     args[pos+3],
					}
				case "system":
					constr.System = &somaproto.TreePropertySystem{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				case "custom":
					constr.Custom = &somaproto.TreePropertyCustom{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				case "oncall":
					constr.Oncall = &somaproto.TreePropertyOncall{}
					switch args[pos+2] {
					case "id":
						constr.Oncall.OncallId = args[pos+3]
					case "name":
						constr.Oncall.Name = args[pos+3]
					default:
						u.Abort(fmt.Sprintf("Syntax error, can not constraint oncall to %s",
							args[pos+2]))
					}
				case "native":
					constr.Native = &somaproto.TreePropertyNative{
						Name:  args[pos+2],
						Value: args[pos+3],
					}
				}
				constraints = append(constraints, constr)
				skip = true
				skipcount = 3
				continue argloop
			case "on":
				result["on/type"] = append(result["on/type"], args[pos+1])
				result["on/object"] = append(result["on/object"], args[pos+2])
				result[val] = append(result[val], fmt.Sprintf("%s::%s", args[pos+1], args[pos+2]))
				skip = true
				skipcount = 2
				continue argloop
			default:
				// regular key/value keyword
				result[val] = append(result[val], args[pos+1])
				skip = true
				skipcount = 1
				continue argloop
			}
		}
		// error is reached if argument was not skipped and not a
		// recognized keyword
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if all required keywords were collected
	for _, key := range required {
		if _, ok := result[key]; !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specuified once
	for _, key := range unique {
		// check ok since unique may still be optional
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result, constraints, thresholds
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package vault

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/hashicorp/vault/helper/logformat"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/physical"
	log "github.com/mgutz/logxi/v1"
)

var (
	clusterTestPausePeriod = 2 * time.Second
)

func TestClusterFetching(t *testing.T) {
	c, _, _ := TestCoreUnsealed(t)

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

	cluster, err := c.Cluster()
	if err != nil {
		t.Fatal(err)
	}
	// Test whether expected values are found
	if cluster == nil || cluster.Name == "" || cluster.ID == "" {
		t.Fatalf("cluster information missing: cluster: %#v", cluster)
	}
}

func TestClusterHAFetching(t *testing.T) {
	logger := logformat.NewVaultLogger(log.LevelTrace)

	redirect := "http://127.0.0.1:8200"

	c, err := NewCore(&CoreConfig{
		Physical:     physical.NewInmemHA(logger),
		HAPhysical:   physical.NewInmemHA(logger),
		RedirectAddr: redirect,
		DisableMlock: true,
	})
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key, _ := TestCoreInit(t, c)
	if _, err := TestCoreUnseal(c, TestKeyCopy(key)); err != nil {
		t.Fatalf("unseal err: %s", err)
	}

	// Verify unsealed
	sealed, err := c.Sealed()
	if err != nil {
		t.Fatalf("err checking seal status: %s", err)
	}
	if sealed {
		t.Fatal("should not be sealed")
	}

	// Wait for core to become active
	TestWaitActive(t, c)

	cluster, err := c.Cluster()
	if err != nil {
		t.Fatal(err)
	}
	// Test whether expected values are found
	if cluster == nil || cluster.Name == "" || cluster.ID == "" {
		t.Fatalf("cluster information missing: cluster:%#v", cluster)
	}
}

func TestCluster_ListenForRequests(t *testing.T) {
	// Make this nicer for tests
	manualStepDownSleepPeriod = 5 * time.Second

	cores := TestCluster(t, []http.Handler{nil, nil, nil}, nil, false)
	for _, core := range cores {
		defer core.CloseListeners()
	}

	root := cores[0].Root

	// Wait for core to become active
	TestWaitActive(t, cores[0].Core)

	// Use this to have a valid config after sealing since ClusterTLSConfig returns nil
	var lastTLSConfig *tls.Config
	checkListenersFunc := func(expectFail bool) {
		tlsConfig, err := cores[0].ClusterTLSConfig()
		if err != nil {
			if err.Error() != ErrSealed.Error() {
				t.Fatal(err)
			}
			tlsConfig = lastTLSConfig
		} else {
			tlsConfig.NextProtos = []string{"h2"}
			lastTLSConfig = tlsConfig
		}

		for _, ln := range cores[0].Listeners {
			tcpAddr, ok := ln.Addr().(*net.TCPAddr)
			if !ok {
				t.Fatal("%s not a TCP port", tcpAddr.String())
			}

			conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", tcpAddr.IP.String(), tcpAddr.Port+1), tlsConfig)
			if err != nil {
				if expectFail {
					t.Logf("testing %s:%d unsuccessful as expected", tcpAddr.IP.String(), tcpAddr.Port+1)
					continue
				}
				t.Fatalf("error: %v\nlisteners are\n%#v\n%#v\n", err, cores[0].Listeners[0], cores[0].Listeners[1])
			}
			if expectFail {
				t.Fatalf("testing %s:%d not unsuccessful as expected", tcpAddr.IP.String(), tcpAddr.Port+1)
			}
			err = conn.Handshake()
			if err != nil {
				t.Fatal(err)
			}
			connState := conn.ConnectionState()
			switch {
			case connState.Version != tls.VersionTLS12:
				t.Fatal("version mismatch")
			case connState.NegotiatedProtocol != "h2" || !connState.NegotiatedProtocolIsMutual:
				t.Fatal("bad protocol negotiation")
			}
			t.Logf("testing %s:%d successful", tcpAddr.IP.String(), tcpAddr.Port+1)
		}
	}

	time.Sleep(clusterTestPausePeriod)
	checkListenersFunc(false)

	err := cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}

	// StepDown doesn't wait during actual preSeal so give time for listeners
	// to close
	time.Sleep(clusterTestPausePeriod)
	checkListenersFunc(true)

	// After this period it should be active again
	time.Sleep(manualStepDownSleepPeriod)
	checkListenersFunc(false)

	err = cores[0].Seal(root)
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	// After sealing it should be inactive again
	checkListenersFunc(true)
}

func TestCluster_ForwardRequests(t *testing.T) {
	// Make this nicer for tests
	manualStepDownSleepPeriod = 5 * time.Second

	testCluster_ForwardRequestsCommon(t, false)
	testCluster_ForwardRequestsCommon(t, true)
	os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "")
}

func testCluster_ForwardRequestsCommon(t *testing.T, rpc bool) {
	if rpc {
		os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "1")
	} else {
		os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "")
	}

	handler1 := http.NewServeMux()
	handler1.HandleFunc("/core1", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(201)
		w.Write([]byte("core1"))
	})
	handler2 := http.NewServeMux()
	handler2.HandleFunc("/core2", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(202)
		w.Write([]byte("core2"))
	})
	handler3 := http.NewServeMux()
	handler3.HandleFunc("/core3", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(203)
		w.Write([]byte("core3"))
	})

	cores := TestCluster(t, []http.Handler{handler1, handler2, handler3}, nil, true)
	for _, core := range cores {
		defer core.CloseListeners()
	}

	root := cores[0].Root

	// Wait for core to become active
	TestWaitActive(t, cores[0].Core)

	// Test forwarding a request. Since we're going directly from core to core
	// with no fallback we know that if it worked, request handling is working
	testCluster_ForwardRequests(t, cores[1], "core1")
	testCluster_ForwardRequests(t, cores[2], "core1")

	//
	// Now we do a bunch of round-robining. The point is to make sure that as
	// nodes come and go, we can always successfully forward to the active
	// node.
	//

	// Ensure active core is cores[1] and test
	err := cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[1].Core)
	testCluster_ForwardRequests(t, cores[0], "core2")
	testCluster_ForwardRequests(t, cores[2], "core2")

	// Ensure active core is cores[2] and test
	err = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[2].Core)
	testCluster_ForwardRequests(t, cores[0], "core3")
	testCluster_ForwardRequests(t, cores[1], "core3")

	// Ensure active core is cores[0] and test
	err = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[0].Core)
	testCluster_ForwardRequests(t, cores[1], "core1")
	testCluster_ForwardRequests(t, cores[2], "core1")

	// Ensure active core is cores[1] and test
	err = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[1].Core)
	testCluster_ForwardRequests(t, cores[0], "core2")
	testCluster_ForwardRequests(t, cores[2], "core2")

	// Ensure active core is cores[2] and test
	err = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[2].Core)
	testCluster_ForwardRequests(t, cores[0], "core3")
	testCluster_ForwardRequests(t, cores[1], "core3")
}

func testCluster_ForwardRequests(t *testing.T, c *TestClusterCore, remoteCoreID string) {
	standby, err := c.Standby()
	if err != nil {
		t.Fatal(err)
	}
	if !standby {
		t.Fatal("expected core to be standby")
	}

	// We need to call Leader as that refreshes the connection info
	isLeader, _, err := c.Leader()
	if err != nil {
		t.Fatal(err)
	}
	if isLeader {
		t.Fatal("core should not be leader")
	}

	bodBuf := bytes.NewReader([]byte(`{ "foo": "bar", "zip": "zap" }`))
	req, err := http.NewRequest("PUT", "https://pushit.real.good:9281/"+remoteCoreID, bodBuf)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("X-Vault-Token", c.Root)

	statusCode, header, respBytes, err := c.ForwardRequest(req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if header == nil {
		t.Fatal("err: expected at least a content-type header")
	}
	if header.Get("Content-Type") != "application/json" {
		t.Fatalf("bad content-type: %s", header.Get("Content-Type"))
	}

	body := string(respBytes)

	if body != remoteCoreID {
		t.Fatalf("expected %s, got %s", remoteCoreID, body)
	}
	switch body {
	case "core1":
		if statusCode != 201 {
			t.Fatal("bad response")
		}
	case "core2":
		if statusCode != 202 {
			t.Fatal("bad response")
		}
	case "core3":
		if statusCode != 203 {
			t.Fatal("bad response")
		}
	}
}
<commit_msg>Fix other clustering tests on OSX<commit_after>package vault

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/hashicorp/vault/helper/logformat"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/physical"
	log "github.com/mgutz/logxi/v1"
)

var (
	clusterTestPausePeriod = 2 * time.Second
)

func TestClusterFetching(t *testing.T) {
	c, _, _ := TestCoreUnsealed(t)

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

	cluster, err := c.Cluster()
	if err != nil {
		t.Fatal(err)
	}
	// Test whether expected values are found
	if cluster == nil || cluster.Name == "" || cluster.ID == "" {
		t.Fatalf("cluster information missing: cluster: %#v", cluster)
	}
}

func TestClusterHAFetching(t *testing.T) {
	logger := logformat.NewVaultLogger(log.LevelTrace)

	redirect := "http://127.0.0.1:8200"

	c, err := NewCore(&CoreConfig{
		Physical:     physical.NewInmemHA(logger),
		HAPhysical:   physical.NewInmemHA(logger),
		RedirectAddr: redirect,
		DisableMlock: true,
	})
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	key, _ := TestCoreInit(t, c)
	if _, err := TestCoreUnseal(c, TestKeyCopy(key)); err != nil {
		t.Fatalf("unseal err: %s", err)
	}

	// Verify unsealed
	sealed, err := c.Sealed()
	if err != nil {
		t.Fatalf("err checking seal status: %s", err)
	}
	if sealed {
		t.Fatal("should not be sealed")
	}

	// Wait for core to become active
	TestWaitActive(t, c)

	cluster, err := c.Cluster()
	if err != nil {
		t.Fatal(err)
	}
	// Test whether expected values are found
	if cluster == nil || cluster.Name == "" || cluster.ID == "" {
		t.Fatalf("cluster information missing: cluster:%#v", cluster)
	}
}

func TestCluster_ListenForRequests(t *testing.T) {
	// Make this nicer for tests
	manualStepDownSleepPeriod = 5 * time.Second

	cores := TestCluster(t, []http.Handler{nil, nil, nil}, nil, false)
	for _, core := range cores {
		defer core.CloseListeners()
	}

	root := cores[0].Root

	// Wait for core to become active
	TestWaitActive(t, cores[0].Core)

	// Use this to have a valid config after sealing since ClusterTLSConfig returns nil
	var lastTLSConfig *tls.Config
	checkListenersFunc := func(expectFail bool) {
		tlsConfig, err := cores[0].ClusterTLSConfig()
		if err != nil {
			if err.Error() != ErrSealed.Error() {
				t.Fatal(err)
			}
			tlsConfig = lastTLSConfig
		} else {
			tlsConfig.NextProtos = []string{"h2"}
			lastTLSConfig = tlsConfig
		}

		for _, ln := range cores[0].Listeners {
			tcpAddr, ok := ln.Addr().(*net.TCPAddr)
			if !ok {
				t.Fatal("%s not a TCP port", tcpAddr.String())
			}

			conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", tcpAddr.IP.String(), tcpAddr.Port+10), tlsConfig)
			if err != nil {
				if expectFail {
					t.Logf("testing %s:%d unsuccessful as expected", tcpAddr.IP.String(), tcpAddr.Port+10)
					continue
				}
				t.Fatalf("error: %v\nlisteners are\n%#v\n%#v\n", err, cores[0].Listeners[0], cores[0].Listeners[1])
			}
			if expectFail {
				t.Fatalf("testing %s:%d not unsuccessful as expected", tcpAddr.IP.String(), tcpAddr.Port+10)
			}
			err = conn.Handshake()
			if err != nil {
				t.Fatal(err)
			}
			connState := conn.ConnectionState()
			switch {
			case connState.Version != tls.VersionTLS12:
				t.Fatal("version mismatch")
			case connState.NegotiatedProtocol != "h2" || !connState.NegotiatedProtocolIsMutual:
				t.Fatal("bad protocol negotiation")
			}
			t.Logf("testing %s:%d successful", tcpAddr.IP.String(), tcpAddr.Port+10)
		}
	}

	time.Sleep(clusterTestPausePeriod)
	checkListenersFunc(false)

	err := cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}

	// StepDown doesn't wait during actual preSeal so give time for listeners
	// to close
	time.Sleep(clusterTestPausePeriod)
	checkListenersFunc(true)

	// After this period it should be active again
	time.Sleep(manualStepDownSleepPeriod)
	checkListenersFunc(false)

	err = cores[0].Seal(root)
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	// After sealing it should be inactive again
	checkListenersFunc(true)
}

func TestCluster_ForwardRequests(t *testing.T) {
	// Make this nicer for tests
	manualStepDownSleepPeriod = 5 * time.Second

	testCluster_ForwardRequestsCommon(t, false)
	testCluster_ForwardRequestsCommon(t, true)
	os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "")
}

func testCluster_ForwardRequestsCommon(t *testing.T, rpc bool) {
	if rpc {
		os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "1")
	} else {
		os.Setenv("VAULT_USE_GRPC_REQUEST_FORWARDING", "")
	}

	handler1 := http.NewServeMux()
	handler1.HandleFunc("/core1", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(201)
		w.Write([]byte("core1"))
	})
	handler2 := http.NewServeMux()
	handler2.HandleFunc("/core2", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(202)
		w.Write([]byte("core2"))
	})
	handler3 := http.NewServeMux()
	handler3.HandleFunc("/core3", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Add("Content-Type", "application/json")
		w.WriteHeader(203)
		w.Write([]byte("core3"))
	})

	cores := TestCluster(t, []http.Handler{handler1, handler2, handler3}, nil, true)
	for _, core := range cores {
		defer core.CloseListeners()
	}

	root := cores[0].Root

	// Wait for core to become active
	TestWaitActive(t, cores[0].Core)

	// Test forwarding a request. Since we're going directly from core to core
	// with no fallback we know that if it worked, request handling is working
	testCluster_ForwardRequests(t, cores[1], "core1")
	testCluster_ForwardRequests(t, cores[2], "core1")

	//
	// Now we do a bunch of round-robining. The point is to make sure that as
	// nodes come and go, we can always successfully forward to the active
	// node.
	//

	// Ensure active core is cores[1] and test
	err := cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[1].Core)
	testCluster_ForwardRequests(t, cores[0], "core2")
	testCluster_ForwardRequests(t, cores[2], "core2")

	// Ensure active core is cores[2] and test
	err = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[2].Core)
	testCluster_ForwardRequests(t, cores[0], "core3")
	testCluster_ForwardRequests(t, cores[1], "core3")

	// Ensure active core is cores[0] and test
	err = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[0].Core)
	testCluster_ForwardRequests(t, cores[1], "core1")
	testCluster_ForwardRequests(t, cores[2], "core1")

	// Ensure active core is cores[1] and test
	err = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[2].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[1].Core)
	testCluster_ForwardRequests(t, cores[0], "core2")
	testCluster_ForwardRequests(t, cores[2], "core2")

	// Ensure active core is cores[2] and test
	err = cores[1].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(clusterTestPausePeriod)
	_ = cores[0].StepDown(&logical.Request{
		Operation:   logical.UpdateOperation,
		Path:        "sys/step-down",
		ClientToken: root,
	})
	time.Sleep(clusterTestPausePeriod)
	TestWaitActive(t, cores[2].Core)
	testCluster_ForwardRequests(t, cores[0], "core3")
	testCluster_ForwardRequests(t, cores[1], "core3")
}

func testCluster_ForwardRequests(t *testing.T, c *TestClusterCore, remoteCoreID string) {
	standby, err := c.Standby()
	if err != nil {
		t.Fatal(err)
	}
	if !standby {
		t.Fatal("expected core to be standby")
	}

	// We need to call Leader as that refreshes the connection info
	isLeader, _, err := c.Leader()
	if err != nil {
		t.Fatal(err)
	}
	if isLeader {
		t.Fatal("core should not be leader")
	}

	bodBuf := bytes.NewReader([]byte(`{ "foo": "bar", "zip": "zap" }`))
	req, err := http.NewRequest("PUT", "https://pushit.real.good:9281/"+remoteCoreID, bodBuf)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("X-Vault-Token", c.Root)

	statusCode, header, respBytes, err := c.ForwardRequest(req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if header == nil {
		t.Fatal("err: expected at least a content-type header")
	}
	if header.Get("Content-Type") != "application/json" {
		t.Fatalf("bad content-type: %s", header.Get("Content-Type"))
	}

	body := string(respBytes)

	if body != remoteCoreID {
		t.Fatalf("expected %s, got %s", remoteCoreID, body)
	}
	switch body {
	case "core1":
		if statusCode != 201 {
			t.Fatal("bad response")
		}
	case "core2":
		if statusCode != 202 {
			t.Fatal("bad response")
		}
	case "core3":
		if statusCode != 203 {
			t.Fatal("bad response")
		}
	}
}
<|endoftext|>


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

// +build race

package safe

import (
	"launchpad.net/gocheck"
	"sync"
)

func (s *S) TestSafeBufferIsThreadSafe(c *gocheck.C) {
	var buf Buffer
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		buf.Write([]byte("something"))
		wg.Done()
	}()
	go func() {
		var p [4]byte
		buf.Read(p[:])
		wg.Done()
	}()
	buf.Reset()
	wg.Wait()
}

func (s *S) TestSafeWriterIsThreadSafe(c *gocheck.C) {
	var buf Buffer
	writer := NewWriter(&buf)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		writer.Write([]byte("something"))
		wg.Done()
	}()
	go func() {
		writer.Write([]byte("otherthing"))
		wg.Done()
	}()
	wg.Wait()
}
<commit_msg>safe: add race test to ensure that Reader is safe<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build race

package safe

import (
	"bytes"
	"launchpad.net/gocheck"
	"sync"
)

func (s *S) TestSafeBufferIsThreadSafe(c *gocheck.C) {
	var buf Buffer
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		buf.Write([]byte("something"))
		wg.Done()
	}()
	go func() {
		var p [4]byte
		buf.Read(p[:])
		wg.Done()
	}()
	buf.Reset()
	wg.Wait()
}

func (s *S) TestSafeWriterIsThreadSafe(c *gocheck.C) {
	var buf bytes.Buffer
	writer := NewWriter(&buf)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		writer.Write([]byte("something"))
		wg.Done()
	}()
	go func() {
		writer.Write([]byte("otherthing"))
		wg.Done()
	}()
	wg.Wait()
}

func (s *S) TestSafeReaderIsThreadSafe(c *gocheck.C) {
	var buf [8]byte
	buffer := bytes.NewBufferString("hello world something")
	reader := NewReader(buffer)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		reader.Read(buf[:])
		wg.Done()
	}()
	go func() {
		reader.Read(buf[:])
		wg.Done()
	}()
	wg.Wait()
}
<|endoftext|>