<commit_before>// Package swarm implements a connection muxer with a pair of channels
// to synchronize all network communication.
package swarm

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	pst "github.com/jbenet/go-stream-muxer"
	"github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	addrutil "github.com/libp2p/go-addr-util"
	conn "github.com/libp2p/go-libp2p-conn"
	ci "github.com/libp2p/go-libp2p-crypto"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	metrics "github.com/libp2p/go-libp2p-metrics"
	mconn "github.com/libp2p/go-libp2p-metrics/conn"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	transport "github.com/libp2p/go-libp2p-transport"
	filter "github.com/libp2p/go-maddr-filter"
	ps "github.com/libp2p/go-peerstream"
	tcpt "github.com/libp2p/go-tcp-transport"
	ma "github.com/multiformats/go-multiaddr"
	psmss "github.com/whyrusleeping/go-smux-multistream"
	spdy "github.com/whyrusleeping/go-smux-spdystream"
	yamux "github.com/whyrusleeping/go-smux-yamux"
	mafilter "github.com/whyrusleeping/multiaddr-filter"
	ws "github.com/whyrusleeping/ws-transport"
)

var log = logging.Logger("swarm2")

// PSTransport is the default peerstream transport that will be used by
// any libp2p swarms.
var PSTransport pst.Transport

func init() {
	msstpt := psmss.NewBlankTransport()

	ymxtpt := &yamux.Transport{
		AcceptBacklog:          8192,
		ConnectionWriteTimeout: time.Second * 10,
		KeepAliveInterval:      time.Second * 30,
		EnableKeepAlive:        true,
		MaxStreamWindowSize:    uint32(1024 * 512),
		LogOutput:              ioutil.Discard,
	}

	msstpt.AddTransport("/yamux/1.0.0", ymxtpt)
	msstpt.AddTransport("/spdy/3.1.0", spdy.Transport)

	// allow overriding of muxer preferences
	if prefs := os.Getenv("LIBP2P_MUX_PREFS"); prefs != "" {
		msstpt.OrderPreference = strings.Fields(prefs)
	}

	PSTransport = msstpt
}

// Swarm is a connection muxer, allowing connections to other peers to
// be opened and closed, while still using the same Chan for all
// communication. The Chan sends/receives Messages, which note the
// destination or source Peer.
//
// Uses peerstream.Swarm
type Swarm struct {
	swarm *ps.Swarm
	local peer.ID
	peers pstore.Peerstore
	connh ConnHandler

	dsync *DialSync
	backf dialbackoff
	dialT time.Duration // mainly for tests

	dialer *conn.Dialer

	notifmu sync.RWMutex
	notifs  map[inet.Notifiee]ps.Notifiee

	transports []transport.Transport

	// filters for addresses that shouldnt be dialed
	Filters *filter.Filters

	// file descriptor rate limited
	fdRateLimit chan struct{}

	proc goprocess.Process
	ctx  context.Context
	bwc  metrics.Reporter

	limiter *dialLimiter

	protec ipnet.Protector
}

func NewSwarm(ctx context.Context, listenAddrs []ma.Multiaddr, local peer.ID,
	peers pstore.Peerstore, bwc metrics.Reporter) (*Swarm, error) {
	return NewSwarmWithProtector(ctx, listenAddrs, local, peers, nil, PSTransport, bwc)
}

// NewSwarm constructs a Swarm, with a Chan.
func NewSwarmWithProtector(ctx context.Context, listenAddrs []ma.Multiaddr, local peer.ID,
	peers pstore.Peerstore, protec ipnet.Protector, tpt pst.Transport, bwc metrics.Reporter) (*Swarm, error) {

	listenAddrs, err := filterAddrs(listenAddrs)
	if err != nil {
		return nil, err
	}

	var wrap func(c transport.Conn) transport.Conn
	if bwc != nil {
		wrap = func(c transport.Conn) transport.Conn {
			return mconn.WrapConn(bwc, c)
		}
	}

	s := &Swarm{
		swarm:  ps.NewSwarm(tpt),
		local:  local,
		peers:  peers,
		ctx:    ctx,
		dialT:  DialTimeout,
		notifs: make(map[inet.Notifiee]ps.Notifiee),
		transports: []transport.Transport{
			tcpt.NewTCPTransport(),
			new(ws.WebsocketTransport),
		},
		bwc:         bwc,
		fdRateLimit: make(chan struct{}, concurrentFdDials),
		Filters:     filter.NewFilters(),
		dialer:      conn.NewDialer(local, peers.PrivKey(local), wrap),
		protec:      protec,
	}
	s.dialer.Protector = protec

	s.dsync = NewDialSync(s.doDial)
	s.limiter = newDialLimiter(s.dialAddr)

	// configure Swarm
	s.proc = goprocessctx.WithContextAndTeardown(ctx, s.teardown)
	s.SetConnHandler(nil) // make sure to setup our own conn handler.

	err = s.setupInterfaces(listenAddrs)
	if err != nil {
		return nil, err
	}

	return s, nil
}

func NewBlankSwarm(ctx context.Context, id peer.ID, privkey ci.PrivKey, pstpt pst.Transport) *Swarm {
	s := &Swarm{
		swarm:       ps.NewSwarm(pstpt),
		local:       id,
		peers:       pstore.NewPeerstore(),
		ctx:         ctx,
		dialT:       DialTimeout,
		notifs:      make(map[inet.Notifiee]ps.Notifiee),
		fdRateLimit: make(chan struct{}, concurrentFdDials),
		Filters:     filter.NewFilters(),
		dialer:      conn.NewDialer(id, privkey, nil),
	}

	// configure Swarm
	s.limiter = newDialLimiter(s.dialAddr)
	s.proc = goprocessctx.WithContextAndTeardown(ctx, s.teardown)
	s.SetConnHandler(nil) // make sure to setup our own conn handler.

	return s
}

func (s *Swarm) AddTransport(t transport.Transport) {
	s.transports = append(s.transports, t)
}

func (s *Swarm) teardown() error {
	return s.swarm.Close()
}

// AddAddrFilter adds a multiaddr filter to the set of filters the swarm will
// use to determine which addresses not to dial to.
func (s *Swarm) AddAddrFilter(f string) error {
	m, err := mafilter.NewMask(f)
	if err != nil {
		return err
	}

	s.Filters.AddDialFilter(m)
	return nil
}

func filterAddrs(listenAddrs []ma.Multiaddr) ([]ma.Multiaddr, error) {
	if len(listenAddrs) > 0 {
		filtered := addrutil.FilterUsableAddrs(listenAddrs)
		if len(filtered) < 1 {
			return nil, fmt.Errorf("swarm cannot use any addr in: %s", listenAddrs)
		}
		listenAddrs = filtered
	}

	return listenAddrs, nil
}

// Listen sets up listeners for all of the given addresses
func (s *Swarm) Listen(addrs ...ma.Multiaddr) error {
	addrs, err := filterAddrs(addrs)
	if err != nil {
		return err
	}

	return s.setupInterfaces(addrs)
}

// Process returns the Process of the swarm
func (s *Swarm) Process() goprocess.Process {
	return s.proc
}

// Context returns the context of the swarm
func (s *Swarm) Context() context.Context {
	return s.ctx
}

// Close stops the Swarm.
func (s *Swarm) Close() error {
	return s.proc.Close()
}

// StreamSwarm returns the underlying peerstream.Swarm
func (s *Swarm) StreamSwarm() *ps.Swarm {
	return s.swarm
}

// SetConnHandler assigns the handler for new connections.
// See peerstream. You will rarely use this. See SetStreamHandler
func (s *Swarm) SetConnHandler(handler ConnHandler) {

	// handler is nil if user wants to clear the old handler.
	if handler == nil {
		s.swarm.SetConnHandler(func(psconn *ps.Conn) {
			s.connHandler(psconn)
		})
		return
	}

	s.swarm.SetConnHandler(func(psconn *ps.Conn) {
		// sc is nil if closed in our handler.
		if sc := s.connHandler(psconn); sc != nil {
			// call the user's handler. in a goroutine for sync safety.
			go handler(sc)
		}
	})
}

// SetStreamHandler assigns the handler for new streams.
// See peerstream.
func (s *Swarm) SetStreamHandler(handler inet.StreamHandler) {
	s.swarm.SetStreamHandler(func(s *ps.Stream) {
		handler((*Stream)(s))
	})
}

// NewStreamWithPeer creates a new stream on any available connection to p
func (s *Swarm) NewStreamWithPeer(ctx context.Context, p peer.ID) (*Stream, error) {
	// if we have no connections, try connecting.
	if len(s.ConnectionsToPeer(p)) == 0 {
		log.Debug("Swarm: NewStreamWithPeer no connections. Attempting to connect...")
		if _, err := s.Dial(ctx, p); err != nil {
			return nil, err
		}
	}
	log.Debug("Swarm: NewStreamWithPeer...")

	// TODO: think about passing a context down to NewStreamWithGroup
	st, err := s.swarm.NewStreamWithGroup(p)
	return (*Stream)(st), err
}

// ConnectionsToPeer returns all the live connections to p
func (s *Swarm) ConnectionsToPeer(p peer.ID) []*Conn {
	return wrapConns(ps.ConnsWithGroup(p, s.swarm.Conns()))
}

func (s *Swarm) HaveConnsToPeer(p peer.ID) bool {
	for _, c := range s.swarm.Conns() {
		if c.InGroup(p) {
			return true
		}
	}
	return false
}

// Connections returns a slice of all connections.
func (s *Swarm) Connections() []*Conn {
	return wrapConns(s.swarm.Conns())
}

// CloseConnection removes a given peer from swarm + closes the connection
func (s *Swarm) CloseConnection(p peer.ID) error {
	conns := s.swarm.ConnsWithGroup(p) // boom.
	for _, c := range conns {
		c.Close()
	}
	return nil
}

// Peers returns a copy of the set of peers swarm is connected to.
func (s *Swarm) Peers() []peer.ID {
	conns := s.Connections()

	seen := make(map[peer.ID]struct{})
	peers := make([]peer.ID, 0, len(conns))
	for _, c := range conns {
		p := c.RemotePeer()
		if _, found := seen[p]; found {
			continue
		}

		seen[p] = struct{}{}
		peers = append(peers, p)
	}
	return peers
}

// LocalPeer returns the local peer swarm is associated to.
func (s *Swarm) LocalPeer() peer.ID {
	return s.local
}

// Backoff returns the dialbackoff object for this swarm.
func (s *Swarm) Backoff() *dialbackoff {
	return &s.backf
}

// notifyAll sends a signal to all Notifiees
func (s *Swarm) notifyAll(notify func(inet.Notifiee)) {
	s.notifmu.RLock()
	for f := range s.notifs {
		go notify(f)
	}
	s.notifmu.RUnlock()
}

// Notify signs up Notifiee to receive signals when events happen
func (s *Swarm) Notify(f inet.Notifiee) {
	// wrap with our notifiee, to translate function calls
	n := &ps2netNotifee{net: (*Network)(s), not: f}

	s.notifmu.Lock()
	s.notifs[f] = n
	s.notifmu.Unlock()

	// register for notifications in the peer swarm.
	s.swarm.Notify(n)
}

// StopNotify unregisters Notifiee fromr receiving signals
func (s *Swarm) StopNotify(f inet.Notifiee) {
	s.notifmu.Lock()
	n, found := s.notifs[f]
	if found {
		delete(s.notifs, f)
	}
	s.notifmu.Unlock()

	if found {
		s.swarm.StopNotify(n)
	}
}

type ps2netNotifee struct {
	net *Network
	not inet.Notifiee
}

func (n *ps2netNotifee) Connected(c *ps.Conn) {
	n.not.Connected(n.net, inet.Conn((*Conn)(c)))
}

func (n *ps2netNotifee) Disconnected(c *ps.Conn) {
	n.not.Disconnected(n.net, inet.Conn((*Conn)(c)))
}

func (n *ps2netNotifee) OpenedStream(s *ps.Stream) {
	n.not.OpenedStream(n.net, (*Stream)(s))
}

func (n *ps2netNotifee) ClosedStream(s *ps.Stream) {
	n.not.ClosedStream(n.net, (*Stream)(s))
}
<commit_msg>Fix go-ws-transport package name<commit_after>// Package swarm implements a connection muxer with a pair of channels
// to synchronize all network communication.
package swarm

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	pst "github.com/jbenet/go-stream-muxer"
	"github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	addrutil "github.com/libp2p/go-addr-util"
	conn "github.com/libp2p/go-libp2p-conn"
	ci "github.com/libp2p/go-libp2p-crypto"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	metrics "github.com/libp2p/go-libp2p-metrics"
	mconn "github.com/libp2p/go-libp2p-metrics/conn"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	transport "github.com/libp2p/go-libp2p-transport"
	filter "github.com/libp2p/go-maddr-filter"
	ps "github.com/libp2p/go-peerstream"
	tcpt "github.com/libp2p/go-tcp-transport"
	ws "github.com/libp2p/go-ws-transport"
	ma "github.com/multiformats/go-multiaddr"
	psmss "github.com/whyrusleeping/go-smux-multistream"
	spdy "github.com/whyrusleeping/go-smux-spdystream"
	yamux "github.com/whyrusleeping/go-smux-yamux"
	mafilter "github.com/whyrusleeping/multiaddr-filter"
)

var log = logging.Logger("swarm2")

// PSTransport is the default peerstream transport that will be used by
// any libp2p swarms.
var PSTransport pst.Transport

func init() {
	msstpt := psmss.NewBlankTransport()

	ymxtpt := &yamux.Transport{
		AcceptBacklog:          8192,
		ConnectionWriteTimeout: time.Second * 10,
		KeepAliveInterval:      time.Second * 30,
		EnableKeepAlive:        true,
		MaxStreamWindowSize:    uint32(1024 * 512),
		LogOutput:              ioutil.Discard,
	}

	msstpt.AddTransport("/yamux/1.0.0", ymxtpt)
	msstpt.AddTransport("/spdy/3.1.0", spdy.Transport)

	// allow overriding of muxer preferences
	if prefs := os.Getenv("LIBP2P_MUX_PREFS"); prefs != "" {
		msstpt.OrderPreference = strings.Fields(prefs)
	}

	PSTransport = msstpt
}

// Swarm is a connection muxer, allowing connections to other peers to
// be opened and closed, while still using the same Chan for all
// communication. The Chan sends/receives Messages, which note the
// destination or source Peer.
//
// Uses peerstream.Swarm
type Swarm struct {
	swarm *ps.Swarm
	local peer.ID
	peers pstore.Peerstore
	connh ConnHandler

	dsync *DialSync
	backf dialbackoff
	dialT time.Duration // mainly for tests

	dialer *conn.Dialer

	notifmu sync.RWMutex
	notifs  map[inet.Notifiee]ps.Notifiee

	transports []transport.Transport

	// filters for addresses that shouldnt be dialed
	Filters *filter.Filters

	// file descriptor rate limited
	fdRateLimit chan struct{}

	proc goprocess.Process
	ctx  context.Context
	bwc  metrics.Reporter

	limiter *dialLimiter

	protec ipnet.Protector
}

func NewSwarm(ctx context.Context, listenAddrs []ma.Multiaddr, local peer.ID,
	peers pstore.Peerstore, bwc metrics.Reporter) (*Swarm, error) {
	return NewSwarmWithProtector(ctx, listenAddrs, local, peers, nil, PSTransport, bwc)
}

// NewSwarm constructs a Swarm, with a Chan.
func NewSwarmWithProtector(ctx context.Context, listenAddrs []ma.Multiaddr, local peer.ID,
	peers pstore.Peerstore, protec ipnet.Protector, tpt pst.Transport, bwc metrics.Reporter) (*Swarm, error) {

	listenAddrs, err := filterAddrs(listenAddrs)
	if err != nil {
		return nil, err
	}

	var wrap func(c transport.Conn) transport.Conn
	if bwc != nil {
		wrap = func(c transport.Conn) transport.Conn {
			return mconn.WrapConn(bwc, c)
		}
	}

	s := &Swarm{
		swarm:  ps.NewSwarm(tpt),
		local:  local,
		peers:  peers,
		ctx:    ctx,
		dialT:  DialTimeout,
		notifs: make(map[inet.Notifiee]ps.Notifiee),
		transports: []transport.Transport{
			tcpt.NewTCPTransport(),
			new(ws.WebsocketTransport),
		},
		bwc:         bwc,
		fdRateLimit: make(chan struct{}, concurrentFdDials),
		Filters:     filter.NewFilters(),
		dialer:      conn.NewDialer(local, peers.PrivKey(local), wrap),
		protec:      protec,
	}
	s.dialer.Protector = protec

	s.dsync = NewDialSync(s.doDial)
	s.limiter = newDialLimiter(s.dialAddr)

	// configure Swarm
	s.proc = goprocessctx.WithContextAndTeardown(ctx, s.teardown)
	s.SetConnHandler(nil) // make sure to setup our own conn handler.

	err = s.setupInterfaces(listenAddrs)
	if err != nil {
		return nil, err
	}

	return s, nil
}

func NewBlankSwarm(ctx context.Context, id peer.ID, privkey ci.PrivKey, pstpt pst.Transport) *Swarm {
	s := &Swarm{
		swarm:       ps.NewSwarm(pstpt),
		local:       id,
		peers:       pstore.NewPeerstore(),
		ctx:         ctx,
		dialT:       DialTimeout,
		notifs:      make(map[inet.Notifiee]ps.Notifiee),
		fdRateLimit: make(chan struct{}, concurrentFdDials),
		Filters:     filter.NewFilters(),
		dialer:      conn.NewDialer(id, privkey, nil),
	}

	// configure Swarm
	s.limiter = newDialLimiter(s.dialAddr)
	s.proc = goprocessctx.WithContextAndTeardown(ctx, s.teardown)
	s.SetConnHandler(nil) // make sure to setup our own conn handler.

	return s
}

func (s *Swarm) AddTransport(t transport.Transport) {
	s.transports = append(s.transports, t)
}

func (s *Swarm) teardown() error {
	return s.swarm.Close()
}

// AddAddrFilter adds a multiaddr filter to the set of filters the swarm will
// use to determine which addresses not to dial to.
func (s *Swarm) AddAddrFilter(f string) error {
	m, err := mafilter.NewMask(f)
	if err != nil {
		return err
	}

	s.Filters.AddDialFilter(m)
	return nil
}

func filterAddrs(listenAddrs []ma.Multiaddr) ([]ma.Multiaddr, error) {
	if len(listenAddrs) > 0 {
		filtered := addrutil.FilterUsableAddrs(listenAddrs)
		if len(filtered) < 1 {
			return nil, fmt.Errorf("swarm cannot use any addr in: %s", listenAddrs)
		}
		listenAddrs = filtered
	}

	return listenAddrs, nil
}

// Listen sets up listeners for all of the given addresses
func (s *Swarm) Listen(addrs ...ma.Multiaddr) error {
	addrs, err := filterAddrs(addrs)
	if err != nil {
		return err
	}

	return s.setupInterfaces(addrs)
}

// Process returns the Process of the swarm
func (s *Swarm) Process() goprocess.Process {
	return s.proc
}

// Context returns the context of the swarm
func (s *Swarm) Context() context.Context {
	return s.ctx
}

// Close stops the Swarm.
func (s *Swarm) Close() error {
	return s.proc.Close()
}

// StreamSwarm returns the underlying peerstream.Swarm
func (s *Swarm) StreamSwarm() *ps.Swarm {
	return s.swarm
}

// SetConnHandler assigns the handler for new connections.
// See peerstream. You will rarely use this. See SetStreamHandler
func (s *Swarm) SetConnHandler(handler ConnHandler) {

	// handler is nil if user wants to clear the old handler.
	if handler == nil {
		s.swarm.SetConnHandler(func(psconn *ps.Conn) {
			s.connHandler(psconn)
		})
		return
	}

	s.swarm.SetConnHandler(func(psconn *ps.Conn) {
		// sc is nil if closed in our handler.
		if sc := s.connHandler(psconn); sc != nil {
			// call the user's handler. in a goroutine for sync safety.
			go handler(sc)
		}
	})
}

// SetStreamHandler assigns the handler for new streams.
// See peerstream.
func (s *Swarm) SetStreamHandler(handler inet.StreamHandler) {
	s.swarm.SetStreamHandler(func(s *ps.Stream) {
		handler((*Stream)(s))
	})
}

// NewStreamWithPeer creates a new stream on any available connection to p
func (s *Swarm) NewStreamWithPeer(ctx context.Context, p peer.ID) (*Stream, error) {
	// if we have no connections, try connecting.
	if len(s.ConnectionsToPeer(p)) == 0 {
		log.Debug("Swarm: NewStreamWithPeer no connections. Attempting to connect...")
		if _, err := s.Dial(ctx, p); err != nil {
			return nil, err
		}
	}
	log.Debug("Swarm: NewStreamWithPeer...")

	// TODO: think about passing a context down to NewStreamWithGroup
	st, err := s.swarm.NewStreamWithGroup(p)
	return (*Stream)(st), err
}

// ConnectionsToPeer returns all the live connections to p
func (s *Swarm) ConnectionsToPeer(p peer.ID) []*Conn {
	return wrapConns(ps.ConnsWithGroup(p, s.swarm.Conns()))
}

func (s *Swarm) HaveConnsToPeer(p peer.ID) bool {
	for _, c := range s.swarm.Conns() {
		if c.InGroup(p) {
			return true
		}
	}
	return false
}

// Connections returns a slice of all connections.
func (s *Swarm) Connections() []*Conn {
	return wrapConns(s.swarm.Conns())
}

// CloseConnection removes a given peer from swarm + closes the connection
func (s *Swarm) CloseConnection(p peer.ID) error {
	conns := s.swarm.ConnsWithGroup(p) // boom.
	for _, c := range conns {
		c.Close()
	}
	return nil
}

// Peers returns a copy of the set of peers swarm is connected to.
func (s *Swarm) Peers() []peer.ID {
	conns := s.Connections()

	seen := make(map[peer.ID]struct{})
	peers := make([]peer.ID, 0, len(conns))
	for _, c := range conns {
		p := c.RemotePeer()
		if _, found := seen[p]; found {
			continue
		}

		seen[p] = struct{}{}
		peers = append(peers, p)
	}
	return peers
}

// LocalPeer returns the local peer swarm is associated to.
func (s *Swarm) LocalPeer() peer.ID {
	return s.local
}

// Backoff returns the dialbackoff object for this swarm.
func (s *Swarm) Backoff() *dialbackoff {
	return &s.backf
}

// notifyAll sends a signal to all Notifiees
func (s *Swarm) notifyAll(notify func(inet.Notifiee)) {
	s.notifmu.RLock()
	for f := range s.notifs {
		go notify(f)
	}
	s.notifmu.RUnlock()
}

// Notify signs up Notifiee to receive signals when events happen
func (s *Swarm) Notify(f inet.Notifiee) {
	// wrap with our notifiee, to translate function calls
	n := &ps2netNotifee{net: (*Network)(s), not: f}

	s.notifmu.Lock()
	s.notifs[f] = n
	s.notifmu.Unlock()

	// register for notifications in the peer swarm.
	s.swarm.Notify(n)
}

// StopNotify unregisters Notifiee fromr receiving signals
func (s *Swarm) StopNotify(f inet.Notifiee) {
	s.notifmu.Lock()
	n, found := s.notifs[f]
	if found {
		delete(s.notifs, f)
	}
	s.notifmu.Unlock()

	if found {
		s.swarm.StopNotify(n)
	}
}

type ps2netNotifee struct {
	net *Network
	not inet.Notifiee
}

func (n *ps2netNotifee) Connected(c *ps.Conn) {
	n.not.Connected(n.net, inet.Conn((*Conn)(c)))
}

func (n *ps2netNotifee) Disconnected(c *ps.Conn) {
	n.not.Disconnected(n.net, inet.Conn((*Conn)(c)))
}

func (n *ps2netNotifee) OpenedStream(s *ps.Stream) {
	n.not.OpenedStream(n.net, (*Stream)(s))
}

func (n *ps2netNotifee) ClosedStream(s *ps.Stream) {
	n.not.ClosedStream(n.net, (*Stream)(s))
}
<|endoftext|>


<commit_before>// Copyright 2017-2020 The Usacloud 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 disk

import (
	"github.com/sacloud/libsacloud/v2/sacloud/types"
	"github.com/sacloud/usacloud/pkg/cmd/base"
)

var createCommand = &base.Command{
	Name:     "create",
	Category: "basics",
	Order:    20,

	ParameterInitializer: func() interface{} {
		return newCreateParameter()
	},
}

type createParameter struct {
	base.ZoneParameter `cli:",squash" mapconv:",squash"`

	Name            string     `cli:",category=disk"`
	Description     string     `cli:",category=disk"`
	Tags            []string   `cli:",category=disk"`
	IconID          types.ID   `cli:",category=disk"`
	DiskPlan        string     `cli:",category=disk,options=disk_plan" mapconv:"DiskPlanID,filters=disk_plan_to_id" validate:"oneof=ssd hdd"`
	Connection      string     `cli:",category=disk,options=disk_connection" validate:"oneof=virtio ide"`
	SourceDiskID    types.ID   `cli:",category=disk"`
	SourceArchiveID types.ID   `cli:",category=disk"`
	ServerID        types.ID   `cli:",category=disk"`
	SizeGB          int        `cli:"size,category=disk"`
	DistantFrom     []types.ID `cli:",category=disk"`
	OSType          string     `cli:",category=disk,options=os_type" mapconv:",filters=os_type"`

	base.ConfirmParameter `cli:",squash" mapconv:"-"`
	base.OutputParameter  `cli:",squash" mapconv:"-"`
}

func newCreateParameter() *createParameter {
	return &createParameter{
		// TODO デフォルト値はここで設定する
	}
}

func init() {
	Resource.AddCommand(createCommand)
}
<commit_msg>Default value<commit_after>// Copyright 2017-2020 The Usacloud 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 disk

import (
	"github.com/sacloud/libsacloud/v2/sacloud/types"
	"github.com/sacloud/usacloud/pkg/cmd/base"
)

var createCommand = &base.Command{
	Name:     "create",
	Category: "basics",
	Order:    20,

	ParameterInitializer: func() interface{} {
		return newCreateParameter()
	},
}

type createParameter struct {
	base.ZoneParameter `cli:",squash" mapconv:",squash"`

	Name            string     `cli:",category=disk"`
	Description     string     `cli:",category=disk"`
	Tags            []string   `cli:",category=disk"`
	IconID          types.ID   `cli:",category=disk"`
	DiskPlan        string     `cli:",category=disk,options=disk_plan" mapconv:"DiskPlanID,filters=disk_plan_to_id" validate:"oneof=ssd hdd"`
	Connection      string     `cli:",category=disk,options=disk_connection" validate:"oneof=virtio ide"`
	SourceDiskID    types.ID   `cli:",category=disk"`
	SourceArchiveID types.ID   `cli:",category=disk"`
	ServerID        types.ID   `cli:",category=disk"`
	SizeGB          int        `cli:"size,category=disk"`
	DistantFrom     []types.ID `cli:",category=disk"`
	OSType          string     `cli:",category=disk,options=os_type" mapconv:",filters=os_type"`

	base.ConfirmParameter `cli:",squash" mapconv:"-"`
	base.OutputParameter  `cli:",squash" mapconv:"-"`
}

func newCreateParameter() *createParameter {
	return &createParameter{
		DiskPlan:   "ssd",
		Connection: "virtio",
		SizeGB:     20,
	}
}

func init() {
	Resource.AddCommand(createCommand)
}
<|endoftext|>


<commit_before>package cmd

import (
	"errors"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/cloudnativelabs/kube-router/pkg/controllers/netpol"
	"github.com/cloudnativelabs/kube-router/pkg/controllers/proxy"
	"github.com/cloudnativelabs/kube-router/pkg/controllers/routing"
	"github.com/cloudnativelabs/kube-router/pkg/healthcheck"
	"github.com/cloudnativelabs/kube-router/pkg/metrics"
	"github.com/cloudnativelabs/kube-router/pkg/options"
	"github.com/cloudnativelabs/kube-router/pkg/version"
	"k8s.io/klog/v2"

	"time"

	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// KubeRouter holds the information needed to run server
type KubeRouter struct {
	Client kubernetes.Interface
	Config *options.KubeRouterConfig
}

// NewKubeRouterDefault returns a KubeRouter object
func NewKubeRouterDefault(config *options.KubeRouterConfig) (*KubeRouter, error) {

	var clientconfig *rest.Config
	var err error
	version.PrintVersion(true)
	// Use out of cluster config if the URL or kubeconfig have been specified. Otherwise use incluster config.
	if len(config.Master) != 0 || len(config.Kubeconfig) != 0 {
		clientconfig, err = clientcmd.BuildConfigFromFlags(config.Master, config.Kubeconfig)
		if err != nil {
			return nil, errors.New("Failed to build configuration from CLI: " + err.Error())
		}
	} else {
		clientconfig, err = rest.InClusterConfig()
		if err != nil {
			return nil, errors.New("unable to initialize inclusterconfig: " + err.Error())
		}
	}

	clientset, err := kubernetes.NewForConfig(clientconfig)
	if err != nil {
		return nil, errors.New("Failed to create Kubernetes client: " + err.Error())
	}

	return &KubeRouter{Client: clientset, Config: config}, nil
}

// CleanupConfigAndExit performs Cleanup on all three controllers
func CleanupConfigAndExit() {
	npc := netpol.NetworkPolicyController{}
	npc.Cleanup()

	nsc := proxy.NetworkServicesController{}
	nsc.Cleanup()

	nrc := routing.NetworkRoutingController{}
	nrc.Cleanup()
}

// Run starts the controllers and waits forever till we get SIGINT or SIGTERM
func (kr *KubeRouter) Run() error {
	var err error
	var ipsetMutex sync.Mutex
	var wg sync.WaitGroup
	healthChan := make(chan *healthcheck.ControllerHeartbeat, 10)
	defer close(healthChan)
	stopCh := make(chan struct{})

	if !(kr.Config.RunFirewall || kr.Config.RunServiceProxy || kr.Config.RunRouter) {
		klog.Info("Router, Firewall or Service proxy functionality must be specified. Exiting!")
		os.Exit(0)
	}

	hc, err := healthcheck.NewHealthController(kr.Config)
	if err != nil {
		return errors.New("Failed to create health controller: " + err.Error())
	}
	wg.Add(1)
	go hc.RunServer(stopCh, &wg)

	informerFactory := informers.NewSharedInformerFactory(kr.Client, 0)
	svcInformer := informerFactory.Core().V1().Services().Informer()
	epInformer := informerFactory.Core().V1().Endpoints().Informer()
	podInformer := informerFactory.Core().V1().Pods().Informer()
	nodeInformer := informerFactory.Core().V1().Nodes().Informer()
	nsInformer := informerFactory.Core().V1().Namespaces().Informer()
	npInformer := informerFactory.Networking().V1().NetworkPolicies().Informer()
	informerFactory.Start(stopCh)

	err = kr.CacheSyncOrTimeout(informerFactory, stopCh)
	if err != nil {
		return errors.New("Failed to synchronize cache: " + err.Error())
	}

	hc.SetAlive()
	wg.Add(1)
	go hc.RunCheck(healthChan, stopCh, &wg)

	if kr.Config.MetricsPort > 0 {
		kr.Config.MetricsEnabled = true
		mc, err := metrics.NewMetricsController(kr.Config)
		if err != nil {
			return errors.New("Failed to create metrics controller: " + err.Error())
		}
		wg.Add(1)
		go mc.Run(healthChan, stopCh, &wg)

	} else if kr.Config.MetricsPort > 65535 {
		klog.Errorf("Metrics port must be over 0 and under 65535, given port: %d", kr.Config.MetricsPort)
		kr.Config.MetricsEnabled = false
	} else {
		kr.Config.MetricsEnabled = false
	}

	if kr.Config.BGPGracefulRestart {
		if kr.Config.BGPGracefulRestartTime > time.Second*4095 {
			return errors.New("BGPGracefuleRestartTime should be less than 4095 seconds")
		}
		if kr.Config.BGPGracefulRestartTime <= 0 {
			return errors.New("BGPGracefuleRestartTime must be positive")
		}

		if kr.Config.BGPGracefulRestartDeferralTime > time.Hour*18 {
			return errors.New("BGPGracefuleRestartDeferralTime should be less than 18 hours")
		}
		if kr.Config.BGPGracefulRestartDeferralTime <= 0 {
			return errors.New("BGPGracefuleRestartDeferralTime must be positive")
		}
	}

	if kr.Config.RunRouter {
		nrc, err := routing.NewNetworkRoutingController(kr.Client, kr.Config,
			nodeInformer, svcInformer, epInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network routing controller: " + err.Error())
		}

		nodeInformer.AddEventHandler(nrc.NodeEventHandler)
		svcInformer.AddEventHandler(nrc.ServiceEventHandler)
		epInformer.AddEventHandler(nrc.EndpointsEventHandler)

		wg.Add(1)
		go nrc.Run(healthChan, stopCh, &wg)

		// wait for the pod networking related firewall rules to be setup before network policies
		if kr.Config.RunFirewall {
			nrc.CNIFirewallSetup.L.Lock()
			nrc.CNIFirewallSetup.Wait()
			nrc.CNIFirewallSetup.L.Unlock()
		}
	}

	if kr.Config.RunServiceProxy {
		nsc, err := proxy.NewNetworkServicesController(kr.Client, kr.Config,
			svcInformer, epInformer, podInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network services controller: " + err.Error())
		}

		svcInformer.AddEventHandler(nsc.ServiceEventHandler)
		epInformer.AddEventHandler(nsc.EndpointsEventHandler)

		wg.Add(1)
		go nsc.Run(healthChan, stopCh, &wg)

		// wait for the proxy firewall rules to be setup before network policies
		if kr.Config.RunFirewall {
			nsc.ProxyFirewallSetup.L.Lock()
			nsc.ProxyFirewallSetup.Wait()
			nsc.ProxyFirewallSetup.L.Unlock()
		}
	}

	if kr.Config.RunFirewall {
		npc, err := netpol.NewNetworkPolicyController(kr.Client,
			kr.Config, podInformer, npInformer, nsInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network policy controller: " + err.Error())
		}

		podInformer.AddEventHandler(npc.PodEventHandler)
		nsInformer.AddEventHandler(npc.NamespaceEventHandler)
		npInformer.AddEventHandler(npc.NetworkPolicyEventHandler)

		wg.Add(1)
		go npc.Run(healthChan, stopCh, &wg)
	}

	// Handle SIGINT and SIGTERM
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	<-ch

	klog.Infof("Shutting down the controllers")
	close(stopCh)

	wg.Wait()
	return nil
}

// CacheSyncOrTimeout performs cache synchronization under timeout limit
func (kr *KubeRouter) CacheSyncOrTimeout(informerFactory informers.SharedInformerFactory, stopCh <-chan struct{}) error {
	syncOverCh := make(chan struct{})
	go func() {
		informerFactory.WaitForCacheSync(stopCh)
		close(syncOverCh)
	}()

	select {
	case <-time.After(kr.Config.CacheSyncTimeout):
		return errors.New(kr.Config.CacheSyncTimeout.String() + " timeout")
	case <-syncOverCh:
		return nil
	}
}
<commit_msg>fix typo Gracefule -> Graceful<commit_after>package cmd

import (
	"errors"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/cloudnativelabs/kube-router/pkg/controllers/netpol"
	"github.com/cloudnativelabs/kube-router/pkg/controllers/proxy"
	"github.com/cloudnativelabs/kube-router/pkg/controllers/routing"
	"github.com/cloudnativelabs/kube-router/pkg/healthcheck"
	"github.com/cloudnativelabs/kube-router/pkg/metrics"
	"github.com/cloudnativelabs/kube-router/pkg/options"
	"github.com/cloudnativelabs/kube-router/pkg/version"
	"k8s.io/klog/v2"

	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// KubeRouter holds the information needed to run server
type KubeRouter struct {
	Client kubernetes.Interface
	Config *options.KubeRouterConfig
}

// NewKubeRouterDefault returns a KubeRouter object
func NewKubeRouterDefault(config *options.KubeRouterConfig) (*KubeRouter, error) {

	var clientconfig *rest.Config
	var err error
	version.PrintVersion(true)
	// Use out of cluster config if the URL or kubeconfig have been specified. Otherwise use incluster config.
	if len(config.Master) != 0 || len(config.Kubeconfig) != 0 {
		clientconfig, err = clientcmd.BuildConfigFromFlags(config.Master, config.Kubeconfig)
		if err != nil {
			return nil, errors.New("Failed to build configuration from CLI: " + err.Error())
		}
	} else {
		clientconfig, err = rest.InClusterConfig()
		if err != nil {
			return nil, errors.New("unable to initialize inclusterconfig: " + err.Error())
		}
	}

	clientset, err := kubernetes.NewForConfig(clientconfig)
	if err != nil {
		return nil, errors.New("Failed to create Kubernetes client: " + err.Error())
	}

	return &KubeRouter{Client: clientset, Config: config}, nil
}

// CleanupConfigAndExit performs Cleanup on all three controllers
func CleanupConfigAndExit() {
	npc := netpol.NetworkPolicyController{}
	npc.Cleanup()

	nsc := proxy.NetworkServicesController{}
	nsc.Cleanup()

	nrc := routing.NetworkRoutingController{}
	nrc.Cleanup()
}

// Run starts the controllers and waits forever till we get SIGINT or SIGTERM
func (kr *KubeRouter) Run() error {
	var err error
	var ipsetMutex sync.Mutex
	var wg sync.WaitGroup
	healthChan := make(chan *healthcheck.ControllerHeartbeat, 10)
	defer close(healthChan)
	stopCh := make(chan struct{})

	if !(kr.Config.RunFirewall || kr.Config.RunServiceProxy || kr.Config.RunRouter) {
		klog.Info("Router, Firewall or Service proxy functionality must be specified. Exiting!")
		os.Exit(0)
	}

	hc, err := healthcheck.NewHealthController(kr.Config)
	if err != nil {
		return errors.New("Failed to create health controller: " + err.Error())
	}
	wg.Add(1)
	go hc.RunServer(stopCh, &wg)

	informerFactory := informers.NewSharedInformerFactory(kr.Client, 0)
	svcInformer := informerFactory.Core().V1().Services().Informer()
	epInformer := informerFactory.Core().V1().Endpoints().Informer()
	podInformer := informerFactory.Core().V1().Pods().Informer()
	nodeInformer := informerFactory.Core().V1().Nodes().Informer()
	nsInformer := informerFactory.Core().V1().Namespaces().Informer()
	npInformer := informerFactory.Networking().V1().NetworkPolicies().Informer()
	informerFactory.Start(stopCh)

	err = kr.CacheSyncOrTimeout(informerFactory, stopCh)
	if err != nil {
		return errors.New("Failed to synchronize cache: " + err.Error())
	}

	hc.SetAlive()
	wg.Add(1)
	go hc.RunCheck(healthChan, stopCh, &wg)

	if kr.Config.MetricsPort > 0 {
		kr.Config.MetricsEnabled = true
		mc, err := metrics.NewMetricsController(kr.Config)
		if err != nil {
			return errors.New("Failed to create metrics controller: " + err.Error())
		}
		wg.Add(1)
		go mc.Run(healthChan, stopCh, &wg)

	} else if kr.Config.MetricsPort > 65535 {
		klog.Errorf("Metrics port must be over 0 and under 65535, given port: %d", kr.Config.MetricsPort)
		kr.Config.MetricsEnabled = false
	} else {
		kr.Config.MetricsEnabled = false
	}

	if kr.Config.BGPGracefulRestart {
		if kr.Config.BGPGracefulRestartTime > time.Second*4095 {
			return errors.New("BGPGracefulRestartTime should be less than 4095 seconds")
		}
		if kr.Config.BGPGracefulRestartTime <= 0 {
			return errors.New("BGPGracefulRestartTime must be positive")
		}

		if kr.Config.BGPGracefulRestartDeferralTime > time.Hour*18 {
			return errors.New("BGPGracefulRestartDeferralTime should be less than 18 hours")
		}
		if kr.Config.BGPGracefulRestartDeferralTime <= 0 {
			return errors.New("BGPGracefulRestartDeferralTime must be positive")
		}
	}

	if kr.Config.RunRouter {
		nrc, err := routing.NewNetworkRoutingController(kr.Client, kr.Config,
			nodeInformer, svcInformer, epInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network routing controller: " + err.Error())
		}

		nodeInformer.AddEventHandler(nrc.NodeEventHandler)
		svcInformer.AddEventHandler(nrc.ServiceEventHandler)
		epInformer.AddEventHandler(nrc.EndpointsEventHandler)

		wg.Add(1)
		go nrc.Run(healthChan, stopCh, &wg)

		// wait for the pod networking related firewall rules to be setup before network policies
		if kr.Config.RunFirewall {
			nrc.CNIFirewallSetup.L.Lock()
			nrc.CNIFirewallSetup.Wait()
			nrc.CNIFirewallSetup.L.Unlock()
		}
	}

	if kr.Config.RunServiceProxy {
		nsc, err := proxy.NewNetworkServicesController(kr.Client, kr.Config,
			svcInformer, epInformer, podInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network services controller: " + err.Error())
		}

		svcInformer.AddEventHandler(nsc.ServiceEventHandler)
		epInformer.AddEventHandler(nsc.EndpointsEventHandler)

		wg.Add(1)
		go nsc.Run(healthChan, stopCh, &wg)

		// wait for the proxy firewall rules to be setup before network policies
		if kr.Config.RunFirewall {
			nsc.ProxyFirewallSetup.L.Lock()
			nsc.ProxyFirewallSetup.Wait()
			nsc.ProxyFirewallSetup.L.Unlock()
		}
	}

	if kr.Config.RunFirewall {
		npc, err := netpol.NewNetworkPolicyController(kr.Client,
			kr.Config, podInformer, npInformer, nsInformer, &ipsetMutex)
		if err != nil {
			return errors.New("Failed to create network policy controller: " + err.Error())
		}

		podInformer.AddEventHandler(npc.PodEventHandler)
		nsInformer.AddEventHandler(npc.NamespaceEventHandler)
		npInformer.AddEventHandler(npc.NetworkPolicyEventHandler)

		wg.Add(1)
		go npc.Run(healthChan, stopCh, &wg)
	}

	// Handle SIGINT and SIGTERM
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	<-ch

	klog.Infof("Shutting down the controllers")
	close(stopCh)

	wg.Wait()
	return nil
}

// CacheSyncOrTimeout performs cache synchronization under timeout limit
func (kr *KubeRouter) CacheSyncOrTimeout(informerFactory informers.SharedInformerFactory, stopCh <-chan struct{}) error {
	syncOverCh := make(chan struct{})
	go func() {
		informerFactory.WaitForCacheSync(stopCh)
		close(syncOverCh)
	}()

	select {
	case <-time.After(kr.Config.CacheSyncTimeout):
		return errors.New(kr.Config.CacheSyncTimeout.String() + " timeout")
	case <-syncOverCh:
		return nil
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 cmd

import (
	"fmt"
	"io"

	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	cmdconfig "k8s.io/kubernetes/pkg/kubectl/cmd/config"
	"k8s.io/kubernetes/pkg/kubectl/cmd/rollout"
	"k8s.io/kubernetes/pkg/kubectl/cmd/set"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/util/flag"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	bash_completion_func = `# call kubectl get $1,
__kubectl_override_flag_list=(kubeconfig cluster user context namespace server)
__kubectl_override_flags()
{
    local ${__kubectl_override_flag_list[*]} two_word_of of
    for w in "${words[@]}"; do
        if [ -n "${two_word_of}" ]; then
            eval "${two_word_of}=\"--${two_word_of}=\${w}\""
            two_word_of=
            continue
        fi
        for of in "${__kubectl_override_flag_list[@]}"; do
            case "${w}" in
                --${of}=*)
                    eval "${of}=\"${w}\""
                    ;;
                --${of})
                    two_word_of="${of}"
                    ;;
            esac
        done
        if [ "${w}" == "--all-namespaces" ]; then
            namespace="--all-namespaces"
        fi
    done
    for of in "${__kubectl_override_flag_list[@]}"; do
        if eval "test -n \"\$${of}\""; then
            eval "echo \${${of}}"
        fi
    done
}

__kubectl_get_namespaces()
{
    local template kubectl_out
    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
    if kubectl_out=$(kubectl get -o template --template="${template}" namespace 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

__kubectl_parse_get()
{
    local template
    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
    local kubectl_out
    if kubectl_out=$(kubectl get $(__kubectl_override_flags) -o template --template="${template}" "$1" 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

__kubectl_get_resource()
{
    if [[ ${#nouns[@]} -eq 0 ]]; then
        return 1
    fi
    __kubectl_parse_get "${nouns[${#nouns[@]} -1]}"
}

__kubectl_get_resource_pod()
{
    __kubectl_parse_get "pod"
}

__kubectl_get_resource_rc()
{
    __kubectl_parse_get "rc"
}

# $1 is the name of the pod we want to get the list of containers inside
__kubectl_get_containers()
{
    local template
    template="{{ range .spec.containers  }}{{ .name }} {{ end }}"
    __debug "${FUNCNAME} nouns are ${nouns[*]}"

    local len="${#nouns[@]}"
    if [[ ${len} -ne 1 ]]; then
        return
    fi
    local last=${nouns[${len} -1]}
    local kubectl_out
    if kubectl_out=$(kubectl get $(__kubectl_override_flags) -o template --template="${template}" pods "${last}" 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

# Require both a pod and a container to be specified
__kubectl_require_pod_and_container()
{
    if [[ ${#nouns[@]} -eq 0 ]]; then
        __kubectl_parse_get pods
        return 0
    fi;
    __kubectl_get_containers
    return 0
}

__custom_func() {
    case ${last_command} in
        kubectl_get | kubectl_describe | kubectl_delete | kubectl_label | kubectl_stop | kubectl_edit | kubectl_patch |\
        kubectl_annotate | kubectl_expose)
            __kubectl_get_resource
            return
            ;;
        kubectl_logs)
            __kubectl_require_pod_and_container
            return
            ;;
        kubectl_exec)
            __kubectl_get_resource_pod
            return
            ;;
        kubectl_rolling-update)
            __kubectl_get_resource_rc
            return
            ;;
        *)
            ;;
    esac
}
`

	// If you add a resource to this list, please also take a look at pkg/kubectl/kubectl.go
	// and add a short forms entry in expandResourceShortcut() when appropriate.
	valid_resources = `Valid resource types include:
   * componentstatuses (aka 'cs')
   * configmaps (aka 'cm')
   * daemonsets (aka 'ds')
   * deployments (aka 'deploy')
   * events (aka 'ev')
   * endpoints (aka 'ep')
   * horizontalpodautoscalers (aka 'hpa')
   * ingress (aka 'ing')
   * jobs
   * limitranges (aka 'limits')
   * nodes (aka 'no')
   * namespaces (aka 'ns')
   * petsets (alpha feature, may be unstable)
   * pods (aka 'po')
   * persistentvolumes (aka 'pv')
   * persistentvolumeclaims (aka 'pvc')
   * quota
   * resourcequotas (aka 'quota')
   * replicasets (aka 'rs')
   * replicationcontrollers (aka 'rc')
   * secrets
   * serviceaccounts (aka 'sa')
   * services (aka 'svc')
`
	usage_template = `{{if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}{{end}}{{if .HasExample}}

Examples:
{{ .Example }}{{end}}{{ if .HasAvailableSubCommands}}

Available Sub-commands:{{range .Commands}}{{if .IsAvailableCommand}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{ if .HasLocalFlags}}

Flags:
{{.LocalFlags.FlagUsages | trimRightSpace}}{{end}}{{ if .HasInheritedFlags}}

Global Flags:
{{.InheritedFlags.FlagUsages | trimRightSpace}}{{end}}{{if .HasHelpSubCommands}}

Additional help topics:{{range .Commands}}{{if .IsHelpCommand}}
  {{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}

Usage:{{if .Runnable}}
  {{if .HasFlags}}{{appendIfNotPresent .UseLine "[flags]"}}{{else}}{{.UseLine}}{{end}}{{end}}{{ if .HasSubCommands }}
  {{ .CommandPath}} [command]

Use "{{.CommandPath}} [command] --help" for more information about a command.{{end}}
`
	help_template = `{{with or .Long .Short }}{{. | trim}}{{end}}{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
)

// NewKubectlCommand creates the `kubectl` command and its nested children.
func NewKubectlCommand(f *cmdutil.Factory, in io.Reader, out, err io.Writer) *cobra.Command {
	// Parent command to which all subcommands are added.
	cmds := &cobra.Command{
		Use:   "kubectl",
		Short: "kubectl controls the Kubernetes cluster manager",
		Long: `kubectl controls the Kubernetes cluster manager.

Find more information at https://github.com/kubernetes/kubernetes.`,
		Run: runHelp,
		BashCompletionFunction: bash_completion_func,
	}

	f.BindFlags(cmds.PersistentFlags())
	f.BindExternalFlags(cmds.PersistentFlags())

	// From this point and forward we get warnings on flags that contain "_" separators
	cmds.SetGlobalNormalizationFunc(flag.WarnWordSepNormalizeFunc)

	groups := templates.CommandGroups{
		{
			Message: "Basic Commands (Beginner):",
			Commands: []*cobra.Command{
				NewCmdCreate(f, out),
				NewCmdExposeService(f, out),
				NewCmdRun(f, in, out, err),
				set.NewCmdSet(f, out),
			},
		},
		{
			Message: "Basic Commands (Intermediate):",
			Commands: []*cobra.Command{
				NewCmdGet(f, out, err),
				NewCmdExplain(f, out, err),
				NewCmdEdit(f, out, err),
				NewCmdDelete(f, out),
			},
		},
		{
			Message: "Deploy Commands:",
			Commands: []*cobra.Command{
				rollout.NewCmdRollout(f, out),
				NewCmdRollingUpdate(f, out),
				NewCmdScale(f, out),
				NewCmdAutoscale(f, out),
			},
		},
		{
			Message: "Cluster Management Commands:",
			Commands: []*cobra.Command{
				NewCmdClusterInfo(f, out),
				NewCmdTop(f, out),
				NewCmdCordon(f, out),
				NewCmdUncordon(f, out),
				NewCmdDrain(f, out),
				NewCmdTaint(f, out),
			},
		},
		{
			Message: "Troubleshooting and Debugging Commands:",
			Commands: []*cobra.Command{
				NewCmdDescribe(f, out, err),
				NewCmdLogs(f, out),
				NewCmdAttach(f, in, out, err),
				NewCmdExec(f, in, out, err),
				NewCmdPortForward(f, out, err),
				NewCmdProxy(f, out),
			},
		},
		{
			Message: "Advanced Commands:",
			Commands: []*cobra.Command{
				NewCmdApply(f, out),
				NewCmdPatch(f, out),
				NewCmdReplace(f, out),
				NewCmdConvert(f, out),
			},
		},
		{
			Message: "Settings Commands:",
			Commands: []*cobra.Command{
				NewCmdLabel(f, out),
				NewCmdAnnotate(f, out),
				NewCmdCompletion(f, out),
			},
		},
	}
	groups.Add(cmds)

	filters := []string{
		"options",
		Deprecated("kubectl", "delete", cmds, NewCmdStop(f, out)),
		Deprecated("kubectl", "config set-context", cmds, NewCmdNamespace(out)),
	}
	templates.ActsAsRootCommand(cmds, filters, groups...)

	if cmds.Flag("namespace") != nil {
		if cmds.Flag("namespace").Annotations == nil {
			cmds.Flag("namespace").Annotations = map[string][]string{}
		}
		cmds.Flag("namespace").Annotations[cobra.BashCompCustom] = append(
			cmds.Flag("namespace").Annotations[cobra.BashCompCustom],
			"__kubectl_get_namespaces",
		)
	}

	cmds.AddCommand(cmdconfig.NewCmdConfig(clientcmd.NewDefaultPathOptions(), out))
	cmds.AddCommand(NewCmdVersion(f, out))
	cmds.AddCommand(NewCmdApiVersions(f, out))
	cmds.AddCommand(NewCmdOptions(out))

	return cmds
}

func runHelp(cmd *cobra.Command, args []string) {
	cmd.Help()
}

func printDeprecationWarning(command, alias string) {
	glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", alias, command)
}

func Deprecated(baseName, to string, parent, cmd *cobra.Command) string {
	cmd.Long = fmt.Sprintf("Deprecated: This command is deprecated, all its functionalities are covered by \"%s %s\"", baseName, to)
	cmd.Short = fmt.Sprintf("Deprecated: %s", to)
	parent.AddCommand(cmd)
	return cmd.Name()
}
<commit_msg>Adding clusters to the list of valid resources printed in kubectl help<commit_after>/*
Copyright 2014 The Kubernetes 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 cmd

import (
	"fmt"
	"io"

	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	cmdconfig "k8s.io/kubernetes/pkg/kubectl/cmd/config"
	"k8s.io/kubernetes/pkg/kubectl/cmd/rollout"
	"k8s.io/kubernetes/pkg/kubectl/cmd/set"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/util/flag"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	bash_completion_func = `# call kubectl get $1,
__kubectl_override_flag_list=(kubeconfig cluster user context namespace server)
__kubectl_override_flags()
{
    local ${__kubectl_override_flag_list[*]} two_word_of of
    for w in "${words[@]}"; do
        if [ -n "${two_word_of}" ]; then
            eval "${two_word_of}=\"--${two_word_of}=\${w}\""
            two_word_of=
            continue
        fi
        for of in "${__kubectl_override_flag_list[@]}"; do
            case "${w}" in
                --${of}=*)
                    eval "${of}=\"${w}\""
                    ;;
                --${of})
                    two_word_of="${of}"
                    ;;
            esac
        done
        if [ "${w}" == "--all-namespaces" ]; then
            namespace="--all-namespaces"
        fi
    done
    for of in "${__kubectl_override_flag_list[@]}"; do
        if eval "test -n \"\$${of}\""; then
            eval "echo \${${of}}"
        fi
    done
}

__kubectl_get_namespaces()
{
    local template kubectl_out
    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
    if kubectl_out=$(kubectl get -o template --template="${template}" namespace 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

__kubectl_parse_get()
{
    local template
    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
    local kubectl_out
    if kubectl_out=$(kubectl get $(__kubectl_override_flags) -o template --template="${template}" "$1" 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

__kubectl_get_resource()
{
    if [[ ${#nouns[@]} -eq 0 ]]; then
        return 1
    fi
    __kubectl_parse_get "${nouns[${#nouns[@]} -1]}"
}

__kubectl_get_resource_pod()
{
    __kubectl_parse_get "pod"
}

__kubectl_get_resource_rc()
{
    __kubectl_parse_get "rc"
}

# $1 is the name of the pod we want to get the list of containers inside
__kubectl_get_containers()
{
    local template
    template="{{ range .spec.containers  }}{{ .name }} {{ end }}"
    __debug "${FUNCNAME} nouns are ${nouns[*]}"

    local len="${#nouns[@]}"
    if [[ ${len} -ne 1 ]]; then
        return
    fi
    local last=${nouns[${len} -1]}
    local kubectl_out
    if kubectl_out=$(kubectl get $(__kubectl_override_flags) -o template --template="${template}" pods "${last}" 2>/dev/null); then
        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
    fi
}

# Require both a pod and a container to be specified
__kubectl_require_pod_and_container()
{
    if [[ ${#nouns[@]} -eq 0 ]]; then
        __kubectl_parse_get pods
        return 0
    fi;
    __kubectl_get_containers
    return 0
}

__custom_func() {
    case ${last_command} in
        kubectl_get | kubectl_describe | kubectl_delete | kubectl_label | kubectl_stop | kubectl_edit | kubectl_patch |\
        kubectl_annotate | kubectl_expose)
            __kubectl_get_resource
            return
            ;;
        kubectl_logs)
            __kubectl_require_pod_and_container
            return
            ;;
        kubectl_exec)
            __kubectl_get_resource_pod
            return
            ;;
        kubectl_rolling-update)
            __kubectl_get_resource_rc
            return
            ;;
        *)
            ;;
    esac
}
`

	// If you add a resource to this list, please also take a look at pkg/kubectl/kubectl.go
	// and add a short forms entry in expandResourceShortcut() when appropriate.
	// TODO: This should be populated using the discovery information from apiserver.
	valid_resources = `Valid resource types include:
   * clusters (valid only for federation apiservers)
   * componentstatuses (aka 'cs')
   * configmaps (aka 'cm')
   * daemonsets (aka 'ds')
   * deployments (aka 'deploy')
   * events (aka 'ev')
   * endpoints (aka 'ep')
   * horizontalpodautoscalers (aka 'hpa')
   * ingress (aka 'ing')
   * jobs
   * limitranges (aka 'limits')
   * nodes (aka 'no')
   * namespaces (aka 'ns')
   * petsets (alpha feature, may be unstable)
   * pods (aka 'po')
   * persistentvolumes (aka 'pv')
   * persistentvolumeclaims (aka 'pvc')
   * quota
   * resourcequotas (aka 'quota')
   * replicasets (aka 'rs')
   * replicationcontrollers (aka 'rc')
   * secrets
   * serviceaccounts (aka 'sa')
   * services (aka 'svc')
`
	usage_template = `{{if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}{{end}}{{if .HasExample}}

Examples:
{{ .Example }}{{end}}{{ if .HasAvailableSubCommands}}

Available Sub-commands:{{range .Commands}}{{if .IsAvailableCommand}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{ if .HasLocalFlags}}

Flags:
{{.LocalFlags.FlagUsages | trimRightSpace}}{{end}}{{ if .HasInheritedFlags}}

Global Flags:
{{.InheritedFlags.FlagUsages | trimRightSpace}}{{end}}{{if .HasHelpSubCommands}}

Additional help topics:{{range .Commands}}{{if .IsHelpCommand}}
  {{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}

Usage:{{if .Runnable}}
  {{if .HasFlags}}{{appendIfNotPresent .UseLine "[flags]"}}{{else}}{{.UseLine}}{{end}}{{end}}{{ if .HasSubCommands }}
  {{ .CommandPath}} [command]

Use "{{.CommandPath}} [command] --help" for more information about a command.{{end}}
`
	help_template = `{{with or .Long .Short }}{{. | trim}}{{end}}{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
)

// NewKubectlCommand creates the `kubectl` command and its nested children.
func NewKubectlCommand(f *cmdutil.Factory, in io.Reader, out, err io.Writer) *cobra.Command {
	// Parent command to which all subcommands are added.
	cmds := &cobra.Command{
		Use:   "kubectl",
		Short: "kubectl controls the Kubernetes cluster manager",
		Long: `kubectl controls the Kubernetes cluster manager.

Find more information at https://github.com/kubernetes/kubernetes.`,
		Run: runHelp,
		BashCompletionFunction: bash_completion_func,
	}

	f.BindFlags(cmds.PersistentFlags())
	f.BindExternalFlags(cmds.PersistentFlags())

	// From this point and forward we get warnings on flags that contain "_" separators
	cmds.SetGlobalNormalizationFunc(flag.WarnWordSepNormalizeFunc)

	groups := templates.CommandGroups{
		{
			Message: "Basic Commands (Beginner):",
			Commands: []*cobra.Command{
				NewCmdCreate(f, out),
				NewCmdExposeService(f, out),
				NewCmdRun(f, in, out, err),
				set.NewCmdSet(f, out),
			},
		},
		{
			Message: "Basic Commands (Intermediate):",
			Commands: []*cobra.Command{
				NewCmdGet(f, out, err),
				NewCmdExplain(f, out, err),
				NewCmdEdit(f, out, err),
				NewCmdDelete(f, out),
			},
		},
		{
			Message: "Deploy Commands:",
			Commands: []*cobra.Command{
				rollout.NewCmdRollout(f, out),
				NewCmdRollingUpdate(f, out),
				NewCmdScale(f, out),
				NewCmdAutoscale(f, out),
			},
		},
		{
			Message: "Cluster Management Commands:",
			Commands: []*cobra.Command{
				NewCmdClusterInfo(f, out),
				NewCmdTop(f, out),
				NewCmdCordon(f, out),
				NewCmdUncordon(f, out),
				NewCmdDrain(f, out),
				NewCmdTaint(f, out),
			},
		},
		{
			Message: "Troubleshooting and Debugging Commands:",
			Commands: []*cobra.Command{
				NewCmdDescribe(f, out, err),
				NewCmdLogs(f, out),
				NewCmdAttach(f, in, out, err),
				NewCmdExec(f, in, out, err),
				NewCmdPortForward(f, out, err),
				NewCmdProxy(f, out),
			},
		},
		{
			Message: "Advanced Commands:",
			Commands: []*cobra.Command{
				NewCmdApply(f, out),
				NewCmdPatch(f, out),
				NewCmdReplace(f, out),
				NewCmdConvert(f, out),
			},
		},
		{
			Message: "Settings Commands:",
			Commands: []*cobra.Command{
				NewCmdLabel(f, out),
				NewCmdAnnotate(f, out),
				NewCmdCompletion(f, out),
			},
		},
	}
	groups.Add(cmds)

	filters := []string{
		"options",
		Deprecated("kubectl", "delete", cmds, NewCmdStop(f, out)),
		Deprecated("kubectl", "config set-context", cmds, NewCmdNamespace(out)),
	}
	templates.ActsAsRootCommand(cmds, filters, groups...)

	if cmds.Flag("namespace") != nil {
		if cmds.Flag("namespace").Annotations == nil {
			cmds.Flag("namespace").Annotations = map[string][]string{}
		}
		cmds.Flag("namespace").Annotations[cobra.BashCompCustom] = append(
			cmds.Flag("namespace").Annotations[cobra.BashCompCustom],
			"__kubectl_get_namespaces",
		)
	}

	cmds.AddCommand(cmdconfig.NewCmdConfig(clientcmd.NewDefaultPathOptions(), out))
	cmds.AddCommand(NewCmdVersion(f, out))
	cmds.AddCommand(NewCmdApiVersions(f, out))
	cmds.AddCommand(NewCmdOptions(out))

	return cmds
}

func runHelp(cmd *cobra.Command, args []string) {
	cmd.Help()
}

func printDeprecationWarning(command, alias string) {
	glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", alias, command)
}

func Deprecated(baseName, to string, parent, cmd *cobra.Command) string {
	cmd.Long = fmt.Sprintf("Deprecated: This command is deprecated, all its functionalities are covered by \"%s %s\"", baseName, to)
	cmd.Short = fmt.Sprintf("Deprecated: %s", to)
	parent.AddCommand(cmd)
	return cmd.Name()
}
<|endoftext|>


<commit_before>package kring

import (
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"syscall"

	"github.com/awnumar/memguard"
	"github.com/jsipprell/keyctl"
	"github.com/oklog/ulid"
	"github.com/stephane-martin/skewer/sys/semaphore"
	"golang.org/x/crypto/ed25519"
)

func joinSessionKeyRing() error {
	_, _, errno := syscall.Syscall6(syscall_keyctl, 1, 0, 0, 0, 0, 0)
	if errno != 0 {
		return errno
	}
	return nil
}

type ring struct {
	creds RingCreds
}

func GetRing(creds RingCreds) Ring {
	return &ring{creds: creds}
}

func (r *ring) Destroy() {
	r.creds.Secret.Destroy()
	destroySem(r.creds.SessionID)
}

func NewRing() (r Ring, err error) {
	err = joinSessionKeyRing()
	if err != nil {
		return nil, err
	}
	creds, err := NewCreds()
	if err != nil {
		return nil, err
	}
	return GetRing(creds), nil
}

func (r *ring) WriteRingPass(w io.Writer) (err error) {
	_, err = w.Write(r.creds.Secret.Buffer())
	return err
}

func (r *ring) GetSessionID() ulid.ULID {
	return r.creds.SessionID
}

func getSecret(session ulid.ULID, label string) (pubkey *memguard.LockedBuffer, err error) {
	sessionStr := session.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		fmt.Fprintln(os.Stderr, "new semaphore error", err)
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		fmt.Fprintln(os.Stderr, "semaphore lock error", err)
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-%s-%s", label, sessionStr))
	if err != nil {
		return nil, err
	}
	data, err := key.Get()
	if err != nil {
		return nil, err
	}
	secret, err := memguard.NewImmutableFromBytes(data)
	if err != nil {
		return nil, err
	}
	return secret, nil
}

func (r *ring) GetSignaturePubkey() (pubkey *memguard.LockedBuffer, err error) {
	return getSecret(r.creds.SessionID, "sigpubkey")
}

func (r *ring) NewSignaturePubkey() (privkey *memguard.LockedBuffer, err error) {
	sessionStr := r.creds.SessionID.String()
	pub, priv, err := ed25519.GenerateKey(nil)
	if err != nil {
		return nil, err
	}

	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	pubkey, err := memguard.NewImmutableFromBytes(pub)
	if err != nil {
		return nil, err
	}
	privkey, err = memguard.NewImmutableFromBytes(priv)
	if err != nil {
		return nil, err
	}
	_, err = keyring.Add(fmt.Sprintf("skewer-sigpubkey-%s", sessionStr), pubkey.Buffer())
	pubkey.Destroy()
	if err != nil {
		privkey.Destroy()
		return nil, err
	}
	return privkey, nil
}

func (r *ring) NewBoxSecret() (secret *memguard.LockedBuffer, err error) {
	sessionStr := r.creds.SessionID.String()
	secretKey := make([]byte, 32)
	_, err = rand.Read(secretKey)
	if err != nil {
		return nil, err
	}

	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	secret, err = memguard.NewImmutableFromBytes(secretKey)
	if err != nil {
		return nil, err
	}
	_, err = keyring.Add(fmt.Sprintf("skewer-boxsecret-%s", sessionStr), secret.Buffer())
	if err != nil {
		secret.Destroy()
		return nil, err
	}
	return secret, nil
}

func (r *ring) GetBoxSecret() (secret *memguard.LockedBuffer, err error) {
	return getSecret(r.creds.SessionID, "boxsecret")
}

func (r *ring) DeleteBoxSecret() error {
	sessionStr := r.creds.SessionID.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return err
	}
	err = sem.Lock()
	if err != nil {
		return err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-boxsecret-%s", sessionStr))
	if err != nil {
		return err
	}
	return key.Unlink()
}

func (r *ring) DeleteSignaturePubKey() error {
	sessionStr := r.creds.SessionID.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return err
	}
	err = sem.Lock()
	if err != nil {
		return err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-sigpubkey-%s", sessionStr))
	if err != nil {
		return err
	}
	return key.Unlink()
}
<commit_msg>encrypt secrets into the linux kernel ring<commit_after>package kring

import (
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"syscall"

	"github.com/awnumar/memguard"
	"github.com/jsipprell/keyctl"
	"github.com/oklog/ulid"
	"github.com/stephane-martin/skewer/sys/semaphore"
	"github.com/stephane-martin/skewer/utils/sbox"
	"golang.org/x/crypto/ed25519"
)

func joinSessionKeyRing() error {
	_, _, errno := syscall.Syscall6(syscall_keyctl, 1, 0, 0, 0, 0, 0)
	if errno != 0 {
		return errno
	}
	return nil
}

type ring struct {
	creds RingCreds
}

func GetRing(creds RingCreds) Ring {
	return &ring{creds: creds}
}

func (r *ring) Destroy() {
	r.creds.Secret.Destroy()
	destroySem(r.creds.SessionID)
}

func NewRing() (r Ring, err error) {
	err = joinSessionKeyRing()
	if err != nil {
		return nil, err
	}
	creds, err := NewCreds()
	if err != nil {
		return nil, err
	}
	return GetRing(creds), nil
}

func (r *ring) WriteRingPass(w io.Writer) (err error) {
	_, err = w.Write(r.creds.Secret.Buffer())
	return err
}

func (r *ring) GetSessionID() ulid.ULID {
	return r.creds.SessionID
}

func getSecret(creds RingCreds, label string) (pubkey *memguard.LockedBuffer, err error) {
	sessionStr := creds.SessionID.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		fmt.Fprintln(os.Stderr, "new semaphore error", err)
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		fmt.Fprintln(os.Stderr, "semaphore lock error", err)
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-%s-%s", label, sessionStr))
	if err != nil {
		return nil, err
	}
	encrypted, err := key.Get()
	if err != nil {
		return nil, err
	}
	decrypted, err := sbox.Decrypt(encrypted, creds.Secret)
	if err != nil {
		return nil, err
	}
	secret, err := memguard.NewImmutableFromBytes(decrypted)
	if err != nil {
		return nil, err
	}
	return secret, nil
}

func (r *ring) GetSignaturePubkey() (pubkey *memguard.LockedBuffer, err error) {
	return getSecret(r.creds, "sigpubkey")
}

func (r *ring) NewSignaturePubkey() (privkey *memguard.LockedBuffer, err error) {
	sessionStr := r.creds.SessionID.String()
	pub, priv, err := ed25519.GenerateKey(nil)
	if err != nil {
		return nil, err
	}

	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	pubkey, err := memguard.NewImmutableFromBytes(pub)
	if err != nil {
		return nil, err
	}
	privkey, err = memguard.NewImmutableFromBytes(priv)
	if err != nil {
		return nil, err
	}
	encryptedPubKey, err := sbox.Encrypt(pubkey.Buffer(), r.creds.Secret)
	if err != nil {
		return nil, err
	}
	_, err = keyring.Add(fmt.Sprintf("skewer-sigpubkey-%s", sessionStr), encryptedPubKey)
	pubkey.Destroy()
	if err != nil {
		privkey.Destroy()
		return nil, err
	}
	return privkey, nil
}

func (r *ring) NewBoxSecret() (secret *memguard.LockedBuffer, err error) {
	sessionStr := r.creds.SessionID.String()
	secretKey := make([]byte, 32)
	_, err = rand.Read(secretKey)
	if err != nil {
		return nil, err
	}

	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return nil, err
	}
	err = sem.Lock()
	if err != nil {
		return nil, err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return nil, err
	}
	secret, err = memguard.NewImmutableFromBytes(secretKey)
	if err != nil {
		return nil, err
	}
	encryptedSecret, err := sbox.Encrypt(secret.Buffer(), r.creds.Secret)
	if err != nil {
		return nil, err
	}
	_, err = keyring.Add(fmt.Sprintf("skewer-boxsecret-%s", sessionStr), encryptedSecret)
	if err != nil {
		secret.Destroy()
		return nil, err
	}
	return secret, nil
}

func (r *ring) GetBoxSecret() (secret *memguard.LockedBuffer, err error) {
	return getSecret(r.creds, "boxsecret")
}

func (r *ring) DeleteBoxSecret() error {
	sessionStr := r.creds.SessionID.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return err
	}
	err = sem.Lock()
	if err != nil {
		return err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-boxsecret-%s", sessionStr))
	if err != nil {
		return err
	}
	return key.Unlink()
}

func (r *ring) DeleteSignaturePubKey() error {
	sessionStr := r.creds.SessionID.String()
	sem, err := semaphore.New(fmt.Sprintf("skw%s", sessionStr))
	if err != nil {
		return err
	}
	err = sem.Lock()
	if err != nil {
		return err
	}
	defer func() {
		sem.Unlock()
		sem.Close()
	}()

	keyring, err := keyctl.SessionKeyring()
	if err != nil {
		return err
	}
	key, err := keyring.Search(fmt.Sprintf("skewer-sigpubkey-%s", sessionStr))
	if err != nil {
		return err
	}
	return key.Unlink()
}
<|endoftext|>


<commit_before>package gosher

import (
	"bufio"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func (s *SshClient) download(remotePath string, localPath string) (*SshResponse, error) {
	localPathInfo, err := os.Stat(localPath)
	destinationDirectory := localPath
	var useSpecifiedFilename bool
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		} else {
			//OK - create file/dir
			useSpecifiedFilename = true
		}
	} else if localPathInfo.IsDir() {
		//ok - use name of remotePath
		//localPath = filepath.Join(localPath, filepath.Base(remotePath))
		destinationDirectory = localPath
		useSpecifiedFilename = false
	} else {
		destinationDirectory = filepath.Dir(localPath)
		useSpecifiedFilename = true
	}
	//from-scp
	response := NewSshResponse(s.address, s.session.Stdout, s.session.Stderr)

	if err != nil {
		return response, err
	}
	defer s.session.Close()
	ce := make(chan error)
	go func() {
		cw, err := s.session.StdinPipe()
		if err != nil {
			ce <- err
			return
		}
		defer cw.Close()
		r, err := s.session.StdoutPipe()
		if err != nil {
			ce <- err
			return
		}
		err = sendByte(cw, 0)
		if err != nil {
			ce <- err
			return
		}
		//defer r.Close()
		//use a scanner for processing individual commands, but not files themselves
		scanner := bufio.NewScanner(r)
		more := true
		first := true
		for more {
			cmdArr := make([]byte, 1)
			n, err := r.Read(cmdArr)
			if err != nil {
				if err == io.EOF {
					//no problem.
				} else {
					ce <- err
				}
				return
			}
			if n < 1 {
				ce <- errors.New("Error reading next byte from standard input")
				return
			}
			cmd := cmdArr[0]
			switch cmd {
			case 0x0:
				//continue
			case 'E':
				//E command: go back out of dir
				destinationDirectory = filepath.Dir(destinationDirectory)
				err = sendByte(cw, 0)
				if err != nil {
					ce <- err
					return
				}
			case 0xA:
				//0xA command: end?
				err = sendByte(cw, 0)
				if err != nil {
					ce <- err
					return
				}

				return
			default:
				scanner.Scan()
				err = scanner.Err()
				if err != nil {
					if err == io.EOF {
						//no problem.
					} else {
						ce <- err
					}
					return
				}
				//first line
				cmdFull := scanner.Text()
				//remainder, split by spaces
				parts := strings.SplitN(cmdFull, " ", 3)

				switch cmd {
				case 0x1:
					ce <- errors.New(cmdFull[1:])
					return
				case 'D', 'C':
					mode, err := strconv.ParseInt(parts[0], 8, 32)
					if err != nil {
						ce <- err
						return
					}
					sizeUint, err := strconv.ParseUint(parts[1], 10, 64)
					size := int64(sizeUint)
					if err != nil {
						ce <- err
						return
					}
					rcvFilename := parts[2]
					var filename string
					//use the specified filename from the destination (only for top-level item)
					if useSpecifiedFilename && first {
						filename = filepath.Base(localPath)
					} else {
						filename = rcvFilename
					}
					err = sendByte(cw, 0)
					if err != nil {
						ce <- err
						return
					}
					if cmd == 'C' {
						//C command - file
						thisLocalPath := filepath.Join(destinationDirectory, filename)
						tot := int64(0)

						fw, err := os.Create(thisLocalPath)
						if err != nil {
							ce <- err
							return
						}
						defer fw.Close()

						//buffered by 4096 bytes
						bufferSize := int64(4096)
						for tot < size {
							if bufferSize > size-tot {
								bufferSize = size - tot
							}
							b := make([]byte, bufferSize)
							n, err = r.Read(b)
							if err != nil {
								ce <- err
								return
							}
							tot += int64(n)
							//write to file
							_, err = fw.Write(b[:n])
							if err != nil {
								ce <- err
								return
							}
						}
						//close file writer & check error
						err = fw.Close()
						if err != nil {
							ce <- err
							return
						}
						//get next byte from channel reader
						nb := make([]byte, 1)
						_, err = r.Read(nb)
						if err != nil {
							ce <- err
							return
						}
						//TODO check value received in nb
						//send null-byte back
						_, err = cw.Write([]byte{0})
						if err != nil {
							ce <- err
							return
						}
					} else {
						//D command (directory)
						thisDstFile := filepath.Join(destinationDirectory, filename)
						fileMode := os.FileMode(uint32(mode))
						err = os.MkdirAll(thisDstFile, fileMode)
						if err != nil {
							ce <- err
							return
						}
						destinationDirectory = thisDstFile
					}
				default:
					return
				}
			}
			first = false
		}
		err = cw.Close()
		if err != nil {
			ce <- err
			return
		}
	}()
	remoteOpts := "-fr"
	err = s.session.Run("/usr/bin/scp " + remoteOpts + " " + remotePath)
	return response, err
}

func sendByte(w io.Writer, val byte) error {
	_, err := w.Write([]byte{val})
	return err
}
<commit_msg>Refactoring the download part<commit_after>package gosher

import (
	"bufio"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func (s *SshClient) download(remotePath string, localPath string) (*SshResponse, error) {
	localPathInfo, err := os.Stat(localPath)
	destinationDirectory := localPath
	var useSpecifiedFilename bool
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		} else {
			//OK - create file/dir
			useSpecifiedFilename = true
		}
	} else if localPathInfo.IsDir() {
		//ok - use name of remotePath
		destinationDirectory = localPath
		useSpecifiedFilename = false
	} else {
		destinationDirectory = filepath.Dir(localPath)
		useSpecifiedFilename = true
	}
	//from-scp
	response := NewSshResponse(s.address, s.session.Stdout, s.session.Stderr)

	if err != nil {
		return response, err
	}
	defer s.session.Close()
	errorChannel := make(chan error)
	go s.manageDownloads(errorChannel, destinationDirectory, useSpecifiedFilename, localPath)
	remoteOpts := "-fr"
	err = s.session.Run("/usr/bin/scp " + remoteOpts + " " + remotePath)
	return response, err
}

func (s *SshClient) manageDownloads(errorChannel chan error, destinationDirectory string,
	useSpecifiedFilename bool, localPath string) {
	inPipe, err := s.session.StdinPipe()
	if err != nil {
		errorChannel <- err
		return
	}
	defer inPipe.Close()
	outPipe, err := s.session.StdoutPipe()
	if err != nil {
		errorChannel <- err
		return
	}
	err = sendByte(inPipe, 0)
	if err != nil {
		errorChannel <- err
		return
	}
	//use a scanner for processing individual commands, but not files themselves
	scanner := bufio.NewScanner(outPipe)
	more := true
	isFirstCommand := true
	for more {
		commandArray := make([]byte, 1)
		commandLength, err := outPipe.Read(commandArray)
		if err != nil {
			if err == io.EOF {
				//no problem.
			} else {
				errorChannel <- err
			}
			return
		}
		if commandLength < 1 {
			errorChannel <- errors.New("Error reading next byte from standard input")
			return
		}
		command := commandArray[0]
		switch command {
		case 0x0:
			//continue
		case 'E':
			//E command: go back out of dir
			destinationDirectory = filepath.Dir(destinationDirectory)
			err = sendByte(inPipe, 0)
			if err != nil {
				errorChannel <- err
				return
			}
		case 0xA:
			//0xA command: end?
			err = sendByte(inPipe, 0)
			if err != nil {
				errorChannel <- err
				return
			}
			return
		default:
			scanner.Scan()
			err = scanner.Err()
			if err != nil {
				if err == io.EOF {
					// no problem.
				} else {
					errorChannel <- err
				}
				return
			}
			// first line
			fullCommand := scanner.Text()
			// remainder, split by spaces
			splitCommands := strings.SplitN(fullCommand, " ", 3)

			switch command {
			case 0x1:
				errorChannel <- errors.New(fullCommand[1:])
				return
			case 'D', 'C':
				if err = s.manageWrites(splitCommands, inPipe, command, isFirstCommand,
					outPipe, destinationDirectory, useSpecifiedFilename, localPath, commandLength); err != nil {
					errorChannel <- err
					return
				}
			default:
				return
			}
		}
		isFirstCommand = false
	}
	err = inPipe.Close()
	if err != nil {
		errorChannel <- err
		return
	}
}

func (s *SshClient) manageWrites(splitCommands []string, inPipe io.WriteCloser, command byte, isFirstCommand bool,
	outPipe io.Reader, destinationDirectory string, useSpecifiedFilename bool, localPath string, commandLength int) error {
	mode, err := strconv.ParseInt(splitCommands[0], 8, 32)
	if err != nil {
		return err
	}
	sizeUint, err := strconv.ParseUint(splitCommands[1], 10, 64)
	commandSize := int64(sizeUint)
	if err != nil {
		return err
	}
	rcvFilename := splitCommands[2]
	var filename string
	//use the specified filename from the destination (only for top-level item)
	if useSpecifiedFilename && isFirstCommand {
		filename = filepath.Base(localPath)
	} else {
		filename = rcvFilename
	}
	err = sendByte(inPipe, 0)
	if err != nil {
		return err
	}
	if command == 'C' {
		//C command - file
		if err = s.writeFileFromPipe(destinationDirectory, filename, inPipe,
			commandLength, commandSize, outPipe); err != nil {
			return err
		}
	} else {
		//D command (directory)
		thisDstFile := filepath.Join(destinationDirectory, filename)
		fileMode := os.FileMode(uint32(mode))
		err = os.MkdirAll(thisDstFile, fileMode)
		if err != nil {
			return err
		}
		destinationDirectory = thisDstFile
	}
	return nil
}

func (s *SshClient) writeFileFromPipe(destinationDirectory string, filename string,
	inPipe io.WriteCloser, commandLength int, commandSize int64, outPipe io.Reader) error {
	thisLocalPath := filepath.Join(destinationDirectory, filename)
	tot := int64(0)

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

	//buffered by 4096 bytes
	bufferSize := int64(4096)
	for tot < commandSize {
		if bufferSize > commandSize-tot {
			bufferSize = commandSize - tot
		}
		b := make([]byte, bufferSize)
		commandLength, err = outPipe.Read(b)
		if err != nil {
			return err
		}
		tot += int64(commandLength)
		//write to file
		_, err = fileWriter.Write(b[:commandLength])
		if err != nil {
			return err
		}
	}
	//close file writer & check error
	err = fileWriter.Close()
	if err != nil {
		return err
	}
	//get next byte from channel reader
	nextByte := make([]byte, 1)
	_, err = outPipe.Read(nextByte)
	if err != nil {
		return err
	}
	//send null-byte back
	_, err = inPipe.Write([]byte{0})
	if err != nil {
		return err
	}
	return nil
}

func sendByte(w io.Writer, val byte) error {
	_, err := w.Write([]byte{val})
	return err
}
<|endoftext|>


<commit_before>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/google/syzkaller/pkg/compiler"
	"github.com/google/syzkaller/pkg/osutil"
)

func extract(info *compiler.ConstInfo, cc string, args []string, addSource string, declarePrintf bool) (
	map[string]uint64, map[string]bool, error) {
	data := &CompileData{
		AddSource:     addSource,
		Defines:       info.Defines,
		Includes:      info.Includes,
		Values:        info.Consts,
		DeclarePrintf: declarePrintf,
	}
	undeclared := make(map[string]bool)
	bin, out, err := compile(cc, args, data)
	if err != nil {
		// Some consts and syscall numbers are not defined on some archs.
		// Figure out from compiler output undefined consts,
		// and try to compile again without them.
		valMap := make(map[string]bool)
		for _, val := range info.Consts {
			valMap[val] = true
		}
		for _, errMsg := range []string{
			"error: ‘([a-zA-Z0-9_]+)’ undeclared",
			"error: '([a-zA-Z0-9_]+)' undeclared",
			"note: in expansion of macro ‘([a-zA-Z0-9_]+)’",
			"error: use of undeclared identifier '([a-zA-Z0-9_]+)'",
		} {
			re := regexp.MustCompile(errMsg)
			matches := re.FindAllSubmatch(out, -1)
			for _, match := range matches {
				val := string(match[1])
				if valMap[val] {
					undeclared[val] = true
				}
			}
		}
		data.Values = nil
		for _, v := range info.Consts {
			if undeclared[v] {
				continue
			}
			data.Values = append(data.Values, v)
		}
		bin, out, err = compile(cc, args, data)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to run compiler: %v\n%v", err, string(out))
		}
	}
	defer os.Remove(bin)

	out, err = osutil.Command(bin).CombinedOutput()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to run flags binary: %v\n%v", err, string(out))
	}
	flagVals := strings.Split(string(out), " ")
	if len(out) == 0 {
		flagVals = nil
	}
	if len(flagVals) != len(data.Values) {
		return nil, nil, fmt.Errorf("fetched wrong number of values %v, want != %v",
			len(flagVals), len(data.Values))
	}
	res := make(map[string]uint64)
	for i, name := range data.Values {
		val := flagVals[i]
		n, err := strconv.ParseUint(val, 10, 64)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to parse value: %v (%v)", err, val)
		}
		res[name] = n
	}
	return res, undeclared, nil
}

type CompileData struct {
	AddSource     string
	Defines       map[string]string
	Includes      []string
	Values        []string
	DeclarePrintf bool
}

func compile(cc string, args []string, data *CompileData) (bin string, out []byte, err error) {
	src := new(bytes.Buffer)
	if err := srcTemplate.Execute(src, data); err != nil {
		return "", nil, fmt.Errorf("failed to generate source: %v", err)
	}
	binFile, err := osutil.TempFile("syz-extract-bin")
	if err != nil {
		return "", nil, err
	}
	args = append(args, []string{
		"-x", "c", "-",
		"-o", binFile,
		"-w",
	}...)
	cmd := osutil.Command(cc, args...)
	cmd.Stdin = src
	if out, err := cmd.CombinedOutput(); err != nil {
		os.Remove(binFile)
		return "", out, err
	}
	return binFile, nil, nil
}

var srcTemplate = template.Must(template.New("").Parse(`
#define __asm__(...)

{{range $incl := $.Includes}}
#include <{{$incl}}>
{{end}}

{{range $name, $val := $.Defines}}
#ifndef {{$name}}
#	define {{$name}} {{$val}}
#endif
{{end}}

{{.AddSource}}

{{if .DeclarePrintf}}
int printf(const char *format, ...);
{{end}}

int main() {
	int i;
	unsigned long long vals[] = {
		{{range $val := $.Values}}(unsigned long long){{$val}},
		{{end}}
	};
	for (i = 0; i < sizeof(vals)/sizeof(vals[0]); i++) {
		if (i != 0)
			printf(" ");
		printf("%llu", vals[i]);
	}
	return 0;
}
`))
<commit_msg>sys/syz-extract: support more identifier quoting options<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/google/syzkaller/pkg/compiler"
	"github.com/google/syzkaller/pkg/osutil"
)

func extract(info *compiler.ConstInfo, cc string, args []string, addSource string, declarePrintf bool) (
	map[string]uint64, map[string]bool, error) {
	data := &CompileData{
		AddSource:     addSource,
		Defines:       info.Defines,
		Includes:      info.Includes,
		Values:        info.Consts,
		DeclarePrintf: declarePrintf,
	}
	undeclared := make(map[string]bool)
	bin, out, err := compile(cc, args, data)
	if err != nil {
		// Some consts and syscall numbers are not defined on some archs.
		// Figure out from compiler output undefined consts,
		// and try to compile again without them.
		valMap := make(map[string]bool)
		for _, val := range info.Consts {
			valMap[val] = true
		}
		for _, errMsg := range []string{
			`error: [‘']([a-zA-Z0-9_]+)[’'] undeclared`,
			`note: in expansion of macro [‘']([a-zA-Z0-9_]+)[’']`,
			`error: use of undeclared identifier [‘']([a-zA-Z0-9_]+)[’']`,
		} {
			re := regexp.MustCompile(errMsg)
			matches := re.FindAllSubmatch(out, -1)
			for _, match := range matches {
				val := string(match[1])
				if valMap[val] {
					undeclared[val] = true
				}
			}
		}
		data.Values = nil
		for _, v := range info.Consts {
			if undeclared[v] {
				continue
			}
			data.Values = append(data.Values, v)
		}
		bin, out, err = compile(cc, args, data)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to run compiler: %v\n%v", err, string(out))
		}
	}
	defer os.Remove(bin)

	out, err = osutil.Command(bin).CombinedOutput()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to run flags binary: %v\n%v", err, string(out))
	}
	flagVals := strings.Split(string(out), " ")
	if len(out) == 0 {
		flagVals = nil
	}
	if len(flagVals) != len(data.Values) {
		return nil, nil, fmt.Errorf("fetched wrong number of values %v, want != %v",
			len(flagVals), len(data.Values))
	}
	res := make(map[string]uint64)
	for i, name := range data.Values {
		val := flagVals[i]
		n, err := strconv.ParseUint(val, 10, 64)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to parse value: %v (%v)", err, val)
		}
		res[name] = n
	}
	return res, undeclared, nil
}

type CompileData struct {
	AddSource     string
	Defines       map[string]string
	Includes      []string
	Values        []string
	DeclarePrintf bool
}

func compile(cc string, args []string, data *CompileData) (bin string, out []byte, err error) {
	src := new(bytes.Buffer)
	if err := srcTemplate.Execute(src, data); err != nil {
		return "", nil, fmt.Errorf("failed to generate source: %v", err)
	}
	binFile, err := osutil.TempFile("syz-extract-bin")
	if err != nil {
		return "", nil, err
	}
	args = append(args, []string{
		"-x", "c", "-",
		"-o", binFile,
		"-w",
	}...)
	cmd := osutil.Command(cc, args...)
	cmd.Stdin = src
	if out, err := cmd.CombinedOutput(); err != nil {
		os.Remove(binFile)
		return "", out, err
	}
	return binFile, nil, nil
}

var srcTemplate = template.Must(template.New("").Parse(`
#define __asm__(...)

{{range $incl := $.Includes}}
#include <{{$incl}}>
{{end}}

{{range $name, $val := $.Defines}}
#ifndef {{$name}}
#	define {{$name}} {{$val}}
#endif
{{end}}

{{.AddSource}}

{{if .DeclarePrintf}}
int printf(const char *format, ...);
{{end}}

int main() {
	int i;
	unsigned long long vals[] = {
		{{range $val := $.Values}}(unsigned long long){{$val}},
		{{end}}
	};
	for (i = 0; i < sizeof(vals)/sizeof(vals[0]); i++) {
		if (i != 0)
			printf(" ");
		printf("%llu", vals[i]);
	}
	return 0;
}
`))
<|endoftext|>