<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cezarsa/form"
	"github.com/ghodss/yaml"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/event"
)

type EventList struct {
	fs     *gnuflag.FlagSet
	filter eventFilter
}

type eventFilter struct {
	KindName    string
	Target      string
	TargetValue string
	OwnerName   string
	Running     bool
}

func (f *eventFilter) queryString(client *cmd.Client) (url.Values, error) {
	values, err := form.EncodeToValues(f)
	if err != nil {
		return nil, err
	}
	for k, v := range values {
		values.Del(k)
		values[strings.ToLower(k)] = v
	}
	if !f.Running {
		values.Del("running")
	}
	return values, nil
}

func (f *eventFilter) flags(fs *gnuflag.FlagSet) {
	name := "Filter events by kind name"
	fs.StringVar(&f.KindName, "kind", "", name)
	fs.StringVar(&f.KindName, "k", "", name)
	name = "Filter events by target name"
	fs.StringVar(&f.Target, "target", "", name)
	fs.StringVar(&f.Target, "t", "", name)
	name = "Filter events by target value"
	fs.StringVar(&f.TargetValue, "target-value", "", name)
	fs.StringVar(&f.TargetValue, "v", "", name)
	name = "Filter events by owner name"
	fs.StringVar(&f.OwnerName, "owner", "", name)
	fs.StringVar(&f.OwnerName, "o", "", name)
	name = "Shows only currently running events"
	fs.BoolVar(&f.Running, "running", false, name)
	fs.BoolVar(&f.Running, "r", false, name)
}

func (c *EventList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "event-list",
		Usage: "event-list [-k kindName]",
		Desc:  `Lists events possibly filtering them.`,
	}
}

func (c *EventList) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.filter.flags(c.fs)
	}
	return c.fs
}

func (c *EventList) Run(context *cmd.Context, client *cmd.Client) error {
	qs, err := c.filter.queryString(client)
	if err != nil {
		return err
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events?%s", qs.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusNoContent {
		return nil
	}
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evts []event.Event
	err = json.Unmarshal(result, &evts)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(evts, context)
}

var reEmailShort = regexp.MustCompile(`@.*$`)

func (c *EventList) Show(evts []event.Event, context *cmd.Context) error {
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"ID", "Start (duration)", "Success", "Owner", "Kind", "Target"}
	for i := range evts {
		evt := &evts[i]
		if evt.Target.Type == "container" {
			evt.Target.Value = evt.Target.Value[:12]
		}
		fullTarget := fmt.Sprintf("%s: %s", evt.Target.Type, evt.Target.Value)
		startFmt := evt.StartTime.Format(time.RFC822Z)
		owner := reEmailShort.ReplaceAllString(evt.Owner.Name, "@…")
		var ts, success string
		if evt.Running {
			ts = fmt.Sprintf("%s (…)", startFmt)
			success = "…"
		} else {
			ts = fmt.Sprintf("%s (%v)", startFmt, evt.EndTime.Sub(evt.StartTime))
			success = fmt.Sprintf("%v", evt.Error == "")
			if evt.CancelInfo.Canceled {
				success += " ✗"
			}
		}
		row := cmd.Row{evt.UniqueID.Hex(), ts, success, owner, evt.Kind.Name, fullTarget}
		var color string
		if evt.Running {
			color = "yellow"
		} else if evt.CancelInfo.Canceled {
			color = "magenta"
		} else if evt.Error != "" {
			color = "red"
		}
		if color != "" {
			for i, v := range row {
				if v != "" {
					row[i] = cmd.Colorfy(v, color, "", "")
				}
			}
		}
		tbl.AddRow(row)
	}
	fmt.Fprintf(context.Stdout, "%s", tbl.String())
	return nil
}

type EventInfo struct{}

func (c *EventInfo) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-info",
		Usage:   "event-info <event-id>",
		Desc:    `Show detailed information about one single event.`,
		MinArgs: 1,
		MaxArgs: 1,
	}
}

func (c *EventInfo) Run(context *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s", context.Args[0]))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evt event.Event
	err = json.Unmarshal(result, &evt)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(&evt, context)
}

func (c *EventInfo) Show(evt *event.Event, context *cmd.Context) error {
	type item struct {
		label string
		value string
	}
	startFmt := evt.StartTime.Format(time.RFC822Z)
	var endFmt string
	if evt.Running {
		endFmt = fmt.Sprintf("running (%v)", time.Since(evt.StartTime))
	} else {
		endFmt = fmt.Sprintf("%s (%v)", evt.EndTime.Format(time.RFC822Z), evt.EndTime.Sub(evt.StartTime))
	}
	items := []item{
		{"ID", evt.UniqueID.Hex()},
		{"Start", startFmt},
		{"End", endFmt},
		{"Target", fmt.Sprintf("%s(%s)", evt.Target.Type, evt.Target.Value)},
		{"Kind", fmt.Sprintf("%s(%s)", evt.Kind.Type, evt.Kind.Name)},
		{"Owner", fmt.Sprintf("%s(%s)", evt.Owner.Type, evt.Owner.Name)},
	}
	successful := evt.Error == ""
	successfulStr := strconv.FormatBool(successful)
	if successful {
		if evt.Running {
			successfulStr = "…"
		}
		items = append(items, item{"Success", successfulStr})
	} else {
		redError := cmd.Colorfy(fmt.Sprintf("%q", evt.Error), "red", "", "")
		redSuccess := cmd.Colorfy(successfulStr, "red", "", "")
		items = append(items, []item{
			{"Success", redSuccess},
			{"Error", redError},
		}...)
	}
	items = append(items, []item{
		{"Cancelable", strconv.FormatBool(evt.Cancelable)},
		{"Canceled", strconv.FormatBool(evt.CancelInfo.Canceled)},
	}...)
	if evt.CancelInfo.Canceled {
		items = append(items, []item{
			{"  Reason", evt.CancelInfo.Reason},
			{"  By", evt.CancelInfo.Owner},
			{"  At", evt.CancelInfo.AckTime.Format(time.RFC822Z)},
		}...)
	}
	labels := []string{"Start", "End", "Other"}
	for i, fn := range []func(interface{}) error{evt.StartData, evt.EndData, evt.OtherData} {
		var data interface{}
		err := fn(&data)
		if err == nil && data != nil {
			str, err := yaml.Marshal(data)
			if err == nil {
				padded := padLines(string(str), "    ")
				items = append(items, item{fmt.Sprintf("%s Custom Data", labels[i]), "\n" + padded})
			}
		}
	}
	if evt.Log != "" {
		items = append(items, item{"Log", "\n" + padLines(evt.Log, "    ")})
	}
	var maxSz int
	for _, item := range items {
		sz := len(item.label)
		if len(item.value) > 0 && item.value[0] != '\n' && sz > maxSz {
			maxSz = sz
		}
	}
	for _, item := range items {
		count := (maxSz - len(item.label)) + 1
		var pad string
		if count > 0 && len(item.value) > 0 && item.value[0] != '\n' {
			pad = strings.Repeat(" ", count)
		}
		label := cmd.Colorfy(item.label+":", "cyan", "", "")
		fmt.Fprintf(context.Stdout, "%s%s%s\n", label, pad, item.value)
	}
	return nil
}

var rePadLines = regexp.MustCompile(`(?m)^(.+)`)

func padLines(s string, pad string) string {
	return rePadLines.ReplaceAllString(s, pad+`$1`)
}

type EventCancel struct {
	cmd.ConfirmationCommand
}

func (c *EventCancel) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-cancel",
		Usage:   "event-cancel <event-id> <reason> [-y]",
		Desc:    `Cancel a running event.`,
		MinArgs: 2,
	}
}

func (c *EventCancel) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(context, "Are you sure you want to cancel this event?") {
		return nil
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s/cancel", context.Args[0]))
	if err != nil {
		return err
	}
	v := url.Values{}
	v.Set("reason", strings.Join(context.Args[1:], " "))
	request, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintln(context.Stdout, "Event successfully canceled.")
	return nil
}
<commit_msg>client: update event filter serializing<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cezarsa/form"
	"github.com/ghodss/yaml"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/event"
)

type EventList struct {
	fs     *gnuflag.FlagSet
	filter eventFilter
}

type eventFilter struct {
	filter  event.Filter
	running bool
}

func (f *eventFilter) queryString(client *cmd.Client) (url.Values, error) {
	if f.running {
		f.filter.Running = &f.running
	}
	values, err := form.EncodeToValues(f.filter)
	if err != nil {
		return nil, err
	}
	for k, v := range values {
		values.Del(k)
		values[strings.ToLower(k)] = v
	}
	if f.filter.Running == nil {
		values.Del("running")
	}
	return values, nil
}

func (f *eventFilter) flags(fs *gnuflag.FlagSet) {
	name := "Filter events by kind name"
	fs.StringVar(&f.filter.KindName, "kind", "", name)
	fs.StringVar(&f.filter.KindName, "k", "", name)
	name = "Filter events by target type"
	ptr := (*string)(&f.filter.Target.Type)
	fs.StringVar(ptr, "target", "", name)
	fs.StringVar(ptr, "t", "", name)
	name = "Filter events by target value"
	fs.StringVar(&f.filter.Target.Value, "target-value", "", name)
	fs.StringVar(&f.filter.Target.Value, "v", "", name)
	name = "Filter events by owner name"
	fs.StringVar(&f.filter.OwnerName, "owner", "", name)
	fs.StringVar(&f.filter.OwnerName, "o", "", name)
	name = "Shows only currently running events"
	fs.BoolVar(&f.running, "running", false, name)
	fs.BoolVar(&f.running, "r", false, name)
}

func (c *EventList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "event-list",
		Usage: "event-list [-k kindName]",
		Desc:  `Lists events possibly filtering them.`,
	}
}

func (c *EventList) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.filter.flags(c.fs)
	}
	return c.fs
}

func (c *EventList) Run(context *cmd.Context, client *cmd.Client) error {
	qs, err := c.filter.queryString(client)
	if err != nil {
		return err
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events?%s", qs.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusNoContent {
		return nil
	}
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evts []event.Event
	err = json.Unmarshal(result, &evts)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(evts, context)
}

var reEmailShort = regexp.MustCompile(`@.*$`)

func (c *EventList) Show(evts []event.Event, context *cmd.Context) error {
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"ID", "Start (duration)", "Success", "Owner", "Kind", "Target"}
	for i := range evts {
		evt := &evts[i]
		if evt.Target.Type == "container" {
			evt.Target.Value = evt.Target.Value[:12]
		}
		fullTarget := fmt.Sprintf("%s: %s", evt.Target.Type, evt.Target.Value)
		startFmt := evt.StartTime.Format(time.RFC822Z)
		owner := reEmailShort.ReplaceAllString(evt.Owner.Name, "@…")
		var ts, success string
		if evt.Running {
			ts = fmt.Sprintf("%s (…)", startFmt)
			success = "…"
		} else {
			ts = fmt.Sprintf("%s (%v)", startFmt, evt.EndTime.Sub(evt.StartTime))
			success = fmt.Sprintf("%v", evt.Error == "")
			if evt.CancelInfo.Canceled {
				success += " ✗"
			}
		}
		row := cmd.Row{evt.UniqueID.Hex(), ts, success, owner, evt.Kind.Name, fullTarget}
		var color string
		if evt.Running {
			color = "yellow"
		} else if evt.CancelInfo.Canceled {
			color = "magenta"
		} else if evt.Error != "" {
			color = "red"
		}
		if color != "" {
			for i, v := range row {
				if v != "" {
					row[i] = cmd.Colorfy(v, color, "", "")
				}
			}
		}
		tbl.AddRow(row)
	}
	fmt.Fprintf(context.Stdout, "%s", tbl.String())
	return nil
}

type EventInfo struct{}

func (c *EventInfo) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-info",
		Usage:   "event-info <event-id>",
		Desc:    `Show detailed information about one single event.`,
		MinArgs: 1,
		MaxArgs: 1,
	}
}

func (c *EventInfo) Run(context *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s", context.Args[0]))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	result, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	var evt event.Event
	err = json.Unmarshal(result, &evt)
	if err != nil {
		return fmt.Errorf("unable to unmarshal %q: %s", string(result), err)
	}
	return c.Show(&evt, context)
}

func (c *EventInfo) Show(evt *event.Event, context *cmd.Context) error {
	type item struct {
		label string
		value string
	}
	startFmt := evt.StartTime.Format(time.RFC822Z)
	var endFmt string
	if evt.Running {
		endFmt = fmt.Sprintf("running (%v)", time.Since(evt.StartTime))
	} else {
		endFmt = fmt.Sprintf("%s (%v)", evt.EndTime.Format(time.RFC822Z), evt.EndTime.Sub(evt.StartTime))
	}
	items := []item{
		{"ID", evt.UniqueID.Hex()},
		{"Start", startFmt},
		{"End", endFmt},
		{"Target", fmt.Sprintf("%s(%s)", evt.Target.Type, evt.Target.Value)},
		{"Kind", fmt.Sprintf("%s(%s)", evt.Kind.Type, evt.Kind.Name)},
		{"Owner", fmt.Sprintf("%s(%s)", evt.Owner.Type, evt.Owner.Name)},
	}
	successful := evt.Error == ""
	successfulStr := strconv.FormatBool(successful)
	if successful {
		if evt.Running {
			successfulStr = "…"
		}
		items = append(items, item{"Success", successfulStr})
	} else {
		redError := cmd.Colorfy(fmt.Sprintf("%q", evt.Error), "red", "", "")
		redSuccess := cmd.Colorfy(successfulStr, "red", "", "")
		items = append(items, []item{
			{"Success", redSuccess},
			{"Error", redError},
		}...)
	}
	items = append(items, []item{
		{"Cancelable", strconv.FormatBool(evt.Cancelable)},
		{"Canceled", strconv.FormatBool(evt.CancelInfo.Canceled)},
	}...)
	if evt.CancelInfo.Canceled {
		items = append(items, []item{
			{"  Reason", evt.CancelInfo.Reason},
			{"  By", evt.CancelInfo.Owner},
			{"  At", evt.CancelInfo.AckTime.Format(time.RFC822Z)},
		}...)
	}
	labels := []string{"Start", "End", "Other"}
	for i, fn := range []func(interface{}) error{evt.StartData, evt.EndData, evt.OtherData} {
		var data interface{}
		err := fn(&data)
		if err == nil && data != nil {
			str, err := yaml.Marshal(data)
			if err == nil {
				padded := padLines(string(str), "    ")
				items = append(items, item{fmt.Sprintf("%s Custom Data", labels[i]), "\n" + padded})
			}
		}
	}
	if evt.Log != "" {
		items = append(items, item{"Log", "\n" + padLines(evt.Log, "    ")})
	}
	var maxSz int
	for _, item := range items {
		sz := len(item.label)
		if len(item.value) > 0 && item.value[0] != '\n' && sz > maxSz {
			maxSz = sz
		}
	}
	for _, item := range items {
		count := (maxSz - len(item.label)) + 1
		var pad string
		if count > 0 && len(item.value) > 0 && item.value[0] != '\n' {
			pad = strings.Repeat(" ", count)
		}
		label := cmd.Colorfy(item.label+":", "cyan", "", "")
		fmt.Fprintf(context.Stdout, "%s%s%s\n", label, pad, item.value)
	}
	return nil
}

var rePadLines = regexp.MustCompile(`(?m)^(.+)`)

func padLines(s string, pad string) string {
	return rePadLines.ReplaceAllString(s, pad+`$1`)
}

type EventCancel struct {
	cmd.ConfirmationCommand
}

func (c *EventCancel) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "event-cancel",
		Usage:   "event-cancel <event-id> <reason> [-y]",
		Desc:    `Cancel a running event.`,
		MinArgs: 2,
	}
}

func (c *EventCancel) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(context, "Are you sure you want to cancel this event?") {
		return nil
	}
	u, err := cmd.GetURLVersion("1.1", fmt.Sprintf("/events/%s/cancel", context.Args[0]))
	if err != nil {
		return err
	}
	v := url.Values{}
	v.Set("reason", strings.Join(context.Args[1:], " "))
	request, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintln(context.Stdout, "Event successfully canceled.")
	return nil
}
<|endoftext|>


<commit_before>package model

import (
	"github.com/guregu/null"
)

// StateCD
// 0 = pending
// 1 = fetching
// 2 = analyzing
// 3 = done
// 4 = error

// Project has uploaded repository information.
type Project struct {
	UUID   string   `json:"uuid"`
	UserID null.Int `json:"user_id"`

	CartfileContent       string `json:"cartfile_content"`
	PodfileLockContent    string `json:"podfile_content"`
	PBXprojectContent     string `json:"pbxporj_content"`
	SupportedSwiftVersion string `json:"supported_swift_version"`
	Repository            string `json:"repository"`
	LastFetchedAt         Time   `json:"last_fetched_at"`
	StateCD               int    `json:"status_cd"`
}
<commit_msg>Set primary key and foreign key to projects<commit_after>package model

import (
	"github.com/guregu/null"
)

// StateCD
// 0 = pending
// 1 = fetching
// 2 = analyzing
// 3 = done
// 4 = error

// Project has uploaded repository information.
type Project struct {
	UUID   string   `json:"uuid" gorm:"primary_key"`
	UserID null.Int `json:"user_id" gorm:"ForeignKey:ID"`

	CartfileContent       string `json:"cartfile_content"`
	PodfileLockContent    string `json:"podfile_content"`
	PBXprojectContent     string `json:"pbxporj_content"`
	SupportedSwiftVersion string `json:"supported_swift_version"`
	Repository            string `json:"repository"`
	LastFetchedAt         Time   `json:"last_fetched_at"`
	StateCD               int    `json:"status_cd"`
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

const (
	workerEtcdPrefix = "workers"
	maxBackoff       = 5 * time.Second
)

type datum struct {
	files   []*pfs.FileInfo
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	// Send datums to this channel to be processed
	DataCh() chan<- *datum
	// Receive datums that failed to be processed from this channel
	FailCh() <-chan *datum
	// Receive hashtrees of the outputs of successfully processing datums
	SuccessCh() <-chan hashtree.HashTree
}

type workerPool struct {
	// When this context is canceled, the worker pool should clean up all
	// its resources.
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// workersMap is a map from a worker's address to the function that
	// can be used to release its resources.
	workersMap     map[string]worker
	workersMapLock sync.Mutex
	// objClient is the client for Pachyderm's object store
	objClient pfs.ObjectAPIClient
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// Used to delete worker pods
	kubeClient *kube.Client
	namespace  string
	// The job that spawned the worker pool
	jobID string
	// workers get datums from this channel.
	dataCh chan *datum
	// workers send datums to this channel when they fail to process
	// the datums.
	failCh chan *datum
	// workers send the hashtrees of the outputs of processing datums to
	// this channel.
	successCh chan hashtree.HashTree
}

func (w *workerPool) discoverWorkers() {
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		protolion.Infof("watching `%s` for workers", w.workerDir)
		watcher, err := watch.NewWatcher(w.ctx, w.etcdClient, w.workerDir)
		if err != nil {
			return err
		}
		defer watcher.Close()
		for {
			resp, ok := <-watcher.Watch()
			if !ok {
				return fmt.Errorf("watcher closed for unknown reasons")
			}
			if err := resp.Err; err != nil {
				return err
			}
			addr := path.Base(string(resp.Key))
			podName := string(resp.Value)
			switch resp.Type {
			case watch.EventPut:
				if err := w.addWorker(addr, podName); err != nil {
					return err
				}
			case watch.EventDelete:
				if err := w.delWorker(addr); err != nil {
					return err
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Errorf("error discovering workers for %v: %v; retrying in %v", w.workerDir, err, d)
		return nil
	})
}

type worker struct {
	cancel  context.CancelFunc
	podName string
}

func (w *workerPool) addWorker(addr string, podName string) error {
	w.workersMapLock.Lock()
	defer w.workersMapLock.Unlock()

	if cancel, ok := w.workersMap[addr]; ok {
		cancel()
	}

	workerCtx, cancelFn := context.WithCancel(w.ctx)
	w.workersMap[addr] = worker{
		cancel:  cancelFn,
		podName: podName,
	}

	protolion.Infof("launching new worker for %s at %v", w.workerDir, addr)
	go w.runWorker(workerCtx, addr)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	w.workersMapLock.Lock()
	defer w.workersMapLock.Unlock()

	worker, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}

	worker.cancel()
	if err := w.kubeClient.Pods(w.namespace).Delete(worker.podName); err != nil {
		return err
	}
	protolion.Infof("deleting worker for %s at %v", w.workerDir, addr)

	return nil
}

func (w *workerPool) runWorker(ctx context.Context, addr string) {
	defer func() {
		protolion.Infof("goro for worker %s for job %s is exiting", addr, w.jobID)
	}()

	var workerClient workerpkg.WorkerClient
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		conn, err := grpc.DialContext(ctx, fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure())
		if err != nil {
			return err
		}
		workerClient = workerpkg.NewWorkerClient(conn)
		return nil
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Infof("error establishing connection with worker %s; retrying in %v", addr, d)
		return nil
	})

	for true {
		var dt *datum
		var ok bool
		b := backoff.NewExponentialBackOff()
		if err := backoff.RetryNotify(func() error {
			select {
			case <-ctx.Done():
				return
			case dt, ok = <-w.dataCh:
				if !ok {
					return
				}
			}
			resp, err := workerClient.Process(ctx, &workerpkg.ProcessRequest{
				JobID: w.jobID,
				Data:  dt.files,
			})
		}, b, func(err error, d time.Duration) {
			if d > maxBackoff {
				return err
			}
			protolion.Errorf("worker %s for job %s failed to process datum %v with error %s; retrying in %s", addr, w.jobID, dt.files, err, d)
		}); err != nil {
			// If this worker keeps failing to process the datum (note that
			// failing to process a datum is different than if the user code
			// ran and returned a non-zero exit code), we eventually give up
			// and delete the worker pod.
			select {
			case w.failCh <- dt:
			case <-ctx.Done():
			}
			if err := w.delWorker(addr); err != nil {
				// If we can't delete the worker for some reason, we will
				// just have to carry on.
				protolion.Errorf("error deleting worker: %v", addr)
				continue
			}
			return
		}
		func() (retErr error) {
			defer func() {
				if retErr != nil {
					protolion.Errorf("datum error in job %s: %v", w.jobID, retErr)
					select {
					case w.failCh <- dt:
					case <-ctx.Done():
					}
				}
			}()
			if resp.Tag != nil {
				var buffer bytes.Buffer
				getTagClient, err := w.objClient.GetTag(ctx, &pfs.Tag{resp.Tag.Name})
				if err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				if err := grpcutil.WriteFromStreamingBytesClient(getTagClient, &buffer); err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				tree, err := hashtree.Deserialize(buffer.Bytes())
				if err != nil {
					return fmt.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", addr, dt.files, err)
				}
				w.successCh <- tree
			} else if resp.Failed {
				return fmt.Errorf("user code failed to process datum %v", dt.files)
			} else {
				return fmt.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", addr, dt.files)
			}
			return nil
		}()
	}
}

func (w *workerPool) DataCh() chan<- *datum {
	return w.dataCh
}

func status(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]*pps.WorkerStatus, error) {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return nil, err
	}
	var result []*pps.WorkerStatus
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return nil, err
		}
		result = append(result, status)
	}
	return result, nil
}

func cancel(ctx context.Context, id string, etcdClient *etcd.Client,
	etcdPrefix string, jobID string, dataFilter []string) error {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return err
	}
	success := false
	for _, workerClient := range workerClients {
		resp, err := workerClient.Cancel(ctx, &workerpkg.CancelRequest{
			JobID:       jobID,
			DataFilters: dataFilter,
		})
		if err != nil {
			return err
		}
		if resp.Success {
			success = true
		}
	}
	if !success {
		return fmt.Errorf("datum matching filter %+v could not be found for jobID %s", dataFilter, jobID)
	}
	return nil
}

func (w *workerPool) FailCh() <-chan *datum {
	return w.failCh
}

func (w *workerPool) SuccessCh() <-chan hashtree.HashTree {
	return w.successCh
}

// workerPool generates a new worker pool that talks to the replication
// controller identified by rcName.
// Each workerPool is supposed to be owned by a single job, identified
// by jobID.
func (a *apiServer) newWorkerPool(ctx context.Context, rcName string, jobID string) (WorkerPool, error) {
	objClient, err := a.getObjectClient()
	if err != nil {
		return nil, err
	}
	wp := &workerPool{
		ctx:        ctx,
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, rcName),
		workersMap: make(map[string]context.CancelFunc),
		objClient:  objClient,
		etcdClient: a.etcdClient,
		kubeClient: a.kubeClient,
		namespace:  a.namespace,
		jobID:      jobID,
		dataCh:     make(chan *datum),
		failCh:     make(chan *datum),
		successCh:  make(chan hashtree.HashTree),
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers()
	return wp, nil
}

func workerClients(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]workerpkg.WorkerClient, error) {
	resp, err := etcdClient.Get(ctx, path.Join(etcdPrefix, workerEtcdPrefix, id), etcd.WithPrefix())
	if err != nil {
		return nil, err
	}
	protolion.Printf("resp: %+v\n", resp)

	var result []workerpkg.WorkerClient
	for _, kv := range resp.Kvs {
		conn, err := grpc.Dial(fmt.Sprintf("%s:%d", path.Base(string(kv.Key)), client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithBlock(), grpc.WithTimeout(5*time.Second))
		if err != nil {
			return nil, err
		}
		result = append(result, workerpkg.NewWorkerClient(conn))
	}
	return result, nil
}
<commit_msg>Make code compile<commit_after>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
	"k8s.io/kubernetes/pkg/api"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

const (
	workerEtcdPrefix = "workers"
	maxBackoff       = 5 * time.Second
)

type datum struct {
	files   []*pfs.FileInfo
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	// Send datums to this channel to be processed
	DataCh() chan<- *datum
	// Receive datums that failed to be processed from this channel
	FailCh() <-chan *datum
	// Receive hashtrees of the outputs of successfully processing datums
	SuccessCh() <-chan hashtree.HashTree
}

type workerPool struct {
	// When this context is canceled, the worker pool should clean up all
	// its resources.
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// workersMap is a map from a worker's address to the function that
	// can be used to release its resources.
	workersMap     map[string]worker
	workersMapLock sync.Mutex
	// objClient is the client for Pachyderm's object store
	objClient pfs.ObjectAPIClient
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// Used to delete worker pods
	kubeClient *kube.Client
	namespace  string
	// The job that spawned the worker pool
	jobID string
	// workers get datums from this channel.
	dataCh chan *datum
	// workers send datums to this channel when they fail to process
	// the datums.
	failCh chan *datum
	// workers send the hashtrees of the outputs of processing datums to
	// this channel.
	successCh chan hashtree.HashTree
}

func (w *workerPool) discoverWorkers() {
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		protolion.Infof("watching `%s` for workers", w.workerDir)
		watcher, err := watch.NewWatcher(w.ctx, w.etcdClient, w.workerDir)
		if err != nil {
			return err
		}
		defer watcher.Close()
		for {
			resp, ok := <-watcher.Watch()
			if !ok {
				return fmt.Errorf("watcher closed for unknown reasons")
			}
			if err := resp.Err; err != nil {
				return err
			}
			addr := path.Base(string(resp.Key))
			podName := string(resp.Value)
			switch resp.Type {
			case watch.EventPut:
				if err := w.addWorker(addr, podName); err != nil {
					return err
				}
			case watch.EventDelete:
				if err := w.delWorker(addr); err != nil {
					return err
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Errorf("error discovering workers for %v: %v; retrying in %v", w.workerDir, err, d)
		return nil
	})
}

type worker struct {
	cancel  context.CancelFunc
	podName string
}

func (w *workerPool) addWorker(addr string, podName string) error {
	w.workersMapLock.Lock()
	defer w.workersMapLock.Unlock()

	if worker, ok := w.workersMap[addr]; ok {
		worker.cancel()
	}

	workerCtx, cancelFn := context.WithCancel(w.ctx)
	w.workersMap[addr] = worker{
		cancel:  cancelFn,
		podName: podName,
	}

	protolion.Infof("launching new worker for %s at %v", w.workerDir, addr)
	go w.runWorker(workerCtx, addr)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	w.workersMapLock.Lock()
	defer w.workersMapLock.Unlock()

	worker, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}

	worker.cancel()
	zeroVal := int64(0)
	if err := w.kubeClient.Pods(w.namespace).Delete(worker.podName, &api.DeleteOptions{
		GracePeriodSeconds: &zeroVal,
	}); err != nil {
		return err
	}
	protolion.Infof("deleting worker for %s at %v", w.workerDir, addr)

	return nil
}

func (w *workerPool) runWorker(ctx context.Context, addr string) {
	defer func() {
		protolion.Infof("goro for worker %s for job %s is exiting", addr, w.jobID)
	}()

	var workerClient workerpkg.WorkerClient
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		conn, err := grpc.DialContext(ctx, fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure())
		if err != nil {
			return err
		}
		workerClient = workerpkg.NewWorkerClient(conn)
		return nil
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Infof("error establishing connection with worker %s; retrying in %v", addr, d)
		return nil
	})

	for true {
		var dt *datum
		var ok bool
		select {
		case <-ctx.Done():
			return
		case dt, ok = <-w.dataCh:
			if !ok {
				return
			}
		}
		var resp *workerpkg.ProcessResponse
		var err error
		b := backoff.NewExponentialBackOff()
		if err := backoff.RetryNotify(func() error {
			resp, err = workerClient.Process(ctx, &workerpkg.ProcessRequest{
				JobID: w.jobID,
				Data:  dt.files,
			})
			return err
		}, b, func(err error, d time.Duration) error {
			if d > maxBackoff {
				return err
			}
			protolion.Errorf("worker %s for job %s failed to process datum %v with error %s; retrying in %s", addr, w.jobID, dt.files, err, d)
			return nil
		}); err != nil {
			select {
			case w.failCh <- dt:
			case <-ctx.Done():
			}
			// If this worker keeps failing to process the datum (note that
			// failing to process a datum is different than if the user code
			// ran and returned a non-zero exit code), we eventually give up
			// and delete the worker pod.
			if err := w.delWorker(addr); err != nil {
				// If we can't delete the worker for some reason, we will
				// just have to carry on.
				protolion.Errorf("error deleting worker: %v", addr)
				continue
			}
			return
		}
		func() (retErr error) {
			defer func() {
				if retErr != nil {
					protolion.Errorf("datum error in job %s: %v", w.jobID, retErr)
					select {
					case w.failCh <- dt:
					case <-ctx.Done():
					}
				}
			}()
			if resp.Tag != nil {
				var buffer bytes.Buffer
				getTagClient, err := w.objClient.GetTag(ctx, &pfs.Tag{resp.Tag.Name})
				if err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				if err := grpcutil.WriteFromStreamingBytesClient(getTagClient, &buffer); err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				tree, err := hashtree.Deserialize(buffer.Bytes())
				if err != nil {
					return fmt.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", addr, dt.files, err)
				}
				w.successCh <- tree
			} else if resp.Failed {
				return fmt.Errorf("user code failed to process datum %v", dt.files)
			} else {
				return fmt.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", addr, dt.files)
			}
			return nil
		}()
	}
}

func (w *workerPool) DataCh() chan<- *datum {
	return w.dataCh
}

func status(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]*pps.WorkerStatus, error) {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return nil, err
	}
	var result []*pps.WorkerStatus
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return nil, err
		}
		result = append(result, status)
	}
	return result, nil
}

func cancel(ctx context.Context, id string, etcdClient *etcd.Client,
	etcdPrefix string, jobID string, dataFilter []string) error {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return err
	}
	success := false
	for _, workerClient := range workerClients {
		resp, err := workerClient.Cancel(ctx, &workerpkg.CancelRequest{
			JobID:       jobID,
			DataFilters: dataFilter,
		})
		if err != nil {
			return err
		}
		if resp.Success {
			success = true
		}
	}
	if !success {
		return fmt.Errorf("datum matching filter %+v could not be found for jobID %s", dataFilter, jobID)
	}
	return nil
}

func (w *workerPool) FailCh() <-chan *datum {
	return w.failCh
}

func (w *workerPool) SuccessCh() <-chan hashtree.HashTree {
	return w.successCh
}

// workerPool generates a new worker pool that talks to the replication
// controller identified by rcName.
// Each workerPool is supposed to be owned by a single job, identified
// by jobID.
func (a *apiServer) newWorkerPool(ctx context.Context, rcName string, jobID string) (WorkerPool, error) {
	objClient, err := a.getObjectClient()
	if err != nil {
		return nil, err
	}
	wp := &workerPool{
		ctx:        ctx,
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, rcName),
		workersMap: make(map[string]worker),
		objClient:  objClient,
		etcdClient: a.etcdClient,
		kubeClient: a.kubeClient,
		namespace:  a.namespace,
		jobID:      jobID,
		dataCh:     make(chan *datum),
		failCh:     make(chan *datum),
		successCh:  make(chan hashtree.HashTree),
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers()
	return wp, nil
}

func workerClients(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]workerpkg.WorkerClient, error) {
	resp, err := etcdClient.Get(ctx, path.Join(etcdPrefix, workerEtcdPrefix, id), etcd.WithPrefix())
	if err != nil {
		return nil, err
	}
	protolion.Printf("resp: %+v\n", resp)

	var result []workerpkg.WorkerClient
	for _, kv := range resp.Kvs {
		conn, err := grpc.Dial(fmt.Sprintf("%s:%d", path.Base(string(kv.Key)), client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithBlock(), grpc.WithTimeout(5*time.Second))
		if err != nil {
			return nil, err
		}
		result = append(result, workerpkg.NewWorkerClient(conn))
	}
	return result, nil
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package state

import (
	"fmt"
	"strings"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/avalanche"
	"github.com/ava-labs/avalanchego/snow/consensus/snowstorm"
	"github.com/ava-labs/avalanchego/utils/formatting"
	"github.com/ava-labs/avalanchego/utils/hashing"
)

// uniqueVertex acts as a cache for vertices in the database.
//
// If a vertex is loaded, it will have one canonical uniqueVertex. The vertex
// will eventually be evicted from memory, when the uniqueVertex is evicted from
// the cache. If the uniqueVertex has a function called again afther this
// eviction, the vertex will be re-loaded from the database.
type uniqueVertex struct {
	serializer *Serializer

	vtxID ids.ID
	v     *vertexState
}

// newUniqueVertex returns a uniqueVertex instance from [b] by checking the cache
// and then parsing the vertex bytes on a cache miss.
func newUniqueVertex(s *Serializer, b []byte) (*uniqueVertex, error) {
	vtx := &uniqueVertex{
		vtxID:      ids.NewID(hashing.ComputeHash256Array(b)),
		serializer: s,
		v:          &vertexState{},
	}

	unique := vtx.serializer.state.UniqueVertex(vtx)
	if unique != vtx {
		// If the vertex is known and in the cache,
		// then we have no new information to add
		// so we return it as is.
		if unique.v.status != choices.Unknown {
			return unique, nil
		}

		// The status was Unknown, so we need to parse and persist the vertex
		innerVertex, err := s.parseVertex(b)
		if err != nil {
			return nil, err
		}
		unique.v.vtx = innerVertex
		return unique, unique.persist()
	}

	// The vertex was not in the cache, so mark it as unique and
	// attempt to parse the vertex bytes.
	vtx.v.unique = true
	innerVertex, err := s.parseVertex(b)
	if err != nil {
		return nil, err
	}

	// Persist the vertex if necessary
	vtx.v.vtx = innerVertex
	return vtx, vtx.persist()
}

func (vtx *uniqueVertex) refresh() {
	if vtx.v == nil {
		vtx.v = &vertexState{}
	}
	if vtx.v.unique {
		return
	}

	unique := vtx.serializer.state.UniqueVertex(vtx)
	prevVtx := vtx.v.vtx
	if unique == vtx {
		vtx.v.status = vtx.serializer.state.Status(vtx.ID())
		vtx.v.unique = true
	} else {
		// If someone is in the cache, they must be up to date
		*vtx = *unique
	}

	switch {
	case vtx.v.vtx == nil && prevVtx == nil:
		vtx.v.vtx = vtx.serializer.state.Vertex(vtx.ID())
	case vtx.v.vtx == nil:
		vtx.v.vtx = prevVtx
	}
}

func (vtx *uniqueVertex) Evict() {
	if vtx.v != nil {
		vtx.v.unique = false
		// make sure the parents can be garbage collected
		vtx.v.parents = nil
	}
}

// persist writes the vertex and status to the database if necessary
// and also updates the current status of [vtx]
// Assumes the inner vertex is non-nil
func (vtx *uniqueVertex) persist() error {
	if vtx.v.status != choices.Unknown {
		return nil
	}

	status := vtx.serializer.state.Status(vtx.ID())
	if status == choices.Unknown {
		if err := vtx.serializer.state.SetStatus(vtx.ID(), choices.Processing); err != nil {
			return err
		}
		if err := vtx.serializer.state.SetVertex(vtx.v.vtx); err != nil {
			return err
		}
		vtx.v.status = choices.Processing
	} else {
		vtx.v.status = status
	}

	return nil
}

func (vtx *uniqueVertex) setVertex(innerVtx *innerVertex) error {
	vtx.refresh()
	if vtx.v.vtx != nil {
		return nil
	}
	vtx.v.vtx = innerVtx
	if err := vtx.serializer.state.SetVertex(innerVtx); err != nil {
		return err
	}
	return vtx.setStatus(choices.Processing)
}

func (vtx *uniqueVertex) setStatus(status choices.Status) error {
	vtx.refresh()
	if vtx.v.status == status {
		return nil
	}
	vtx.v.status = status
	return vtx.serializer.state.SetStatus(vtx.ID(), status)
}

func (vtx *uniqueVertex) ID() ids.ID { return vtx.vtxID }

func (vtx *uniqueVertex) Accept() error {
	if err := vtx.setStatus(choices.Accepted); err != nil {
		return err
	}

	vtx.serializer.edge.Add(vtx.vtxID)
	parents, err := vtx.Parents()
	if err != nil {
		return err
	}

	for _, parent := range parents {
		vtx.serializer.edge.Remove(parent.ID())
	}

	if err := vtx.serializer.state.SetEdge(vtx.serializer.edge.List()); err != nil {
		return fmt.Errorf("failed to set edge while accepting vertex %s due to %w", vtx.vtxID, err)
	}

	// Should never traverse into parents of a decided vertex. Allows for the
	// parents to be garbage collected
	vtx.v.parents = nil

	return vtx.serializer.db.Commit()
}

func (vtx *uniqueVertex) Reject() error {
	if err := vtx.setStatus(choices.Rejected); err != nil {
		return err
	}

	// Should never traverse into parents of a decided vertex. Allows for the
	// parents to be garbage collected
	vtx.v.parents = nil

	return vtx.serializer.db.Commit()
}

func (vtx *uniqueVertex) Status() choices.Status { vtx.refresh(); return vtx.v.status }

func (vtx *uniqueVertex) Parents() ([]avalanche.Vertex, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return nil, fmt.Errorf("failed to get parents for vertex with status: %s", vtx.v.status)
	}

	if len(vtx.v.parents) != len(vtx.v.vtx.parentIDs) {
		vtx.v.parents = make([]avalanche.Vertex, len(vtx.v.vtx.parentIDs))
		for i, parentID := range vtx.v.vtx.parentIDs {
			vtx.v.parents[i] = &uniqueVertex{
				serializer: vtx.serializer,
				vtxID:      parentID,
			}
		}
	}

	return vtx.v.parents, nil
}

func (vtx *uniqueVertex) Height() (uint64, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return 0, fmt.Errorf("failed to get height for vertex with status: %s", vtx.v.status)
	}

	return vtx.v.vtx.height, nil
}

func (vtx *uniqueVertex) Txs() ([]snowstorm.Tx, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return nil, fmt.Errorf("failed to get txs for vertex with status: %s", vtx.v.status)
	}

	if len(vtx.v.vtx.txs) != len(vtx.v.txs) {
		vtx.v.txs = make([]snowstorm.Tx, len(vtx.v.vtx.txs))
		for i, tx := range vtx.v.vtx.txs {
			vtx.v.txs[i] = tx
		}
	}

	return vtx.v.txs, nil
}

func (vtx *uniqueVertex) Bytes() []byte { return vtx.v.vtx.Bytes() }

func (vtx *uniqueVertex) Verify() error { return vtx.v.vtx.Verify() }

func (vtx *uniqueVertex) String() string {
	sb := strings.Builder{}

	parents, err := vtx.Parents()
	if err != nil {
		sb.WriteString(fmt.Sprintf("Vertex(ID = %s, Error=error while retrieving vertex parents: %s)", vtx.ID(), err))
		return sb.String()
	}
	txs, err := vtx.Txs()
	if err != nil {
		sb.WriteString(fmt.Sprintf("Vertex(ID = %s, Error=error while retrieving vertex txs: %s)", vtx.ID(), err))
		return sb.String()
	}

	sb.WriteString(fmt.Sprintf(
		"Vertex(ID = %s, Status = %s, Number of Dependencies = %d, Number of Transactions = %d)",
		vtx.ID(),
		vtx.Status(),
		len(parents),
		len(txs),
	))

	parentFormat := fmt.Sprintf("\n    Parent[%s]: ID = %%s, Status = %%s",
		formatting.IntFormat(len(parents)-1))
	for i, parent := range parents {
		sb.WriteString(fmt.Sprintf(parentFormat, i, parent.ID(), parent.Status()))
	}

	txFormat := fmt.Sprintf("\n    Transaction[%s]: ID = %%s, Status = %%s",
		formatting.IntFormat(len(txs)-1))
	for i, tx := range txs {
		sb.WriteString(fmt.Sprintf(txFormat, i, tx.ID(), tx.Status()))
	}

	return sb.String()
}

type vertexState struct {
	unique bool

	vtx    *innerVertex
	status choices.Status

	parents []avalanche.Vertex
	txs     []snowstorm.Tx
}
<commit_msg>Implement shallow refresh function to allow skipping db reads when possible<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package state

import (
	"fmt"
	"strings"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/avalanche"
	"github.com/ava-labs/avalanchego/snow/consensus/snowstorm"
	"github.com/ava-labs/avalanchego/utils/formatting"
	"github.com/ava-labs/avalanchego/utils/hashing"
)

// uniqueVertex acts as a cache for vertices in the database.
//
// If a vertex is loaded, it will have one canonical uniqueVertex. The vertex
// will eventually be evicted from memory, when the uniqueVertex is evicted from
// the cache. If the uniqueVertex has a function called again afther this
// eviction, the vertex will be re-loaded from the database.
//
// Invariant: a uniqueVertex should never have a vertex state
// where the innerVertex is non-nil and the
// status is Unknown
// ie v.vtx != nil ==> v.status != choices.Unknown at all times
// This is ensured by updating the status first.
type uniqueVertex struct {
	serializer *Serializer

	vtxID ids.ID
	v     *vertexState
}

// newUniqueVertex returns a uniqueVertex instance from [b] by checking the cache
// and then parsing the vertex bytes on a cache miss.
func newUniqueVertex(s *Serializer, b []byte) (*uniqueVertex, error) {
	vtx := &uniqueVertex{
		vtxID:      ids.NewID(hashing.ComputeHash256Array(b)),
		serializer: s,
	}
	vtx.shallowRefresh()

	if vtx.v.vtx != nil {
		return vtx, nil
	}

	innerVertex, err := s.parseVertex(b)
	if err != nil {
		return nil, err
	}
	vtx.v.vtx = innerVertex

	// If the vertex is already known, skip writing it
	// to the database
	if vtx.v.status != choices.Unknown {
		return vtx, nil
	}

	// Otherwise, set the status to Processing and write
	// to the database before returning the vertex
	if err := vtx.serializer.state.SetStatus(vtx.ID(), choices.Processing); err != nil {
		return nil, err
	}
	if err := vtx.serializer.state.SetVertex(vtx.v.vtx); err != nil {
		return nil, err
	}
	vtx.v.status = choices.Processing

	return vtx, nil
}

func (vtx *uniqueVertex) refresh() {
	vtx.shallowRefresh()

	if vtx.v.vtx == nil && vtx.v.status != choices.Unknown {
		vtx.v.vtx = vtx.serializer.state.Vertex(vtx.ID())
	}
}

// shallowRefresh checks the cache for the uniqueVertex and gets the
// most up to date status for [vtx]
// ensures that the status is up to date for this vertex
// inner vertex may be nil after calling shallowRefresh
func (vtx *uniqueVertex) shallowRefresh() {
	if vtx.v == nil {
		vtx.v = &vertexState{}
	}
	if vtx.v.unique {
		return
	}

	unique := vtx.serializer.state.UniqueVertex(vtx)
	prevVtx := vtx.v.vtx
	if unique == vtx {
		vtx.v.status = vtx.serializer.state.Status(vtx.ID())
		vtx.v.unique = true
	} else {
		// If someone is in the cache, they must be up to date
		*vtx = *unique
	}

	if vtx.v.vtx == nil {
		vtx.v.vtx = prevVtx
	}
}

func (vtx *uniqueVertex) Evict() {
	if vtx.v != nil {
		vtx.v.unique = false
		// make sure the parents can be garbage collected
		vtx.v.parents = nil
	}
}

func (vtx *uniqueVertex) setVertex(innerVtx *innerVertex) error {
	vtx.shallowRefresh()
	if vtx.v.vtx != nil {
		return nil
	}

	vtx.v.vtx = innerVtx
	if err := vtx.serializer.state.SetVertex(innerVtx); err != nil {
		return err
	}
	return vtx.setStatus(choices.Processing)
}

func (vtx *uniqueVertex) setStatus(status choices.Status) error {
	vtx.refresh()
	if vtx.v.status == status {
		return nil
	}
	vtx.v.status = status
	return vtx.serializer.state.SetStatus(vtx.ID(), status)
}

func (vtx *uniqueVertex) ID() ids.ID { return vtx.vtxID }

func (vtx *uniqueVertex) Accept() error {
	if err := vtx.setStatus(choices.Accepted); err != nil {
		return err
	}

	vtx.serializer.edge.Add(vtx.vtxID)
	parents, err := vtx.Parents()
	if err != nil {
		return err
	}

	for _, parent := range parents {
		vtx.serializer.edge.Remove(parent.ID())
	}

	if err := vtx.serializer.state.SetEdge(vtx.serializer.edge.List()); err != nil {
		return fmt.Errorf("failed to set edge while accepting vertex %s due to %w", vtx.vtxID, err)
	}

	// Should never traverse into parents of a decided vertex. Allows for the
	// parents to be garbage collected
	vtx.v.parents = nil

	return vtx.serializer.db.Commit()
}

func (vtx *uniqueVertex) Reject() error {
	if err := vtx.setStatus(choices.Rejected); err != nil {
		return err
	}

	// Should never traverse into parents of a decided vertex. Allows for the
	// parents to be garbage collected
	vtx.v.parents = nil

	return vtx.serializer.db.Commit()
}

func (vtx *uniqueVertex) Status() choices.Status { vtx.refresh(); return vtx.v.status }

func (vtx *uniqueVertex) Parents() ([]avalanche.Vertex, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return nil, fmt.Errorf("failed to get parents for vertex with status: %s", vtx.v.status)
	}

	if len(vtx.v.parents) != len(vtx.v.vtx.parentIDs) {
		vtx.v.parents = make([]avalanche.Vertex, len(vtx.v.vtx.parentIDs))
		for i, parentID := range vtx.v.vtx.parentIDs {
			vtx.v.parents[i] = &uniqueVertex{
				serializer: vtx.serializer,
				vtxID:      parentID,
			}
		}
	}

	return vtx.v.parents, nil
}

func (vtx *uniqueVertex) Height() (uint64, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return 0, fmt.Errorf("failed to get height for vertex with status: %s", vtx.v.status)
	}

	return vtx.v.vtx.height, nil
}

func (vtx *uniqueVertex) Txs() ([]snowstorm.Tx, error) {
	vtx.refresh()

	if vtx.v.vtx == nil {
		return nil, fmt.Errorf("failed to get txs for vertex with status: %s", vtx.v.status)
	}

	if len(vtx.v.vtx.txs) != len(vtx.v.txs) {
		vtx.v.txs = make([]snowstorm.Tx, len(vtx.v.vtx.txs))
		for i, tx := range vtx.v.vtx.txs {
			vtx.v.txs[i] = tx
		}
	}

	return vtx.v.txs, nil
}

func (vtx *uniqueVertex) Bytes() []byte { return vtx.v.vtx.Bytes() }

func (vtx *uniqueVertex) Verify() error { return vtx.v.vtx.Verify() }

func (vtx *uniqueVertex) String() string {
	sb := strings.Builder{}

	parents, err := vtx.Parents()
	if err != nil {
		sb.WriteString(fmt.Sprintf("Vertex(ID = %s, Error=error while retrieving vertex parents: %s)", vtx.ID(), err))
		return sb.String()
	}
	txs, err := vtx.Txs()
	if err != nil {
		sb.WriteString(fmt.Sprintf("Vertex(ID = %s, Error=error while retrieving vertex txs: %s)", vtx.ID(), err))
		return sb.String()
	}

	sb.WriteString(fmt.Sprintf(
		"Vertex(ID = %s, Status = %s, Number of Dependencies = %d, Number of Transactions = %d)",
		vtx.ID(),
		vtx.Status(),
		len(parents),
		len(txs),
	))

	parentFormat := fmt.Sprintf("\n    Parent[%s]: ID = %%s, Status = %%s",
		formatting.IntFormat(len(parents)-1))
	for i, parent := range parents {
		sb.WriteString(fmt.Sprintf(parentFormat, i, parent.ID(), parent.Status()))
	}

	txFormat := fmt.Sprintf("\n    Transaction[%s]: ID = %%s, Status = %%s",
		formatting.IntFormat(len(txs)-1))
	for i, tx := range txs {
		sb.WriteString(fmt.Sprintf(txFormat, i, tx.ID(), tx.Status()))
	}

	return sb.String()
}

type vertexState struct {
	unique bool

	vtx    *innerVertex
	status choices.Status

	parents []avalanche.Vertex
	txs     []snowstorm.Tx
}
<|endoftext|>


<commit_before>package brain

import (
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/lib/prettyprint"
	"io"
	"strings"
)

// PrivilegeLevel is a type to represent different privilege levels.
// since privilege levels in the brain are just strings, they're just a string type here too.
type PrivilegeLevel string

const (
	// ClusterAdminPrivilege allows a user to administer the cluster managed by the brain, and do things like create/delete VMs on accounts they have no explicit right on, grant others AccountAdminPrivilege, and set disc iops_limit
	ClusterAdminPrivilege PrivilegeLevel = "cluster_admin"
	// AccountAdminPrivilege allows a user to create, modify & delete groups and servers in an account.
	AccountAdminPrivilege = "account_admin"
	// GroupAdminPrivilege allows a user to create, modify & delete servers in a specific group.
	GroupAdminPrivilege = "group_admin"
	// VMAdminPrivilege allows a user to modify & administer a server, including increasing the performance (and hence the price on the uk0 cluster) and accessing the console.
	VMAdminPrivilege = "vm_admin"
	// VMConsolePrivilege allows a user to access the console for a particular server.
	VMConsolePrivilege = "vm_console"
)

// String returns the privilege level cast to a string.
func (pl PrivilegeLevel) String() string {
	return string(pl)
}

const (
	PrivilegeTargetTypeVM      = "vm"
	PrivilegeTargetTypeGroup   = "group"
	PrivilegeTargetTypeAccount = "account"
	PrivilegeTargetTypeCluster = "cluster"
)

// Privilege represents a privilege on the brain.
// A user may have multiple privileges, and multiple privileges may be granted on the same object.
// At the moment we're not worried about the extra fields that privileges have on the brain (IP restrictions) because they're unused
type Privilege struct {
	// ID is the numeric ID used mostly by the brain
	ID int `json:"id,omitempty"`
	// Username is the user who the privilege is granted to
	Username string `json:"username,omitempty"`
	// VirtualMachineID is the ID of the virtual machine the privilege is granted on
	VirtualMachineID int `json:"virtual_machine_id,omitempty"`
	// AccountID is the ID of the account the privilege is granted on
	AccountID int `json:"account_id,omitempty"`
	// GroupID is the ID of the group the privilege is granted on
	GroupID int `json:"group_id,omitempty"`
	// Level is the PrivilegeLevel they have
	Level PrivilegeLevel `json:"level,omitempty"`
	// YubikeyRequired is true if the user should have to authenticate with a yubikey in order to use this privilege. Only set it to true if you're sure the user has a yubikey set up on their account, and that they know where it is!
	YubikeyRequired bool `json:"yubikey_required"`
	// YubikeyOTPMaxAge should set how long (in seconds) a yubikey one-time-password would be accepted for, but it might not be used?
	YubikeyOTPMaxAge int `json:"yubikey_otp_max_age,omitempty"`
}

func (p Privilege) TargetType() string {
	return strings.Split(string(p.Level), "_")[0]
}

// String returns a string representation of the Privilege in English.
// Privileges are a little tricky to represent in English because the Privilege itself doesn't know if it exists on a user or if it has just been removed from a user, nor does it now anything about the target it's been granted on/revoked from other than a numerical ID. So we do the best we can.
func (p Privilege) String() string {
	requiresYubikey := ""
	if p.YubikeyRequired {
		requiresYubikey = " (requires yubikey)"
	}
	switch p.TargetType() {
	case PrivilegeTargetTypeVM:
		return fmt.Sprintf("%s on VM #%d for %s%s", p.Level, p.VirtualMachineID, p.Username, requiresYubikey)
	case PrivilegeTargetTypeGroup:
		return fmt.Sprintf("%s on group #%d for %s%s", p.Level, p.GroupID, p.Username, requiresYubikey)
	case PrivilegeTargetTypeAccount:
		return fmt.Sprintf("%s on account #%d for %s%s", p.Level, p.AccountID, p.Username, requiresYubikey)
	}
	return fmt.Sprintf("%s for %s%s", p.Level, p.Username, requiresYubikey)
}

// PrettyPrint nicely formats the Privilege and sends it to the given writer.
// At the moment, the detail parameter is ignored.
func (p Privilege) PrettyPrint(wr io.Writer, detail prettyprint.DetailLevel) (err error) {
	_, err = wr.Write([]byte(p.String()))
	return
}

// Privileges is used to allow API consumers to use IndexOf on the array of privileges.
type Privileges []*Privilege

// IndexOf finds the privilege given in the list of privileges, ignoring the Privilege ID and returns the index. If it couldn't find it, returns -1.
func (ps Privileges) IndexOf(priv Privilege) int {
	if priv.Username == "" || priv.Level == "" {
		return -1
	}
	for i, p := range ps {
		if p.VirtualMachineID == priv.VirtualMachineID &&
			p.GroupID == priv.GroupID && p.AccountID == priv.AccountID &&
			p.YubikeyRequired == priv.YubikeyRequired &&
			p.Level == priv.Level && p.Username == priv.Username {
			return i
		}
	}
	return -1
}
<commit_msg>Add documentation comments to PrivilegeTargetType constants & to Privilege.TargetType<commit_after>package brain

import (
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/lib/prettyprint"
	"io"
	"strings"
)

// PrivilegeLevel is a type to represent different privilege levels.
// since privilege levels in the brain are just strings, they're just a string type here too.
type PrivilegeLevel string

const (
	// ClusterAdminPrivilege allows a user to administer the cluster managed by the brain, and do things like create/delete VMs on accounts they have no explicit right on, grant others AccountAdminPrivilege, and set disc iops_limit
	ClusterAdminPrivilege PrivilegeLevel = "cluster_admin"
	// AccountAdminPrivilege allows a user to create, modify & delete groups and servers in an account.
	AccountAdminPrivilege = "account_admin"
	// GroupAdminPrivilege allows a user to create, modify & delete servers in a specific group.
	GroupAdminPrivilege = "group_admin"
	// VMAdminPrivilege allows a user to modify & administer a server, including increasing the performance (and hence the price on the uk0 cluster) and accessing the console.
	VMAdminPrivilege = "vm_admin"
	// VMConsolePrivilege allows a user to access the console for a particular server.
	VMConsolePrivilege = "vm_console"
)

// String returns the privilege level cast to a string.
func (pl PrivilegeLevel) String() string {
	return string(pl)
}

const (
	// PrivilegeTargetTypeVM is the prefix for all privilege levels that affect VMs
	PrivilegeTargetTypeVM = "vm"
	// PrivilegeTargetTypeGroup is the prefix for all privilege levels that affect Groups
	PrivilegeTargetTypeGroup = "group"
	// PrivilegeTargetTypeAccount is the prefix for all privilege levels that affect Accounts
	PrivilegeTargetTypeAccount = "account"
	// PrivilegeTargetTypeCluster is the prefix for all privilege levels that affect the whole cluster.
	PrivilegeTargetTypeCluster = "cluster"
)

// Privilege represents a privilege on the brain.
// A user may have multiple privileges, and multiple privileges may be granted on the same object.
// At the moment we're not worried about the extra fields that privileges have on the brain (IP restrictions) because they're unused
type Privilege struct {
	// ID is the numeric ID used mostly by the brain
	ID int `json:"id,omitempty"`
	// Username is the user who the privilege is granted to
	Username string `json:"username,omitempty"`
	// VirtualMachineID is the ID of the virtual machine the privilege is granted on
	VirtualMachineID int `json:"virtual_machine_id,omitempty"`
	// AccountID is the ID of the account the privilege is granted on
	AccountID int `json:"account_id,omitempty"`
	// GroupID is the ID of the group the privilege is granted on
	GroupID int `json:"group_id,omitempty"`
	// Level is the PrivilegeLevel they have
	Level PrivilegeLevel `json:"level,omitempty"`
	// YubikeyRequired is true if the user should have to authenticate with a yubikey in order to use this privilege. Only set it to true if you're sure the user has a yubikey set up on their account, and that they know where it is!
	YubikeyRequired bool `json:"yubikey_required"`
	// YubikeyOTPMaxAge should set how long (in seconds) a yubikey one-time-password would be accepted for, but it might not be used?
	YubikeyOTPMaxAge int `json:"yubikey_otp_max_age,omitempty"`
}

// TargetType returns the prefix of the PrivilegeLevel, which should be one of the PrivilegeTargetType* constants.
func (p Privilege) TargetType() string {
	return strings.Split(string(p.Level), "_")[0]
}

// String returns a string representation of the Privilege in English.
// Privileges are a little tricky to represent in English because the Privilege itself doesn't know if it exists on a user or if it has just been removed from a user, nor does it now anything about the target it's been granted on/revoked from other than a numerical ID. So we do the best we can.
func (p Privilege) String() string {
	requiresYubikey := ""
	if p.YubikeyRequired {
		requiresYubikey = " (requires yubikey)"
	}
	switch p.TargetType() {
	case PrivilegeTargetTypeVM:
		return fmt.Sprintf("%s on VM #%d for %s%s", p.Level, p.VirtualMachineID, p.Username, requiresYubikey)
	case PrivilegeTargetTypeGroup:
		return fmt.Sprintf("%s on group #%d for %s%s", p.Level, p.GroupID, p.Username, requiresYubikey)
	case PrivilegeTargetTypeAccount:
		return fmt.Sprintf("%s on account #%d for %s%s", p.Level, p.AccountID, p.Username, requiresYubikey)
	}
	return fmt.Sprintf("%s for %s%s", p.Level, p.Username, requiresYubikey)
}

// PrettyPrint nicely formats the Privilege and sends it to the given writer.
// At the moment, the detail parameter is ignored.
func (p Privilege) PrettyPrint(wr io.Writer, detail prettyprint.DetailLevel) (err error) {
	_, err = wr.Write([]byte(p.String()))
	return
}

// Privileges is used to allow API consumers to use IndexOf on the array of privileges.
type Privileges []*Privilege

// IndexOf finds the privilege given in the list of privileges, ignoring the Privilege ID and returns the index. If it couldn't find it, returns -1.
func (ps Privileges) IndexOf(priv Privilege) int {
	if priv.Username == "" || priv.Level == "" {
		return -1
	}
	for i, p := range ps {
		if p.VirtualMachineID == priv.VirtualMachineID &&
			p.GroupID == priv.GroupID && p.AccountID == priv.AccountID &&
			p.YubikeyRequired == priv.YubikeyRequired &&
			p.Level == priv.Level && p.Username == priv.Username {
			return i
		}
	}
	return -1
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"code.google.com/p/go.crypto/nacl/box"
	protobuf "code.google.com/p/gogoprotobuf/proto"
	"crypto/rand"
	"fmt"
	"github.com/andres-erbsen/chatterbox/proto"
	"github.com/andres-erbsen/chatterbox/ratchet"
	"github.com/andres-erbsen/dename/client"
	testutil2 "github.com/andres-erbsen/dename/server/testutil" //TODO: Move MakeToken to TestUtil
	"github.com/andres-erbsen/dename/testutil"
	"testing"
	"time"
)

func handleError(err error, t *testing.T) {
	if err != nil {
		t.Error(err)
	}
}

func TestMessageEncryptionAuthentication(t *testing.T) {
	t.Skip("apparently the tests can only run a single binary involving a dename server. this may have something to do with TCP TIME_WAIT. I (andres) am disabling this test to make the other one pass.")
	config, f := testutil.SingleServer(t)
	defer f()
	time.Sleep(100)

	ska, dnmca := createNewUser("Alice", t, config)
	skb, dnmcb := createNewUser("Bob", t, config)

	ratchA := &ratchet.Ratchet{
		FillAuth:  FillAuthWith(ska),
		CheckAuth: CheckAuthWith(dnmca),
		Rand:      nil,
		Now:       nil,
	}
	ratchB := &ratchet.Ratchet{
		FillAuth:  FillAuthWith(skb),
		CheckAuth: CheckAuthWith(dnmcb),
		Rand:      nil,
		Now:       nil,
	}

	//pka0, ska0, err := box.GenerateKey(rand.Reader)
	//handleError(err, t)
	pkb0, skb0, err := box.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	msg, err := protobuf.Marshal(&proto.Message{
		Subject:  "",
		Contents: []byte("Message"),
		Dename:   "Alice",
	})
	if err != nil {
		t.Fatal(err)
	}
	out := append([]byte{}, (*pkb0)[:]...)

	out = ratchA.EncryptFirst(out, msg, pkb0)
	msg2, err := ratchB.DecryptFirst(out[32:], skb0)
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(msg, msg2) {
		t.Error("Original and decrypted message not the same.")
	}
}

func createNewUser(name string, t *testing.T, config *client.Config) (*[32]byte, *client.Client) {
	newClient, err := client.NewClient(config, nil, nil)
	if err != nil {
		t.Fatal(err)
	}

	//TODO: All these names are horrible, please change them
	pkAuth, skAuth, err := box.GenerateKey(rand.Reader)

	chatProfile := &proto.Profile{
		ServerAddressTCP:  "",
		ServerPortTCP:     -1,
		ServerTransportPK: (proto.Byte32)([32]byte{}),
		UserIDAtServer:    (proto.Byte32)([32]byte{}),
		KeySigningKey:     (proto.Byte32)([32]byte{}),
		MessageAuthKey:    (proto.Byte32)(*pkAuth),
	}

	chatProfileBytes, err := protobuf.Marshal(chatProfile)
	if err != nil {
		t.Fatal(err)
	}

	profile, sk, err := client.NewProfile(nil, nil)
	if err != nil {
		t.Fatal(err)
	}

	client.SetProfileField(profile, PROFILE_FIELD_ID, chatProfileBytes)

	err = newClient.Register(sk, name, profile, testutil2.MakeToken())
	if err != nil {
		t.Fatal(err)
	}

	//Remove this outside of the test
	profile2, err := newClient.Lookup(name)
	if !profile.Equal(profile2) {
		t.Error("Correct profile not added to server.")
		fmt.Printf("profile: %v\n", profile)
		fmt.Printf("profile2: %v\n", profile2)
	}

	return skAuth, newClient
}
<commit_msg>profile verification test stub<commit_after>package client

import (
	"bytes"
	"code.google.com/p/go.crypto/nacl/box"
	protobuf "code.google.com/p/gogoprotobuf/proto"
	"crypto/rand"
	"fmt"
	"github.com/andres-erbsen/chatterbox/proto"
	"github.com/andres-erbsen/chatterbox/ratchet"
	"github.com/andres-erbsen/dename/client"
	dename "github.com/andres-erbsen/dename/protocol"
	testutil2 "github.com/andres-erbsen/dename/server/testutil" //TODO: Move MakeToken to TestUtil
	"github.com/andres-erbsen/dename/testutil"
	"testing"
	"time"
)

func handleError(err error, t *testing.T) {
	if err != nil {
		t.Error(err)
	}
}

func noProfileRatchet(c *client.Client) func(name string, reply *dename.ClientReply) (*dename.Profile, error) {
	return func(name string, reply *dename.ClientReply) (*dename.Profile, error) {
		if reply != nil {
			if profile, err := c.LookupFromReply(name, reply); err == nil {
				return profile, nil
			}
		}
		return c.Lookup(name)
	}
}

func TestMessageEncryptionAuthentication(t *testing.T) {
	t.Skip("apparently the tests can only run a single binary involving a dename server. this may have something to do with TCP TIME_WAIT. I (andres) am disabling this test to make the other one pass.")
	config, f := testutil.SingleServer(t)
	defer f()
	time.Sleep(100)

	ska, dnmca := createNewUser("Alice", t, config)
	skb, dnmcb := createNewUser("Bob", t, config)

	ratchA := &ratchet.Ratchet{
		FillAuth:  FillAuthWith(ska),
		CheckAuth: CheckAuthWith(noProfileRatchet(dnmca)),
		Rand:      nil,
		Now:       nil,
	}
	ratchB := &ratchet.Ratchet{
		FillAuth:  FillAuthWith(skb),
		CheckAuth: CheckAuthWith(noProfileRatchet(dnmcb)),
		Rand:      nil,
		Now:       nil,
	}

	//pka0, ska0, err := box.GenerateKey(rand.Reader)
	//handleError(err, t)
	pkb0, skb0, err := box.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	msg, err := protobuf.Marshal(&proto.Message{
		Subject:  "",
		Contents: []byte("Message"),
		Dename:   "Alice",
	})
	if err != nil {
		t.Fatal(err)
	}
	out := append([]byte{}, (*pkb0)[:]...)

	out = ratchA.EncryptFirst(out, msg, pkb0)
	msg2, err := ratchB.DecryptFirst(out[32:], skb0)
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(msg, msg2) {
		t.Error("Original and decrypted message not the same.")
	}
}

func createNewUser(name string, t *testing.T, config *client.Config) (*[32]byte, *client.Client) {
	newClient, err := client.NewClient(config, nil, nil)
	if err != nil {
		t.Fatal(err)
	}

	//TODO: All these names are horrible, please change them
	pkAuth, skAuth, err := box.GenerateKey(rand.Reader)

	chatProfile := &proto.Profile{
		ServerAddressTCP:  "",
		ServerPortTCP:     -1,
		ServerTransportPK: (proto.Byte32)([32]byte{}),
		UserIDAtServer:    (proto.Byte32)([32]byte{}),
		KeySigningKey:     (proto.Byte32)([32]byte{}),
		MessageAuthKey:    (proto.Byte32)(*pkAuth),
	}

	chatProfileBytes, err := protobuf.Marshal(chatProfile)
	if err != nil {
		t.Fatal(err)
	}

	profile, sk, err := client.NewProfile(nil, nil)
	if err != nil {
		t.Fatal(err)
	}

	client.SetProfileField(profile, PROFILE_FIELD_ID, chatProfileBytes)

	err = newClient.Register(sk, name, profile, testutil2.MakeToken())
	if err != nil {
		t.Fatal(err)
	}

	//Remove this outside of the test
	profile2, err := newClient.Lookup(name)
	if !profile.Equal(profile2) {
		t.Error("Correct profile not added to server.")
		fmt.Printf("profile: %v\n", profile)
		fmt.Printf("profile2: %v\n", profile2)
	}

	return skAuth, newClient
}
<|endoftext|>


<commit_before>/*
  Package certgen id set of utilities used to generate ssh certificates
*/
package certgen

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"golang.org/x/crypto/ssh"
	"math/big"
	//"os"
	"os/exec"
	"time"

	"fmt"
)

const numValidHours = 24

// GetUserPubKeyFromSSSD user authorized keys content based on the running sssd configuration
func GetUserPubKeyFromSSSD(username string) (string, error) {
	cmd := exec.Command("/usr/bin/sss_ssh_authorizedkeys", username)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

func goCertToFileString(c ssh.Certificate, username string) (string, error) {
	certBytes := c.Marshal()
	encoded := base64.StdEncoding.EncodeToString(certBytes)
	fileComment := "/tmp/" + username + "-cert.pub"
	return "ssh-rsa-cert-v01@openssh.com " + encoded + " " + fileComment, nil
}

// gen_user_cert a username and key, returns a short lived cert for that user
func GenSSHCertFileString(username string, userPubKey string, signer ssh.Signer, host_identity string) (string, error) {
	//const numValidHours = 24

	userKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(userPubKey))
	if err != nil {
		return "", err
	}
	keyIdentity := host_identity + "_" + username

	currentEpoch := uint64(time.Now().Unix())
	expireEpoch := currentEpoch + (3600 * numValidHours)

	// The values of the permissions are taken from the default values used
	// by ssh-keygen
	cert := ssh.Certificate{
		Key:             userKey,
		CertType:        ssh.UserCert,
		SignatureKey:    signer.PublicKey(),
		ValidPrincipals: []string{username},
		KeyId:           keyIdentity,
		ValidAfter:      currentEpoch,
		ValidBefore:     expireEpoch,
		Permissions: ssh.Permissions{Extensions: map[string]string{
			"permit-X11-forwarding":   "",
			"permit-agent-forwarding": "",
			"permit-port-forwarding":  "",
			"permit-pty":              "",
			"permit-user-rc":          ""}}}

	err = cert.SignCert(bytes.NewReader(cert.Marshal()), signer)
	if err != nil {
		return "", err
	}
	certString, err := goCertToFileString(cert, username)
	if err != nil {
		return "", err
	}
	return certString, nil
}

func GenSSHCertFileStringFromSSSDPublicKey(userName string, signer ssh.Signer, hostIdentity string) (string, error) {

	userPubKey, err := GetUserPubKeyFromSSSD(userName)
	if err != nil {
		return "", err
	}
	cert, err := GenSSHCertFileString(userName, userPubKey, signer, hostIdentity)
	if err != nil {
		return "", err
	}
	return cert, err
}

/// X509 section
func getPubKeyFromPem(pubkey string) (pub interface{}, err error) {
	block, rest := pem.Decode([]byte(pubkey))
	if block == nil || block.Type != "PUBLIC KEY" {
		err := errors.New(fmt.Sprintf("Cannot decode user public Key '%s' rest='%s'", pubkey, string(rest)))
		if block != nil {
			err = errors.New(fmt.Sprintf("public key bad type %s", block.Type))
		}
		return "", err
	}
	return x509.ParsePKIXPublicKey(block.Bytes)
}

func getPrivateKeyFromPem(privateKey string) (pub interface{}, err error) {
	//TODO handle ecdsa and other non-rsa keys
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		err := errors.New("Cannot decode Private Key")
		return "", err
	}
	switch block.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	default:
		err := errors.New("Cannot process that key")
		return pub, err
	}
}

//copied from https://golang.org/src/crypto/tls/generate_cert.go
func publicKey(priv interface{}) interface{} {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &k.PublicKey
	//case *ecdsa.PrivateKey:
	//	return &k.PublicKey
	default:
		return nil
	}
}

func derBytesCertToCertAndPem(derBytes []byte) (*x509.Certificate, string, error) {
	cert, err := x509.ParseCertificate(derBytes)
	if err != nil {
		return nil, "", err
	}
	pemCert := string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}))
	fmt.Printf("pem:\n%s\n", pemCert)
	return cert, pemCert, nil
}

// return both an internal representation an the pem representation of the string
// As long as the issuer value matches THEN the serial number can be different every time
func GenSelfSignedCACert(commonName string, organization string, caPriv interface{}) (*x509.Certificate, string, error) {
	//// Now do the actual work...
	notBefore := time.Now()
	notAfter := notBefore.Add(24 * 365 * 8 * time.Hour)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, "", err
	}
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   commonName,
			Organization: []string{organization},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,
		KeyUsage:  x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageCertSign,
		//ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		//BasicConstraintsValid: true,
		IsCA: true,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey(caPriv), caPriv)
	if err != nil {

		//log.Fatalf("Failed to create certificate: %s", err)
		return nil, "", err
	}
	return derBytesCertToCertAndPem(derBytes)

}

// From RFC 4120 section 5.2.2 (https://tools.ietf.org/html/rfc4120)
type KerberosPrincipal struct {
	Len       int      `asn1:"explicit,tag:0"`
	Principal []string `asn1:"explicit,tag:1"`
}

// From RFC 4556 section 3.2.2 (https://tools.ietf.org/html/rfc4556.html)
type KRB5PrincipalName struct {
	Realm string `asn1:"explicit,tag:0"`
	//`asn1:"ia5"`
	//`asn1:"ia5"`
	Principal KerberosPrincipal `asn1:"explicit,tag:1"`
	//`asn1:"tag:16"`
}

type PKInitSANAnotherName struct {
	Id    asn1.ObjectIdentifier
	Value KRB5PrincipalName `asn1:"explicit,tag:0"`
}

// This is the m
func changePrintableStringToGeneralString(kerberosRealm string, inString []byte) []byte {
	/*	position := 12
		inString[position] = 27
		position = position + 1 + len(kerberosRealm) + 8
		inString[position] = 27
	*/
	return inString
}

func genSANExtension(userName string, kerberosRealm *string) (*pkix.Extension, error) {
	krbRealm := "EXAMPLE.COM"
	if kerberosRealm != nil {
		krbRealm = *kerberosRealm
	}
	/*
		krbSanExt := KRB5PrincipalName{
			Realm:     krbRealm,
			Principal: KerberosPrincipal{Len: 1, Principal: []string{userName}},
		}
		krbSanExtDer, err := asn1.Marshal(krbSanExt)
		if err != nil {
			return nil, "", err
		}
		fmt.Printf("ext: %+x", krbSanExtDer)
	*/
	//1.3.6.1.5.2.2
	krbSanAnotherName := PKInitSANAnotherName{
		Id: []int{1, 3, 6, 1, 5, 2, 2},
		Value: KRB5PrincipalName{
			Realm:     krbRealm,
			Principal: KerberosPrincipal{Len: 1, Principal: []string{userName}},
		},
	}
	krbSanAnotherNameDer, err := asn1.Marshal(krbSanAnotherName)
	if err != nil {
		return nil, err
	}
	fmt.Printf("ext: %+x\n", krbSanAnotherNameDer)
	krbSanAnotherNameDer = changePrintableStringToGeneralString(krbRealm, krbSanAnotherNameDer)
	fmt.Printf("ext: %+x\n", krbSanAnotherNameDer)
	//Apply fix HERE!!!!

	// inspired by marshalSANs in x509.go
	var rawValues []asn1.RawValue
	rawValues = append(rawValues, asn1.RawValue{Tag: 0, Class: 2, IsCompound: true, Bytes: krbSanAnotherNameDer})

	rawSan, err := asn1.Marshal(rawValues)
	if err != nil {
		return nil, err
	}
	sanExtension := pkix.Extension{
		Id:    []int{2, 5, 29, 17},
		Value: rawSan,
	}

	return &sanExtension, nil
}

// returns an x509 cert that is with the username in the common name
func GenUserX509Cert(userName string, userPub interface{}, caCert *x509.Certificate, caPriv interface{}, kerberosRealm *string) (*x509.Certificate, string, error) {
	//// Now do the actual work...
	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(numValidHours) * time.Hour)

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

	sanExtension, err := genSANExtension(userName, kerberosRealm)
	if err != nil {
		return nil, "", err
	}

	// need to add the extended key usage... that is special for kerberos
	//and also the client key usage
	kerberosClientExtKeyUsage := []int{1, 3, 6, 1, 5, 2, 3, 4}
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   userName,
			Organization: []string{"Acme Co"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		UnknownExtKeyUsage:    []asn1.ObjectIdentifier{kerberosClientExtKeyUsage},
		BasicConstraintsValid: true,
		IsCA: false,
	}
	template.ExtraExtensions = []pkix.Extension{*sanExtension}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, caCert, userPub, caPriv)
	if err != nil {

		//log.Fatalf("Failed to create certificate: %s", err)
		return nil, "", err
	}
	return derBytesCertToCertAndPem(derBytes)

	//pemCert := string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}))

	//return pemCert, nil
}
<commit_msg>dixing change the printiable to general<commit_after>/*
  Package certgen id set of utilities used to generate ssh certificates
*/
package certgen

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"golang.org/x/crypto/ssh"
	"math/big"
	//"os"
	"os/exec"
	"time"

	"fmt"
)

const numValidHours = 24

// GetUserPubKeyFromSSSD user authorized keys content based on the running sssd configuration
func GetUserPubKeyFromSSSD(username string) (string, error) {
	cmd := exec.Command("/usr/bin/sss_ssh_authorizedkeys", username)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

func goCertToFileString(c ssh.Certificate, username string) (string, error) {
	certBytes := c.Marshal()
	encoded := base64.StdEncoding.EncodeToString(certBytes)
	fileComment := "/tmp/" + username + "-cert.pub"
	return "ssh-rsa-cert-v01@openssh.com " + encoded + " " + fileComment, nil
}

// gen_user_cert a username and key, returns a short lived cert for that user
func GenSSHCertFileString(username string, userPubKey string, signer ssh.Signer, host_identity string) (string, error) {
	//const numValidHours = 24

	userKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(userPubKey))
	if err != nil {
		return "", err
	}
	keyIdentity := host_identity + "_" + username

	currentEpoch := uint64(time.Now().Unix())
	expireEpoch := currentEpoch + (3600 * numValidHours)

	// The values of the permissions are taken from the default values used
	// by ssh-keygen
	cert := ssh.Certificate{
		Key:             userKey,
		CertType:        ssh.UserCert,
		SignatureKey:    signer.PublicKey(),
		ValidPrincipals: []string{username},
		KeyId:           keyIdentity,
		ValidAfter:      currentEpoch,
		ValidBefore:     expireEpoch,
		Permissions: ssh.Permissions{Extensions: map[string]string{
			"permit-X11-forwarding":   "",
			"permit-agent-forwarding": "",
			"permit-port-forwarding":  "",
			"permit-pty":              "",
			"permit-user-rc":          ""}}}

	err = cert.SignCert(bytes.NewReader(cert.Marshal()), signer)
	if err != nil {
		return "", err
	}
	certString, err := goCertToFileString(cert, username)
	if err != nil {
		return "", err
	}
	return certString, nil
}

func GenSSHCertFileStringFromSSSDPublicKey(userName string, signer ssh.Signer, hostIdentity string) (string, error) {

	userPubKey, err := GetUserPubKeyFromSSSD(userName)
	if err != nil {
		return "", err
	}
	cert, err := GenSSHCertFileString(userName, userPubKey, signer, hostIdentity)
	if err != nil {
		return "", err
	}
	return cert, err
}

/// X509 section
func getPubKeyFromPem(pubkey string) (pub interface{}, err error) {
	block, rest := pem.Decode([]byte(pubkey))
	if block == nil || block.Type != "PUBLIC KEY" {
		err := errors.New(fmt.Sprintf("Cannot decode user public Key '%s' rest='%s'", pubkey, string(rest)))
		if block != nil {
			err = errors.New(fmt.Sprintf("public key bad type %s", block.Type))
		}
		return "", err
	}
	return x509.ParsePKIXPublicKey(block.Bytes)
}

func getPrivateKeyFromPem(privateKey string) (pub interface{}, err error) {
	//TODO handle ecdsa and other non-rsa keys
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		err := errors.New("Cannot decode Private Key")
		return "", err
	}
	switch block.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	default:
		err := errors.New("Cannot process that key")
		return pub, err
	}
}

//copied from https://golang.org/src/crypto/tls/generate_cert.go
func publicKey(priv interface{}) interface{} {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &k.PublicKey
	//case *ecdsa.PrivateKey:
	//	return &k.PublicKey
	default:
		return nil
	}
}

func derBytesCertToCertAndPem(derBytes []byte) (*x509.Certificate, string, error) {
	cert, err := x509.ParseCertificate(derBytes)
	if err != nil {
		return nil, "", err
	}
	pemCert := string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}))
	fmt.Printf("pem:\n%s\n", pemCert)
	return cert, pemCert, nil
}

// return both an internal representation an the pem representation of the string
// As long as the issuer value matches THEN the serial number can be different every time
func GenSelfSignedCACert(commonName string, organization string, caPriv interface{}) (*x509.Certificate, string, error) {
	//// Now do the actual work...
	notBefore := time.Now()
	notAfter := notBefore.Add(24 * 365 * 8 * time.Hour)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, "", err
	}
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   commonName,
			Organization: []string{organization},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,
		KeyUsage:  x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageCertSign,
		//ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		//BasicConstraintsValid: true,
		IsCA: true,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey(caPriv), caPriv)
	if err != nil {

		//log.Fatalf("Failed to create certificate: %s", err)
		return nil, "", err
	}
	return derBytesCertToCertAndPem(derBytes)

}

// From RFC 4120 section 5.2.2 (https://tools.ietf.org/html/rfc4120)
type KerberosPrincipal struct {
	Len       int      `asn1:"explicit,tag:0"`
	Principal []string `asn1:"explicit,tag:1"`
}

// From RFC 4556 section 3.2.2 (https://tools.ietf.org/html/rfc4556.html)
type KRB5PrincipalName struct {
	Realm string `asn1:"explicit,tag:0"`
	//`asn1:"ia5"`
	//`asn1:"ia5"`
	Principal KerberosPrincipal `asn1:"explicit,tag:1"`
	//`asn1:"tag:16"`
}

type PKInitSANAnotherName struct {
	Id    asn1.ObjectIdentifier
	Value KRB5PrincipalName `asn1:"explicit,tag:0"`
}

// This is the m
func changePrintableStringToGeneralString(kerberosRealm string, inString []byte) []byte {
	position := 16
	inString[position] = 27

	position = position + 1 + len(kerberosRealm) + 14
	inString[position] = 27

	return inString
}

func genSANExtension(userName string, kerberosRealm *string) (*pkix.Extension, error) {
	krbRealm := "EXAMPLE.COM"
	if kerberosRealm != nil {
		krbRealm = *kerberosRealm
	}
	//1.3.6.1.5.2.2
	krbSanAnotherName := PKInitSANAnotherName{
		Id: []int{1, 3, 6, 1, 5, 2, 2},
		Value: KRB5PrincipalName{
			Realm:     krbRealm,
			Principal: KerberosPrincipal{Len: 1, Principal: []string{userName}},
		},
	}
	krbSanAnotherNameDer, err := asn1.Marshal(krbSanAnotherName)
	if err != nil {
		return nil, err
	}
	fmt.Printf("ext: %+x\n", krbSanAnotherNameDer)
	krbSanAnotherNameDer = changePrintableStringToGeneralString(krbRealm, krbSanAnotherNameDer)
	fmt.Printf("ext: %+x\n", krbSanAnotherNameDer)
	//Apply fix HERE!!!!

	// inspired by marshalSANs in x509.go
	var rawValues []asn1.RawValue
	rawValues = append(rawValues, asn1.RawValue{Tag: 0, Class: 2, IsCompound: true, Bytes: krbSanAnotherNameDer})

	rawSan, err := asn1.Marshal(rawValues)
	if err != nil {
		return nil, err
	}
	sanExtension := pkix.Extension{
		Id:    []int{2, 5, 29, 17},
		Value: rawSan,
	}

	return &sanExtension, nil
}

// returns an x509 cert that is with the username in the common name
func GenUserX509Cert(userName string, userPub interface{}, caCert *x509.Certificate, caPriv interface{}, kerberosRealm *string) (*x509.Certificate, string, error) {
	//// Now do the actual work...
	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(numValidHours) * time.Hour)

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

	sanExtension, err := genSANExtension(userName, kerberosRealm)
	if err != nil {
		return nil, "", err
	}

	// need to add the extended key usage... that is special for kerberos
	//and also the client key usage
	kerberosClientExtKeyUsage := []int{1, 3, 6, 1, 5, 2, 3, 4}
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   userName,
			Organization: []string{"Acme Co"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		UnknownExtKeyUsage:    []asn1.ObjectIdentifier{kerberosClientExtKeyUsage},
		BasicConstraintsValid: true,
		IsCA: false,
	}
	template.ExtraExtensions = []pkix.Extension{*sanExtension}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, caCert, userPub, caPriv)
	if err != nil {

		//log.Fatalf("Failed to create certificate: %s", err)
		return nil, "", err
	}
	return derBytesCertToCertAndPem(derBytes)

	//pemCert := string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}))

	//return pemCert, nil
}
<|endoftext|>


<commit_before>package lib

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/url"
	"regexp"
	"strconv"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/pkg/units"
	"github.com/docker/docker/runconfig"
)

var headerRegexp = regexp.MustCompile(`\ADocker/.+\s\((.+)\)\z`)

// ImageBuild sends request to the daemon to build images.
// The Body in the response implement an io.ReadCloser and it's up to the caller to
// close it.
func (cli *Client) ImageBuild(options types.ImageBuildOptions) (types.ImageBuildResponse, error) {
	query, err := imageBuildOptionsToQuery(options)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}

	headers := http.Header(make(map[string][]string))
	buf, err := json.Marshal(options.AuthConfigs)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}
	headers.Add("X-Registry-Config", base64.URLEncoding.EncodeToString(buf))
	headers.Set("Content-Type", "application/tar")

	serverResp, err := cli.postRaw("/build", query, options.Context, headers)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}

	osType := getDockerOS(serverResp.header.Get("Server"))

	return types.ImageBuildResponse{
		Body:   serverResp.body,
		OSType: osType,
	}, nil
}

func imageBuildOptionsToQuery(options types.ImageBuildOptions) (url.Values, error) {
	query := url.Values{
		"t": options.Tags,
	}
	if options.SuppressOutput {
		query.Set("q", "1")
	}
	if options.RemoteContext != "" {
		query.Set("remote", options.RemoteContext)
	}
	if options.NoCache {
		query.Set("nocache", "1")
	}
	if options.Remove {
		query.Set("rm", "1")
	} else {
		query.Set("rm", "0")
	}

	if options.ForceRemove {
		query.Set("forcerm", "1")
	}

	if options.PullParent {
		query.Set("pull", "1")
	}

	if !runconfig.IsolationLevel.IsDefault(runconfig.IsolationLevel(options.Isolation)) {
		query.Set("isolation", options.Isolation)
	}

	query.Set("cpusetcpus", options.CPUSetCPUs)
	query.Set("cpusetmems", options.CPUSetMems)
	query.Set("cpushares", strconv.FormatInt(options.CPUShares, 10))
	query.Set("cpuquota", strconv.FormatInt(options.CPUQuota, 10))
	query.Set("cpuperiod", strconv.FormatInt(options.CPUPeriod, 10))
	query.Set("memory", strconv.FormatInt(options.Memory, 10))
	query.Set("memswap", strconv.FormatInt(options.MemorySwap, 10))
	query.Set("cgroupparent", options.CgroupParent)

	if options.ShmSize != "" {
		parsedShmSize, err := units.RAMInBytes(options.ShmSize)
		if err != nil {
			return query, err
		}
		query.Set("shmsize", strconv.FormatInt(parsedShmSize, 10))
	}

	query.Set("dockerfile", options.Dockerfile)

	ulimitsJSON, err := json.Marshal(options.Ulimits)
	if err != nil {
		return query, err
	}
	query.Set("ulimits", string(ulimitsJSON))

	buildArgs := runconfig.ConvertKVStringsToMap(options.BuildArgs)
	buildArgsJSON, err := json.Marshal(buildArgs)
	if err != nil {
		return query, err
	}
	query.Set("buildargs", string(buildArgsJSON))

	return query, nil
}

func getDockerOS(serverHeader string) string {
	var osType string
	matches := headerRegexp.FindStringSubmatch(serverHeader)
	if len(matches) > 0 {
		osType = matches[1]
	}
	return osType
}
<commit_msg>Replace pkg/units with docker/go-units.<commit_after>package lib

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/url"
	"regexp"
	"strconv"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/runconfig"
	"github.com/docker/go-units"
)

var headerRegexp = regexp.MustCompile(`\ADocker/.+\s\((.+)\)\z`)

// ImageBuild sends request to the daemon to build images.
// The Body in the response implement an io.ReadCloser and it's up to the caller to
// close it.
func (cli *Client) ImageBuild(options types.ImageBuildOptions) (types.ImageBuildResponse, error) {
	query, err := imageBuildOptionsToQuery(options)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}

	headers := http.Header(make(map[string][]string))
	buf, err := json.Marshal(options.AuthConfigs)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}
	headers.Add("X-Registry-Config", base64.URLEncoding.EncodeToString(buf))
	headers.Set("Content-Type", "application/tar")

	serverResp, err := cli.postRaw("/build", query, options.Context, headers)
	if err != nil {
		return types.ImageBuildResponse{}, err
	}

	osType := getDockerOS(serverResp.header.Get("Server"))

	return types.ImageBuildResponse{
		Body:   serverResp.body,
		OSType: osType,
	}, nil
}

func imageBuildOptionsToQuery(options types.ImageBuildOptions) (url.Values, error) {
	query := url.Values{
		"t": options.Tags,
	}
	if options.SuppressOutput {
		query.Set("q", "1")
	}
	if options.RemoteContext != "" {
		query.Set("remote", options.RemoteContext)
	}
	if options.NoCache {
		query.Set("nocache", "1")
	}
	if options.Remove {
		query.Set("rm", "1")
	} else {
		query.Set("rm", "0")
	}

	if options.ForceRemove {
		query.Set("forcerm", "1")
	}

	if options.PullParent {
		query.Set("pull", "1")
	}

	if !runconfig.IsolationLevel.IsDefault(runconfig.IsolationLevel(options.Isolation)) {
		query.Set("isolation", options.Isolation)
	}

	query.Set("cpusetcpus", options.CPUSetCPUs)
	query.Set("cpusetmems", options.CPUSetMems)
	query.Set("cpushares", strconv.FormatInt(options.CPUShares, 10))
	query.Set("cpuquota", strconv.FormatInt(options.CPUQuota, 10))
	query.Set("cpuperiod", strconv.FormatInt(options.CPUPeriod, 10))
	query.Set("memory", strconv.FormatInt(options.Memory, 10))
	query.Set("memswap", strconv.FormatInt(options.MemorySwap, 10))
	query.Set("cgroupparent", options.CgroupParent)

	if options.ShmSize != "" {
		parsedShmSize, err := units.RAMInBytes(options.ShmSize)
		if err != nil {
			return query, err
		}
		query.Set("shmsize", strconv.FormatInt(parsedShmSize, 10))
	}

	query.Set("dockerfile", options.Dockerfile)

	ulimitsJSON, err := json.Marshal(options.Ulimits)
	if err != nil {
		return query, err
	}
	query.Set("ulimits", string(ulimitsJSON))

	buildArgs := runconfig.ConvertKVStringsToMap(options.BuildArgs)
	buildArgsJSON, err := json.Marshal(buildArgs)
	if err != nil {
		return query, err
	}
	query.Set("buildargs", string(buildArgsJSON))

	return query, nil
}

func getDockerOS(serverHeader string) string {
	var osType string
	matches := headerRegexp.FindStringSubmatch(serverHeader)
	if len(matches) > 0 {
		osType = matches[1]
	}
	return osType
}
<|endoftext|>


<commit_before>package main

import (
	"sync"

	elektra "go.libelektra.org/kdb"
)

type handle struct {
	kdb    elektra.KDB
	keySet elektra.KeySet
}

type handlePool struct {
	mut     sync.Mutex
	handles []*handle
}

func initPool(initialSize int) *handlePool {
	pool := &handlePool{
		handles: make([]*handle, initialSize+initialSize*1/4),
	}

	var err error

	for i := 0; i < initialSize; i++ {
		pool.handles[i], err = newHandle()

		if err != nil {
			panic("asd")
		}
	}

	return pool
}

func newHandle() (*handle, error) {
	kdb := elektra.New()

	err := kdb.Open()

	if err != nil {
		return nil, err
	}

	parentKey, err := elektra.NewKey("/")

	if err != nil {
		return nil, err
	}

	ks := elektra.NewKeySet()

	if _, err = kdb.Get(ks, parentKey); err != nil {
		return nil, err
	}

	return &handle{
		kdb:    kdb,
		keySet: ks,
	}, nil
}

func (p *handlePool) pop() *handle {
	var h *handle

	p.mut.Lock()
	size := len(p.handles)

	h, p.handles = p.handles[size-1], p.handles[:size-1]
	p.mut.Unlock()

	return h
}

func (p *handlePool) refill() {
	h, err := newHandle()

	if err != nil {
		panic("could not create new handle")
	}

	p.mut.Lock()
	p.handles = append(p.handles, h)
	defer p.mut.Unlock()
}

func (p *handlePool) Get() *handle {
	p.refill()

	return p.pop()
}
<commit_msg>tool: elektra - use channels for handlepool<commit_after>package main

import (
	elektra "go.libelektra.org/kdb"
)

type handle struct {
	kdb    elektra.KDB
	keySet elektra.KeySet
}

type handlePool struct {
	handles  chan *handle
	doRefill chan int
	size     int
}

func initPool(size int) *handlePool {
	pool := &handlePool{
		handles:  make(chan *handle, size),
		doRefill: make(chan int, 1),
		size:     size,
	}

	go pool.refillLoop()

	pool.refill()

	return pool
}

func newHandle() (*handle, error) {
	kdb := elektra.New()

	err := kdb.Open()

	if err != nil {
		return nil, err
	}

	parentKey, err := elektra.NewKey("/")

	if err != nil {
		return nil, err
	}

	ks := elektra.NewKeySet()

	if _, err = kdb.Get(ks, parentKey); err != nil {
		return nil, err
	}

	return &handle{
		kdb:    kdb,
		keySet: ks,
	}, nil
}

func (p *handlePool) refill() {
	select {
	case p.doRefill <- 1:
	default:
	}
}

func (p *handlePool) pop() *handle {
	return <-p.handles
}

func (p *handlePool) refillLoop() {
	for range p.doRefill {
		for l := len(p.handles); l < p.size; {
			// t := time.Now()
			h, err := newHandle()
			// log.Printf("time to create new handle: %v", time.Now().Sub(t))

			if err != nil {
				panic("could not create new handle")
			}

			p.handles <- h
		}
	}
}

func (p *handlePool) Get() *handle {
	p.refill()

	return p.pop()
}
<|endoftext|>