<commit_before>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"github.com/hashicorp/raft"
	"github.com/nats-io/go-nats-streaming/pb"

	"github.com/nats-io/nats-streaming-server/spb"
)

const batchSize = 200

var (
	fragmentPool    = &sync.Pool{New: func() interface{} { return &spb.RaftSnapshotFragment{} }}
	batchPool       = &sync.Pool{New: func() interface{} { return &spb.Batch{} }}
	subSnapshotPool = &sync.Pool{New: func() interface{} { return &spb.SubSnapshot{} }}
)

// channelSnapshot implements the raft.FSMSnapshot interface.
type channelSnapshot struct {
	*channel
}

func newChannelSnapshot(c *channel) raft.FSMSnapshot {
	return &channelSnapshot{channel: c}
}

// Persist should dump all necessary state to the WriteCloser 'sink',
// and call sink.Close() when finished or call sink.Cancel() on error.
func (c *channelSnapshot) Persist(sink raft.SnapshotSink) (err error) {
	defer func() {
		if err != nil {
			sink.Cancel()
		}
	}()

	if err := c.snapshotMessages(sink); err != nil {
		return err
	}

	if err := c.snapshotSubscriptions(sink); err != nil {
		return err
	}

	if err := c.snapshotAcks(sink); err != nil {
		return err
	}

	return sink.Close()
}

// Release is a no-op.
func (c *channelSnapshot) Release() {}

func (c *channelSnapshot) snapshotMessages(sink raft.SnapshotSink) error {
	// TODO: this is very expensive and might repeatedly fail if messages are
	// constantly being truncated. Is there a way we can optimize this, e.g.
	// handling Restore() out-of-band from Raft?

	first, last, err := c.store.Msgs.FirstAndLastSequence()
	if err != nil {
		return err
	}

	var (
		buf   [4]byte
		batch = batchPool.Get().(*spb.Batch)
	)
	batch.Messages = make([]*pb.MsgProto, 0, batchSize)

	for seq := first; seq <= last; seq++ {
		msg, err := c.store.Msgs.Lookup(seq)
		if err != nil {
			batchPool.Put(batch)
			return err
		}
		// If msg is nil, channel truncation has occurred while snapshotting.
		if msg == nil {
			// Channel truncation has occurred while snapshotting.
			batchPool.Put(batch)
			return fmt.Errorf("channel %q was truncated while snapshotting", c.name)
		}

		// Previous batch is full, ship it.
		if len(batch.Messages) == batchSize {
			fragment := newFragment()
			fragment.FragmentType = spb.RaftSnapshotFragment_Messages
			fragment.MessageBatch = batch
			if err := writeFragment(sink, fragment, buf); err != nil {
				return err
			}

			// Create a new batch.
			batch = batchPool.Get().(*spb.Batch)
			batch.Messages = make([]*pb.MsgProto, 0, batchSize)
		}

		batch.Messages = append(batch.Messages, msg)
	}

	// Ship any partial batch.
	if len(batch.Messages) > 0 {
		fragment := newFragment()
		fragment.FragmentType = spb.RaftSnapshotFragment_Messages
		fragment.MessageBatch = batch
		if err := writeFragment(sink, fragment, buf); err != nil {
			return err
		}
	}

	return nil
}

func (c *channelSnapshot) snapshotSubscriptions(sink raft.SnapshotSink) error {
	var buf [4]byte
	for _, sub := range c.ss.getAllSubs() {
		fragment := newFragment()
		fragment.FragmentType = spb.RaftSnapshotFragment_Subscription
		subSnap := subSnapshotPool.Get().(*spb.SubSnapshot)
		subSnap.Subject = sub.subject
		subSnap.State = &sub.SubState
		fragment.Sub = subSnap
		if err := writeFragment(sink, fragment, buf); err != nil {
			return err
		}
	}
	return nil
}

func (c *channelSnapshot) snapshotAcks(sink raft.SnapshotSink) error {
	// TODO
	return nil
}

// writeFragment writes the marshalled RaftSnapshotFragment to the
// SnapshotSink. The fragment should not be used after this is called.
func writeFragment(sink raft.SnapshotSink, fragment *spb.RaftSnapshotFragment, sizeBuf [4]byte) error {
	data, err := fragment.Marshal()
	if fragment.MessageBatch != nil {
		batchPool.Put(fragment.MessageBatch)
	}
	if fragment.Sub != nil {
		subSnapshotPool.Put(fragment.Sub)
	}
	fragmentPool.Put(fragment)
	if err != nil {
		return err
	}
	binary.BigEndian.PutUint32(sizeBuf[:], uint32(len(data)))
	_, err = sink.Write(sizeBuf[:])
	if err != nil {
		return err
	}
	_, err = sink.Write(data)
	return err
}

func newFragment() *spb.RaftSnapshotFragment {
	fragment := fragmentPool.Get().(*spb.RaftSnapshotFragment)
	fragment.MessageBatch = nil
	return fragment
}

// restoreFromSnapshot restores a channel from a snapshot. This is not called
// concurrently with any other Raft commands.
func (c *channel) restoreFromSnapshot(snapshot io.ReadCloser) error {
	defer snapshot.Close()

	// TODO: this needs to be fully implemented to support restoring the
	// entire Raft log, not just channel messages.
	// TODO: restore acks

	// Drop all existing subs. These will be restored from the snapshot.
	for _, sub := range c.ss.getAllSubs() {
		c.stan.unsubscribe(c, sub.ClientID, sub, true)
	}

	sizeBuf := make([]byte, 4)
	for {
		// Read the fragment size.
		if _, err := io.ReadFull(snapshot, sizeBuf); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		// Read the fragment.
		size := binary.BigEndian.Uint32(sizeBuf)
		buf := make([]byte, size)
		if _, err := io.ReadFull(snapshot, buf); err != nil {
			return err
		}

		// Unmarshal the fragment.
		fragment := newFragment()
		if err := fragment.Unmarshal(buf); err != nil {
			return err
		}

		// Apply the fragment.
		switch fragment.FragmentType {
		case spb.RaftSnapshotFragment_Messages:
			// Channel messages.
			for _, msg := range fragment.MessageBatch.Messages {
				if _, err := c.store.Msgs.Store(msg); err != nil {
					batchPool.Put(fragment.MessageBatch)
					fragmentPool.Put(fragment)
					return err
				}
			}
			batchPool.Put(fragment.MessageBatch)
			fragmentPool.Put(fragment)
		case spb.RaftSnapshotFragment_Subscription:
			// Channel subscription.
			sub := &subState{
				SubState:    *fragment.Sub.State,
				subject:     fragment.Sub.Subject,
				ackWait:     computeAckWait(fragment.Sub.State.AckWaitInSecs),
				acksPending: make(map[uint64]int64),
				store:       c.store.Subs,
			}
			if err := c.stan.addSubscription(c.ss, sub); err != nil {
				subSnapshotPool.Put(fragment.Sub)
				fragmentPool.Put(fragment)
				return err
			}
			subSnapshotPool.Put(fragment.Sub)
			fragmentPool.Put(fragment)
		default:
			panic(fmt.Sprintf("unknown snapshot fragment type %s", fragment.FragmentType))
		}
	}
	return c.store.Msgs.Flush()
}
<commit_msg>Fix spelling mistake<commit_after>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"github.com/hashicorp/raft"
	"github.com/nats-io/go-nats-streaming/pb"

	"github.com/nats-io/nats-streaming-server/spb"
)

const batchSize = 200

var (
	fragmentPool    = &sync.Pool{New: func() interface{} { return &spb.RaftSnapshotFragment{} }}
	batchPool       = &sync.Pool{New: func() interface{} { return &spb.Batch{} }}
	subSnapshotPool = &sync.Pool{New: func() interface{} { return &spb.SubSnapshot{} }}
)

// channelSnapshot implements the raft.FSMSnapshot interface.
type channelSnapshot struct {
	*channel
}

func newChannelSnapshot(c *channel) raft.FSMSnapshot {
	return &channelSnapshot{channel: c}
}

// Persist should dump all necessary state to the WriteCloser 'sink',
// and call sink.Close() when finished or call sink.Cancel() on error.
func (c *channelSnapshot) Persist(sink raft.SnapshotSink) (err error) {
	defer func() {
		if err != nil {
			sink.Cancel()
		}
	}()

	if err := c.snapshotMessages(sink); err != nil {
		return err
	}

	if err := c.snapshotSubscriptions(sink); err != nil {
		return err
	}

	if err := c.snapshotAcks(sink); err != nil {
		return err
	}

	return sink.Close()
}

// Release is a no-op.
func (c *channelSnapshot) Release() {}

func (c *channelSnapshot) snapshotMessages(sink raft.SnapshotSink) error {
	// TODO: this is very expensive and might repeatedly fail if messages are
	// constantly being truncated. Is there a way we can optimize this, e.g.
	// handling Restore() out-of-band from Raft?

	first, last, err := c.store.Msgs.FirstAndLastSequence()
	if err != nil {
		return err
	}

	var (
		buf   [4]byte
		batch = batchPool.Get().(*spb.Batch)
	)
	batch.Messages = make([]*pb.MsgProto, 0, batchSize)

	for seq := first; seq <= last; seq++ {
		msg, err := c.store.Msgs.Lookup(seq)
		if err != nil {
			batchPool.Put(batch)
			return err
		}
		// If msg is nil, channel truncation has occurred while snapshotting.
		if msg == nil {
			// Channel truncation has occurred while snapshotting.
			batchPool.Put(batch)
			return fmt.Errorf("channel %q was truncated while snapshotting", c.name)
		}

		// Previous batch is full, ship it.
		if len(batch.Messages) == batchSize {
			fragment := newFragment()
			fragment.FragmentType = spb.RaftSnapshotFragment_Messages
			fragment.MessageBatch = batch
			if err := writeFragment(sink, fragment, buf); err != nil {
				return err
			}

			// Create a new batch.
			batch = batchPool.Get().(*spb.Batch)
			batch.Messages = make([]*pb.MsgProto, 0, batchSize)
		}

		batch.Messages = append(batch.Messages, msg)
	}

	// Ship any partial batch.
	if len(batch.Messages) > 0 {
		fragment := newFragment()
		fragment.FragmentType = spb.RaftSnapshotFragment_Messages
		fragment.MessageBatch = batch
		if err := writeFragment(sink, fragment, buf); err != nil {
			return err
		}
	}

	return nil
}

func (c *channelSnapshot) snapshotSubscriptions(sink raft.SnapshotSink) error {
	var buf [4]byte
	for _, sub := range c.ss.getAllSubs() {
		fragment := newFragment()
		fragment.FragmentType = spb.RaftSnapshotFragment_Subscription
		subSnap := subSnapshotPool.Get().(*spb.SubSnapshot)
		subSnap.Subject = sub.subject
		subSnap.State = &sub.SubState
		fragment.Sub = subSnap
		if err := writeFragment(sink, fragment, buf); err != nil {
			return err
		}
	}
	return nil
}

func (c *channelSnapshot) snapshotAcks(sink raft.SnapshotSink) error {
	// TODO
	return nil
}

// writeFragment writes the marshaled RaftSnapshotFragment to the
// SnapshotSink. The fragment should not be used after this is called.
func writeFragment(sink raft.SnapshotSink, fragment *spb.RaftSnapshotFragment, sizeBuf [4]byte) error {
	data, err := fragment.Marshal()
	if fragment.MessageBatch != nil {
		batchPool.Put(fragment.MessageBatch)
	}
	if fragment.Sub != nil {
		subSnapshotPool.Put(fragment.Sub)
	}
	fragmentPool.Put(fragment)
	if err != nil {
		return err
	}
	binary.BigEndian.PutUint32(sizeBuf[:], uint32(len(data)))
	_, err = sink.Write(sizeBuf[:])
	if err != nil {
		return err
	}
	_, err = sink.Write(data)
	return err
}

func newFragment() *spb.RaftSnapshotFragment {
	fragment := fragmentPool.Get().(*spb.RaftSnapshotFragment)
	fragment.MessageBatch = nil
	return fragment
}

// restoreFromSnapshot restores a channel from a snapshot. This is not called
// concurrently with any other Raft commands.
func (c *channel) restoreFromSnapshot(snapshot io.ReadCloser) error {
	defer snapshot.Close()

	// TODO: this needs to be fully implemented to support restoring the
	// entire Raft log, not just channel messages.
	// TODO: restore acks

	// Drop all existing subs. These will be restored from the snapshot.
	for _, sub := range c.ss.getAllSubs() {
		c.stan.unsubscribe(c, sub.ClientID, sub, true)
	}

	sizeBuf := make([]byte, 4)
	for {
		// Read the fragment size.
		if _, err := io.ReadFull(snapshot, sizeBuf); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		// Read the fragment.
		size := binary.BigEndian.Uint32(sizeBuf)
		buf := make([]byte, size)
		if _, err := io.ReadFull(snapshot, buf); err != nil {
			return err
		}

		// Unmarshal the fragment.
		fragment := newFragment()
		if err := fragment.Unmarshal(buf); err != nil {
			return err
		}

		// Apply the fragment.
		switch fragment.FragmentType {
		case spb.RaftSnapshotFragment_Messages:
			// Channel messages.
			for _, msg := range fragment.MessageBatch.Messages {
				if _, err := c.store.Msgs.Store(msg); err != nil {
					batchPool.Put(fragment.MessageBatch)
					fragmentPool.Put(fragment)
					return err
				}
			}
			batchPool.Put(fragment.MessageBatch)
			fragmentPool.Put(fragment)
		case spb.RaftSnapshotFragment_Subscription:
			// Channel subscription.
			sub := &subState{
				SubState:    *fragment.Sub.State,
				subject:     fragment.Sub.Subject,
				ackWait:     computeAckWait(fragment.Sub.State.AckWaitInSecs),
				acksPending: make(map[uint64]int64),
				store:       c.store.Subs,
			}
			if err := c.stan.addSubscription(c.ss, sub); err != nil {
				subSnapshotPool.Put(fragment.Sub)
				fragmentPool.Put(fragment)
				return err
			}
			subSnapshotPool.Put(fragment.Sub)
			fragmentPool.Put(fragment)
		default:
			panic(fmt.Sprintf("unknown snapshot fragment type %s", fragment.FragmentType))
		}
	}
	return c.store.Msgs.Flush()
}
<|endoftext|>


<commit_before>// Copyright 2010 Jeremy Wall (jeremy@marzhillstudios.com)
// Use of this source code is governed by the Artistic License 2.0.
// That License is included in the LICENSE file.

package transform

import (
	"code.google.com/p/go-html-transform/css/selector"
	"code.google.com/p/go-html-transform/h5"

	"exp/html"
	"log"
)

// Collector defines an interface for html node collectors.
type Collector interface {
	// Find searches a tree rooted at n and returns a slice of nodes
	// that match a criteria.
	Find(n *html.Node) []*html.Node
}

// The TransformFunc type is the type of a html.Node transformation function.
type TransformFunc func(*html.Node)

// Transformer encapsulates a document under transformation.
type Transformer struct {
	doc h5.Tree
}

// Constructor for a Transformer. It makes a copy of the document
// and transforms that instead of the original.
func NewTransformer(t h5.Tree) *Transformer {
	return newTransformer(t.Clone())
}

func newTransformer(t h5.Tree) *Transformer {
	return &Transformer{doc: t}
}

// The Doc method returns the document under transformation.
func (t *Transformer) Doc() *html.Node {
	return t.doc.Top()
}

func (t *Transformer) String() string {
	return t.doc.String()
}

func (t *Transformer) Clone() *Transformer {
	return NewTransformer(t.doc)
}

func applyFuncToCollector(f TransformFunc, n *html.Node, sel Collector) {
	for _, nn := range sel.Find(n) {
		f(nn)
	}
}

// The ApplyWithSelector method applies a TransformFunc to the nodes matched
// by the CSS3 Selector.
func (t *Transformer) Apply(f TransformFunc, sel string) error {
	sq, err := selector.Selector(sel)
	t.ApplyWithCollector(f, sq)
	return err
}

func (t *Transformer) ApplyWithCollector(f TransformFunc, coll Collector) {
	// TODO come up with a way to walk tree once?
	applyFuncToCollector(f, t.Doc(), coll)
}

// Transform is a bundle of selectors and a transform func. It forms a
// self contained Transfrom on an html document that can be reused.
type Transform struct {
	coll Collector
	f    TransformFunc
}

// Trans creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a valid CSS3 Selector.
// It returns a *Transform or an error if the selector wasn't valid
func Trans(f TransformFunc, sel string) (*Transform, error) {
	sq, err := selector.Selector(sel)
	return TransCollector(f, sq), err
}

// TransCollector creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a Collector
func TransCollector(f TransformFunc, coll Collector) *Transform {
	return &Transform{f: f, coll: coll}
}

// ApplyAll applies a series of Transforms to a document.
//     t.ApplyAll(Trans(f, sel1, sel2), Trans(f2, sel3, sel4))
func (t *Transformer) ApplyAll(ts ...*Transform) {
	for _, spec := range ts {
		t.ApplyWithCollector(spec.f, spec.coll)
	}
}

// AppendChildren creates a TransformFunc that appends the Children passed in.
func AppendChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			if c.Parent != nil {
				c.Parent.RemoveChild(c)
			}
			n.AppendChild(c)
		}
	}
}

// PrependChildren creates a TransformFunc that prepends the Children passed in.
func PrependChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			n.InsertBefore(c, n.FirstChild)
		}
	}
}

// RemoveChildren creates a TransformFunc that removes the Children of the node
// it operates on.
func RemoveChildren() TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
	}
}

func removeChildren(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		defer n.RemoveChild(c)
	}
}

// ReplaceChildren creates a TransformFunc that replaces the Children of the
// node it operates on with the Children passed in.
func ReplaceChildren(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
		for _, c := range ns {
			n.AppendChild(c)
		}
	}
}

func nodeToString(n *html.Node) string {
	t := h5.NewTree(n)
	return t.String()
}

// Replace constructs a TransformFunc that replaces a node with the nodes passed
// in.
func Replace(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		p := n.Parent
		switch p {
		case nil:
			log.Panicf("Attempt to replace Root node: %s", h5.RenderNodesToString([]*html.Node{n}))
		default:
			for _, nc := range ns {
				p.InsertBefore(nc, n)
			}
			p.RemoveChild(n)
		}
	}
}

// DoAll returns a TransformFunc that combines all the TransformFuncs that are
// passed in. Doing each transform in order.
func DoAll(fs ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, f := range fs {
			f(n)
		}
	}
}

// CopyAnd will construct a TransformFunc that will
// make a copy of the node for each passed in TransformFunc
// and replace the passed in node with the resulting transformed
// html.Nodes.
func CopyAnd(fns ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, fn := range fns {
			node := h5.CloneNode(n)
			n.Parent.InsertBefore(node, n)
			fn(node)
		}
		n.Parent.RemoveChild(n)
	}
}

// SubTransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func SubTransform(f TransformFunc, sel string) (TransformFunc, error) {
	sq, err := selector.Selector(sel)
	return SubTransformCollector(f, sq), err
}

// SubTransformSelector constructs a TransformFunc that runs a TransformFunc on
// any nodes collected, using the passed in collector, from the subtree the
// TransformFunc is run on.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func SubTransformCollector(f TransformFunc, coll Collector) TransformFunc {
	return func(n *html.Node) {
		applyFuncToCollector(f, n, coll)
	}
}

// ModifyAttrb creates a TransformFunc that modifies the attributes
// of the node it operates on. If an Attribute with the same name
// as the key doesn't exist it creates it.
func ModifyAttrib(key string, val string) TransformFunc {
	return func(n *html.Node) {
		found := false
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = val
				found = true
			}
		}
		if !found {
			n.Attr = append(n.Attr, html.Attribute{Key: key, Val: val})
		}
	}
}

// TransformAttrib returns a TransformFunc that transforms an attribute on
// the node it operates on using the provided func. It only transforms
// the attribute if it exists.
func TransformAttrib(key string, f func(string) string) TransformFunc {
	return func(n *html.Node) {
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = f(n.Attr[i].Val)
			}
		}
	}
}

// Trace is a debugging wrapper for transform funcs.
// It prints debugging information before and after the TransformFunc
// is applied.
func Trace(f TransformFunc, msg string, args ...interface{}) TransformFunc {
	return func(n *html.Node) {
		log.Printf("TRACE: "+msg, args...)
		p := n.Parent
		if p == nil {
			p = n
		}
		log.Printf("TRACE: Before: %s", h5.NewTree(p).String())
		f(n)
		log.Printf("TRACE: After: %s", h5.NewTree(p).String())
	}
}
<commit_msg>Add Must* versions of functions that can return errors.<commit_after>// Copyright 2010 Jeremy Wall (jeremy@marzhillstudios.com)
// Use of this source code is governed by the Artistic License 2.0.
// That License is included in the LICENSE file.

package transform

import (
	"code.google.com/p/go-html-transform/css/selector"
	"code.google.com/p/go-html-transform/h5"

	"exp/html"
	"log"
)

// Collector defines an interface for html node collectors.
type Collector interface {
	// Find searches a tree rooted at n and returns a slice of nodes
	// that match a criteria.
	Find(n *html.Node) []*html.Node
}

// The TransformFunc type is the type of a html.Node transformation function.
type TransformFunc func(*html.Node)

// Transformer encapsulates a document under transformation.
type Transformer struct {
	doc h5.Tree
}

// Constructor for a Transformer. It makes a copy of the document
// and transforms that instead of the original.
func NewTransformer(t h5.Tree) *Transformer {
	return newTransformer(t.Clone())
}

func newTransformer(t h5.Tree) *Transformer {
	return &Transformer{doc: t}
}

// The Doc method returns the document under transformation.
func (t *Transformer) Doc() *html.Node {
	return t.doc.Top()
}

func (t *Transformer) String() string {
	return t.doc.String()
}

func (t *Transformer) Clone() *Transformer {
	return NewTransformer(t.doc)
}

func applyFuncToCollector(f TransformFunc, n *html.Node, sel Collector) {
	for _, nn := range sel.Find(n) {
		f(nn)
	}
}

// The ApplyWithSelector method applies a TransformFunc to the nodes matched
// by the CSS3 Selector.
func (t *Transformer) Apply(f TransformFunc, sel string) error {
	sq, err := selector.Selector(sel)
	t.ApplyWithCollector(f, sq)
	return err
}

// ApplyWithCollector applies a TransformFunc to the tree using a Collector.
func (t *Transformer) ApplyWithCollector(f TransformFunc, coll Collector) {
	// TODO come up with a way to walk tree once?
	applyFuncToCollector(f, t.Doc(), coll)
}

// Transform is a bundle of selectors and a transform func. It forms a
// self contained Transfrom on an html document that can be reused.
type Transform struct {
	coll Collector
	f    TransformFunc
}

// Trans creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a valid CSS3 Selector.
// It returns a *Transform or an error if the selector wasn't valid
func Trans(f TransformFunc, sel string) (*Transform, error) {
	sq, err := selector.Selector(sel)
	return TransCollector(f, sq), err
}

// MustTrans creates a Transform.
// Panics if the selector wasn't valid.
func MustTrans(f TransformFunc, sel string) *Transform {
	t, err := Trans(f, sel)
	if err != nil { panic(err) }
	return t
}

// TransCollector creates a Transform that you can apply using ApplyAll.
// It takes a TransformFunc and a Collector
func TransCollector(f TransformFunc, coll Collector) *Transform {
	return &Transform{f: f, coll: coll}
}

// ApplyAll applies a series of Transforms to a document.
//     t.ApplyAll(Trans(f, sel1, sel2), Trans(f2, sel3, sel4))
func (t *Transformer) ApplyAll(ts ...*Transform) {
	for _, spec := range ts {
		t.ApplyWithCollector(spec.f, spec.coll)
	}
}

// AppendChildren creates a TransformFunc that appends the Children passed in.
func AppendChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			if c.Parent != nil {
				c.Parent.RemoveChild(c)
			}
			n.AppendChild(c)
		}
	}
}

// PrependChildren creates a TransformFunc that prepends the Children passed in.
func PrependChildren(cs ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		for _, c := range cs {
			n.InsertBefore(c, n.FirstChild)
		}
	}
}

// RemoveChildren creates a TransformFunc that removes the Children of the node
// it operates on.
func RemoveChildren() TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
	}
}

func removeChildren(n *html.Node) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		defer n.RemoveChild(c)
	}
}

// ReplaceChildren creates a TransformFunc that replaces the Children of the
// node it operates on with the Children passed in.
func ReplaceChildren(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		removeChildren(n)
		for _, c := range ns {
			n.AppendChild(c)
		}
	}
}

func nodeToString(n *html.Node) string {
	t := h5.NewTree(n)
	return t.String()
}

// Replace constructs a TransformFunc that replaces a node with the nodes passed
// in.
func Replace(ns ...*html.Node) TransformFunc {
	return func(n *html.Node) {
		p := n.Parent
		switch p {
		case nil:
			log.Panicf("Attempt to replace Root node: %s", h5.RenderNodesToString([]*html.Node{n}))
		default:
			for _, nc := range ns {
				p.InsertBefore(nc, n)
			}
			p.RemoveChild(n)
		}
	}
}

// DoAll returns a TransformFunc that combines all the TransformFuncs that are
// passed in. Doing each transform in order.
func DoAll(fs ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, f := range fs {
			f(n)
		}
	}
}

// CopyAnd will construct a TransformFunc that will
// make a copy of the node for each passed in TransformFunc
// and replace the passed in node with the resulting transformed
// html.Nodes.
func CopyAnd(fns ...TransformFunc) TransformFunc {
	return func(n *html.Node) {
		for _, fn := range fns {
			node := h5.CloneNode(n)
			n.Parent.InsertBefore(node, n)
			fn(node)
		}
		n.Parent.RemoveChild(n)
	}
}

// SubTransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func Subtransform(f TransformFunc, sel string) (TransformFunc, error) {
	sq, err := selector.Selector(sel)
	return SubtransformCollector(f, sq), err
}

// SubTransform constructs a TransformFunc that runs a TransformFunc on
// any nodes in the tree rooted by the node the the TransformFunc is run
// against.
// Panics if the selector string is malformed.
func MustSubtransform(f TransformFunc, sel string) TransformFunc {
	t, err := Subtransform(f, sel)
	if err != nil { panic(err) }
	return t
}

// SubTransformSelector constructs a TransformFunc that runs a TransformFunc on
// any nodes collected, using the passed in collector, from the subtree the
// TransformFunc is run on.
// This is useful for creating self contained Transforms that are
// meant to work on subtrees of the html document.
func SubtransformCollector(f TransformFunc, coll Collector) TransformFunc {
	return func(n *html.Node) {
		applyFuncToCollector(f, n, coll)
	}
}

// ModifyAttrb creates a TransformFunc that modifies the attributes
// of the node it operates on. If an Attribute with the same name
// as the key doesn't exist it creates it.
func ModifyAttrib(key string, val string) TransformFunc {
	return func(n *html.Node) {
		found := false
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = val
				found = true
			}
		}
		if !found {
			n.Attr = append(n.Attr, html.Attribute{Key: key, Val: val})
		}
	}
}

// TransformAttrib returns a TransformFunc that transforms an attribute on
// the node it operates on using the provided func. It only transforms
// the attribute if it exists.
func TransformAttrib(key string, f func(string) string) TransformFunc {
	return func(n *html.Node) {
		for i, attr := range n.Attr {
			if attr.Key == key {
				n.Attr[i].Val = f(n.Attr[i].Val)
			}
		}
	}
}

// Trace is a debugging wrapper for transform funcs.
// It prints debugging information before and after the TransformFunc
// is applied.
func Trace(f TransformFunc, msg string, args ...interface{}) TransformFunc {
	return func(n *html.Node) {
		log.Printf("TRACE: "+msg, args...)
		p := n.Parent
		if p == nil {
			p = n
		}
		log.Printf("TRACE: Before: %s", h5.NewTree(p).String())
		f(n)
		log.Printf("TRACE: After: %s", h5.NewTree(p).String())
	}
}
<|endoftext|>


<commit_before>package helpers

import (
	"context"
	"fmt"
	"os"

	chunk "github.com/ipfs/go-ipfs/importer/chunk"
	dag "github.com/ipfs/go-ipfs/merkledag"
	pi "github.com/ipfs/go-ipfs/thirdparty/posinfo"
	ft "github.com/ipfs/go-ipfs/unixfs"

	node "gx/ipfs/QmZx42H5khbVQhV5odp66TApShV4XCujYazcvYduZ4TroB/go-ipld-node"
)

// BlockSizeLimit specifies the maximum size an imported block can have.
var BlockSizeLimit = 1048576 // 1 MB

// rough estimates on expected sizes
var roughDataBlockSize = chunk.DefaultBlockSize
var roughLinkBlockSize = 1 << 13 // 8KB
var roughLinkSize = 34 + 8 + 5   // sha256 multihash + size + no name + protobuf framing

// DefaultLinksPerBlock governs how the importer decides how many links there
// will be per block. This calculation is based on expected distributions of:
//  * the expected distribution of block sizes
//  * the expected distribution of link sizes
//  * desired access speed
// For now, we use:
//
//   var roughLinkBlockSize = 1 << 13 // 8KB
//   var roughLinkSize = 288          // sha256 + framing + name
//   var DefaultLinksPerBlock = (roughLinkBlockSize / roughLinkSize)
//
// See calc_test.go
var DefaultLinksPerBlock = (roughLinkBlockSize / roughLinkSize)

// ErrSizeLimitExceeded signals that a block is larger than BlockSizeLimit.
var ErrSizeLimitExceeded = fmt.Errorf("object size limit exceeded")

// UnixfsNode is a struct created to aid in the generation
// of unixfs DAG trees
type UnixfsNode struct {
	raw     bool
	rawnode *dag.RawNode
	node    *dag.ProtoNode
	ufmt    *ft.FSNode
	posInfo *pi.PosInfo
}

// NewUnixfsNode creates a new Unixfs node to represent a file
func NewUnixfsNode() *UnixfsNode {
	return &UnixfsNode{
		node: new(dag.ProtoNode),
		ufmt: &ft.FSNode{Type: ft.TFile},
	}
}

// NewUnixfsBlock creates a new Unixfs node to represent a raw data block
func NewUnixfsBlock() *UnixfsNode {
	return &UnixfsNode{
		node: new(dag.ProtoNode),
		ufmt: &ft.FSNode{Type: ft.TRaw},
	}
}

// NewUnixfsNodeFromDag reconstructs a Unixfs node from a given dag node
func NewUnixfsNodeFromDag(nd *dag.ProtoNode) (*UnixfsNode, error) {
	mb, err := ft.FSNodeFromBytes(nd.Data())
	if err != nil {
		return nil, err
	}

	return &UnixfsNode{
		node: nd,
		ufmt: mb,
	}, nil
}

func (n *UnixfsNode) NumChildren() int {
	return n.ufmt.NumChildren()
}

func (n *UnixfsNode) Set(other *UnixfsNode) {
	n.node = other.node
	n.raw = other.raw
	n.rawnode = other.rawnode
	if other.ufmt != nil {
		n.ufmt.Data = other.ufmt.Data
	}
}

func (n *UnixfsNode) GetChild(ctx context.Context, i int, ds dag.DAGService) (*UnixfsNode, error) {
	nd, err := n.node.Links()[i].GetNode(ctx, ds)
	if err != nil {
		return nil, err
	}

	pbn, ok := nd.(*dag.ProtoNode)
	if !ok {
		return nil, dag.ErrNotProtobuf
	}

	return NewUnixfsNodeFromDag(pbn)
}

// addChild will add the given UnixfsNode as a child of the receiver.
// the passed in DagBuilderHelper is used to store the child node an
// pin it locally so it doesnt get lost
func (n *UnixfsNode) AddChild(child *UnixfsNode, db *DagBuilderHelper) error {
	n.ufmt.AddBlockSize(child.DataSize())

	childnode, err := child.GetDagNode()
	if err != nil {
		return err
	}

	// Add a link to this node without storing a reference to the memory
	// This way, we avoid nodes building up and consuming all of our RAM
	err = n.node.AddNodeLinkClean("", childnode)
	if err != nil {
		return err
	}

	_, err = db.batch.Add(childnode)
	if err != nil {
		return err
	}

	return nil
}

// Removes the child node at the given index
func (n *UnixfsNode) RemoveChild(index int, dbh *DagBuilderHelper) {
	n.ufmt.RemoveBlockSize(index)
	n.node.SetLinks(append(n.node.Links()[:index], n.node.Links()[index+1:]...))
}

func (n *UnixfsNode) FileSize() uint64 {
	return n.ufmt.FileSize()
}

func (n *UnixfsNode) SetData(data []byte) {
	n.ufmt.Data = data
}

func (n *UnixfsNode) DataSize() uint64 {
	if n.raw {
		return uint64(len(n.rawnode.RawData()))
	}
	return n.ufmt.FileSize()
}

func (n *UnixfsNode) SetPosInfo(offset uint64, fullPath string, stat os.FileInfo) {
	n.posInfo = &pi.PosInfo{offset, fullPath, stat}
}

// getDagNode fills out the proper formatting for the unixfs node
// inside of a DAG node and returns the dag node
func (n *UnixfsNode) GetDagNode() (node.Node, error) {
	nd, err := n.getBaseDagNode()
	if err != nil {
		return nil, err
	}

	if n.posInfo != nil {
		return &pi.FilestoreNode{
			Node:    nd,
			PosInfo: n.posInfo,
		}, nil
	}

	return nd, nil
}

func (n *UnixfsNode) getBaseDagNode() (node.Node, error) {
	if n.raw {
		return n.rawnode, nil
	}

	data, err := n.ufmt.GetBytes()
	if err != nil {
		return nil, err
	}
	n.node.SetData(data)
	return n.node, nil
}
<commit_msg>Filestore: Remove unused FileSize() method.<commit_after>package helpers

import (
	"context"
	"fmt"
	"os"

	chunk "github.com/ipfs/go-ipfs/importer/chunk"
	dag "github.com/ipfs/go-ipfs/merkledag"
	pi "github.com/ipfs/go-ipfs/thirdparty/posinfo"
	ft "github.com/ipfs/go-ipfs/unixfs"

	node "gx/ipfs/QmZx42H5khbVQhV5odp66TApShV4XCujYazcvYduZ4TroB/go-ipld-node"
)

// BlockSizeLimit specifies the maximum size an imported block can have.
var BlockSizeLimit = 1048576 // 1 MB

// rough estimates on expected sizes
var roughDataBlockSize = chunk.DefaultBlockSize
var roughLinkBlockSize = 1 << 13 // 8KB
var roughLinkSize = 34 + 8 + 5   // sha256 multihash + size + no name + protobuf framing

// DefaultLinksPerBlock governs how the importer decides how many links there
// will be per block. This calculation is based on expected distributions of:
//  * the expected distribution of block sizes
//  * the expected distribution of link sizes
//  * desired access speed
// For now, we use:
//
//   var roughLinkBlockSize = 1 << 13 // 8KB
//   var roughLinkSize = 288          // sha256 + framing + name
//   var DefaultLinksPerBlock = (roughLinkBlockSize / roughLinkSize)
//
// See calc_test.go
var DefaultLinksPerBlock = (roughLinkBlockSize / roughLinkSize)

// ErrSizeLimitExceeded signals that a block is larger than BlockSizeLimit.
var ErrSizeLimitExceeded = fmt.Errorf("object size limit exceeded")

// UnixfsNode is a struct created to aid in the generation
// of unixfs DAG trees
type UnixfsNode struct {
	raw     bool
	rawnode *dag.RawNode
	node    *dag.ProtoNode
	ufmt    *ft.FSNode
	posInfo *pi.PosInfo
}

// NewUnixfsNode creates a new Unixfs node to represent a file
func NewUnixfsNode() *UnixfsNode {
	return &UnixfsNode{
		node: new(dag.ProtoNode),
		ufmt: &ft.FSNode{Type: ft.TFile},
	}
}

// NewUnixfsBlock creates a new Unixfs node to represent a raw data block
func NewUnixfsBlock() *UnixfsNode {
	return &UnixfsNode{
		node: new(dag.ProtoNode),
		ufmt: &ft.FSNode{Type: ft.TRaw},
	}
}

// NewUnixfsNodeFromDag reconstructs a Unixfs node from a given dag node
func NewUnixfsNodeFromDag(nd *dag.ProtoNode) (*UnixfsNode, error) {
	mb, err := ft.FSNodeFromBytes(nd.Data())
	if err != nil {
		return nil, err
	}

	return &UnixfsNode{
		node: nd,
		ufmt: mb,
	}, nil
}

func (n *UnixfsNode) NumChildren() int {
	return n.ufmt.NumChildren()
}

func (n *UnixfsNode) Set(other *UnixfsNode) {
	n.node = other.node
	n.raw = other.raw
	n.rawnode = other.rawnode
	if other.ufmt != nil {
		n.ufmt.Data = other.ufmt.Data
	}
}

func (n *UnixfsNode) GetChild(ctx context.Context, i int, ds dag.DAGService) (*UnixfsNode, error) {
	nd, err := n.node.Links()[i].GetNode(ctx, ds)
	if err != nil {
		return nil, err
	}

	pbn, ok := nd.(*dag.ProtoNode)
	if !ok {
		return nil, dag.ErrNotProtobuf
	}

	return NewUnixfsNodeFromDag(pbn)
}

// addChild will add the given UnixfsNode as a child of the receiver.
// the passed in DagBuilderHelper is used to store the child node an
// pin it locally so it doesnt get lost
func (n *UnixfsNode) AddChild(child *UnixfsNode, db *DagBuilderHelper) error {
	n.ufmt.AddBlockSize(child.DataSize())

	childnode, err := child.GetDagNode()
	if err != nil {
		return err
	}

	// Add a link to this node without storing a reference to the memory
	// This way, we avoid nodes building up and consuming all of our RAM
	err = n.node.AddNodeLinkClean("", childnode)
	if err != nil {
		return err
	}

	_, err = db.batch.Add(childnode)
	if err != nil {
		return err
	}

	return nil
}

// Removes the child node at the given index
func (n *UnixfsNode) RemoveChild(index int, dbh *DagBuilderHelper) {
	n.ufmt.RemoveBlockSize(index)
	n.node.SetLinks(append(n.node.Links()[:index], n.node.Links()[index+1:]...))
}

func (n *UnixfsNode) SetData(data []byte) {
	n.ufmt.Data = data
}

func (n *UnixfsNode) DataSize() uint64 {
	if n.raw {
		return uint64(len(n.rawnode.RawData()))
	}
	return n.ufmt.FileSize()
}

func (n *UnixfsNode) SetPosInfo(offset uint64, fullPath string, stat os.FileInfo) {
	n.posInfo = &pi.PosInfo{offset, fullPath, stat}
}

// getDagNode fills out the proper formatting for the unixfs node
// inside of a DAG node and returns the dag node
func (n *UnixfsNode) GetDagNode() (node.Node, error) {
	nd, err := n.getBaseDagNode()
	if err != nil {
		return nil, err
	}

	if n.posInfo != nil {
		return &pi.FilestoreNode{
			Node:    nd,
			PosInfo: n.posInfo,
		}, nil
	}

	return nd, nil
}

func (n *UnixfsNode) getBaseDagNode() (node.Node, error) {
	if n.raw {
		return n.rawnode, nil
	}

	data, err := n.ufmt.GetBytes()
	if err != nil {
		return nil, err
	}
	n.node.SetData(data)
	return n.node, nil
}
<|endoftext|>


<commit_before>package service

import (
	"database/sql"
	"log"
	"net"
	"os"
	"time"

	pb "github.com/otoolep/go-grpc-pg/proto"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Service represents a gRPC service that communicates with a database backend.
type Service struct {
	grpc *grpc.Server
	db   *sql.DB

	ln   net.Listener
	addr string

	logger *log.Logger
}

// New returns an instantiated service.
func New(addr string, db *sql.DB) *Service {
	s := Service{
		grpc:   grpc.NewServer(),
		db:     db,
		addr:   addr,
		logger: log.New(os.Stderr, "[service] ", log.LstdFlags),
	}

	pb.RegisterDBProviderServer(s.grpc, (*gprcService)(&s))
	return &s
}

// Addr returns the address on which the service is listening.
func (s *Service) Addr() string {
	return s.ln.Addr().String()
}

// Open opens the service, starting it listening on the configured address.
func (s *Service) Open() error {
	ln, err := net.Listen("tcp", s.addr)
	if err != nil {
		return err
	}
	s.ln = ln
	s.logger.Println("listening on", s.ln.Addr().String())

	go func() {
		err := s.grpc.Serve(s.ln)
		if err != nil {
			s.logger.Println("gRPC Serve() returned:", err.Error())
		}
	}()

	return nil
}

// Close closes the service.
func (s *Service) Close() error {
	s.grpc.GracefulStop()
	s.ln = nil
	s.logger.Println("gRPC server stopped")
	return nil
}

// gprcService is an unexported type, that is the same type as Service.
//
// Having the methods that the gRPC service requires on this type means that even though
// the methods are exported, since the type is not, these methods are not visible outside
// this package.
type gprcService Service

// Query implements the Query interface of the gRPC service.
func (g *gprcService) Query(c context.Context, q *pb.QueryRequest) (*pb.QueryResponse, error) {
	start := time.Now()
	rows, err := g.db.Query(q.Stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get the column names.
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	response := pb.QueryResponse{
		Columns: cols,
	}

	// Iterate through each row returned by the query.
	for rows.Next() {
		row := make([]string, len(cols))
		// Get a set of pointers to the strings allocated above.
		rowI := make([]interface{}, len(cols))
		for i := range row {
			rowI[i] = &row[i]
		}

		if err := rows.Scan(rowI...); err != nil {
			return nil, err
		}

		// Add the latest rows to existing rows.
		response.Rows = append(response.Rows, &pb.Row{Values: row})
	}

	g.logger.Printf(`query '%s' completed in %s, %d %s returned`,
		q.Stmt, time.Since(start), len(response.Rows), prettyRows(int64(len(response.Rows))))
	return &response, nil
}

// Exec implements the Exec interface of the gRPC service.
func (g *gprcService) Exec(c context.Context, e *pb.ExecRequest) (*pb.ExecResponse, error) {
	start := time.Now()
	r, err := g.db.Exec(e.Stmt)
	if err != nil {
		return nil, err
	}

	lid, err := r.LastInsertId()
	if err != nil {
		// Not all databases support LastInsertId()
		lid = -1
	}
	ra, err := r.RowsAffected()
	if err != nil {
		return nil, err
	}

	g.logger.Printf(`exec '%s' completed in %s, %d %s affected`,
		e.Stmt, time.Since(start), ra, prettyRows(ra))
	return &pb.ExecResponse{
		LastInsertId: lid,
		RowsAffected: ra,
	}, nil
}

func prettyRows(n int64) string {
	if n == 1 {
		return "row"
	}
	return "rows"
}
<commit_msg>Update service.go<commit_after>package service

import (
	"database/sql"
	"log"
	"net"
	"os"
	"time"

	pb "github.com/otoolep/go-grpc-pg/proto"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Service represents a gRPC service that communicates with a database backend.
type Service struct {
	grpc *grpc.Server
	db   *sql.DB

	ln   net.Listener
	addr string

	logger *log.Logger
}

// New returns an instantiated service.
func New(addr string, db *sql.DB) *Service {
	s := Service{
		grpc:   grpc.NewServer(),
		db:     db,
		addr:   addr,
		logger: log.New(os.Stderr, "[service] ", log.LstdFlags),
	}

	pb.RegisterDBProviderServer(s.grpc, (*gprcService)(&s))
	return &s
}

// Addr returns the address on which the service is listening.
func (s *Service) Addr() string {
	return s.ln.Addr().String()
}

// Open opens the service, starting it listening on the configured address.
func (s *Service) Open() error {
	ln, err := net.Listen("tcp", s.addr)
	if err != nil {
		return err
	}
	s.ln = ln
	s.logger.Println("listening on", s.ln.Addr().String())

	go func() {
		err := s.grpc.Serve(s.ln)
		if err != nil {
			s.logger.Println("gRPC Serve() returned:", err.Error())
		}
	}()

	return nil
}

// Close closes the service.
func (s *Service) Close() error {
	s.grpc.GracefulStop()
	s.ln = nil
	s.logger.Println("gRPC server stopped")
	return nil
}

// gprcService is an unexported type, that is the same type as Service.
//
// Having the methods that the gRPC service requires on this type means that even though
// the methods are exported, since the type is not, these methods are not visible outside
// this package.
type gprcService Service

// Query implements the Query interface of the gRPC service.
func (g *gprcService) Query(c context.Context, q *pb.QueryRequest) (*pb.QueryResponse, error) {
	start := time.Now()
	rows, err := g.db.Query(q.Stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get the column names.
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	response := pb.QueryResponse{
		Columns: cols,
	}

	// Iterate through each row returned by the query.
	for rows.Next() {
		row := make([]string, len(cols))
		// Get a set of pointers to the strings allocated above.
		rowI := make([]interface{}, len(cols))
		for i := range row {
			rowI[i] = &row[i]
		}

		if err := rows.Scan(rowI...); err != nil {
			return nil, err
		}

		// Add the latest rows to existing rows.
		response.Rows = append(response.Rows, &pb.Row{Values: row})
	}

	g.logger.Printf(`query '%s' completed in %s, %d %s returned`,
		q.Stmt, time.Since(start), len(response.Rows), prettyRows(int64(len(response.Rows))))
	return &response, nil
}

// Exec implements the Exec interface of the gRPC service.
func (g *gprcService) Exec(c context.Context, e *pb.ExecRequest) (*pb.ExecResponse, error) {
	start := time.Now()
	r, err := g.db.Exec(e.Stmt)
	if err != nil {
		return nil, err
	}

	lid, err := r.LastInsertId()
	if err != nil {
		// Not all databases support LastInsertId()
		lid = -1
	}
	ra, err := r.RowsAffected()
	if err != nil {
		return nil, err
	}

	g.logger.Printf(`exec '%s' completed in %s, %d %s affected`,
		e.Stmt, time.Since(start), ra, prettyRows(ra))
	return &pb.ExecResponse{
		LastInsertId: lid,
		RowsAffected: ra,
	}, nil
}

// prettyRows returns a singular or plural form of "row", depending on n.
func prettyRows(n int64) string {
	if n == 1 {
		return "row"
	}
	return "rows"
}
<|endoftext|>


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

package shadowsocks

import (
	"errors"
	"fmt"
	"net"
	"runtime/debug"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	logging "github.com/op/go-logging"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"sync"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

const udpBufSize = 64 * 1024

// Wrapper for logger.Debugf during UDP access key searches.
func debugUDP(cipherID, template string, val interface{}) {
	// This is an optimization to reduce unnecessary allocations due to an interaction
	// between Go's inlining/escape analysis and varargs functions like logger.Debugf.
	if logger.IsEnabledFor(logging.DEBUG) {
		logger.Debugf("UDP(%s): "+template, cipherID, val)
	}
}

// upack decrypts src into dst. It tries each cipher until it finds one that authenticates
// correctly. dst and src must not overlap.
func unpack(clientIP net.IP, dst, src []byte, cipherList CipherList) ([]byte, string, shadowaead.Cipher, error) {
	// Try each cipher until we find one that authenticates successfully. This assumes that all ciphers are AEAD.
	// We snapshot the list because it may be modified while we use it.
	for ci, entry := range cipherList.SnapshotForClientIP(clientIP) {
		id, cipher := entry.Value.(*CipherEntry).ID, entry.Value.(*CipherEntry).Cipher
		buf, err := shadowaead.Unpack(dst, src, cipher)
		if err != nil {
			debugUDP(id, "Failed to unpack: %v", err)
			continue
		}
		debugUDP(id, "Found cipher at index %d", ci)
		// Move the active cipher to the front, so that the search is quicker next time.
		cipherList.MarkUsedByClientIP(entry, clientIP)
		return buf, id, cipher, nil
	}
	return nil, "", nil, errors.New("could not find valid cipher")
}

type udpService struct {
	clientConn     net.PacketConn
	natTimeout     time.Duration
	ciphers        CipherList
	m              metrics.ShadowsocksMetrics
	isRunning      bool
	checkAllowedIP func(net.IP) *onet.ConnectionError
}

// NewUDPService creates a UDPService
func NewUDPService(clientConn net.PacketConn, natTimeout time.Duration, cipherList CipherList, m metrics.ShadowsocksMetrics) UDPService {
	return &udpService{clientConn: clientConn, natTimeout: natTimeout, ciphers: cipherList, m: m, checkAllowedIP: onet.RequirePublicIP}
}

// UDPService is a UDP shadowsocks service that can be started and stopped.
type UDPService interface {
	Start()
	Stop() error
}

// Listen on addr for encrypted packets and basically do UDP NAT.
// We take the ciphers as a pointer because it gets replaced on config updates.
func (s *udpService) Start() {
	defer s.clientConn.Close()

	nm := newNATmap(s.natTimeout, s.m)
	cipherBuf := make([]byte, udpBufSize)
	textBuf := make([]byte, udpBufSize)

	s.isRunning = true
	for s.isRunning {
		func() (connError *onet.ConnectionError) {
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in UDP loop: %v", r)
					debug.PrintStack()
				}
			}()
			clientLocation := ""
			keyID := ""
			var clientProxyBytes, proxyTargetBytes int
			var timeToCipher time.Duration
			defer func() {
				status := "OK"
				if connError != nil {
					logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				s.m.AddUDPPacketFromClient(clientLocation, keyID, status, clientProxyBytes, proxyTargetBytes, timeToCipher)
			}()
			clientProxyBytes, clientAddr, err := s.clientConn.ReadFrom(cipherBuf)
			if err != nil {
				if !s.isRunning {
					return nil
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from client", err)
			}
			clientLocation, locErr := s.m.GetLocation(clientAddr)
			if locErr != nil {
				logger.Warningf("Failed location lookup: %v", locErr)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientAddr.String())
			defer logger.Debugf("UDP done with %v", clientAddr.String())
			logger.Debugf("UDP Request from %v with %v bytes", clientAddr, clientProxyBytes)
			unpackStart := time.Now()
			ip := clientAddr.(*net.UDPAddr).IP
			buf, keyID, cipher, err := unpack(ip, textBuf, cipherBuf[:clientProxyBytes], s.ciphers)
			timeToCipher = time.Now().Sub(unpackStart)

			if err != nil {
				return onet.NewConnectionError("ERR_CIPHER", "Failed to upack data from client", err)
			}

			tgtAddr := socks.SplitAddr(buf)
			if tgtAddr == nil {
				return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
			}

			tgtUDPAddr, err := net.ResolveUDPAddr("udp", tgtAddr.String())
			if err != nil {
				return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
			}
			if err := s.checkAllowedIP(tgtUDPAddr.IP); err != nil {
				return err
			}

			payload := buf[len(tgtAddr):]

			targetConn := nm.Get(clientAddr.String())
			if targetConn == nil {
				targetConn, err = net.ListenPacket("udp", "")
				if err != nil {
					return onet.NewConnectionError("ERR_CREATE_SOCKET", "Failed to create UDP socket", err)
				}
				nm.Add(clientAddr, s.clientConn, cipher, targetConn, clientLocation, keyID)
			}
			logger.Debugf("UDP NAT: client %v <-> proxy exit %v", clientAddr, targetConn.LocalAddr())

			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
			}
			return nil
		}()
	}
}

func (s *udpService) Stop() error {
	s.isRunning = false
	return s.clientConn.Close()
}

// Packet NAT table
type natmap struct {
	sync.RWMutex
	keyConn map[string]net.PacketConn
	timeout time.Duration
	metrics metrics.ShadowsocksMetrics
}

func newNATmap(timeout time.Duration, sm metrics.ShadowsocksMetrics) *natmap {
	m := &natmap{metrics: sm}
	m.keyConn = make(map[string]net.PacketConn)
	m.timeout = timeout
	return m
}

func (m *natmap) Get(key string) net.PacketConn {
	m.RLock()
	defer m.RUnlock()
	return m.keyConn[key]
}

func (m *natmap) set(key string, pc net.PacketConn) {
	m.Lock()
	defer m.Unlock()

	m.keyConn[key] = pc
}

func (m *natmap) del(key string) net.PacketConn {
	m.Lock()
	defer m.Unlock()

	pc, ok := m.keyConn[key]
	if ok {
		delete(m.keyConn, key)
		return pc
	}
	return nil
}

func (m *natmap) Add(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn, clientLocation, keyID string) {
	m.set(clientAddr.String(), targetConn)

	m.metrics.AddUDPNatEntry()
	go func() {
		timedCopy(clientAddr, clientConn, cipher, targetConn, m.timeout, clientLocation, keyID, m.metrics)
		m.metrics.RemoveUDPNatEntry()
		if pc := m.del(clientAddr.String()); pc != nil {
			pc.Close()
		}
	}()
}

// copy from src to dst at target with read timeout
func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn,
	timeout time.Duration, clientLocation, keyID string, sm metrics.ShadowsocksMetrics) {
	textBuf := make([]byte, udpBufSize)
	cipherBuf := make([]byte, udpBufSize)

	expired := false
	for !expired {
		var targetProxyBytes, proxyClientBytes int
		connError := func() (connError *onet.ConnectionError) {
			var (
				raddr net.Addr
				err   error
			)
			targetConn.SetReadDeadline(time.Now().Add(timeout))
			targetProxyBytes, raddr, err = targetConn.ReadFrom(textBuf)
			if err != nil {
				if netErr, ok := err.(net.Error); ok {
					if netErr.Timeout() {
						expired = true
						return nil
					}
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from target", err)
			}

			srcAddr := socks.ParseAddr(raddr.String())
			logger.Debugf("UDP response from %v to %v", srcAddr, clientAddr)
			// Shift data buffer to prepend with srcAddr.
			copy(textBuf[len(srcAddr):], textBuf[:targetProxyBytes])
			copy(textBuf, srcAddr)

			buf, err := shadowaead.Pack(cipherBuf, textBuf[:len(srcAddr)+targetProxyBytes], cipher)
			if err != nil {
				return onet.NewConnectionError("ERR_PACK", "Failed to pack data to client", err)
			}
			proxyClientBytes, err = clientConn.WriteTo(buf, clientAddr)
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to client", err)
			}
			return nil
		}()
		status := "OK"
		if connError != nil {
			logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
			status = connError.Status
		}
		sm.AddUDPPacketFromTarget(clientLocation, keyID, status, targetProxyBytes, proxyClientBytes)
	}
}
<commit_msg>Store client location in the NAT map<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package shadowsocks

import (
	"errors"
	"fmt"
	"net"
	"runtime/debug"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	logging "github.com/op/go-logging"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"sync"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

const udpBufSize = 64 * 1024

// Wrapper for logger.Debugf during UDP access key searches.
func debugUDP(cipherID, template string, val interface{}) {
	// This is an optimization to reduce unnecessary allocations due to an interaction
	// between Go's inlining/escape analysis and varargs functions like logger.Debugf.
	if logger.IsEnabledFor(logging.DEBUG) {
		logger.Debugf("UDP(%s): "+template, cipherID, val)
	}
}

// upack decrypts src into dst. It tries each cipher until it finds one that authenticates
// correctly. dst and src must not overlap.
func unpack(clientIP net.IP, dst, src []byte, cipherList CipherList) ([]byte, string, shadowaead.Cipher, error) {
	// Try each cipher until we find one that authenticates successfully. This assumes that all ciphers are AEAD.
	// We snapshot the list because it may be modified while we use it.
	for ci, entry := range cipherList.SnapshotForClientIP(clientIP) {
		id, cipher := entry.Value.(*CipherEntry).ID, entry.Value.(*CipherEntry).Cipher
		buf, err := shadowaead.Unpack(dst, src, cipher)
		if err != nil {
			debugUDP(id, "Failed to unpack: %v", err)
			continue
		}
		debugUDP(id, "Found cipher at index %d", ci)
		// Move the active cipher to the front, so that the search is quicker next time.
		cipherList.MarkUsedByClientIP(entry, clientIP)
		return buf, id, cipher, nil
	}
	return nil, "", nil, errors.New("could not find valid cipher")
}

type udpService struct {
	clientConn     net.PacketConn
	natTimeout     time.Duration
	ciphers        CipherList
	m              metrics.ShadowsocksMetrics
	isRunning      bool
	checkAllowedIP func(net.IP) *onet.ConnectionError
}

// NewUDPService creates a UDPService
func NewUDPService(clientConn net.PacketConn, natTimeout time.Duration, cipherList CipherList, m metrics.ShadowsocksMetrics) UDPService {
	return &udpService{clientConn: clientConn, natTimeout: natTimeout, ciphers: cipherList, m: m, checkAllowedIP: onet.RequirePublicIP}
}

// UDPService is a UDP shadowsocks service that can be started and stopped.
type UDPService interface {
	Start()
	Stop() error
}

// Listen on addr for encrypted packets and basically do UDP NAT.
// We take the ciphers as a pointer because it gets replaced on config updates.
func (s *udpService) Start() {
	defer s.clientConn.Close()

	nm := newNATmap(s.natTimeout, s.m)
	cipherBuf := make([]byte, udpBufSize)
	textBuf := make([]byte, udpBufSize)

	s.isRunning = true
	for s.isRunning {
		func() (connError *onet.ConnectionError) {
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in UDP loop: %v", r)
					debug.PrintStack()
				}
			}()
			clientLocation := ""
			keyID := ""
			var clientProxyBytes, proxyTargetBytes int
			var timeToCipher time.Duration
			defer func() {
				status := "OK"
				if connError != nil {
					logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				s.m.AddUDPPacketFromClient(clientLocation, keyID, status, clientProxyBytes, proxyTargetBytes, timeToCipher)
			}()
			clientProxyBytes, clientAddr, err := s.clientConn.ReadFrom(cipherBuf)
			if err != nil {
				if !s.isRunning {
					return nil
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from client", err)
			}
			defer logger.Debugf("UDP done with %v", clientAddr.String())
			logger.Debugf("UDP Request from %v with %v bytes", clientAddr, clientProxyBytes)
			unpackStart := time.Now()
			ip := clientAddr.(*net.UDPAddr).IP
			buf, keyID, cipher, err := unpack(ip, textBuf, cipherBuf[:clientProxyBytes], s.ciphers)
			timeToCipher = time.Now().Sub(unpackStart)

			if err != nil {
				return onet.NewConnectionError("ERR_CIPHER", "Failed to upack data from client", err)
			}

			tgtAddr := socks.SplitAddr(buf)
			if tgtAddr == nil {
				return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
			}

			tgtUDPAddr, err := net.ResolveUDPAddr("udp", tgtAddr.String())
			if err != nil {
				return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
			}
			if err := s.checkAllowedIP(tgtUDPAddr.IP); err != nil {
				return err
			}

			payload := buf[len(tgtAddr):]

			targetConn, clientLocation := nm.Get(clientAddr.String())
			if targetConn == nil {
				targetConn, err = net.ListenPacket("udp", "")
				if err != nil {
					return onet.NewConnectionError("ERR_CREATE_SOCKET", "Failed to create UDP socket", err)
				}
				clientLocation, locErr := s.m.GetLocation(clientAddr)
				if locErr != nil {
					logger.Warningf("Failed location lookup: %v", locErr)
				}
				logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientAddr.String())
				nm.Add(clientAddr, s.clientConn, cipher, targetConn, clientLocation, keyID)
			}
			logger.Debugf("UDP NAT: client %v <-> proxy exit %v", clientAddr, targetConn.LocalAddr())

			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
			}
			return nil
		}()
	}
}

func (s *udpService) Stop() error {
	s.isRunning = false
	return s.clientConn.Close()
}

type natentry struct {
	conn           net.PacketConn
	clientLocation string
}

// Packet NAT table
type natmap struct {
	sync.RWMutex
	keyConn map[string]natentry
	timeout time.Duration
	metrics metrics.ShadowsocksMetrics
}

func newNATmap(timeout time.Duration, sm metrics.ShadowsocksMetrics) *natmap {
	m := &natmap{metrics: sm}
	m.keyConn = make(map[string]natentry)
	m.timeout = timeout
	return m
}

func (m *natmap) Get(key string) (net.PacketConn, string) {
	m.RLock()
	defer m.RUnlock()
	entry := m.keyConn[key]
	return entry.conn, entry.clientLocation
}

func (m *natmap) set(key string, pc net.PacketConn, clientLocation string) {
	m.Lock()
	defer m.Unlock()

	m.keyConn[key] = natentry{pc, clientLocation}
}

func (m *natmap) del(key string) net.PacketConn {
	m.Lock()
	defer m.Unlock()

	entry, ok := m.keyConn[key]
	if ok {
		delete(m.keyConn, key)
		return entry.conn
	}
	return nil
}

func (m *natmap) Add(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn, clientLocation, keyID string) {
	m.set(clientAddr.String(), targetConn, clientLocation)

	m.metrics.AddUDPNatEntry()
	go func() {
		timedCopy(clientAddr, clientConn, cipher, targetConn, m.timeout, clientLocation, keyID, m.metrics)
		m.metrics.RemoveUDPNatEntry()
		if pc := m.del(clientAddr.String()); pc != nil {
			pc.Close()
		}
	}()
}

// copy from src to dst at target with read timeout
func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn,
	timeout time.Duration, clientLocation, keyID string, sm metrics.ShadowsocksMetrics) {
	textBuf := make([]byte, udpBufSize)
	cipherBuf := make([]byte, udpBufSize)

	expired := false
	for !expired {
		var targetProxyBytes, proxyClientBytes int
		connError := func() (connError *onet.ConnectionError) {
			var (
				raddr net.Addr
				err   error
			)
			targetConn.SetReadDeadline(time.Now().Add(timeout))
			targetProxyBytes, raddr, err = targetConn.ReadFrom(textBuf)
			if err != nil {
				if netErr, ok := err.(net.Error); ok {
					if netErr.Timeout() {
						expired = true
						return nil
					}
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from target", err)
			}

			srcAddr := socks.ParseAddr(raddr.String())
			logger.Debugf("UDP response from %v to %v", srcAddr, clientAddr)
			// Shift data buffer to prepend with srcAddr.
			copy(textBuf[len(srcAddr):], textBuf[:targetProxyBytes])
			copy(textBuf, srcAddr)

			buf, err := shadowaead.Pack(cipherBuf, textBuf[:len(srcAddr)+targetProxyBytes], cipher)
			if err != nil {
				return onet.NewConnectionError("ERR_PACK", "Failed to pack data to client", err)
			}
			proxyClientBytes, err = clientConn.WriteTo(buf, clientAddr)
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to client", err)
			}
			return nil
		}()
		status := "OK"
		if connError != nil {
			logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
			status = connError.Status
		}
		sm.AddUDPPacketFromTarget(clientLocation, keyID, status, targetProxyBytes, proxyClientBytes)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/bgentry/speakeasy"
)

const (
	serviceName        = "aws-vault"
	sessionServiceName = "aws-vault.sessions"
)

type stsClient interface {
	AssumeRole(input *sts.AssumeRoleInput) (*sts.AssumeRoleOutput, error)
	GetSessionToken(input *sts.GetSessionTokenInput) (*sts.GetSessionTokenOutput, error)
}

type VaultProvider struct {
	credentials.Expiry
	Keyring         keyring.Keyring
	Profile         string
	SessionDuration time.Duration
	ExpiryWindow    time.Duration
	profilesConf    profiles
	session         *sts.Credentials
	client          stsClient
}

func NewVaultProvider(k keyring.Keyring, profile string) (*VaultProvider, error) {
	conf, err := parseProfiles()
	if err != nil {
		return nil, err
	}
	return &VaultProvider{
		Keyring:         k,
		Profile:         profile,
		SessionDuration: time.Second * 900, // the shortest AWS will allow
		ExpiryWindow:    time.Second * 90,
		profilesConf:    conf,
	}, nil
}

func (p *VaultProvider) Retrieve() (credentials.Value, error) {
	var session sts.Credentials

	if err := keyring.Unmarshal(p.Keyring, sessionServiceName, p.Profile, &session); err != nil {
		log.Println("Session lookup failed", err)

		session, err = p.getSessionToken(p.SessionDuration)
		if err != nil {
			return credentials.Value{}, err
		}

		if role, ok := p.profilesConf[p.Profile]["role_arn"]; ok {
			session, err = p.assumeRole(session, role)
			if err != nil {
				return credentials.Value{}, err
			}
		}

		keyring.Marshal(p.Keyring, sessionServiceName, p.Profile, session)
	} else {
		log.Printf("Found a cached session token for %s", p.Profile)
	}

	log.Printf("Session token expires in %s", session.Expiration.Sub(time.Now()))
	p.SetExpiration(*session.Expiration, p.ExpiryWindow)

	value := credentials.Value{
		AccessKeyID:     *session.AccessKeyId,
		SecretAccessKey: *session.SecretAccessKey,
		SessionToken:    *session.SessionToken,
	}

	return value, nil
}

func (p *VaultProvider) getSessionToken(length time.Duration) (sts.Credentials, error) {
	source := p.profilesConf.sourceProfile(p.Profile)

	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(length.Seconds())),
	}

	if mfa, ok := p.profilesConf[p.Profile]["mfa_serial"]; ok {
		token, err := speakeasy.Ask(fmt.Sprintf("Enter token for %s: ", mfa))
		if err != nil {
			return sts.Credentials{}, err
		}
		params.SerialNumber = aws.String(mfa)
		params.TokenCode = aws.String(token)
	}

	client := p.client
	if client == nil {
		client = sts.New(&aws.Config{Credentials: credentials.NewChainCredentials(
			p.defaultProviders(source),
		)})
	}

	log.Printf("Getting new session token for profile %s", p.Profile)
	resp, err := client.GetSessionToken(params)
	if err != nil {
		return sts.Credentials{}, err
	}

	return *resp.Credentials, nil
}

func (p *VaultProvider) assumeRole(session sts.Credentials, roleArn string) (sts.Credentials, error) {
	client := p.client
	if client == nil {
		client = sts.New(&aws.Config{Credentials: credentials.NewStaticCredentials(
			*session.AccessKeyId,
			*session.SecretAccessKey,
			*session.SessionToken,
		)})
	}

	// Try to work out a role name that will hopefully end up unique.
	roleSessionName := fmt.Sprintf("%d", time.Now().UTC().UnixNano())

	input := &sts.AssumeRoleInput{
		RoleArn:         aws.String(roleArn),
		RoleSessionName: aws.String(roleSessionName),
		DurationSeconds: aws.Int64(int64(15 * 60)),
	}

	log.Printf("Assuming role %s", roleArn)
	resp, err := client.AssumeRole(input)
	if err != nil {
		return sts.Credentials{}, err
	}

	return *resp.Credentials, nil
}

func (p *VaultProvider) defaultProviders(profile string) []credentials.Provider {
	return []credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{Filename: "", Profile: profile},
		&KeyringProvider{Keyring: p.Keyring, Profile: profile},
	}
}

type KeyringProvider struct {
	Keyring keyring.Keyring
	Profile string
}

func (p *KeyringProvider) IsExpired() bool {
	return false
}

func (p *KeyringProvider) Retrieve() (val credentials.Value, err error) {
	log.Printf("Looking up keyring for %s", p.Profile)
	if err = keyring.Unmarshal(p.Keyring, serviceName, p.Profile, &val); err != nil {
		log.Println("Error looking up keyring", err)
	}
	return
}

func (p *KeyringProvider) Store(val credentials.Value) error {
	return keyring.Marshal(p.Keyring, serviceName, p.Profile, val)
}

func (p *KeyringProvider) Delete() error {
	p.Keyring.Remove(sessionServiceName, p.Profile)
	return p.Keyring.Remove(serviceName, p.Profile)
}
<commit_msg>Purge sessions when adding credentials<commit_after>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/bgentry/speakeasy"
)

const (
	serviceName        = "aws-vault"
	sessionServiceName = "aws-vault.sessions"
)

type stsClient interface {
	AssumeRole(input *sts.AssumeRoleInput) (*sts.AssumeRoleOutput, error)
	GetSessionToken(input *sts.GetSessionTokenInput) (*sts.GetSessionTokenOutput, error)
}

type VaultProvider struct {
	credentials.Expiry
	Keyring         keyring.Keyring
	Profile         string
	SessionDuration time.Duration
	ExpiryWindow    time.Duration
	profilesConf    profiles
	session         *sts.Credentials
	client          stsClient
}

func NewVaultProvider(k keyring.Keyring, profile string) (*VaultProvider, error) {
	conf, err := parseProfiles()
	if err != nil {
		return nil, err
	}
	return &VaultProvider{
		Keyring:         k,
		Profile:         profile,
		SessionDuration: time.Second * 900, // the shortest AWS will allow
		ExpiryWindow:    time.Second * 90,
		profilesConf:    conf,
	}, nil
}

func (p *VaultProvider) Retrieve() (credentials.Value, error) {
	var session sts.Credentials

	if err := keyring.Unmarshal(p.Keyring, sessionServiceName, p.Profile, &session); err != nil {
		log.Println("Session lookup failed", err)

		session, err = p.getSessionToken(p.SessionDuration)
		if err != nil {
			return credentials.Value{}, err
		}

		if role, ok := p.profilesConf[p.Profile]["role_arn"]; ok {
			session, err = p.assumeRole(session, role)
			if err != nil {
				return credentials.Value{}, err
			}
		}

		keyring.Marshal(p.Keyring, sessionServiceName, p.Profile, session)
	} else {
		log.Printf("Found a cached session token for %s", p.Profile)
	}

	log.Printf("Session token expires in %s", session.Expiration.Sub(time.Now()))
	p.SetExpiration(*session.Expiration, p.ExpiryWindow)

	value := credentials.Value{
		AccessKeyID:     *session.AccessKeyId,
		SecretAccessKey: *session.SecretAccessKey,
		SessionToken:    *session.SessionToken,
	}

	return value, nil
}

func (p *VaultProvider) getSessionToken(length time.Duration) (sts.Credentials, error) {
	source := p.profilesConf.sourceProfile(p.Profile)

	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(length.Seconds())),
	}

	if mfa, ok := p.profilesConf[p.Profile]["mfa_serial"]; ok {
		token, err := speakeasy.Ask(fmt.Sprintf("Enter token for %s: ", mfa))
		if err != nil {
			return sts.Credentials{}, err
		}
		params.SerialNumber = aws.String(mfa)
		params.TokenCode = aws.String(token)
	}

	client := p.client
	if client == nil {
		client = sts.New(&aws.Config{Credentials: credentials.NewChainCredentials(
			p.defaultProviders(source),
		)})
	}

	log.Printf("Getting new session token for profile %s", p.Profile)
	resp, err := client.GetSessionToken(params)
	if err != nil {
		return sts.Credentials{}, err
	}

	return *resp.Credentials, nil
}

func (p *VaultProvider) assumeRole(session sts.Credentials, roleArn string) (sts.Credentials, error) {
	client := p.client
	if client == nil {
		client = sts.New(&aws.Config{Credentials: credentials.NewStaticCredentials(
			*session.AccessKeyId,
			*session.SecretAccessKey,
			*session.SessionToken,
		)})
	}

	// Try to work out a role name that will hopefully end up unique.
	roleSessionName := fmt.Sprintf("%d", time.Now().UTC().UnixNano())

	input := &sts.AssumeRoleInput{
		RoleArn:         aws.String(roleArn),
		RoleSessionName: aws.String(roleSessionName),
		DurationSeconds: aws.Int64(int64(15 * 60)),
	}

	log.Printf("Assuming role %s", roleArn)
	resp, err := client.AssumeRole(input)
	if err != nil {
		return sts.Credentials{}, err
	}

	return *resp.Credentials, nil
}

func (p *VaultProvider) defaultProviders(profile string) []credentials.Provider {
	return []credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{Filename: "", Profile: profile},
		&KeyringProvider{Keyring: p.Keyring, Profile: profile},
	}
}

type KeyringProvider struct {
	Keyring keyring.Keyring
	Profile string
}

func (p *KeyringProvider) IsExpired() bool {
	return false
}

func (p *KeyringProvider) Retrieve() (val credentials.Value, err error) {
	log.Printf("Looking up keyring for %s", p.Profile)
	if err = keyring.Unmarshal(p.Keyring, serviceName, p.Profile, &val); err != nil {
		log.Println("Error looking up keyring", err)
	}
	return
}

func (p *KeyringProvider) Store(val credentials.Value) error {
	p.Keyring.Remove(sessionServiceName, p.Profile)
	return keyring.Marshal(p.Keyring, serviceName, p.Profile, val)
}

func (p *KeyringProvider) Delete() error {
	p.Keyring.Remove(sessionServiceName, p.Profile)
	return p.Keyring.Remove(serviceName, p.Profile)
}
<|endoftext|>


<commit_before>package xpi

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"time"

	"github.com/ThalesIgnite/crypto11"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"go.mozilla.org/cose"
)

// populateRsaCache adds an rsa key to the cache every
// XPISigner.rsaCacheSleepDuration, blocks when the cache channel is
// full, and should be run as a goroutine
func (s *XPISigner) populateRsaCache(size int) {
	var (
		err   error
		key   *rsa.PrivateKey
		start time.Time
	)
	for {
		start = time.Now()
		key, err = rsa.GenerateKey(s.rand, size)
		if err != nil {
			log.Fatalf("xpi: error generating RSA key for cache: %s", err)
		}
		if key == nil {
			log.Fatal("xpi: error generated nil RSA key for cache")
		}

		if s.stats != nil {
			s.stats.SendHistogram("xpi.rsa_cache.gen_key_dur", time.Since(start))
		}
		s.rsaCache <- key
		time.Sleep(s.rsaCacheGeneratorSleepDuration)
	}
}

// monitorRsaCacheSize sends the number of cached keys and cache size
// to datadog. It should be run as a goroutine
func (s *XPISigner) monitorRsaCacheSize() {
	if s.stats == nil {
		return
	}
	for {
		s.stats.SendGauge("xpi.rsa_cache.chan_len", len(s.rsaCache))

		// chan capacity should be constant but is useful for
		// knowing % cache filled across deploys
		s.stats.SendGauge("xpi.rsa_cache.chan_cap", cap(s.rsaCache))

		time.Sleep(s.rsaCacheSizeSampleRate)
	}
}

// retrieve a key from the cache or generate one if it takes too long
// or if the size is wrong
func (s *XPISigner) getRsaKey(size int) (*rsa.PrivateKey, error) {
	var (
		err   error
		key   *rsa.PrivateKey
		start time.Time
	)
	start = time.Now()
	select {
	case key = <-s.rsaCache:
		if key.N.BitLen() != size {
			// it's theoritically impossible for this to happen
			// because the end entity has the same key size has
			// the signer, but we're paranoid so handling it
			log.Warnf("WARNING: xpi rsa cache returned a key of size %d when %d was requested", key.N.BitLen(), size)
			key, err = rsa.GenerateKey(s.rand, size)
		}
	case <-time.After(s.rsaCacheFetchTimeout):
		// generate a key if none available
		key, err = rsa.GenerateKey(s.rand, size)
	}

	if s.stats != nil {
		s.stats.SendHistogram("xpi.rsa_cache.get_key", time.Since(start))
	}
	return key, err
}

// makeTemplate returns a pointer to a template for an x509.Certificate EE
func (s *XPISigner) makeTemplate(cn string) *x509.Certificate {
	cndigest := sha256.Sum256([]byte(cn))
	return &x509.Certificate{
		// The maximum length of a serial number per rfc 5280 is 20 bytes / 160 bits
		// https://tools.ietf.org/html/rfc5280#section-4.1.2.2
		// Setting it to nanoseconds guarantees we'll never have two conflicting serials
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		// PKIX requires EE's to have a valid DNS Names when the intermediate has
		// a constraint, so we hash the CN into an fqdn to get something unique enough
		DNSNames: []string{fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])},
		Subject: pkix.Name{
			CommonName:         cn,
			Organization:       []string{"Addons"},
			OrganizationalUnit: []string{s.OU},
			Country:            []string{"US"},
			Province:           []string{"CA"},
			Locality:           []string{"Mountain View"},
		},
		NotBefore:          time.Now(),
		NotAfter:           time.Now().Add(8760 * time.Hour), // one year
		SignatureAlgorithm: s.issuerCert.SignatureAlgorithm,
		KeyUsage:           x509.KeyUsageDigitalSignature,
		ExtKeyUsage:        []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning},
	}
}

// generateIssuerEEKeyPair returns a public and private key pair
// matching the issuer XPISigner issuerKey size and type
func (s *XPISigner) generateIssuerEEKeyPair() (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	switch issuerKey := s.issuerKey.(type) {
	case *rsa.PrivateKey:
		size := issuerKey.N.BitLen()
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate rsa private key of size %d", size)
			return
		}
		if eeKey == nil {
			err = errors.Wrapf(err, "xpi: failed to get rsa private key of size %d", size)
			return
		}

		newKey, ok := eeKey.(*rsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key of size %d to *rsa.PrivateKey", size)
			return
		}
		eePublicKey = newKey.Public()
	case *crypto11.PKCS11PrivateKeyRSA:
		size := issuerKey.PubKey.(*rsa.PublicKey).N.BitLen()
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate rsa private key of size %d", size)
			return
		}
		if eeKey == nil {
			err = errors.Wrapf(err, "xpi: failed to get rsa private key of size %d", size)
			return
		}

		newKey, ok := eeKey.(*rsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key of size %d to *rsa.PrivateKey", size)
			return
		}
		eePublicKey = newKey.Public()
	case *ecdsa.PrivateKey:
		eeKey, err = ecdsa.GenerateKey(issuerKey.Curve, s.rand)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate ecdsa private key on curve %s", issuerKey.Curve.Params().Name)
			return
		}
		newKey, ok := eeKey.(*ecdsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key on curve %s to *ecdsa.PrivateKey", issuerKey.Curve.Params().Name)
			return
		}
		eePublicKey = newKey.Public()
	case *crypto11.PKCS11PrivateKeyECDSA:
		curve := issuerKey.PubKey.(*ecdsa.PublicKey).Curve
		eeKey, err = ecdsa.GenerateKey(curve, s.rand)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate ecdsa private key on curve %s", curve.Params().Name)
			return
		}
		newKey, ok := eeKey.(*ecdsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key on curve %s to *ecdsa.PrivateKey", curve.Params().Name)
			return
		}
		eePublicKey = newKey.Public()
	default:
		err = errors.Errorf("xpi: unrecognized issuer key type for EE: %T", issuerKey)
		return
	}
	return
}

// MakeEndEntity generates a private key and certificate ready to sign a given XPI.
//
// The subject CN of the certificate is taken from the `cn` string argument.
//
// The key type is identical to the key type of the signer that issues
// the certificate when the optional `coseAlg` argument is nil. For
// example, if the signer uses an RSA 2048 key, so will the
// end-entity. When `coseAlg` is not nil, a key type of the COSE
// algorithm is generated.
//
// The signature expiration date is copied over from the issuer.
//
// The signed x509 certificate and private key are returned.
func (s *XPISigner) MakeEndEntity(cn string, coseAlg *cose.Algorithm) (eeCert *x509.Certificate, eeKey crypto.PrivateKey, err error) {
	var (
		eePublicKey crypto.PublicKey
		derCert     []byte
	)

	template := s.makeTemplate(cn)

	if coseAlg == nil {
		eeKey, eePublicKey, err = s.generateIssuerEEKeyPair()
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching issuer")
			return
		}
	} else {
		eeKey, eePublicKey, err = s.generateCOSEKeyPair(coseAlg)
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching COSE Algorithm type %s", coseAlg.Name)
			return
		}
	}

	derCert, err = x509.CreateCertificate(s.rand, template, s.issuerCert, eePublicKey, s.issuerKey)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: failed to create certificate")
		return
	}
	if len(derCert) == 0 {
		err = errors.Errorf("xpi.MakeEndEntity: certificate creation failed for an unknown reason")
		return
	}
	eeCert, err = x509.ParseCertificate(derCert)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: certificate parsing failed")
	}
	return
}
<commit_msg>signer: de-deup xpi pkcs7 key gen cases<commit_after>package xpi

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"time"

	"github.com/ThalesIgnite/crypto11"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"go.mozilla.org/cose"
)

// populateRsaCache adds an rsa key to the cache every
// XPISigner.rsaCacheSleepDuration, blocks when the cache channel is
// full, and should be run as a goroutine
func (s *XPISigner) populateRsaCache(size int) {
	var (
		err   error
		key   *rsa.PrivateKey
		start time.Time
	)
	for {
		start = time.Now()
		key, err = rsa.GenerateKey(s.rand, size)
		if err != nil {
			log.Fatalf("xpi: error generating RSA key for cache: %s", err)
		}
		if key == nil {
			log.Fatal("xpi: error generated nil RSA key for cache")
		}

		if s.stats != nil {
			s.stats.SendHistogram("xpi.rsa_cache.gen_key_dur", time.Since(start))
		}
		s.rsaCache <- key
		time.Sleep(s.rsaCacheGeneratorSleepDuration)
	}
}

// monitorRsaCacheSize sends the number of cached keys and cache size
// to datadog. It should be run as a goroutine
func (s *XPISigner) monitorRsaCacheSize() {
	if s.stats == nil {
		return
	}
	for {
		s.stats.SendGauge("xpi.rsa_cache.chan_len", len(s.rsaCache))

		// chan capacity should be constant but is useful for
		// knowing % cache filled across deploys
		s.stats.SendGauge("xpi.rsa_cache.chan_cap", cap(s.rsaCache))

		time.Sleep(s.rsaCacheSizeSampleRate)
	}
}

// retrieve a key from the cache or generate one if it takes too long
// or if the size is wrong
func (s *XPISigner) getRsaKey(size int) (*rsa.PrivateKey, error) {
	var (
		err   error
		key   *rsa.PrivateKey
		start time.Time
	)
	start = time.Now()
	select {
	case key = <-s.rsaCache:
		if key.N.BitLen() != size {
			// it's theoritically impossible for this to happen
			// because the end entity has the same key size has
			// the signer, but we're paranoid so handling it
			log.Warnf("WARNING: xpi rsa cache returned a key of size %d when %d was requested", key.N.BitLen(), size)
			key, err = rsa.GenerateKey(s.rand, size)
		}
	case <-time.After(s.rsaCacheFetchTimeout):
		// generate a key if none available
		key, err = rsa.GenerateKey(s.rand, size)
	}

	if s.stats != nil {
		s.stats.SendHistogram("xpi.rsa_cache.get_key", time.Since(start))
	}
	return key, err
}

// makeTemplate returns a pointer to a template for an x509.Certificate EE
func (s *XPISigner) makeTemplate(cn string) *x509.Certificate {
	cndigest := sha256.Sum256([]byte(cn))
	return &x509.Certificate{
		// The maximum length of a serial number per rfc 5280 is 20 bytes / 160 bits
		// https://tools.ietf.org/html/rfc5280#section-4.1.2.2
		// Setting it to nanoseconds guarantees we'll never have two conflicting serials
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		// PKIX requires EE's to have a valid DNS Names when the intermediate has
		// a constraint, so we hash the CN into an fqdn to get something unique enough
		DNSNames: []string{fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])},
		Subject: pkix.Name{
			CommonName:         cn,
			Organization:       []string{"Addons"},
			OrganizationalUnit: []string{s.OU},
			Country:            []string{"US"},
			Province:           []string{"CA"},
			Locality:           []string{"Mountain View"},
		},
		NotBefore:          time.Now(),
		NotAfter:           time.Now().Add(8760 * time.Hour), // one year
		SignatureAlgorithm: s.issuerCert.SignatureAlgorithm,
		KeyUsage:           x509.KeyUsageDigitalSignature,
		ExtKeyUsage:        []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning},
	}
}

// getIssuerPubKey returns the public key for ECDSA or RSA keys from
// the crypto stdlib or crypto11
func (s *XPISigner) getIssuerPubKey() (pubKey crypto.PublicKey, err error) {
	switch issuerKey := s.issuerKey.(type) {
	// NB: when these two cases aren't separate golang treats them
	// as crypto.PrivateKey (which doesn't have a .Public method)
	case *rsa.PrivateKey:
		pubKey = issuerKey.Public()
	case *ecdsa.PrivateKey:
		pubKey = issuerKey.Public()
	case *crypto11.PKCS11PrivateKeyRSA:
		pubKey = issuerKey.PubKey
	case *crypto11.PKCS11PrivateKeyECDSA:
		pubKey = issuerKey.PubKey
	default:
		err = errors.Errorf("xpi: cannot get public key for issuer key type %T", issuerKey)
	}
	return
}

// getIssuerRSAKeySize returns the rsa key size in bits for crypto or
// crypto11 issuer keys
func (s *XPISigner) getIssuerRSAKeySize() (size int, err error) {
	pubKey, err := s.getIssuerPubKey()
	if err != nil {
		err = errors.Errorf("xpi: failed to public key to get rsa key size")
		return
	}
	rsaKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		err = errors.Errorf("xpi: failed to cast public key to *rsa.PublicKey to get rsa key size")
		return
	}
	return rsaKey.N.BitLen(), nil
}

// getIssuerECDSACurve returns the ecdsa curve for crypto or crypto11
// issuer keys
func (s *XPISigner) getIssuerECDSACurve() (curve elliptic.Curve, err error) {
	pubKey, err := s.getIssuerPubKey()
	if err != nil {
		err = errors.Errorf("xpi: failed to public key to get ecdsa curve")
		return
	}
	ecKey, ok := pubKey.(*ecdsa.PublicKey)
	if !ok {
		err = errors.Errorf("xpi: failed to cast public key to *ecdsa.PublicKey to get curve")
		return
	}
	return ecKey.Curve, nil
}

// generateIssuerEEKeyPair returns a public and private key pair
// matching the issuer XPISigner issuerKey size and type
func (s *XPISigner) generateIssuerEEKeyPair() (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	switch issuerKey := s.issuerKey.(type) {
	case *rsa.PrivateKey, *crypto11.PKCS11PrivateKeyRSA:
		var size int
		size, err = s.getIssuerRSAKeySize()
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to get rsa key size")
			return
		}
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate rsa private key of size %d", size)
			return
		}
		if eeKey == nil {
			err = errors.Wrapf(err, "xpi: failed to get rsa private key of size %d", size)
			return
		}

		newKey, ok := eeKey.(*rsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key of size %d to *rsa.PrivateKey", size)
			return
		}
		eePublicKey = newKey.Public()
	case *ecdsa.PrivateKey, *crypto11.PKCS11PrivateKeyECDSA:
		var curve elliptic.Curve
		curve, err = s.getIssuerECDSACurve()
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to get ecdsa curve")
			return
		}
		eeKey, err = ecdsa.GenerateKey(curve, s.rand)
		if err != nil {
			err = errors.Wrapf(err, "xpi: failed to generate ecdsa private key on curve %s", curve.Params().Name)
			return
		}

		newKey, ok := eeKey.(*ecdsa.PrivateKey)
		if !ok {
			err = errors.Wrapf(err, "xpi: failed to cast generated key on curve %s to *ecdsa.PrivateKey", curve.Params().Name)
			return
		}
		eePublicKey = newKey.Public()
	default:
		err = errors.Errorf("xpi: unrecognized issuer key type for EE: %T", issuerKey)
		return
	}
	return
}

// MakeEndEntity generates a private key and certificate ready to sign a given XPI.
//
// The subject CN of the certificate is taken from the `cn` string argument.
//
// The key type is identical to the key type of the signer that issues
// the certificate when the optional `coseAlg` argument is nil. For
// example, if the signer uses an RSA 2048 key, so will the
// end-entity. When `coseAlg` is not nil, a key type of the COSE
// algorithm is generated.
//
// The signature expiration date is copied over from the issuer.
//
// The signed x509 certificate and private key are returned.
func (s *XPISigner) MakeEndEntity(cn string, coseAlg *cose.Algorithm) (eeCert *x509.Certificate, eeKey crypto.PrivateKey, err error) {
	var (
		eePublicKey crypto.PublicKey
		derCert     []byte
	)

	template := s.makeTemplate(cn)

	if coseAlg == nil {
		eeKey, eePublicKey, err = s.generateIssuerEEKeyPair()
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching issuer")
			return
		}
	} else {
		eeKey, eePublicKey, err = s.generateCOSEKeyPair(coseAlg)
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching COSE Algorithm type %s", coseAlg.Name)
			return
		}
	}

	derCert, err = x509.CreateCertificate(s.rand, template, s.issuerCert, eePublicKey, s.issuerKey)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: failed to create certificate")
		return
	}
	if len(derCert) == 0 {
		err = errors.Errorf("xpi.MakeEndEntity: certificate creation failed for an unknown reason")
		return
	}
	eeCert, err = x509.ParseCertificate(derCert)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: certificate parsing failed")
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/KIT-MAMID/mamid/msp"
	. "github.com/KIT-MAMID/mamid/slave"
	"golang.org/x/sys/unix"
	"gopkg.in/mgo.v2"
	"os/exec"
)

const MongodExecutableDefaultName = "mongod"
const DefaultMongodSoftShutdownTimeout = "3s" // seconds
const DefaultMongodHardShutdownTimeout = "5s" // seconds

func main() {

	var (
		mongodExecutable, dataDir                                  string
		mongodSoftShutdownTimeoutStr, mongodHardShutdownTimeoutStr string
	)

	flag.StringVar(&dataDir, "data", "", "Persistent data and slave configuration directory")
	mongodExecutableLookupPath, _ := exec.LookPath(MongodExecutableDefaultName)
	flag.StringVar(&mongodExecutable, "mongodExecutable", mongodExecutableLookupPath, "Path to or name of Mongod binary")

	flag.StringVar(&mongodSoftShutdownSeconds, "mongod.shutdownTimeout.soft", DefaultMongodSoftShutdownTimeout,
		"Duration to wait for regular Mongod shutdown call to return. Specify with suffix [ms,s,min,...]")
	flag.StringVar(&mongodHardShutdownSeconds, "mongod.shutdownTimeout.hard", DefaultMongodHardShutdownTimeout,
		"Duration to wait after issuing a shutdown call before the Mongod is killed (SIGKILL). Specify with suffix [ms,s,min,...]")

	flag.Parse()

	// Assert dataDir is valid. TODO should we do this lazyly?

	if dataDir == "" {
		println("No root data directory passed; specify with -data=/path/to/root/dir")
		return
	}

	if err := unix.Access(dataDir, unix.W_OK); err != nil {
		println(fmt.Sprintf("Root data directory %s does not exist or is not writable", dataDir))
		return
	}

	dbDir := fmt.Sprintf("%s/%s", dataDir, DataDBDir) // TODO directory creation should happen in the component that uses the path
	if err := unix.Access(dbDir, unix.R_OK|unix.W_OK|unix.X_OK); err != nil {
		if err := unix.Mkdir(dbDir, 0700); err != nil {
			fmt.Printf("Could not create a readable and writable directory at %s: %s", dbDir, err)
			return
		}
	}

	// Convert timeouts to internal representation

	mongodSoftShutdownTimeout, err := time.ParseDuration(mongodSoftShutdownTimeoutStr)
	if !err {
		fmt.Printf("could not convert soft shutdown timeout to time.Duration: %s", err)
		return
	}

	mongodHardShutdownTimeout, err := time.Duration(mongodHardShutdownTimeoutStr)
	if !err {
		fmt.Printf("could not convert hard shutdown timeout to time.Duration: %s", err)
		return
	}

	processManager := NewProcessManager(mongodExecutable, dataDir)
	configurator := &ConcreteMongodConfigurator{
		dial: mgo.Dial,
		MongodSoftShutdownTimeout: mongodSoftShutdownTimeout,
	}

	controller := NewController(processManager, configurator, mongodHardShutdownTimeout)
	server := msp.NewServer(controller)
	server.Run()
}
<commit_msg>UPD: slave: use logrus<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/KIT-MAMID/mamid/msp"
	. "github.com/KIT-MAMID/mamid/slave"
	"github.com/Sirupsen/logrus"
	"golang.org/x/sys/unix"
	"gopkg.in/mgo.v2"
	"os/exec"
)

var log = logrus.WithField("module", "slave")

const MongodExecutableDefaultName = "mongod"
const DefaultMongodSoftShutdownTimeout = "3s" // seconds
const DefaultMongodHardShutdownTimeout = "5s" // seconds

func main() {

	log.SetLevel(logrus.DebugLevel)

	var (
		mongodExecutable, dataDir                                  string
		mongodSoftShutdownTimeoutStr, mongodHardShutdownTimeoutStr string
	)

	flag.StringVar(&dataDir, "data", "", "Persistent data and slave configuration directory")
	mongodExecutableLookupPath, _ := exec.LookPath(MongodExecutableDefaultName)
	flag.StringVar(&mongodExecutable, "mongodExecutable", mongodExecutableLookupPath, "Path to or name of Mongod binary")

	flag.StringVar(&mongodSoftShutdownSeconds, "mongod.shutdownTimeout.soft", DefaultMongodSoftShutdownTimeout,
		"Duration to wait for regular Mongod shutdown call to return. Specify with suffix [ms,s,min,...]")
	flag.StringVar(&mongodHardShutdownSeconds, "mongod.shutdownTimeout.hard", DefaultMongodHardShutdownTimeout,
		"Duration to wait after issuing a shutdown call before the Mongod is killed (SIGKILL). Specify with suffix [ms,s,min,...]")

	flag.Parse()

	// Assert dataDir is valid. TODO should we do this lazyly?

	if dataDir == "" {
		log.Fatal("No root data directory passed; specify with -data=/path/to/root/dir")
	}

	if err := unix.Access(dataDir, unix.W_OK); err != nil {
		log.Fatal(fmt.Sprintf("Root data directory %s does not exist or is not writable", dataDir))
	}

	dbDir := fmt.Sprintf("%s/%s", dataDir, DataDBDir) // TODO directory creation should happen in the component that uses the path
	if err := unix.Access(dbDir, unix.R_OK|unix.W_OK|unix.X_OK); err != nil {
		if err := unix.Mkdir(dbDir, 0700); err != nil {
			log.Printf("Could not create a readable and writable directory at %s: %s", dbDir, err)
			return
		}
	}

	// Convert timeouts to internal representation

	mongodSoftShutdownTimeout, err := time.ParseDuration(mongodSoftShutdownTimeoutStr)
	if !err {
		log.Fatal("could not convert soft shutdown timeout to time.Duration: %s", err)
	}

	mongodHardShutdownTimeout, err := time.Duration(mongodHardShutdownTimeoutStr)
	if !err {
		log.Fatal("could not convert hard shutdown timeout to time.Duration: %s", err)
	}

	processManager := NewProcessManager(mongodExecutable, dataDir)
	configurator := &ConcreteMongodConfigurator{
		dial: mgo.Dial,
		MongodSoftShutdownTimeout: mongodSoftShutdownTimeout,
	}

	controller := NewController(processManager, configurator, mongodHardShutdownTimeout)
	server := msp.NewServer(controller)
	server.Run()
}
<|endoftext|>


<commit_before>package smartbus

import (
	"net"
	"strings"
        serial "github.com/ivan4th/goserial"
)

// FIXME
const (
	DRIVER_SUBNET = 0x01
	DRIVER_DEVICE_ID = 0x14
	DRIVER_DEVICE_TYPE = 0x0095
	DRIVER_CLIENT_ID = "smartbus"
)

func connect(serialAddress string) (SmartbusIO, error) {
	switch {
	case strings.HasPrefix(serialAddress, "/"):
		if serial, err := serial.OpenPort(&serial.Config{
			Name: serialAddress,
			Baud: 9600,
			Parity: serial.ParityEven,
			Size: serial.Byte8,
		}); err != nil {
			return nil, err
		} else {
			return NewStreamIO(serial), nil
		}
	case serialAddress == "udp":
		if dgramIO, err := NewDatagramIO(); err != nil {
			return nil, err
		} else {
			return dgramIO, nil
		}
	case strings.HasPrefix(serialAddress, "tcp://"):
		if conn, err := net.Dial("tcp", serialAddress[6:]); err != nil {
			return nil, err
		} else {
			return NewStreamIO(conn), nil
		}
	}

	if conn, err := net.Dial("tcp", serialAddress); err != nil {
		return nil, err
	} else {
		return NewStreamIO(conn), nil
	}
}

func NewSmartbusTCPDriver(serialAddress, brokerAddress string) (*Driver, error) {
	model := NewSmartbusModel(func () (SmartbusIO, error) {
		return connect(serialAddress)
	}, DRIVER_SUBNET, DRIVER_DEVICE_ID, DRIVER_DEVICE_TYPE)
	driver := NewDriver(model, func (handler MQTTMessageHandler) MQTTClient {
		return NewPahoMQTTClient(brokerAddress, DRIVER_CLIENT_ID, handler)
	})
	return driver, nil
}
<commit_msg>Use own device id & type.<commit_after>package smartbus

import (
	"net"
	"strings"
        serial "github.com/ivan4th/goserial"
)

// FIXME
const (
	DRIVER_SUBNET = 0x01
	DRIVER_DEVICE_ID = 0x99
	DRIVER_DEVICE_TYPE = 0x1234
	DRIVER_CLIENT_ID = "smartbus"
)

func connect(serialAddress string) (SmartbusIO, error) {
	switch {
	case strings.HasPrefix(serialAddress, "/"):
		if serial, err := serial.OpenPort(&serial.Config{
			Name: serialAddress,
			Baud: 9600,
			Parity: serial.ParityEven,
			Size: serial.Byte8,
		}); err != nil {
			return nil, err
		} else {
			return NewStreamIO(serial), nil
		}
	case serialAddress == "udp":
		if dgramIO, err := NewDatagramIO(); err != nil {
			return nil, err
		} else {
			return dgramIO, nil
		}
	case strings.HasPrefix(serialAddress, "tcp://"):
		if conn, err := net.Dial("tcp", serialAddress[6:]); err != nil {
			return nil, err
		} else {
			return NewStreamIO(conn), nil
		}
	}

	if conn, err := net.Dial("tcp", serialAddress); err != nil {
		return nil, err
	} else {
		return NewStreamIO(conn), nil
	}
}

func NewSmartbusTCPDriver(serialAddress, brokerAddress string) (*Driver, error) {
	model := NewSmartbusModel(func () (SmartbusIO, error) {
		return connect(serialAddress)
	}, DRIVER_SUBNET, DRIVER_DEVICE_ID, DRIVER_DEVICE_TYPE)
	driver := NewDriver(model, func (handler MQTTMessageHandler) MQTTClient {
		return NewPahoMQTTClient(brokerAddress, DRIVER_CLIENT_ID, handler)
	})
	return driver, nil
}
<|endoftext|>


<commit_before>package sockjs

import (
	"sync"
)

type connections struct {
	connections map[string]*conn
	mu          sync.RWMutex
}

type connFactory func() *conn

func newConnections() connections {
	return connections{
		connections: make(map[string]*conn),
	}
}

func (c *connections) get(sessid string) (conn *conn, exists bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	conn, exists = c.connections[sessid]
	return
}

func (c *connections) getOrCreate(sessid string, f connFactory) (conn *conn, exists bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	conn, exists = c.connections[sessid]
	if !exists {
		c.connections[sessid] = f()
		conn = c.connections[sessid]
	}
	return
}

func (c *connections) delete(sessid string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.connections, sessid)
}
<commit_msg>embedded mutex<commit_after>package sockjs

import (
	"sync"
)

type connections struct {
	sync.RWMutex
	connections map[string]*conn
}

type connFactory func() *conn

func newConnections() connections {
	return connections{
		connections: make(map[string]*conn),
	}
}

func (c *connections) get(sessid string) (conn *conn, exists bool) {
	c.RLock()
	defer c.RUnlock()
	conn, exists = c.connections[sessid]
	return
}

func (c *connections) getOrCreate(sessid string, f connFactory) (conn *conn, exists bool) {
	c.Lock()
	defer c.Unlock()
	conn, exists = c.connections[sessid]
	if !exists {
		c.connections[sessid] = f()
		conn = c.connections[sessid]
	}
	return
}

func (c *connections) delete(sessid string) {
	c.Lock()
	defer c.Unlock()
	delete(c.connections, sessid)
}
<|endoftext|>