<commit_before>// package merkledag implements the ipfs Merkle DAG datastructures.
package merkledag

import (
	"fmt"
	"sync"

	blocks "github.com/ipfs/go-ipfs/blocks"
	key "github.com/ipfs/go-ipfs/blocks/key"
	bserv "github.com/ipfs/go-ipfs/blockservice"
	logging "gx/ipfs/QmYtB7Qge8cJpXc4irsEp8zRqfnZMBeB7aTrMEkPk67DRv/go-log"
	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

var log = logging.Logger("merkledag")
var ErrNotFound = fmt.Errorf("merkledag: not found")

// DAGService is an IPFS Merkle DAG service.
type DAGService interface {
	Add(*Node) (key.Key, error)
	Get(context.Context, key.Key) (*Node, error)
	Remove(*Node) error

	// GetDAG returns, in order, all the single leve child
	// nodes of the passed in node.
	GetMany(context.Context, []key.Key) <-chan *NodeOption

	Batch() *Batch
}

func NewDAGService(bs *bserv.BlockService) DAGService {
	return &dagService{bs}
}

// dagService is an IPFS Merkle DAG service.
// - the root is virtual (like a forest)
// - stores nodes' data in a BlockService
// TODO: should cache Nodes that are in memory, and be
//       able to free some of them when vm pressure is high
type dagService struct {
	Blocks *bserv.BlockService
}

// Add adds a node to the dagService, storing the block in the BlockService
func (n *dagService) Add(nd *Node) (key.Key, error) {
	if n == nil { // FIXME remove this assertion. protect with constructor invariant
		return "", fmt.Errorf("dagService is nil")
	}

	d, err := nd.EncodeProtobuf(false)
	if err != nil {
		return "", err
	}

	mh, err := nd.Multihash()
	if err != nil {
		return "", err
	}

	b, _ := blocks.NewBlockWithHash(d, mh)

	return n.Blocks.AddBlock(b)
}

func (n *dagService) Batch() *Batch {
	return &Batch{ds: n, MaxSize: 8 * 1024 * 1024}
}

// Get retrieves a node from the dagService, fetching the block in the BlockService
func (n *dagService) Get(ctx context.Context, k key.Key) (*Node, error) {
	if k == "" {
		return nil, ErrNotFound
	}
	if n == nil {
		return nil, fmt.Errorf("dagService is nil")
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	b, err := n.Blocks.GetBlock(ctx, k)
	if err != nil {
		if err == bserv.ErrNotFound {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("Failed to get block for %s: %v", k.B58String(), err)
	}

	res, err := DecodeProtobuf(b.Data())
	if err != nil {
		return nil, fmt.Errorf("Failed to decode Protocol Buffers: %v", err)
	}
	return res, nil
}

func (n *dagService) Remove(nd *Node) error {
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

// FetchGraph fetches all nodes that are children of the given node
func FetchGraph(ctx context.Context, root *Node, serv DAGService) error {
	return EnumerateChildrenAsync(ctx, serv, root, key.NewKeySet())
}

// FindLinks searches this nodes links for the given key,
// returns the indexes of any links pointing to it
func FindLinks(links []key.Key, k key.Key, start int) []int {
	var out []int
	for i, lnk_k := range links[start:] {
		if k == lnk_k {
			out = append(out, i+start)
		}
	}
	return out
}

type NodeOption struct {
	Node *Node
	Err  error
}

func (ds *dagService) GetMany(ctx context.Context, keys []key.Key) <-chan *NodeOption {
	out := make(chan *NodeOption, len(keys))
	blocks := ds.Blocks.GetBlocks(ctx, keys)
	var count int

	go func() {
		defer close(out)
		for {
			select {
			case b, ok := <-blocks:
				if !ok {
					if count != len(keys) {
						out <- &NodeOption{Err: fmt.Errorf("failed to fetch all nodes")}
					}
					return
				}
				nd, err := DecodeProtobuf(b.Data())
				if err != nil {
					out <- &NodeOption{Err: err}
					return
				}

				// buffered, no need to select
				out <- &NodeOption{Node: nd}
				count++

			case <-ctx.Done():
				out <- &NodeOption{Err: ctx.Err()}
				return
			}
		}
	}()
	return out
}

// GetDAG will fill out all of the links of the given Node.
// It returns a channel of nodes, which the caller can receive
// all the child nodes of 'root' on, in proper order.
func GetDAG(ctx context.Context, ds DAGService, root *Node) []NodeGetter {
	var keys []key.Key
	for _, lnk := range root.Links {
		keys = append(keys, key.Key(lnk.Hash))
	}

	return GetNodes(ctx, ds, keys)
}

// GetNodes returns an array of 'NodeGetter' promises, with each corresponding
// to the key with the same index as the passed in keys
func GetNodes(ctx context.Context, ds DAGService, keys []key.Key) []NodeGetter {

	// Early out if no work to do
	if len(keys) == 0 {
		return nil
	}

	promises := make([]NodeGetter, len(keys))
	for i := range keys {
		promises[i] = newNodePromise(ctx)
	}

	dedupedKeys := dedupeKeys(keys)
	go func() {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		nodechan := ds.GetMany(ctx, dedupedKeys)

		for count := 0; count < len(keys); {
			select {
			case opt, ok := <-nodechan:
				if !ok {
					for _, p := range promises {
						p.Fail(ErrNotFound)
					}
					return
				}

				if opt.Err != nil {
					for _, p := range promises {
						p.Fail(opt.Err)
					}
					return
				}

				nd := opt.Node

				k, err := nd.Key()
				if err != nil {
					log.Error("Failed to get node key: ", err)
					continue
				}

				is := FindLinks(keys, k, 0)
				for _, i := range is {
					count++
					promises[i].Send(nd)
				}
			case <-ctx.Done():
				return
			}
		}
	}()
	return promises
}

// Remove duplicates from a list of keys
func dedupeKeys(ks []key.Key) []key.Key {
	kmap := make(map[key.Key]struct{})
	var out []key.Key
	for _, k := range ks {
		if _, ok := kmap[k]; !ok {
			kmap[k] = struct{}{}
			out = append(out, k)
		}
	}
	return out
}

func newNodePromise(ctx context.Context) NodeGetter {
	return &nodePromise{
		recv: make(chan *Node, 1),
		ctx:  ctx,
		err:  make(chan error, 1),
	}
}

type nodePromise struct {
	cache *Node
	clk   sync.Mutex
	recv  chan *Node
	ctx   context.Context
	err   chan error
}

// NodeGetter provides a promise like interface for a dag Node
// the first call to Get will block until the Node is received
// from its internal channels, subsequent calls will return the
// cached node.
type NodeGetter interface {
	Get(context.Context) (*Node, error)
	Fail(err error)
	Send(*Node)
}

func (np *nodePromise) Fail(err error) {
	np.clk.Lock()
	v := np.cache
	np.clk.Unlock()

	// if promise has a value, don't fail it
	if v != nil {
		return
	}

	np.err <- err
}

func (np *nodePromise) Send(nd *Node) {
	var already bool
	np.clk.Lock()
	if np.cache != nil {
		already = true
	}
	np.cache = nd
	np.clk.Unlock()

	if already {
		panic("sending twice to the same promise is an error!")
	}

	np.recv <- nd
}

func (np *nodePromise) Get(ctx context.Context) (*Node, error) {
	np.clk.Lock()
	c := np.cache
	np.clk.Unlock()
	if c != nil {
		return c, nil
	}

	select {
	case nd := <-np.recv:
		return nd, nil
	case <-np.ctx.Done():
		return nil, np.ctx.Err()
	case <-ctx.Done():
		return nil, ctx.Err()
	case err := <-np.err:
		return nil, err
	}
}

type Batch struct {
	ds *dagService

	blocks  []blocks.Block
	size    int
	MaxSize int
}

func (t *Batch) Add(nd *Node) (key.Key, error) {
	d, err := nd.EncodeProtobuf(false)
	if err != nil {
		return "", err
	}

	mh, err := nd.Multihash()
	if err != nil {
		return "", err
	}

	b, _ := blocks.NewBlockWithHash(d, mh)

	k := key.Key(mh)

	t.blocks = append(t.blocks, b)
	t.size += len(b.Data())
	if t.size > t.MaxSize {
		return k, t.Commit()
	}
	return k, nil
}

func (t *Batch) Commit() error {
	_, err := t.ds.Blocks.AddBlocks(t.blocks)
	t.blocks = nil
	t.size = 0
	return err
}

// EnumerateChildren will walk the dag below the given root node and add all
// unseen children to the passed in set.
// TODO: parallelize to avoid disk latency perf hits?
func EnumerateChildren(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	for _, lnk := range root.Links {
		k := key.Key(lnk.Hash)
		if !set.Has(k) {
			set.Add(k)
			child, err := ds.Get(ctx, k)
			if err != nil {
				return err
			}
			err = EnumerateChildren(ctx, ds, child, set)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func EnumerateChildrenAsync(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	toprocess := make(chan []key.Key, 8)
	nodes := make(chan *NodeOption, 8)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(toprocess)

	go fetchNodes(ctx, ds, toprocess, nodes)

	nodes <- &NodeOption{Node: root}
	live := 1

	for {
		select {
		case opt, ok := <-nodes:
			if !ok {
				return nil
			}

			if opt.Err != nil {
				return opt.Err
			}

			nd := opt.Node

			// a node has been fetched
			live--

			var keys []key.Key
			for _, lnk := range nd.Links {
				k := key.Key(lnk.Hash)
				if !set.Has(k) {
					set.Add(k)
					live++
					keys = append(keys, k)
				}
			}

			if live == 0 {
				return nil
			}

			if len(keys) > 0 {
				select {
				case toprocess <- keys:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func fetchNodes(ctx context.Context, ds DAGService, in <-chan []key.Key, out chan<- *NodeOption) {
	var wg sync.WaitGroup
	defer func() {
		// wait for all 'get' calls to complete so we don't accidentally send
		// on a closed channel
		wg.Wait()
		close(out)
	}()

	get := func(ks []key.Key) {
		defer wg.Done()
		nodes := ds.GetMany(ctx, ks)
		for opt := range nodes {
			select {
			case out <- opt:
			case <-ctx.Done():
				return
			}
		}
	}

	for ks := range in {
		wg.Add(1)
		go get(ks)
	}
}
<commit_msg>return a better error if the ref is not an object<commit_after>// package merkledag implements the ipfs Merkle DAG datastructures.
package merkledag

import (
	"fmt"
	"strings"
	"sync"

	blocks "github.com/ipfs/go-ipfs/blocks"
	key "github.com/ipfs/go-ipfs/blocks/key"
	bserv "github.com/ipfs/go-ipfs/blockservice"
	logging "gx/ipfs/QmYtB7Qge8cJpXc4irsEp8zRqfnZMBeB7aTrMEkPk67DRv/go-log"
	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

var log = logging.Logger("merkledag")
var ErrNotFound = fmt.Errorf("merkledag: not found")

// DAGService is an IPFS Merkle DAG service.
type DAGService interface {
	Add(*Node) (key.Key, error)
	Get(context.Context, key.Key) (*Node, error)
	Remove(*Node) error

	// GetDAG returns, in order, all the single leve child
	// nodes of the passed in node.
	GetMany(context.Context, []key.Key) <-chan *NodeOption

	Batch() *Batch
}

func NewDAGService(bs *bserv.BlockService) DAGService {
	return &dagService{bs}
}

// dagService is an IPFS Merkle DAG service.
// - the root is virtual (like a forest)
// - stores nodes' data in a BlockService
// TODO: should cache Nodes that are in memory, and be
//       able to free some of them when vm pressure is high
type dagService struct {
	Blocks *bserv.BlockService
}

// Add adds a node to the dagService, storing the block in the BlockService
func (n *dagService) Add(nd *Node) (key.Key, error) {
	if n == nil { // FIXME remove this assertion. protect with constructor invariant
		return "", fmt.Errorf("dagService is nil")
	}

	d, err := nd.EncodeProtobuf(false)
	if err != nil {
		return "", err
	}

	mh, err := nd.Multihash()
	if err != nil {
		return "", err
	}

	b, _ := blocks.NewBlockWithHash(d, mh)

	return n.Blocks.AddBlock(b)
}

func (n *dagService) Batch() *Batch {
	return &Batch{ds: n, MaxSize: 8 * 1024 * 1024}
}

// Get retrieves a node from the dagService, fetching the block in the BlockService
func (n *dagService) Get(ctx context.Context, k key.Key) (*Node, error) {
	if k == "" {
		return nil, ErrNotFound
	}
	if n == nil {
		return nil, fmt.Errorf("dagService is nil")
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	b, err := n.Blocks.GetBlock(ctx, k)
	if err != nil {
		if err == bserv.ErrNotFound {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("Failed to get block for %s: %v", k.B58String(), err)
	}

	res, err := DecodeProtobuf(b.Data())
	if err != nil {
		if strings.Contains(err.Error(), "Unmarshal failed") {
			return nil, fmt.Errorf("%s was not a valid merkledag node", k)
		}
		return nil, fmt.Errorf("Failed to decode Protocol Buffers: %v", err)
	}
	return res, nil
}

func (n *dagService) Remove(nd *Node) error {
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

// FetchGraph fetches all nodes that are children of the given node
func FetchGraph(ctx context.Context, root *Node, serv DAGService) error {
	return EnumerateChildrenAsync(ctx, serv, root, key.NewKeySet())
}

// FindLinks searches this nodes links for the given key,
// returns the indexes of any links pointing to it
func FindLinks(links []key.Key, k key.Key, start int) []int {
	var out []int
	for i, lnk_k := range links[start:] {
		if k == lnk_k {
			out = append(out, i+start)
		}
	}
	return out
}

type NodeOption struct {
	Node *Node
	Err  error
}

func (ds *dagService) GetMany(ctx context.Context, keys []key.Key) <-chan *NodeOption {
	out := make(chan *NodeOption, len(keys))
	blocks := ds.Blocks.GetBlocks(ctx, keys)
	var count int

	go func() {
		defer close(out)
		for {
			select {
			case b, ok := <-blocks:
				if !ok {
					if count != len(keys) {
						out <- &NodeOption{Err: fmt.Errorf("failed to fetch all nodes")}
					}
					return
				}
				nd, err := DecodeProtobuf(b.Data())
				if err != nil {
					out <- &NodeOption{Err: err}
					return
				}

				// buffered, no need to select
				out <- &NodeOption{Node: nd}
				count++

			case <-ctx.Done():
				out <- &NodeOption{Err: ctx.Err()}
				return
			}
		}
	}()
	return out
}

// GetDAG will fill out all of the links of the given Node.
// It returns a channel of nodes, which the caller can receive
// all the child nodes of 'root' on, in proper order.
func GetDAG(ctx context.Context, ds DAGService, root *Node) []NodeGetter {
	var keys []key.Key
	for _, lnk := range root.Links {
		keys = append(keys, key.Key(lnk.Hash))
	}

	return GetNodes(ctx, ds, keys)
}

// GetNodes returns an array of 'NodeGetter' promises, with each corresponding
// to the key with the same index as the passed in keys
func GetNodes(ctx context.Context, ds DAGService, keys []key.Key) []NodeGetter {

	// Early out if no work to do
	if len(keys) == 0 {
		return nil
	}

	promises := make([]NodeGetter, len(keys))
	for i := range keys {
		promises[i] = newNodePromise(ctx)
	}

	dedupedKeys := dedupeKeys(keys)
	go func() {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		nodechan := ds.GetMany(ctx, dedupedKeys)

		for count := 0; count < len(keys); {
			select {
			case opt, ok := <-nodechan:
				if !ok {
					for _, p := range promises {
						p.Fail(ErrNotFound)
					}
					return
				}

				if opt.Err != nil {
					for _, p := range promises {
						p.Fail(opt.Err)
					}
					return
				}

				nd := opt.Node

				k, err := nd.Key()
				if err != nil {
					log.Error("Failed to get node key: ", err)
					continue
				}

				is := FindLinks(keys, k, 0)
				for _, i := range is {
					count++
					promises[i].Send(nd)
				}
			case <-ctx.Done():
				return
			}
		}
	}()
	return promises
}

// Remove duplicates from a list of keys
func dedupeKeys(ks []key.Key) []key.Key {
	kmap := make(map[key.Key]struct{})
	var out []key.Key
	for _, k := range ks {
		if _, ok := kmap[k]; !ok {
			kmap[k] = struct{}{}
			out = append(out, k)
		}
	}
	return out
}

func newNodePromise(ctx context.Context) NodeGetter {
	return &nodePromise{
		recv: make(chan *Node, 1),
		ctx:  ctx,
		err:  make(chan error, 1),
	}
}

type nodePromise struct {
	cache *Node
	clk   sync.Mutex
	recv  chan *Node
	ctx   context.Context
	err   chan error
}

// NodeGetter provides a promise like interface for a dag Node
// the first call to Get will block until the Node is received
// from its internal channels, subsequent calls will return the
// cached node.
type NodeGetter interface {
	Get(context.Context) (*Node, error)
	Fail(err error)
	Send(*Node)
}

func (np *nodePromise) Fail(err error) {
	np.clk.Lock()
	v := np.cache
	np.clk.Unlock()

	// if promise has a value, don't fail it
	if v != nil {
		return
	}

	np.err <- err
}

func (np *nodePromise) Send(nd *Node) {
	var already bool
	np.clk.Lock()
	if np.cache != nil {
		already = true
	}
	np.cache = nd
	np.clk.Unlock()

	if already {
		panic("sending twice to the same promise is an error!")
	}

	np.recv <- nd
}

func (np *nodePromise) Get(ctx context.Context) (*Node, error) {
	np.clk.Lock()
	c := np.cache
	np.clk.Unlock()
	if c != nil {
		return c, nil
	}

	select {
	case nd := <-np.recv:
		return nd, nil
	case <-np.ctx.Done():
		return nil, np.ctx.Err()
	case <-ctx.Done():
		return nil, ctx.Err()
	case err := <-np.err:
		return nil, err
	}
}

type Batch struct {
	ds *dagService

	blocks  []blocks.Block
	size    int
	MaxSize int
}

func (t *Batch) Add(nd *Node) (key.Key, error) {
	d, err := nd.EncodeProtobuf(false)
	if err != nil {
		return "", err
	}

	mh, err := nd.Multihash()
	if err != nil {
		return "", err
	}

	b, _ := blocks.NewBlockWithHash(d, mh)

	k := key.Key(mh)

	t.blocks = append(t.blocks, b)
	t.size += len(b.Data())
	if t.size > t.MaxSize {
		return k, t.Commit()
	}
	return k, nil
}

func (t *Batch) Commit() error {
	_, err := t.ds.Blocks.AddBlocks(t.blocks)
	t.blocks = nil
	t.size = 0
	return err
}

// EnumerateChildren will walk the dag below the given root node and add all
// unseen children to the passed in set.
// TODO: parallelize to avoid disk latency perf hits?
func EnumerateChildren(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	for _, lnk := range root.Links {
		k := key.Key(lnk.Hash)
		if !set.Has(k) {
			set.Add(k)
			child, err := ds.Get(ctx, k)
			if err != nil {
				return err
			}
			err = EnumerateChildren(ctx, ds, child, set)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func EnumerateChildrenAsync(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	toprocess := make(chan []key.Key, 8)
	nodes := make(chan *NodeOption, 8)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(toprocess)

	go fetchNodes(ctx, ds, toprocess, nodes)

	nodes <- &NodeOption{Node: root}
	live := 1

	for {
		select {
		case opt, ok := <-nodes:
			if !ok {
				return nil
			}

			if opt.Err != nil {
				return opt.Err
			}

			nd := opt.Node

			// a node has been fetched
			live--

			var keys []key.Key
			for _, lnk := range nd.Links {
				k := key.Key(lnk.Hash)
				if !set.Has(k) {
					set.Add(k)
					live++
					keys = append(keys, k)
				}
			}

			if live == 0 {
				return nil
			}

			if len(keys) > 0 {
				select {
				case toprocess <- keys:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func fetchNodes(ctx context.Context, ds DAGService, in <-chan []key.Key, out chan<- *NodeOption) {
	var wg sync.WaitGroup
	defer func() {
		// wait for all 'get' calls to complete so we don't accidentally send
		// on a closed channel
		wg.Wait()
		close(out)
	}()

	get := func(ks []key.Key) {
		defer wg.Done()
		nodes := ds.GetMany(ctx, ks)
		for opt := range nodes {
			select {
			case out <- opt:
			case <-ctx.Done():
				return
			}
		}
	}

	for ks := range in {
		wg.Add(1)
		go get(ks)
	}
}
<|endoftext|>


<commit_before>package mesh

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"sync"
)

type GossipChannel struct {
	sync.Mutex
	name         string
	ourself      *LocalPeer
	routes       *Routes
	gossiper     Gossiper
	senders      connectionSenders
	broadcasters peerSenders
}

type connectionSenders map[Connection]*GossipSender
type peerSenders map[PeerName]*GossipSender

func NewGossipChannel(channelName string, ourself *LocalPeer, routes *Routes, g Gossiper) *GossipChannel {
	return &GossipChannel{
		name:         channelName,
		ourself:      ourself,
		routes:       routes,
		gossiper:     g,
		senders:      make(connectionSenders),
		broadcasters: make(peerSenders)}
}

func (router *Router) handleGossip(tag ProtocolTag, payload []byte) error {
	decoder := gob.NewDecoder(bytes.NewReader(payload))
	var channelName string
	if err := decoder.Decode(&channelName); err != nil {
		return err
	}
	channel := router.gossipChannel(channelName)
	var srcName PeerName
	if err := decoder.Decode(&srcName); err != nil {
		return err
	}
	switch tag {
	case ProtocolGossipUnicast:
		return channel.deliverUnicast(srcName, payload, decoder)
	case ProtocolGossipBroadcast:
		return channel.deliverBroadcast(srcName, payload, decoder)
	case ProtocolGossip:
		return channel.deliver(srcName, payload, decoder)
	}
	return nil
}

func (c *GossipChannel) deliverUnicast(srcName PeerName, origPayload []byte, dec *gob.Decoder) error {
	var destName PeerName
	if err := dec.Decode(&destName); err != nil {
		return err
	}
	if c.ourself.Name == destName {
		var payload []byte
		if err := dec.Decode(&payload); err != nil {
			return err
		}
		return c.gossiper.OnGossipUnicast(srcName, payload)
	}
	if err := c.relayUnicast(destName, origPayload); err != nil {
		c.log(err)
	}
	return nil
}

func (c *GossipChannel) deliverBroadcast(srcName PeerName, _ []byte, dec *gob.Decoder) error {
	var payload []byte
	if err := dec.Decode(&payload); err != nil {
		return err
	}
	data, err := c.gossiper.OnGossipBroadcast(srcName, payload)
	if err != nil || data == nil {
		return err
	}
	if err := c.relayBroadcast(srcName, data); err != nil {
		c.log(err)
	}
	return nil
}

func (c *GossipChannel) deliver(srcName PeerName, _ []byte, dec *gob.Decoder) error {
	var payload []byte
	if err := dec.Decode(&payload); err != nil {
		return err
	}
	data, err := c.gossiper.OnGossip(payload)
	if err != nil || data == nil {
		return err
	}
	c.Send(srcName, data)
	return nil
}

func (c *GossipChannel) GossipUnicast(dstPeerName PeerName, msg []byte) error {
	return c.relayUnicast(dstPeerName, GobEncode(c.name, c.ourself.Name, dstPeerName, msg))
}

func (c *GossipChannel) GossipBroadcast(update GossipData) error {
	return c.relayBroadcast(c.ourself.Name, update)
}

func (c *GossipChannel) relayUnicast(dstPeerName PeerName, buf []byte) (err error) {
	if relayPeerName, found := c.routes.UnicastAll(dstPeerName); !found {
		err = fmt.Errorf("unknown relay destination: %s", dstPeerName)
	} else if conn, found := c.ourself.ConnectionTo(relayPeerName); !found {
		err = fmt.Errorf("unable to find connection to relay peer %s", relayPeerName)
	} else {
		conn.(ProtocolSender).SendProtocolMsg(ProtocolMsg{ProtocolGossipUnicast, buf})
	}
	return err
}

func (c *GossipChannel) relayBroadcast(srcName PeerName, update GossipData) error {
	names := c.routes.PeerNames() // do this outside the lock so they don't nest
	c.Lock()
	defer c.Unlock()
	// GC - randomly (courtesy of go's map iterator) pick some
	// existing broadcasters and stop&remove them if their source peer
	// is unknown. We stop as soon as we encounter a valid entry; the
	// idea being that when there is little or no garbage then this
	// executes close to O(1)[1], whereas when there is lots of
	// garbage we remove it quickly.
	//
	// [1] TODO Unfortunately, due to the desire to avoid nested
	// locks, instead of simply invoking Peers.Fetch(name) below, we
	// have that Peers.Names() invocation above. That is O(n_peers) at
	// best.
	for name, broadcaster := range c.broadcasters {
		if _, found := names[name]; !found {
			delete(c.broadcasters, name)
			broadcaster.Stop()
		} else {
			break
		}
	}
	broadcaster, found := c.broadcasters[srcName]
	if !found {
		broadcaster = NewGossipSender(func(pending GossipData) { c.sendBroadcast(srcName, pending) })
		c.broadcasters[srcName] = broadcaster
	}
	broadcaster.Send(update)
	return nil
}

func (c *GossipChannel) sendBroadcast(srcName PeerName, update GossipData) {
	c.routes.EnsureRecalculated()
	nextHops := c.routes.BroadcastAll(srcName)
	if len(nextHops) == 0 {
		return
	}
	connections := c.ourself.ConnectionsTo(nextHops)
	for _, msg := range update.Encode() {
		protocolMsg := ProtocolMsg{ProtocolGossipBroadcast, GobEncode(c.name, srcName, msg)}
		// FIXME a single blocked connection can stall us
		for _, conn := range connections {
			conn.(ProtocolSender).SendProtocolMsg(protocolMsg)
		}
	}
}

func (c *GossipChannel) Send(srcName PeerName, data GossipData) {
	c.routes.EnsureRecalculated()
	destinations := c.routes.RandomNeighbours(srcName)
	if len(destinations) == 0 {
		return
	}
	c.sendDown(c.ourself.ConnectionsTo(destinations), data)
}

func (c *GossipChannel) SendDown(conn Connection, data GossipData) {
	c.sendDown([]Connection{conn}, data)
}

func (c *GossipChannel) sendDown(selectedConnections []Connection, data GossipData) {
	if len(selectedConnections) == 0 {
		return
	}
	connections := c.ourself.Connections()
	c.Lock()
	defer c.Unlock()
	// GC - randomly (courtesy of go's map iterator) pick some
	// existing senders and stop&remove them if the associated
	// connection is no longer active.  We stop as soon as we
	// encounter a valid entry; the idea being that when there is
	// little or no garbage then this executes close to O(1)[1],
	// whereas when there is lots of garbage we remove it quickly.
	//
	// [1] TODO Unfortunately, due to the desire to avoid nested
	// locks, instead of simply invoking LocalPeer.ConnectionTo(name),
	// we operate on LocalPeer.Connections(). That is
	// O(n_our_connections) at best.
	for conn, sender := range c.senders {
		if _, found := connections[conn]; !found {
			delete(c.senders, conn)
			sender.Stop()
		} else {
			break
		}
	}
	// start senders, if necessary, and send.
	for _, conn := range selectedConnections {
		sender, found := c.senders[conn]
		if !found {
			sender = c.makeSender(conn)
			c.senders[conn] = sender
		}
		sender.Send(data)
	}
}

// We have seen a couple of failures which suggest a >128GB slice was encountered.
// 100MB should be enough for anyone.
const maxFeasibleMessageLen = 100 * 1024 * 1024

func (c *GossipChannel) makeSender(conn Connection) *GossipSender {
	return NewGossipSender(func(pending GossipData) {
		for _, msg := range pending.Encode() {
			if len(msg) > maxFeasibleMessageLen {
				panic(fmt.Sprintf("Gossip message too large: len=%d bytes; on channel '%s' from %+v", len(msg), c.name, pending))
			}
			protocolMsg := ProtocolMsg{ProtocolGossip, GobEncode(c.name, c.ourself.Name, msg)}
			conn.(ProtocolSender).SendProtocolMsg(protocolMsg)
		}
	})
}

func (c *GossipChannel) log(args ...interface{}) {
	log.Println(append(append([]interface{}{}, "[gossip "+c.name+"]:"), args...)...)
}

func GobEncode(items ...interface{}) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, i := range items {
		checkFatal(enc.Encode(i))
	}
	return buf.Bytes()
}
<commit_msg>cosmetic<commit_after>package mesh

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"sync"
)

type GossipChannel struct {
	sync.Mutex
	name         string
	ourself      *LocalPeer
	routes       *Routes
	gossiper     Gossiper
	senders      connectionSenders
	broadcasters peerSenders
}

type connectionSenders map[Connection]*GossipSender
type peerSenders map[PeerName]*GossipSender

func NewGossipChannel(channelName string, ourself *LocalPeer, routes *Routes, g Gossiper) *GossipChannel {
	return &GossipChannel{
		name:         channelName,
		ourself:      ourself,
		routes:       routes,
		gossiper:     g,
		senders:      make(connectionSenders),
		broadcasters: make(peerSenders)}
}

func (router *Router) handleGossip(tag ProtocolTag, payload []byte) error {
	decoder := gob.NewDecoder(bytes.NewReader(payload))
	var channelName string
	if err := decoder.Decode(&channelName); err != nil {
		return err
	}
	channel := router.gossipChannel(channelName)
	var srcName PeerName
	if err := decoder.Decode(&srcName); err != nil {
		return err
	}
	switch tag {
	case ProtocolGossipUnicast:
		return channel.deliverUnicast(srcName, payload, decoder)
	case ProtocolGossipBroadcast:
		return channel.deliverBroadcast(srcName, payload, decoder)
	case ProtocolGossip:
		return channel.deliver(srcName, payload, decoder)
	}
	return nil
}

func (c *GossipChannel) deliverUnicast(srcName PeerName, origPayload []byte, dec *gob.Decoder) error {
	var destName PeerName
	if err := dec.Decode(&destName); err != nil {
		return err
	}
	if c.ourself.Name == destName {
		var payload []byte
		if err := dec.Decode(&payload); err != nil {
			return err
		}
		return c.gossiper.OnGossipUnicast(srcName, payload)
	}
	if err := c.relayUnicast(destName, origPayload); err != nil {
		c.log(err)
	}
	return nil
}

func (c *GossipChannel) deliverBroadcast(srcName PeerName, _ []byte, dec *gob.Decoder) error {
	var payload []byte
	if err := dec.Decode(&payload); err != nil {
		return err
	}
	data, err := c.gossiper.OnGossipBroadcast(srcName, payload)
	if err != nil || data == nil {
		return err
	}
	if err := c.relayBroadcast(srcName, data); err != nil {
		c.log(err)
	}
	return nil
}

func (c *GossipChannel) deliver(srcName PeerName, _ []byte, dec *gob.Decoder) error {
	var payload []byte
	if err := dec.Decode(&payload); err != nil {
		return err
	}
	update, err := c.gossiper.OnGossip(payload)
	if err != nil || update == nil {
		return err
	}
	c.Send(srcName, update)
	return nil
}

func (c *GossipChannel) GossipUnicast(dstPeerName PeerName, msg []byte) error {
	return c.relayUnicast(dstPeerName, GobEncode(c.name, c.ourself.Name, dstPeerName, msg))
}

func (c *GossipChannel) GossipBroadcast(update GossipData) error {
	return c.relayBroadcast(c.ourself.Name, update)
}

func (c *GossipChannel) relayUnicast(dstPeerName PeerName, buf []byte) (err error) {
	if relayPeerName, found := c.routes.UnicastAll(dstPeerName); !found {
		err = fmt.Errorf("unknown relay destination: %s", dstPeerName)
	} else if conn, found := c.ourself.ConnectionTo(relayPeerName); !found {
		err = fmt.Errorf("unable to find connection to relay peer %s", relayPeerName)
	} else {
		conn.(ProtocolSender).SendProtocolMsg(ProtocolMsg{ProtocolGossipUnicast, buf})
	}
	return err
}

func (c *GossipChannel) relayBroadcast(srcName PeerName, update GossipData) error {
	names := c.routes.PeerNames() // do this outside the lock so they don't nest
	c.Lock()
	defer c.Unlock()
	// GC - randomly (courtesy of go's map iterator) pick some
	// existing broadcasters and stop&remove them if their source peer
	// is unknown. We stop as soon as we encounter a valid entry; the
	// idea being that when there is little or no garbage then this
	// executes close to O(1)[1], whereas when there is lots of
	// garbage we remove it quickly.
	//
	// [1] TODO Unfortunately, due to the desire to avoid nested
	// locks, instead of simply invoking Peers.Fetch(name) below, we
	// have that Peers.Names() invocation above. That is O(n_peers) at
	// best.
	for name, broadcaster := range c.broadcasters {
		if _, found := names[name]; !found {
			delete(c.broadcasters, name)
			broadcaster.Stop()
		} else {
			break
		}
	}
	broadcaster, found := c.broadcasters[srcName]
	if !found {
		broadcaster = NewGossipSender(func(pending GossipData) { c.sendBroadcast(srcName, pending) })
		c.broadcasters[srcName] = broadcaster
	}
	broadcaster.Send(update)
	return nil
}

func (c *GossipChannel) sendBroadcast(srcName PeerName, update GossipData) {
	c.routes.EnsureRecalculated()
	nextHops := c.routes.BroadcastAll(srcName)
	if len(nextHops) == 0 {
		return
	}
	connections := c.ourself.ConnectionsTo(nextHops)
	for _, msg := range update.Encode() {
		protocolMsg := ProtocolMsg{ProtocolGossipBroadcast, GobEncode(c.name, srcName, msg)}
		// FIXME a single blocked connection can stall us
		for _, conn := range connections {
			conn.(ProtocolSender).SendProtocolMsg(protocolMsg)
		}
	}
}

func (c *GossipChannel) Send(srcName PeerName, data GossipData) {
	c.routes.EnsureRecalculated()
	destinations := c.routes.RandomNeighbours(srcName)
	if len(destinations) == 0 {
		return
	}
	c.sendDown(c.ourself.ConnectionsTo(destinations), data)
}

func (c *GossipChannel) SendDown(conn Connection, data GossipData) {
	c.sendDown([]Connection{conn}, data)
}

func (c *GossipChannel) sendDown(selectedConnections []Connection, data GossipData) {
	if len(selectedConnections) == 0 {
		return
	}
	connections := c.ourself.Connections()
	c.Lock()
	defer c.Unlock()
	// GC - randomly (courtesy of go's map iterator) pick some
	// existing senders and stop&remove them if the associated
	// connection is no longer active.  We stop as soon as we
	// encounter a valid entry; the idea being that when there is
	// little or no garbage then this executes close to O(1)[1],
	// whereas when there is lots of garbage we remove it quickly.
	//
	// [1] TODO Unfortunately, due to the desire to avoid nested
	// locks, instead of simply invoking LocalPeer.ConnectionTo(name),
	// we operate on LocalPeer.Connections(). That is
	// O(n_our_connections) at best.
	for conn, sender := range c.senders {
		if _, found := connections[conn]; !found {
			delete(c.senders, conn)
			sender.Stop()
		} else {
			break
		}
	}
	// start senders, if necessary, and send.
	for _, conn := range selectedConnections {
		sender, found := c.senders[conn]
		if !found {
			sender = c.makeSender(conn)
			c.senders[conn] = sender
		}
		sender.Send(data)
	}
}

// We have seen a couple of failures which suggest a >128GB slice was encountered.
// 100MB should be enough for anyone.
const maxFeasibleMessageLen = 100 * 1024 * 1024

func (c *GossipChannel) makeSender(conn Connection) *GossipSender {
	return NewGossipSender(func(pending GossipData) {
		for _, msg := range pending.Encode() {
			if len(msg) > maxFeasibleMessageLen {
				panic(fmt.Sprintf("Gossip message too large: len=%d bytes; on channel '%s' from %+v", len(msg), c.name, pending))
			}
			protocolMsg := ProtocolMsg{ProtocolGossip, GobEncode(c.name, c.ourself.Name, msg)}
			conn.(ProtocolSender).SendProtocolMsg(protocolMsg)
		}
	})
}

func (c *GossipChannel) log(args ...interface{}) {
	log.Println(append(append([]interface{}{}, "[gossip "+c.name+"]:"), args...)...)
}

func GobEncode(items ...interface{}) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, i := range items {
		checkFatal(enc.Encode(i))
	}
	return buf.Bytes()
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"time"
)

var Version = "Development"

var overrideableBazelFlags []string = []string{
	"--test_output=",
	"--config=",
	"--curses=no",
	"--output_groups=",
	"--keep_going",
	"-k",
	"--strategy",
	"--define",
	"--override_repository",
	"--features=",
	"--test_arg",
}

var debounceDuration = flag.Duration("debounce", 100*time.Millisecond, "Debounce duration")
var logToFile = flag.String("log_to_file", "-", "Log iBazel stderr to a file instead of os.Stderr")

func usage() {
	fmt.Fprintf(os.Stderr, `iBazel - Version %s

A file watcher for Bazel. Whenever a source file used in a specified
target, run, build, or test the specified targets.

Usage:

ibazel [flags] build|test|run targets...

Example:

ibazel test //path/to/my/testing:target
ibazel test //path/to/my/testing/targets/...
ibazel run //path/to/my/runnable:target -- --arguments --for_your=binary
ibazel build //path/to/my/buildable:target

Supported Bazel flags:
  %s

To add to this list, edit
https://github.com/bazelbuild/bazel-watcher/blob/master/ibazel/main.go

iBazel flags:
`, Version, strings.Join(overrideableBazelFlags, "\n  "))
	flag.PrintDefaults()
}

func isOverrideableBazelFlag(arg string) bool {
	for _, overrideable := range overrideableBazelFlags {
		if strings.HasPrefix(arg, overrideable) {
			return true
		}
	}
	return false
}

func parseArgs(in []string) (targets, bazelArgs, args []string) {
	afterDoubleDash := false
	for _, arg := range in {
		if afterDoubleDash {
			// Put it in the extra args section if we are after a double dash.
			args = append(args, arg)
		} else {
			// Check to see if this token is a double dash.
			if arg == "--" {
				afterDoubleDash = true
				continue
			}

			// Check to see if this flag is on the bazel whitelist of flags.
			if isOverrideableBazelFlag(arg) {
				bazelArgs = append(bazelArgs, arg)
				continue
			}

			// If none of those things then it's probably a target.
			targets = append(targets, arg)
		}
	}
	return
}

// main entrypoint for IBazel.
func main() {
	flag.Usage = usage
	flag.Parse()

	if *logToFile != "-" {
		var err error
		logFile, err := os.OpenFile(*logToFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			panic(err)
		}
		os.Stderr = logFile
	}

	if len(flag.Args()) < 2 {
		usage()
		return
	}

	command := strings.ToLower(flag.Args()[0])
	args := flag.Args()[1:]

	i, err := New()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating iBazel: %s\n", err)
		os.Exit(1)
	}
	i.SetDebounceDuration(*debounceDuration)
	defer i.Cleanup()

	// increase the number of files that this process can
	// have open.
	err = setUlimit()
	if err != nil {
		fmt.Fprint(os.Stderr, "error setting higher file descriptor limit for this process: ", err)
	}

	handle(i, command, args)
}

func handle(i *IBazel, command string, args []string) {
	targets, bazelArgs, args := parseArgs(args)
	i.SetBazelArgs(bazelArgs)

	switch command {
	case "build":
		i.Build(targets...)
	case "test":
		i.Test(targets...)
	case "run":
		// Run only takes one argument
		i.Run(targets[0], args)
	default:
		fmt.Fprintf(os.Stderr, "Asked me to perform %s. I don't know how to do that.", command)
		usage()
		return
	}
}
<commit_msg>Add --test_filter flag to ibazel (#216)<commit_after>// Copyright 2017 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"time"
)

var Version = "Development"

var overrideableBazelFlags []string = []string{
	"--test_output=",
	"--config=",
	"--curses=no",
	"--output_groups=",
	"--keep_going",
	"-k",
	"--strategy",
	"--define",
	"--override_repository",
	"--features=",
	"--test_arg",
	"--test_filter",
}

var debounceDuration = flag.Duration("debounce", 100*time.Millisecond, "Debounce duration")
var logToFile = flag.String("log_to_file", "-", "Log iBazel stderr to a file instead of os.Stderr")

func usage() {
	fmt.Fprintf(os.Stderr, `iBazel - Version %s

A file watcher for Bazel. Whenever a source file used in a specified
target, run, build, or test the specified targets.

Usage:

ibazel [flags] build|test|run targets...

Example:

ibazel test //path/to/my/testing:target
ibazel test //path/to/my/testing/targets/...
ibazel run //path/to/my/runnable:target -- --arguments --for_your=binary
ibazel build //path/to/my/buildable:target

Supported Bazel flags:
  %s

To add to this list, edit
https://github.com/bazelbuild/bazel-watcher/blob/master/ibazel/main.go

iBazel flags:
`, Version, strings.Join(overrideableBazelFlags, "\n  "))
	flag.PrintDefaults()
}

func isOverrideableBazelFlag(arg string) bool {
	for _, overrideable := range overrideableBazelFlags {
		if strings.HasPrefix(arg, overrideable) {
			return true
		}
	}
	return false
}

func parseArgs(in []string) (targets, bazelArgs, args []string) {
	afterDoubleDash := false
	for _, arg := range in {
		if afterDoubleDash {
			// Put it in the extra args section if we are after a double dash.
			args = append(args, arg)
		} else {
			// Check to see if this token is a double dash.
			if arg == "--" {
				afterDoubleDash = true
				continue
			}

			// Check to see if this flag is on the bazel whitelist of flags.
			if isOverrideableBazelFlag(arg) {
				bazelArgs = append(bazelArgs, arg)
				continue
			}

			// If none of those things then it's probably a target.
			targets = append(targets, arg)
		}
	}
	return
}

// main entrypoint for IBazel.
func main() {
	flag.Usage = usage
	flag.Parse()

	if *logToFile != "-" {
		var err error
		logFile, err := os.OpenFile(*logToFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			panic(err)
		}
		os.Stderr = logFile
	}

	if len(flag.Args()) < 2 {
		usage()
		return
	}

	command := strings.ToLower(flag.Args()[0])
	args := flag.Args()[1:]

	i, err := New()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating iBazel: %s\n", err)
		os.Exit(1)
	}
	i.SetDebounceDuration(*debounceDuration)
	defer i.Cleanup()

	// increase the number of files that this process can
	// have open.
	err = setUlimit()
	if err != nil {
		fmt.Fprint(os.Stderr, "error setting higher file descriptor limit for this process: ", err)
	}

	handle(i, command, args)
}

func handle(i *IBazel, command string, args []string) {
	targets, bazelArgs, args := parseArgs(args)
	i.SetBazelArgs(bazelArgs)

	switch command {
	case "build":
		i.Build(targets...)
	case "test":
		i.Test(targets...)
	case "run":
		// Run only takes one argument
		i.Run(targets[0], args)
	default:
		fmt.Fprintf(os.Stderr, "Asked me to perform %s. I don't know how to do that.", command)
		usage()
		return
	}
}
<|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 (
	"bytes"
	"fmt"
	"io"
	"net"
	"time"

	logging "github.com/op/go-logging"

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

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

func ensureBytes(reader io.Reader, buf []byte, bytesNeeded int) ([]byte, error) {
	if cap(buf) < bytesNeeded {
		return buf, io.ErrShortBuffer
	}
	bytesToRead := bytesNeeded - len(buf)
	if bytesToRead <= 0 {
		return buf, nil
	}
	n, err := io.ReadFull(reader, buf[len(buf):bytesNeeded])
	buf = buf[:len(buf)+n]
	if (err == nil || err == io.EOF) && n < bytesToRead {
		err = io.ErrUnexpectedEOF
	}
	return buf, err
}

func findAccessKey(clientConn onet.DuplexConn, cipherList map[string]shadowaead.Cipher) (string, onet.DuplexConn, error) {
	// This must have enough space to hold the salt + 2 bytes chunk length + AEAD tag (Oeverhead) for any cipher
	replayBytes := make([]byte, 0, 32+2+16)
	// Constant of zeroes to use as the start chunk count. This must be as big as the max NonceSize() across all ciphers.
	zeroCountBuf := make([]byte, 12) // MaxCountSize
	// To hold the decrypted chunk length.
	chunkLenBuf := [2]byte{}
	var err error

	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	// TODO: Reorder list to try previously successful ciphers first for the client IP.
	// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
	for id, cipher := range cipherList {
		replayBytes, err = ensureBytes(clientConn, replayBytes, cipher.SaltSize())
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		salt := replayBytes[:cipher.SaltSize()]
		aead, err := cipher.Decrypter(salt)
		replayBytes, err = ensureBytes(clientConn, replayBytes, cipher.SaltSize()+2+aead.Overhead())
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		cipherText := replayBytes[cipher.SaltSize() : cipher.SaltSize()+2+aead.Overhead()]
		_, err = aead.Open(chunkLenBuf[:0], zeroCountBuf[:aead.NonceSize()], cipherText, nil)
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		if logger.IsEnabledFor(logging.DEBUG) {
			logger.Debugf("Selected TCP cipher %v", id)
		}
		ssr := NewShadowsocksReader(io.MultiReader(bytes.NewReader(replayBytes), clientConn), cipher)
		ssw := NewShadowsocksWriter(clientConn, cipher)
		return id, onet.WrapConn(clientConn, ssr, ssw).(onet.DuplexConn), nil
	}
	return "", nil, fmt.Errorf("Could not find valid TCP cipher")
}

type tcpService struct {
	listener  *net.TCPListener
	ciphers   *map[string]shadowaead.Cipher
	m         metrics.ShadowsocksMetrics
	isRunning bool
}

func NewTCPService(listener *net.TCPListener, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) TCPService {
	return &tcpService{listener: listener, ciphers: ciphers, m: m}
}

type TCPService interface {
	Start()
	Stop() error
}

func (s *tcpService) Start() {
	s.isRunning = true
	for s.isRunning {
		var clientConn onet.DuplexConn
		clientConn, err := s.listener.AcceptTCP()
		if err != nil {
			if !s.isRunning {
				return
			}
			logger.Errorf("Failed to accept: %v", err)
		}

		go func() (connError *onet.ConnectionError) {
			clientLocation, err := s.m.GetLocation(clientConn.RemoteAddr())
			if err != nil {
				logger.Warningf("Failed location lookup: %v", err)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientConn.RemoteAddr().String())
			s.m.AddOpenTCPConnection(clientLocation)
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			connStart := time.Now()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			keyID := ""
			var proxyMetrics metrics.ProxyMetrics
			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
			defer func() {
				connEnd := time.Now()
				connDuration := connEnd.Sub(connStart)
				clientConn.Close()
				status := "OK"
				if connError != nil {
					logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				logger.Debugf("Done with status %v, duration %v", status, connDuration)
				s.m.AddClosedTCPConnection(clientLocation, keyID, status, proxyMetrics, connDuration)
			}()

			keyID, clientConn, err := findAccessKey(clientConn, *s.ciphers)
			if err != nil {
				return &onet.ConnectionError{"ERR_CIPHER", "Failed to find a valid cipher", err}
			}

			tgtAddr, err := socks.ReadAddr(clientConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_READ_ADDRESS", "Failed to get target address", err}
			}
			tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
			if err != nil {
				return &onet.ConnectionError{"ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err}
			}
			if !tgtTCPAddr.IP.IsGlobalUnicast() {
				return &onet.ConnectionError{"ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err}
			}

			tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
			if err != nil {
				return &onet.ConnectionError{"ERR_CONNECT", "Failed to connect to target", err}
			}
			defer tgtTCPConn.Close()
			tgtTCPConn.SetKeepAlive(true)
			tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			// TODO: Disable logging in production. This is sensitive.
			logger.Debugf("proxy %s <-> %s", clientConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
			_, _, err = onet.Relay(clientConn, tgtConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_RELAY", "Failed to relay traffic", err}
			}
			return nil
		}()
	}
}

func (s *tcpService) Stop() error {
	s.isRunning = false
	return s.listener.Close()
}
<commit_msg>Add comment to ensureBytes<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 (
	"bytes"
	"fmt"
	"io"
	"net"
	"time"

	logging "github.com/op/go-logging"

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

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

// Reads bytes from reader and appends to buf to ensure the needed number of bytes.
// The cpacity of buf must be at least bytesNeeded.
func ensureBytes(reader io.Reader, buf []byte, bytesNeeded int) ([]byte, error) {
	if cap(buf) < bytesNeeded {
		return buf, io.ErrShortBuffer
	}
	bytesToRead := bytesNeeded - len(buf)
	if bytesToRead <= 0 {
		return buf, nil
	}
	n, err := io.ReadFull(reader, buf[len(buf):bytesNeeded])
	buf = buf[:len(buf)+n]
	if (err == nil || err == io.EOF) && n < bytesToRead {
		err = io.ErrUnexpectedEOF
	}
	return buf, err
}

func findAccessKey(clientConn onet.DuplexConn, cipherList map[string]shadowaead.Cipher) (string, onet.DuplexConn, error) {
	// This must have enough space to hold the salt + 2 bytes chunk length + AEAD tag (Oeverhead) for any cipher
	replayBytes := make([]byte, 0, 32+2+16)
	// Constant of zeroes to use as the start chunk count. This must be as big as the max NonceSize() across all ciphers.
	zeroCountBuf := make([]byte, 12) // MaxCountSize
	// To hold the decrypted chunk length.
	chunkLenBuf := [2]byte{}
	var err error

	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	// TODO: Reorder list to try previously successful ciphers first for the client IP.
	// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
	for id, cipher := range cipherList {
		replayBytes, err = ensureBytes(clientConn, replayBytes, cipher.SaltSize())
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		salt := replayBytes[:cipher.SaltSize()]
		aead, err := cipher.Decrypter(salt)
		replayBytes, err = ensureBytes(clientConn, replayBytes, cipher.SaltSize()+2+aead.Overhead())
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		cipherText := replayBytes[cipher.SaltSize() : cipher.SaltSize()+2+aead.Overhead()]
		_, err = aead.Open(chunkLenBuf[:0], zeroCountBuf[:aead.NonceSize()], cipherText, nil)
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed TCP ciper %v: %v", id, err)
			}
			continue
		}
		if logger.IsEnabledFor(logging.DEBUG) {
			logger.Debugf("Selected TCP cipher %v", id)
		}
		ssr := NewShadowsocksReader(io.MultiReader(bytes.NewReader(replayBytes), clientConn), cipher)
		ssw := NewShadowsocksWriter(clientConn, cipher)
		return id, onet.WrapConn(clientConn, ssr, ssw).(onet.DuplexConn), nil
	}
	return "", nil, fmt.Errorf("Could not find valid TCP cipher")
}

type tcpService struct {
	listener  *net.TCPListener
	ciphers   *map[string]shadowaead.Cipher
	m         metrics.ShadowsocksMetrics
	isRunning bool
}

func NewTCPService(listener *net.TCPListener, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) TCPService {
	return &tcpService{listener: listener, ciphers: ciphers, m: m}
}

type TCPService interface {
	Start()
	Stop() error
}

func (s *tcpService) Start() {
	s.isRunning = true
	for s.isRunning {
		var clientConn onet.DuplexConn
		clientConn, err := s.listener.AcceptTCP()
		if err != nil {
			if !s.isRunning {
				return
			}
			logger.Errorf("Failed to accept: %v", err)
		}

		go func() (connError *onet.ConnectionError) {
			clientLocation, err := s.m.GetLocation(clientConn.RemoteAddr())
			if err != nil {
				logger.Warningf("Failed location lookup: %v", err)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientConn.RemoteAddr().String())
			s.m.AddOpenTCPConnection(clientLocation)
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			connStart := time.Now()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			keyID := ""
			var proxyMetrics metrics.ProxyMetrics
			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
			defer func() {
				connEnd := time.Now()
				connDuration := connEnd.Sub(connStart)
				clientConn.Close()
				status := "OK"
				if connError != nil {
					logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				logger.Debugf("Done with status %v, duration %v", status, connDuration)
				s.m.AddClosedTCPConnection(clientLocation, keyID, status, proxyMetrics, connDuration)
			}()

			keyID, clientConn, err := findAccessKey(clientConn, *s.ciphers)
			if err != nil {
				return &onet.ConnectionError{"ERR_CIPHER", "Failed to find a valid cipher", err}
			}

			tgtAddr, err := socks.ReadAddr(clientConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_READ_ADDRESS", "Failed to get target address", err}
			}
			tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
			if err != nil {
				return &onet.ConnectionError{"ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err}
			}
			if !tgtTCPAddr.IP.IsGlobalUnicast() {
				return &onet.ConnectionError{"ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err}
			}

			tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
			if err != nil {
				return &onet.ConnectionError{"ERR_CONNECT", "Failed to connect to target", err}
			}
			defer tgtTCPConn.Close()
			tgtTCPConn.SetKeepAlive(true)
			tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			// TODO: Disable logging in production. This is sensitive.
			logger.Debugf("proxy %s <-> %s", clientConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
			_, _, err = onet.Relay(clientConn, tgtConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_RELAY", "Failed to relay traffic", err}
			}
			return nil
		}()
	}
}

func (s *tcpService) Stop() error {
	s.isRunning = false
	return s.listener.Close()
}
<|endoftext|>


<commit_before>package casper

import (
	"fmt"
	"io/ioutil"
	"syscall"
)

func SingleInstane(pidfile string) {
	if e := lockPidFile(pidfile); e != nil {
		pid, _ := ioutil.ReadFile(pidfile)
		panic(fmt.Errorf("process already run: %v", string(pid)))
	}

}

func lockPidFile(pidfile string) error {
	fd, e := syscall.Open(pidfile, syscall.O_CREAT|syscall.O_RDWR, 0777)
	if e != nil {
		return e
	}

	e = syscall.Flock(fd, syscall.LOCK_NB|syscall.LOCK_EX)
	if e != nil {
		return e
	}

	e := syscall.Ftruncate(fd, 0)
	if e != nil {
		return e
	}

	_, e = syscall.Write(fd, []byte(fmt.Sprintf("%d", syscall.Getpid())))
	if e != nil {
		return e
	}

	return nil
}

/*
func main() {
	SingleInstane("/tmp/pid")
	time.Sleep(100 * time.Second)
}
*/
<commit_msg>bug fix<commit_after>package casper

import (
	"fmt"
	"io/ioutil"
	"syscall"
)

func SingleInstane(pidfile string) {
	if e := lockPidFile(pidfile); e != nil {
		pid, _ := ioutil.ReadFile(pidfile)
		panic(fmt.Errorf("process already run: %v", string(pid)))
	}

}

func lockPidFile(pidfile string) error {
	fd, e := syscall.Open(pidfile, syscall.O_CREAT|syscall.O_RDWR, 0777)
	if e != nil {
		return e
	}

	e = syscall.Flock(fd, syscall.LOCK_NB|syscall.LOCK_EX)
	if e != nil {
		return e
	}

	e = syscall.Ftruncate(fd, 0)
	if e != nil {
		return e
	}

	_, e = syscall.Write(fd, []byte(fmt.Sprintf("%d", syscall.Getpid())))
	if e != nil {
		return e
	}

	return nil
}

/*
func main() {
	SingleInstane("/tmp/pid")
	time.Sleep(100 * time.Second)
}
*/
<|endoftext|>


<commit_before>// Code generated by protoc-gen-gohttp. DO NOT EDIT.
// source: httprule/all_pattern.proto

package httprulepb

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// AllPatternHTTPConverter has a function to convert AllPatternServer interface to http.HandlerFunc.
type AllPatternHTTPConverter struct {
	srv AllPatternServer
}

// NewAllPatternHTTPConverter returns AllPatternHTTPConverter.
func NewAllPatternHTTPConverter(srv AllPatternServer) *AllPatternHTTPConverter {
	return &AllPatternHTTPConverter{
		srv: srv,
	}
}

// AllPattern returns AllPatternServer interface's AllPattern converted to http.HandlerFunc.
func (h *AllPatternHTTPConverter) AllPattern(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) http.HandlerFunc {
	if cb == nil {
		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				p := status.New(codes.Unknown, err.Error()).Proto()
				switch r.Header.Get("Content-Type") {
				case "application/protobuf", "application/x-protobuf":
					buf, err := proto.Marshal(p)
					if err != nil {
						return
					}
					if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
						return
					}
				case "application/json":
					if err := json.NewEncoder(w).Encode(p); err != nil {
						return
					}
				default:
				}
			}
		}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		arg := &AllPatternRequest{}
		contentType := r.Header.Get("Content-Type")
		if r.Method != http.MethodGet {
			body, err := ioutil.ReadAll(r.Body)
			if err != nil {
				cb(ctx, w, r, nil, nil, err)
				return
			}

			switch contentType {
			case "application/protobuf", "application/x-protobuf":
				if err := proto.Unmarshal(body, arg); err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
			case "application/json":
				if err := jsonpb.Unmarshal(bytes.NewBuffer(body), arg); err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
			default:
				w.WriteHeader(http.StatusUnsupportedMediaType)
				_, err := fmt.Fprintf(w, "Unsupported Content-Type: %s", contentType)
				cb(ctx, w, r, nil, nil, err)
				return
			}
		}

		ret, err := h.srv.AllPattern(ctx, arg)
		if err != nil {
			cb(ctx, w, r, arg, nil, err)
			return
		}

		accepts := strings.Split(r.Header.Get("Accept"), ",")
		accept := accepts[0]
		if accept == "*/*" || accept == "" {
			if contentType != "" {
				accept = contentType
			} else {
				accept = "application/json"
			}
		}

		switch accept {
		case "application/protobuf", "application/x-protobuf":
			buf, err := proto.Marshal(ret)
			if err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
			if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		case "application/json":
			m := jsonpb.Marshaler{
				EnumsAsInts:  true,
				EmitDefaults: true,
			}
			if err := m.Marshal(w, ret); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		default:
			w.WriteHeader(http.StatusUnsupportedMediaType)
			_, err := fmt.Fprintf(w, "Unsupported Accept: %s", accept)
			cb(ctx, w, r, arg, ret, err)
			return
		}
		cb(ctx, w, r, arg, ret, nil)
	})
}

// AllPatternWithName returns Service name, Method name and AllPatternServer interface's AllPattern converted to http.HandlerFunc.
func (h *AllPatternHTTPConverter) AllPatternWithName(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) (string, string, http.HandlerFunc) {
	return "AllPattern", "AllPattern", h.AllPattern(cb)
}

func (h *AllPatternHTTPConverter) AllPatternHTTPRule(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) (string, string, http.HandlerFunc) {
	if cb == nil {
		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				p := status.New(codes.Unknown, err.Error()).Proto()
				switch r.Header.Get("Content-Type") {
				case "application/protobuf", "application/x-protobuf":
					buf, err := proto.Marshal(p)
					if err != nil {
						return
					}
					if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
						return
					}
				case "application/json":
					if err := json.NewEncoder(w).Encode(p); err != nil {
						return
					}
				default:
				}
			}
		}
	}
	return http.MethodGet, "/all/pattern", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		arg := &AllPatternRequest{}
		contentType := r.Header.Get("Content-Type")
		if r.Method == http.MethodGet {
			var err error
			arg.String_ = r.URL.Query().Get("string")
		}

		ret, err := h.srv.AllPattern(ctx, arg)
		if err != nil {
			cb(ctx, w, r, arg, nil, err)
			return
		}

		accepts := strings.Split(r.Header.Get("Accept"), ",")
		accept := accepts[0]
		if accept == "*/*" || accept == "" {
			if contentType != "" {
				accept = contentType
			} else {
				accept = "application/json"
			}
		}

		switch accept {
		case "application/protobuf", "application/x-protobuf":
			buf, err := proto.Marshal(ret)
			if err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
			if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		case "application/json":
			m := jsonpb.Marshaler{
				EnumsAsInts:  true,
				EmitDefaults: true,
			}
			if err := m.Marshal(w, ret); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		default:
			w.WriteHeader(http.StatusUnsupportedMediaType)
			_, err := fmt.Fprintf(w, "Unsupported Accept: %s", accept)
			cb(ctx, w, r, arg, ret, err)
			return
		}
		cb(ctx, w, r, arg, ret, nil)
	})
}
<commit_msg>Added single var type test case<commit_after>// Code generated by protoc-gen-gohttp. DO NOT EDIT.
// source: httprule/all_pattern.proto

package httprulepb

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// AllPatternHTTPConverter has a function to convert AllPatternServer interface to http.HandlerFunc.
type AllPatternHTTPConverter struct {
	srv AllPatternServer
}

// NewAllPatternHTTPConverter returns AllPatternHTTPConverter.
func NewAllPatternHTTPConverter(srv AllPatternServer) *AllPatternHTTPConverter {
	return &AllPatternHTTPConverter{
		srv: srv,
	}
}

// AllPattern returns AllPatternServer interface's AllPattern converted to http.HandlerFunc.
func (h *AllPatternHTTPConverter) AllPattern(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) http.HandlerFunc {
	if cb == nil {
		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				p := status.New(codes.Unknown, err.Error()).Proto()
				switch r.Header.Get("Content-Type") {
				case "application/protobuf", "application/x-protobuf":
					buf, err := proto.Marshal(p)
					if err != nil {
						return
					}
					if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
						return
					}
				case "application/json":
					if err := json.NewEncoder(w).Encode(p); err != nil {
						return
					}
				default:
				}
			}
		}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		arg := &AllPatternRequest{}
		contentType := r.Header.Get("Content-Type")
		if r.Method != http.MethodGet {
			body, err := ioutil.ReadAll(r.Body)
			if err != nil {
				cb(ctx, w, r, nil, nil, err)
				return
			}

			switch contentType {
			case "application/protobuf", "application/x-protobuf":
				if err := proto.Unmarshal(body, arg); err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
			case "application/json":
				if err := jsonpb.Unmarshal(bytes.NewBuffer(body), arg); err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
			default:
				w.WriteHeader(http.StatusUnsupportedMediaType)
				_, err := fmt.Fprintf(w, "Unsupported Content-Type: %s", contentType)
				cb(ctx, w, r, nil, nil, err)
				return
			}
		}

		ret, err := h.srv.AllPattern(ctx, arg)
		if err != nil {
			cb(ctx, w, r, arg, nil, err)
			return
		}

		accepts := strings.Split(r.Header.Get("Accept"), ",")
		accept := accepts[0]
		if accept == "*/*" || accept == "" {
			if contentType != "" {
				accept = contentType
			} else {
				accept = "application/json"
			}
		}

		switch accept {
		case "application/protobuf", "application/x-protobuf":
			buf, err := proto.Marshal(ret)
			if err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
			if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		case "application/json":
			m := jsonpb.Marshaler{
				EnumsAsInts:  true,
				EmitDefaults: true,
			}
			if err := m.Marshal(w, ret); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		default:
			w.WriteHeader(http.StatusUnsupportedMediaType)
			_, err := fmt.Fprintf(w, "Unsupported Accept: %s", accept)
			cb(ctx, w, r, arg, ret, err)
			return
		}
		cb(ctx, w, r, arg, ret, nil)
	})
}

// AllPatternWithName returns Service name, Method name and AllPatternServer interface's AllPattern converted to http.HandlerFunc.
func (h *AllPatternHTTPConverter) AllPatternWithName(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) (string, string, http.HandlerFunc) {
	return "AllPattern", "AllPattern", h.AllPattern(cb)
}

func (h *AllPatternHTTPConverter) AllPatternHTTPRule(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) (string, string, http.HandlerFunc) {
	if cb == nil {
		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				p := status.New(codes.Unknown, err.Error()).Proto()
				switch r.Header.Get("Content-Type") {
				case "application/protobuf", "application/x-protobuf":
					buf, err := proto.Marshal(p)
					if err != nil {
						return
					}
					if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
						return
					}
				case "application/json":
					if err := json.NewEncoder(w).Encode(p); err != nil {
						return
					}
				default:
				}
			}
		}
	}
	return http.MethodGet, "/all/pattern", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		arg := &AllPatternRequest{}
		contentType := r.Header.Get("Content-Type")
		if r.Method == http.MethodGet {
			{
				d, err := strconv.ParseFloat(r.URL.Query().Get("double"), 64)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Double = d
			}
			{
				f, err := strconv.ParseFloat(r.URL.Query().Get("float"), 32)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Float = float32(f)
			}
			{
				i32, err := strconv.ParseInt(r.URL.Query().Get("int32"), 10, 32)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Int32 = int32(i32)
			}
			{
				i64, err := strconv.ParseInt(r.URL.Query().Get("int64"), 10, 64)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Int64 = i64
			}
			{
				ui32, err := strconv.ParseUint(r.URL.Query().Get("uint32"), 10, 32)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Uint32 = uint32(ui32)
			}
			{
				ui64, err := strconv.ParseUint(r.URL.Query().Get("uint64"), 10, 64)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Uint64 = uint64(ui64)
			}
			{
				f32, err := strconv.ParseUint(r.URL.Query().Get("fixed32"), 10, 32)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Fixed32 = uint32(f32)
			}
			{
				f64, err := strconv.ParseUint(r.URL.Query().Get("fixed64"), 10, 64)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Fixed64 = uint64(f64)
			}
			{
				sf32, err := strconv.ParseInt(r.URL.Query().Get("sfixed32"), 10, 32)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Sfixed32 = int32(sf32)
			}
			{
				sf64, err := strconv.ParseInt(r.URL.Query().Get("sfixed64"), 10, 64)
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Sfixed64 = int64(sf64)
			}
			{
				b, err := strconv.ParseBool(r.URL.Query().Get("bool"))
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Bool = b
			}
			arg.String_ = r.URL.Query().Get("string")
			{
				b, err := strconv.ParseBool(r.URL.Query().Get("bool"))
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Bool = b
			}
			{
				b, err := base64.StdEncoding.DecodeString(r.URL.Query().Get("bytes"))
				if err != nil {
					cb(ctx, w, r, nil, nil, err)
					return
				}
				arg.Bytes = b
			}
		}

		ret, err := h.srv.AllPattern(ctx, arg)
		if err != nil {
			cb(ctx, w, r, arg, nil, err)
			return
		}

		accepts := strings.Split(r.Header.Get("Accept"), ",")
		accept := accepts[0]
		if accept == "*/*" || accept == "" {
			if contentType != "" {
				accept = contentType
			} else {
				accept = "application/json"
			}
		}

		switch accept {
		case "application/protobuf", "application/x-protobuf":
			buf, err := proto.Marshal(ret)
			if err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
			if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		case "application/json":
			m := jsonpb.Marshaler{
				EnumsAsInts:  true,
				EmitDefaults: true,
			}
			if err := m.Marshal(w, ret); err != nil {
				cb(ctx, w, r, arg, ret, err)
				return
			}
		default:
			w.WriteHeader(http.StatusUnsupportedMediaType)
			_, err := fmt.Fprintf(w, "Unsupported Accept: %s", accept)
			cb(ctx, w, r, arg, ret, err)
			return
		}
		cb(ctx, w, r, arg, ret, nil)
	})
}
<|endoftext|>


<commit_before>// +build windows

/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package windows

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/Microsoft/hcsshim"
	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/api/types"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/events"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/plugin"
	"github.com/containerd/containerd/runtime"
	"github.com/containerd/containerd/windows/hcsshimtypes"
	"github.com/containerd/typeurl"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
)

const (
	runtimeName              = "windows"
	hcsshimOwner             = "containerd"
	defaultTerminateDuration = 5 * time.Minute
)

var (
	pluginID = fmt.Sprintf("%s.%s", plugin.RuntimePlugin, runtimeName)
)

var _ = (runtime.Runtime)(&windowsRuntime{})

func init() {
	plugin.Register(&plugin.Registration{
		ID:     runtimeName,
		Type:   plugin.RuntimePlugin,
		InitFn: New,
		Requires: []plugin.Type{
			plugin.MetadataPlugin,
		},
	})
}

// New returns a new Windows runtime
func New(ic *plugin.InitContext) (interface{}, error) {
	if err := os.MkdirAll(ic.Root, 0700); err != nil {
		return nil, errors.Wrapf(err, "could not create state directory at %s", ic.Root)
	}
	r := &windowsRuntime{
		root:    ic.Root,
		pidPool: newPidPool(),

		events:    make(chan interface{}, 4096),
		publisher: ic.Events,
		// TODO(mlaventure): windows needs a stat monitor
		monitor: nil,
		tasks:   runtime.NewTaskList(),
	}

	// Load our existing containers and kill/delete them. We don't support
	// reattaching to them
	r.cleanup(ic.Context)

	return r, nil
}

type windowsRuntime struct {
	sync.Mutex

	root    string
	pidPool *pidPool

	publisher events.Publisher
	events    chan interface{}

	monitor runtime.TaskMonitor
	tasks   *runtime.TaskList
}

func (r *windowsRuntime) ID() string {
	return pluginID
}

func (r *windowsRuntime) Create(ctx context.Context, id string, opts runtime.CreateOpts) (runtime.Task, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}

	s, err := typeurl.UnmarshalAny(opts.Spec)
	if err != nil {
		return nil, err
	}
	spec := s.(*specs.Spec)

	var createOpts *hcsshimtypes.CreateOptions
	if opts.Options != nil {
		o, err := typeurl.UnmarshalAny(opts.Options)
		if err != nil {
			return nil, err
		}
		createOpts = o.(*hcsshimtypes.CreateOptions)
	} else {
		createOpts = &hcsshimtypes.CreateOptions{}
	}

	if createOpts.TerminateDuration == 0 {
		createOpts.TerminateDuration = defaultTerminateDuration
	}

	if len(opts.Rootfs) == 0 {
		return nil, errors.Wrap(errdefs.ErrInvalidArgument, "rootfs was not provided to container create")
	}
	spec.Windows.LayerFolders = append(spec.Windows.LayerFolders, opts.Rootfs[0].Source)
	parentLayerPaths, err := opts.Rootfs[0].GetParentPaths()
	if err != nil {
		return nil, err
	}
	spec.Windows.LayerFolders = append(spec.Windows.LayerFolders, parentLayerPaths...)

	return r.newTask(ctx, namespace, id, opts.Rootfs, spec, opts.IO, createOpts)
}

func (r *windowsRuntime) Get(ctx context.Context, id string) (runtime.Task, error) {
	return r.tasks.Get(ctx, id)
}

func (r *windowsRuntime) Tasks(ctx context.Context) ([]runtime.Task, error) {
	return r.tasks.GetAll(ctx)
}

func (r *windowsRuntime) Delete(ctx context.Context, t runtime.Task) (*runtime.Exit, error) {
	wt, ok := t.(*task)
	if !ok {
		return nil, errors.Wrap(errdefs.ErrInvalidArgument, "no a windows task")
	}

	// TODO(mlaventure): stop monitor on this task

	var (
		err      error
		state, _ = wt.State(ctx)
	)
	switch state.Status {
	case runtime.StoppedStatus:
		fallthrough
	case runtime.CreatedStatus:
		// if it's stopped or in created state, we need to shutdown the
		// container before removing it
		if err = wt.stop(ctx); err != nil {
			return nil, err
		}
	default:
		return nil, errors.Wrap(errdefs.ErrFailedPrecondition,
			"cannot delete a non-stopped task")
	}

	var rtExit *runtime.Exit
	if p := wt.getProcess(t.ID()); p != nil {
		ec, ea, err := p.ExitCode()
		if err != nil {
			return nil, err
		}
		rtExit = &runtime.Exit{
			Pid:       wt.pid,
			Status:    ec,
			Timestamp: ea,
		}
	} else {
		rtExit = &runtime.Exit{
			Pid:       wt.pid,
			Status:    255,
			Timestamp: time.Now().UTC(),
		}
	}

	wt.cleanup()
	r.tasks.Delete(ctx, t.ID())

	r.publisher.Publish(ctx,
		runtime.TaskDeleteEventTopic,
		&eventstypes.TaskDelete{
			ContainerID: wt.id,
			Pid:         wt.pid,
			ExitStatus:  rtExit.Status,
			ExitedAt:    rtExit.Timestamp,
		})

	if err := mount.UnmountAll(wt.rootfs[0].Source, 0); err != nil {
		log.G(ctx).WithError(err).WithField("path", wt.rootfs[0].Source).
			Warn("failed to unmount rootfs on failure")
	}

	// We were never started, return failure
	return rtExit, nil
}

func (r *windowsRuntime) newTask(ctx context.Context, namespace, id string, rootfs []mount.Mount, spec *specs.Spec, io runtime.IO, createOpts *hcsshimtypes.CreateOptions) (*task, error) {
	var (
		err  error
		pset *pipeSet
	)

	if pset, err = newPipeSet(ctx, io); err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			pset.Close()
		}
	}()

	var pid uint32
	if pid, err = r.pidPool.Get(); err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			r.pidPool.Put(pid)
		}
	}()

	if err := mount.All(rootfs, ""); err != nil {
		return nil, errors.Wrap(err, "failed to mount rootfs")
	}
	defer func() {
		if err != nil {
			if err := mount.UnmountAll(rootfs[0].Source, 0); err != nil {
				log.G(ctx).WithError(err).WithField("path", rootfs[0].Source).
					Warn("failed to unmount rootfs on failure")
			}
		}
	}()

	var (
		conf *hcsshim.ContainerConfig
		nsid = namespace + "-" + id
	)
	if conf, err = newWindowsContainerConfig(ctx, hcsshimOwner, nsid, spec); err != nil {
		return nil, err
	}

	ctr, err := hcsshim.CreateContainer(nsid, conf)
	if err != nil {
		return nil, errors.Wrapf(err, "hcsshim failed to create task")
	}
	defer func() {
		if err != nil {
			ctr.Terminate()
			ctr.Wait()
			ctr.Close()
		}
	}()

	if err = ctr.Start(); err != nil {
		return nil, errors.Wrap(err, "hcsshim failed to spawn task")
	}

	t := &task{
		id:                id,
		namespace:         namespace,
		pid:               pid,
		io:                pset,
		status:            runtime.CreatedStatus,
		spec:              spec,
		processes:         make(map[string]*process),
		hyperV:            spec.Windows.HyperV != nil,
		publisher:         r.publisher,
		rwLayer:           conf.LayerFolderPath,
		rootfs:            rootfs,
		pidPool:           r.pidPool,
		hcsContainer:      ctr,
		terminateDuration: createOpts.TerminateDuration,
	}
	// Create the new process but don't start it
	pconf := newWindowsProcessConfig(t.spec.Process, t.io)
	if _, err = t.newProcess(ctx, t.id, pconf, t.io); err != nil {
		return nil, err
	}
	r.tasks.Add(ctx, t)

	var eventRootfs []*types.Mount
	for _, m := range rootfs {
		eventRootfs = append(eventRootfs, &types.Mount{
			Type:    m.Type,
			Source:  m.Source,
			Options: m.Options,
		})
	}

	r.publisher.Publish(ctx,
		runtime.TaskCreateEventTopic,
		&eventstypes.TaskCreate{
			ContainerID: id,
			IO: &eventstypes.TaskIO{
				Stdin:    io.Stdin,
				Stdout:   io.Stdout,
				Stderr:   io.Stderr,
				Terminal: io.Terminal,
			},
			Pid:    t.pid,
			Rootfs: eventRootfs,
			// TODO: what should be in Bundle for windows?
		})

	return t, nil
}

func (r *windowsRuntime) cleanup(ctx context.Context) {
	cp, err := hcsshim.GetContainers(hcsshim.ComputeSystemQuery{
		Types:  []string{"Container"},
		Owners: []string{hcsshimOwner},
	})
	if err != nil {
		log.G(ctx).Warn("failed to retrieve running containers")
		return
	}

	for _, p := range cp {
		container, err := hcsshim.OpenContainer(p.ID)
		if err != nil {
			log.G(ctx).Warnf("failed open container %s", p.ID)
			continue
		}

		err = container.Terminate()
		if err == nil || hcsshim.IsPending(err) || hcsshim.IsAlreadyStopped(err) {
			container.Wait()
		}
		container.Close()
	}
}
<commit_msg>windows: define platform in runtime plugin init<commit_after>// +build windows

/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package windows

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/Microsoft/hcsshim"
	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/api/types"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/events"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/platforms"
	"github.com/containerd/containerd/plugin"
	"github.com/containerd/containerd/runtime"
	"github.com/containerd/containerd/windows/hcsshimtypes"
	"github.com/containerd/typeurl"
	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
	runtimespec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
)

const (
	runtimeName              = "windows"
	hcsshimOwner             = "containerd"
	defaultTerminateDuration = 5 * time.Minute
)

var (
	pluginID = fmt.Sprintf("%s.%s", plugin.RuntimePlugin, runtimeName)
)

var _ = (runtime.Runtime)(&windowsRuntime{})

func init() {
	plugin.Register(&plugin.Registration{
		ID:     runtimeName,
		Type:   plugin.RuntimePlugin,
		InitFn: New,
		Requires: []plugin.Type{
			plugin.MetadataPlugin,
		},
	})
}

// New returns a new Windows runtime
func New(ic *plugin.InitContext) (interface{}, error) {
	ic.Meta.Platforms = []imagespec.Platform{platforms.DefaultSpec()}

	if err := os.MkdirAll(ic.Root, 0700); err != nil {
		return nil, errors.Wrapf(err, "could not create state directory at %s", ic.Root)
	}
	r := &windowsRuntime{
		root:    ic.Root,
		pidPool: newPidPool(),

		events:    make(chan interface{}, 4096),
		publisher: ic.Events,
		// TODO(mlaventure): windows needs a stat monitor
		monitor: nil,
		tasks:   runtime.NewTaskList(),
	}

	// Load our existing containers and kill/delete them. We don't support
	// reattaching to them
	r.cleanup(ic.Context)

	return r, nil
}

type windowsRuntime struct {
	sync.Mutex

	root    string
	pidPool *pidPool

	publisher events.Publisher
	events    chan interface{}

	monitor runtime.TaskMonitor
	tasks   *runtime.TaskList
}

func (r *windowsRuntime) ID() string {
	return pluginID
}

func (r *windowsRuntime) Create(ctx context.Context, id string, opts runtime.CreateOpts) (runtime.Task, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}

	s, err := typeurl.UnmarshalAny(opts.Spec)
	if err != nil {
		return nil, err
	}
	spec := s.(*runtimespec.Spec)

	var createOpts *hcsshimtypes.CreateOptions
	if opts.Options != nil {
		o, err := typeurl.UnmarshalAny(opts.Options)
		if err != nil {
			return nil, err
		}
		createOpts = o.(*hcsshimtypes.CreateOptions)
	} else {
		createOpts = &hcsshimtypes.CreateOptions{}
	}

	if createOpts.TerminateDuration == 0 {
		createOpts.TerminateDuration = defaultTerminateDuration
	}

	if len(opts.Rootfs) == 0 {
		return nil, errors.Wrap(errdefs.ErrInvalidArgument, "rootfs was not provided to container create")
	}
	spec.Windows.LayerFolders = append(spec.Windows.LayerFolders, opts.Rootfs[0].Source)
	parentLayerPaths, err := opts.Rootfs[0].GetParentPaths()
	if err != nil {
		return nil, err
	}
	spec.Windows.LayerFolders = append(spec.Windows.LayerFolders, parentLayerPaths...)

	return r.newTask(ctx, namespace, id, opts.Rootfs, spec, opts.IO, createOpts)
}

func (r *windowsRuntime) Get(ctx context.Context, id string) (runtime.Task, error) {
	return r.tasks.Get(ctx, id)
}

func (r *windowsRuntime) Tasks(ctx context.Context) ([]runtime.Task, error) {
	return r.tasks.GetAll(ctx)
}

func (r *windowsRuntime) Delete(ctx context.Context, t runtime.Task) (*runtime.Exit, error) {
	wt, ok := t.(*task)
	if !ok {
		return nil, errors.Wrap(errdefs.ErrInvalidArgument, "no a windows task")
	}

	// TODO(mlaventure): stop monitor on this task

	var (
		err      error
		state, _ = wt.State(ctx)
	)
	switch state.Status {
	case runtime.StoppedStatus:
		fallthrough
	case runtime.CreatedStatus:
		// if it's stopped or in created state, we need to shutdown the
		// container before removing it
		if err = wt.stop(ctx); err != nil {
			return nil, err
		}
	default:
		return nil, errors.Wrap(errdefs.ErrFailedPrecondition,
			"cannot delete a non-stopped task")
	}

	var rtExit *runtime.Exit
	if p := wt.getProcess(t.ID()); p != nil {
		ec, ea, err := p.ExitCode()
		if err != nil {
			return nil, err
		}
		rtExit = &runtime.Exit{
			Pid:       wt.pid,
			Status:    ec,
			Timestamp: ea,
		}
	} else {
		rtExit = &runtime.Exit{
			Pid:       wt.pid,
			Status:    255,
			Timestamp: time.Now().UTC(),
		}
	}

	wt.cleanup()
	r.tasks.Delete(ctx, t.ID())

	r.publisher.Publish(ctx,
		runtime.TaskDeleteEventTopic,
		&eventstypes.TaskDelete{
			ContainerID: wt.id,
			Pid:         wt.pid,
			ExitStatus:  rtExit.Status,
			ExitedAt:    rtExit.Timestamp,
		})

	if err := mount.UnmountAll(wt.rootfs[0].Source, 0); err != nil {
		log.G(ctx).WithError(err).WithField("path", wt.rootfs[0].Source).
			Warn("failed to unmount rootfs on failure")
	}

	// We were never started, return failure
	return rtExit, nil
}

func (r *windowsRuntime) newTask(ctx context.Context, namespace, id string, rootfs []mount.Mount, spec *runtimespec.Spec, io runtime.IO, createOpts *hcsshimtypes.CreateOptions) (*task, error) {
	var (
		err  error
		pset *pipeSet
	)

	if pset, err = newPipeSet(ctx, io); err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			pset.Close()
		}
	}()

	var pid uint32
	if pid, err = r.pidPool.Get(); err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			r.pidPool.Put(pid)
		}
	}()

	if err := mount.All(rootfs, ""); err != nil {
		return nil, errors.Wrap(err, "failed to mount rootfs")
	}
	defer func() {
		if err != nil {
			if err := mount.UnmountAll(rootfs[0].Source, 0); err != nil {
				log.G(ctx).WithError(err).WithField("path", rootfs[0].Source).
					Warn("failed to unmount rootfs on failure")
			}
		}
	}()

	var (
		conf *hcsshim.ContainerConfig
		nsid = namespace + "-" + id
	)
	if conf, err = newWindowsContainerConfig(ctx, hcsshimOwner, nsid, spec); err != nil {
		return nil, err
	}

	ctr, err := hcsshim.CreateContainer(nsid, conf)
	if err != nil {
		return nil, errors.Wrapf(err, "hcsshim failed to create task")
	}
	defer func() {
		if err != nil {
			ctr.Terminate()
			ctr.Wait()
			ctr.Close()
		}
	}()

	if err = ctr.Start(); err != nil {
		return nil, errors.Wrap(err, "hcsshim failed to spawn task")
	}

	t := &task{
		id:                id,
		namespace:         namespace,
		pid:               pid,
		io:                pset,
		status:            runtime.CreatedStatus,
		spec:              spec,
		processes:         make(map[string]*process),
		hyperV:            spec.Windows.HyperV != nil,
		publisher:         r.publisher,
		rwLayer:           conf.LayerFolderPath,
		rootfs:            rootfs,
		pidPool:           r.pidPool,
		hcsContainer:      ctr,
		terminateDuration: createOpts.TerminateDuration,
	}
	// Create the new process but don't start it
	pconf := newWindowsProcessConfig(t.spec.Process, t.io)
	if _, err = t.newProcess(ctx, t.id, pconf, t.io); err != nil {
		return nil, err
	}
	r.tasks.Add(ctx, t)

	var eventRootfs []*types.Mount
	for _, m := range rootfs {
		eventRootfs = append(eventRootfs, &types.Mount{
			Type:    m.Type,
			Source:  m.Source,
			Options: m.Options,
		})
	}

	r.publisher.Publish(ctx,
		runtime.TaskCreateEventTopic,
		&eventstypes.TaskCreate{
			ContainerID: id,
			IO: &eventstypes.TaskIO{
				Stdin:    io.Stdin,
				Stdout:   io.Stdout,
				Stderr:   io.Stderr,
				Terminal: io.Terminal,
			},
			Pid:    t.pid,
			Rootfs: eventRootfs,
			// TODO: what should be in Bundle for windows?
		})

	return t, nil
}

func (r *windowsRuntime) cleanup(ctx context.Context) {
	cp, err := hcsshim.GetContainers(hcsshim.ComputeSystemQuery{
		Types:  []string{"Container"},
		Owners: []string{hcsshimOwner},
	})
	if err != nil {
		log.G(ctx).Warn("failed to retrieve running containers")
		return
	}

	for _, p := range cp {
		container, err := hcsshim.OpenContainer(p.ID)
		if err != nil {
			log.G(ctx).Warnf("failed open container %s", p.ID)
			continue
		}

		err = container.Terminate()
		if err == nil || hcsshim.IsPending(err) || hcsshim.IsAlreadyStopped(err) {
			container.Wait()
		}
		container.Close()
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package openapi

import (
	"encoding/json"

	yaml "gopkg.in/yaml.v2"

	"github.com/caixw/apidoc/doc"
)

func parse(doc *doc.Doc) (*OpenAPI, error) {
	openapi := &OpenAPI{
		OpenAPI: doc.APIDoc,
		Info: &Info{
			Title:       doc.Title,
			Description: doc.Content,
			Contact:     newContact(doc.Contact),
			License:     newLicense(doc.License),
			Version:     doc.Version,
		},
		Servers: make([]*Server, 0, len(doc.Servers)),
		Tags:    make([]*Tag, 0, len(doc.Tags)),

		// TODO Paths
	}

	for _, srv := range doc.Servers {
		openapi.Servers = append(openapi.Servers, newServer(srv))
	}

	for _, tag := range doc.Tags {
		openapi.Tags = append(openapi.Tags, newTag(tag))
	}

	return openapi, nil
}

// JSON 输出 JSON 格式数据
func JSON(doc *doc.Doc) ([]byte, error) {
	openapi, err := parse(doc)
	if err != nil {
		return nil, err
	}

	return json.Marshal(openapi)
}

// YAML 输出 YAML 格式数据
func YAML(doc *doc.Doc) ([]byte, error) {
	openapi, err := parse(doc)
	if err != nil {
		return nil, err
	}

	return yaml.Marshal(openapi)
}
<commit_msg>完成转换成 openapi 的功能<commit_after>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package openapi

import (
	"encoding/json"
	"strings"

	yaml "gopkg.in/yaml.v2"

	"github.com/caixw/apidoc/doc"
)

func parse(doc *doc.Doc) (*OpenAPI, error) {
	openapi := &OpenAPI{
		OpenAPI: doc.APIDoc,
		Info: &Info{
			Title:       doc.Title,
			Description: doc.Content,
			Contact:     newContact(doc.Contact),
			License:     newLicense(doc.License),
			Version:     doc.Version,
		},
		Servers: make([]*Server, 0, len(doc.Servers)),
		Tags:    make([]*Tag, 0, len(doc.Tags)),
	}

	for _, srv := range doc.Servers {
		openapi.Servers = append(openapi.Servers, newServer(srv))
	}

	for _, tag := range doc.Tags {
		openapi.Tags = append(openapi.Tags, newTag(tag))
	}

	if err := parsePaths(openapi, doc); err != nil {
		return nil, err
	}

	return openapi, nil
}

func parsePaths(openapi *OpenAPI, doc *doc.Doc) *Error {
	for _, api := range doc.Apis {
		p := openapi.Paths[api.Path]
		if p == nil {
			p = &PathItem{}
			openapi.Paths[api.Path] = p
		}

		operation, err := setOperation(p, api.Method)
		if err != nil {
			err.Field = "paths." + err.Field
			return err
		}

		operation.Tags = api.Tags
		operation.Description = api.Description
		operation.Deprecated = api.Deprecated != ""
		setOperationParams(operation, api)

		if len(api.Requests) > 0 {
			content := make(map[string]*MediaType, len(api.Requests))
			for _, r := range api.Requests {
				examples := make(map[string]*Example, len(r.Examples))
				for _, exp := range r.Examples {
					examples[exp.Mimetype] = &Example{
						Summary: exp.Summary,
						Value:   ExampleValue(exp.Value),
					}
				}

				content[r.Mimetype] = &MediaType{
					Schema:   &Schema{Schema: r.Type},
					Examples: examples,
				}
			}

			operation.RequestBody = &RequestBody{
				Content: content,
			}
		}

		operation.Responses = make(map[string]*Response, len(api.Responses))
		for _, resp := range api.Responses {
			r, found := operation.Responses[resp.Status]
			if !found {
				r = &Response{}
				operation.Responses[resp.Status] = r
			}

			if r.Headers == nil {
				r.Headers = make(map[string]*Header, 10)
			}
			for _, h := range resp.Headers {
				r.Headers[h.Name] = &Header{
					Description: Description(h.Summary),
				}
			}

			if r.Content == nil {
				r.Content = make(map[string]*MediaType, 10)
			}
			examples := make(map[string]*Example, len(resp.Examples))
			for _, exp := range resp.Examples {
				examples[exp.Mimetype] = &Example{
					Summary: exp.Summary,
					Value:   ExampleValue(exp.Value),
				}
			}
			r.Content[resp.Mimetype] = &MediaType{
				Schema:   &Schema{Schema: resp.Type},
				Examples: examples,
			}
		}
	} // end for doc.Apis

	return nil
}

func setOperationParams(operation *Operation, api *doc.API) {
	operation.Parameters = make([]*Parameter, 0, len(api.Params)+len(api.Queries)+len(api.Requests[0].Headers))

	for _, param := range api.Params {
		operation.Parameters = append(operation.Parameters, &Parameter{
			Name:        param.Name,
			IN:          ParameterINPath,
			Description: Description(param.Summary),
			Required:    !param.Optional,
			Schema:      &Schema{Schema: param.Type},
		})
	}

	for _, param := range api.Queries {
		operation.Parameters = append(operation.Parameters, &Parameter{
			Name:        param.Name,
			IN:          ParameterINQuery,
			Description: Description(param.Summary),
			Required:    !param.Optional,
			Schema:      &Schema{Schema: param.Type},
		})
	}

	if len(api.Requests) > 0 {
		for _, param := range api.Requests[0].Headers {
			operation.Parameters = append(operation.Parameters, &Parameter{
				Name:        param.Name,
				IN:          ParameterINHeader,
				Description: Description(param.Summary),
				Required:    !param.Optional,
			})
		}
	}
}

func setOperation(path *PathItem, method string) (*Operation, *Error) {
	operation := &Operation{}

	switch strings.ToUpper(method) {
	case "GET":
		if path.Get != nil {
			return nil, &Error{Field: "Get", Message: "已经存在"}
		}
		path.Get = operation
	case "DELETE":
		if path.Delete != nil {
			return nil, &Error{Field: "Delete", Message: "已经存在"}
		}
		path.Delete = operation
	case "POST":
		if path.Post != nil {
			return nil, &Error{Field: "Post", Message: "已经存在"}
		}
		path.Post = operation
	case "PUT":
		if path.Put != nil {
			return nil, &Error{Field: "Put", Message: "已经存在"}
		}
		path.Put = operation
	case "PATCH":
		if path.Patch != nil {
			return nil, &Error{Field: "Patch", Message: "已经存在"}
		}
		path.Patch = operation
	case "OPTIONS":
		if path.Options != nil {
			return nil, &Error{Field: "Options", Message: "已经存在"}
		}
		path.Options = operation
	case "HEAD":
		if path.Head != nil {
			return nil, &Error{Field: "Head", Message: "已经存在"}
		}
		path.Head = operation
	case "TRACE":
		if path.Trace != nil {
			return nil, &Error{Field: "Trace", Message: "已经存在"}
		}
		path.Trace = operation
	}

	return operation, nil
}

// JSON 输出 JSON 格式数据
func JSON(doc *doc.Doc) ([]byte, error) {
	openapi, err := parse(doc)
	if err != nil {
		return nil, err
	}

	return json.Marshal(openapi)
}

// YAML 输出 YAML 格式数据
func YAML(doc *doc.Doc) ([]byte, error) {
	openapi, err := parse(doc)
	if err != nil {
		return nil, err
	}

	return yaml.Marshal(openapi)
}
<|endoftext|>