<commit_before>// Copyright 2020 Google 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 cloudsqlconn

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"cloud.google.com/go/cloudsqlconn/errtypes"
	"cloud.google.com/go/cloudsqlconn/internal/cloudsql"
	"cloud.google.com/go/cloudsqlconn/internal/trace"
	"github.com/google/uuid"
	"golang.org/x/net/proxy"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"
	sqladmin "google.golang.org/api/sqladmin/v1beta4"
)

const (
	// versionString indicates the version of this library.
	versionString = "0.1.0-dev"
	userAgent     = "cloud-sql-go-connector/" + versionString

	// defaultTCPKeepAlive is the default keep alive value used on connections to a Cloud SQL instance.
	defaultTCPKeepAlive = 30 * time.Second
	// serverProxyPort is the port the server-side proxy receives connections on.
	serverProxyPort = "3307"
)

var (
	// defaultKey is the default RSA public/private keypair used by the clients.
	defaultKey    *rsa.PrivateKey
	defaultKeyErr error
	keyOnce       sync.Once
)

func getDefaultKeys() (*rsa.PrivateKey, error) {
	keyOnce.Do(func() {
		defaultKey, defaultKeyErr = rsa.GenerateKey(rand.Reader, 2048)
	})
	return defaultKey, defaultKeyErr
}

// A Dialer is used to create connections to Cloud SQL instances.
//
// Use NewDialer to initialize a Dialer.
type Dialer struct {
	lock sync.RWMutex
	// instances map connection names (e.g., my-project:us-central1:my-instance)
	// to *cloudsql.Instance types.
	instances      map[string]*cloudsql.Instance
	key            *rsa.PrivateKey
	refreshTimeout time.Duration

	sqladmin *sqladmin.Service

	// defaultDialCfg holds the constructor level DialOptions, so that it can
	// be copied and mutated by the Dial function.
	defaultDialCfg dialCfg

	// dialerID uniquely identifies a Dialer. Used for monitoring purposes,
	// *only* when a client has configured OpenCensus exporters.
	dialerID string

	// dialFunc is the function used to connect to the address on the named
	// network. By default it is golang.org/x/net/proxy#Dial.
	dialFunc func(cxt context.Context, network, addr string) (net.Conn, error)

	// iamTokenSource supplies the OAuth2 token used for IAM DB Authn. If IAM DB
	// Authn is not enabled, iamTokenSource will be nil.
	iamTokenSource oauth2.TokenSource
}

// NewDialer creates a new Dialer.
//
// Initial calls to NewDialer make take longer than normal because generation of an
// RSA keypair is performed. Calls with a WithRSAKeyPair DialOption or after a default
// RSA keypair is generated will be faster.
func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
	cfg := &dialerConfig{
		refreshTimeout: 30 * time.Second,
		sqladminOpts:   []option.ClientOption{option.WithUserAgent(userAgent)},
		dialFunc:       proxy.Dial,
	}
	for _, opt := range opts {
		opt(cfg)
		if cfg.err != nil {
			return nil, cfg.err
		}
	}
	// If callers have not provided a token source, either explicitly with
	// WithTokenSource or implicitly with WithCredentialsJSON etc, then use the
	// default token source.
	if cfg.useIAMAuthN && cfg.tokenSource == nil {
		ts, err := google.DefaultTokenSource(ctx, sqladmin.SqlserviceAdminScope)
		if err != nil {
			return nil, fmt.Errorf("failed to create token source: %v", err)
		}
		cfg.tokenSource = ts
	}
	// If IAM Authn is not explicitly enabled, remove the token source.
	if !cfg.useIAMAuthN {
		cfg.tokenSource = nil
	}

	if cfg.rsaKey == nil {
		key, err := getDefaultKeys()
		if err != nil {
			return nil, fmt.Errorf("failed to generate RSA keys: %v", err)
		}
		cfg.rsaKey = key
	}

	client, err := sqladmin.NewService(ctx, cfg.sqladminOpts...)
	if err != nil {
		return nil, fmt.Errorf("failed to create sqladmin client: %v", err)
	}

	dialCfg := dialCfg{
		ipType:       cloudsql.PublicIP,
		tcpKeepAlive: defaultTCPKeepAlive,
	}
	for _, opt := range cfg.dialOpts {
		opt(&dialCfg)
	}

	if err := trace.InitMetrics(); err != nil {
		return nil, err
	}
	d := &Dialer{
		instances:      make(map[string]*cloudsql.Instance),
		key:            cfg.rsaKey,
		refreshTimeout: cfg.refreshTimeout,
		sqladmin:       client,
		defaultDialCfg: dialCfg,
		dialerID:       uuid.New().String(),
		iamTokenSource: cfg.tokenSource,
		dialFunc:       cfg.dialFunc,
	}
	return d, nil
}

// Dial returns a net.Conn connected to the specified Cloud SQL instance. The instance argument must be the
// instance's connection name, which is in the format "project-name:region:instance-name".
func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption) (conn net.Conn, err error) {
	startTime := time.Now()
	var endDial trace.EndSpanFunc
	ctx, endDial = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn.Dial",
		trace.AddInstanceName(instance),
		trace.AddDialerID(d.dialerID),
	)
	defer func() { endDial(err) }()
	cfg := d.defaultDialCfg
	for _, opt := range opts {
		opt(&cfg)
	}

	var endInfo trace.EndSpanFunc
	ctx, endInfo = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.InstanceInfo")
	i, err := d.instance(instance)
	if err != nil {
		endInfo(err)
		return nil, err
	}
	addr, tlsCfg, err := i.ConnectInfo(ctx, cfg.ipType)
	if err != nil {
		endInfo(err)
		return nil, err
	}
	endInfo(err)

	var connectEnd trace.EndSpanFunc
	ctx, connectEnd = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.Connect")
	defer func() { connectEnd(err) }()
	addr = net.JoinHostPort(addr, serverProxyPort)
	conn, err = d.dialFunc(ctx, "tcp", addr)
	if err != nil {
		// refresh the instance info in case it caused the connection failure
		i.ForceRefresh()
		return nil, errtypes.NewDialError("failed to dial", i.String(), err)
	}
	if c, ok := conn.(*net.TCPConn); ok {
		if err := c.SetKeepAlive(true); err != nil {
			return nil, errtypes.NewDialError("failed to set keep-alive", i.String(), err)
		}
		if err := c.SetKeepAlivePeriod(cfg.tcpKeepAlive); err != nil {
			return nil, errtypes.NewDialError("failed to set keep-alive period", i.String(), err)
		}
	}
	tlsConn := tls.Client(conn, tlsCfg)
	if err := tlsConn.Handshake(); err != nil {
		// refresh the instance info in case it caused the handshake failure
		i.ForceRefresh()
		_ = tlsConn.Close() // best effort close attempt
		return nil, errtypes.NewDialError("handshake failed", i.String(), err)
	}
	latency := time.Since(startTime).Milliseconds()
	go func() {
		n := atomic.AddUint64(&i.OpenConns, 1)
		trace.RecordOpenConnections(ctx, int64(n), d.dialerID, i.String())
		trace.RecordDialLatency(ctx, instance, d.dialerID, latency)
	}()

	return newInstrumentedConn(tlsConn, func() {
		n := atomic.AddUint64(&i.OpenConns, ^uint64(0))
		trace.RecordOpenConnections(context.Background(), int64(n), d.dialerID, i.String())
	}), nil
}

// EngineVersion returns the engine type and version for the instance. The value will
// corespond to one of the following types for the instance:
// https://cloud.google.com/sql/docs/mysql/admin-api/rest/v1beta4/SqlDatabaseVersion
func (d *Dialer) EngineVersion(ctx context.Context, instance string) (string, error) {
	i, err := d.instance(instance)
	if err != nil {
		return "", err
	}
	e, err := i.InstanceEngineVersion(ctx)
	if err != nil {
		return "", err
	}
	return e, nil
}

// newInstrumentedConn initializes an instrumentedConn that on closing will
// decrement the number of open connects and record the result.
func newInstrumentedConn(conn net.Conn, closeFunc func()) *instrumentedConn {
	return &instrumentedConn{
		Conn:      conn,
		closeFunc: closeFunc,
	}
}

// instrumentedConn wraps a net.Conn and invokes closeFunc when the connection
// is closed.
type instrumentedConn struct {
	net.Conn
	closeFunc func()
}

// Close delegates to the underylying net.Conn interface and reports the close
// to the provided closeFunc only when Close returns no error.
func (i *instrumentedConn) Close() error {
	err := i.Conn.Close()
	if err != nil {
		return err
	}
	go i.closeFunc()
	return nil
}

// Close closes the Dialer; it prevents the Dialer from refreshing the information
// needed to connect. Additional dial operations may succeed until the information
// expires.
func (d *Dialer) Close() {
	d.lock.Lock()
	defer d.lock.Unlock()
	for _, i := range d.instances {
		i.Close()
	}
}

func (d *Dialer) instance(connName string) (*cloudsql.Instance, error) {
	// Check instance cache
	d.lock.RLock()
	i, ok := d.instances[connName]
	d.lock.RUnlock()
	if !ok {
		d.lock.Lock()
		// Recheck to ensure instance wasn't created between locks
		i, ok = d.instances[connName]
		if !ok {
			// Create a new instance
			var err error
			i, err = cloudsql.NewInstance(connName, d.sqladmin, d.key, d.refreshTimeout, d.iamTokenSource)
			if err != nil {
				d.lock.Unlock()
				return nil, err
			}
			d.instances[connName] = i
		}
		d.lock.Unlock()
	}
	return i, nil
}
<commit_msg>fix!: dialer is a io.Closer (#76)<commit_after>// Copyright 2020 Google 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 cloudsqlconn

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"cloud.google.com/go/cloudsqlconn/errtypes"
	"cloud.google.com/go/cloudsqlconn/internal/cloudsql"
	"cloud.google.com/go/cloudsqlconn/internal/trace"
	"github.com/google/uuid"
	"golang.org/x/net/proxy"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"
	sqladmin "google.golang.org/api/sqladmin/v1beta4"
)

const (
	// versionString indicates the version of this library.
	versionString = "0.1.0-dev"
	userAgent     = "cloud-sql-go-connector/" + versionString

	// defaultTCPKeepAlive is the default keep alive value used on connections to a Cloud SQL instance.
	defaultTCPKeepAlive = 30 * time.Second
	// serverProxyPort is the port the server-side proxy receives connections on.
	serverProxyPort = "3307"
)

var (
	// defaultKey is the default RSA public/private keypair used by the clients.
	defaultKey    *rsa.PrivateKey
	defaultKeyErr error
	keyOnce       sync.Once
)

func getDefaultKeys() (*rsa.PrivateKey, error) {
	keyOnce.Do(func() {
		defaultKey, defaultKeyErr = rsa.GenerateKey(rand.Reader, 2048)
	})
	return defaultKey, defaultKeyErr
}

// A Dialer is used to create connections to Cloud SQL instances.
//
// Use NewDialer to initialize a Dialer.
type Dialer struct {
	lock sync.RWMutex
	// instances map connection names (e.g., my-project:us-central1:my-instance)
	// to *cloudsql.Instance types.
	instances      map[string]*cloudsql.Instance
	key            *rsa.PrivateKey
	refreshTimeout time.Duration

	sqladmin *sqladmin.Service

	// defaultDialCfg holds the constructor level DialOptions, so that it can
	// be copied and mutated by the Dial function.
	defaultDialCfg dialCfg

	// dialerID uniquely identifies a Dialer. Used for monitoring purposes,
	// *only* when a client has configured OpenCensus exporters.
	dialerID string

	// dialFunc is the function used to connect to the address on the named
	// network. By default it is golang.org/x/net/proxy#Dial.
	dialFunc func(cxt context.Context, network, addr string) (net.Conn, error)

	// iamTokenSource supplies the OAuth2 token used for IAM DB Authn. If IAM DB
	// Authn is not enabled, iamTokenSource will be nil.
	iamTokenSource oauth2.TokenSource
}

// NewDialer creates a new Dialer.
//
// Initial calls to NewDialer make take longer than normal because generation of an
// RSA keypair is performed. Calls with a WithRSAKeyPair DialOption or after a default
// RSA keypair is generated will be faster.
func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
	cfg := &dialerConfig{
		refreshTimeout: 30 * time.Second,
		sqladminOpts:   []option.ClientOption{option.WithUserAgent(userAgent)},
		dialFunc:       proxy.Dial,
	}
	for _, opt := range opts {
		opt(cfg)
		if cfg.err != nil {
			return nil, cfg.err
		}
	}
	// If callers have not provided a token source, either explicitly with
	// WithTokenSource or implicitly with WithCredentialsJSON etc, then use the
	// default token source.
	if cfg.useIAMAuthN && cfg.tokenSource == nil {
		ts, err := google.DefaultTokenSource(ctx, sqladmin.SqlserviceAdminScope)
		if err != nil {
			return nil, fmt.Errorf("failed to create token source: %v", err)
		}
		cfg.tokenSource = ts
	}
	// If IAM Authn is not explicitly enabled, remove the token source.
	if !cfg.useIAMAuthN {
		cfg.tokenSource = nil
	}

	if cfg.rsaKey == nil {
		key, err := getDefaultKeys()
		if err != nil {
			return nil, fmt.Errorf("failed to generate RSA keys: %v", err)
		}
		cfg.rsaKey = key
	}

	client, err := sqladmin.NewService(ctx, cfg.sqladminOpts...)
	if err != nil {
		return nil, fmt.Errorf("failed to create sqladmin client: %v", err)
	}

	dialCfg := dialCfg{
		ipType:       cloudsql.PublicIP,
		tcpKeepAlive: defaultTCPKeepAlive,
	}
	for _, opt := range cfg.dialOpts {
		opt(&dialCfg)
	}

	if err := trace.InitMetrics(); err != nil {
		return nil, err
	}
	d := &Dialer{
		instances:      make(map[string]*cloudsql.Instance),
		key:            cfg.rsaKey,
		refreshTimeout: cfg.refreshTimeout,
		sqladmin:       client,
		defaultDialCfg: dialCfg,
		dialerID:       uuid.New().String(),
		iamTokenSource: cfg.tokenSource,
		dialFunc:       cfg.dialFunc,
	}
	return d, nil
}

// Dial returns a net.Conn connected to the specified Cloud SQL instance. The instance argument must be the
// instance's connection name, which is in the format "project-name:region:instance-name".
func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption) (conn net.Conn, err error) {
	startTime := time.Now()
	var endDial trace.EndSpanFunc
	ctx, endDial = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn.Dial",
		trace.AddInstanceName(instance),
		trace.AddDialerID(d.dialerID),
	)
	defer func() { endDial(err) }()
	cfg := d.defaultDialCfg
	for _, opt := range opts {
		opt(&cfg)
	}

	var endInfo trace.EndSpanFunc
	ctx, endInfo = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.InstanceInfo")
	i, err := d.instance(instance)
	if err != nil {
		endInfo(err)
		return nil, err
	}
	addr, tlsCfg, err := i.ConnectInfo(ctx, cfg.ipType)
	if err != nil {
		endInfo(err)
		return nil, err
	}
	endInfo(err)

	var connectEnd trace.EndSpanFunc
	ctx, connectEnd = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.Connect")
	defer func() { connectEnd(err) }()
	addr = net.JoinHostPort(addr, serverProxyPort)
	conn, err = d.dialFunc(ctx, "tcp", addr)
	if err != nil {
		// refresh the instance info in case it caused the connection failure
		i.ForceRefresh()
		return nil, errtypes.NewDialError("failed to dial", i.String(), err)
	}
	if c, ok := conn.(*net.TCPConn); ok {
		if err := c.SetKeepAlive(true); err != nil {
			return nil, errtypes.NewDialError("failed to set keep-alive", i.String(), err)
		}
		if err := c.SetKeepAlivePeriod(cfg.tcpKeepAlive); err != nil {
			return nil, errtypes.NewDialError("failed to set keep-alive period", i.String(), err)
		}
	}
	tlsConn := tls.Client(conn, tlsCfg)
	if err := tlsConn.Handshake(); err != nil {
		// refresh the instance info in case it caused the handshake failure
		i.ForceRefresh()
		_ = tlsConn.Close() // best effort close attempt
		return nil, errtypes.NewDialError("handshake failed", i.String(), err)
	}
	latency := time.Since(startTime).Milliseconds()
	go func() {
		n := atomic.AddUint64(&i.OpenConns, 1)
		trace.RecordOpenConnections(ctx, int64(n), d.dialerID, i.String())
		trace.RecordDialLatency(ctx, instance, d.dialerID, latency)
	}()

	return newInstrumentedConn(tlsConn, func() {
		n := atomic.AddUint64(&i.OpenConns, ^uint64(0))
		trace.RecordOpenConnections(context.Background(), int64(n), d.dialerID, i.String())
	}), nil
}

// EngineVersion returns the engine type and version for the instance. The value will
// corespond to one of the following types for the instance:
// https://cloud.google.com/sql/docs/mysql/admin-api/rest/v1beta4/SqlDatabaseVersion
func (d *Dialer) EngineVersion(ctx context.Context, instance string) (string, error) {
	i, err := d.instance(instance)
	if err != nil {
		return "", err
	}
	e, err := i.InstanceEngineVersion(ctx)
	if err != nil {
		return "", err
	}
	return e, nil
}

// newInstrumentedConn initializes an instrumentedConn that on closing will
// decrement the number of open connects and record the result.
func newInstrumentedConn(conn net.Conn, closeFunc func()) *instrumentedConn {
	return &instrumentedConn{
		Conn:      conn,
		closeFunc: closeFunc,
	}
}

// instrumentedConn wraps a net.Conn and invokes closeFunc when the connection
// is closed.
type instrumentedConn struct {
	net.Conn
	closeFunc func()
}

// Close delegates to the underylying net.Conn interface and reports the close
// to the provided closeFunc only when Close returns no error.
func (i *instrumentedConn) Close() error {
	err := i.Conn.Close()
	if err != nil {
		return err
	}
	go i.closeFunc()
	return nil
}

// Close closes the Dialer; it prevents the Dialer from refreshing the information
// needed to connect. Additional dial operations may succeed until the information
// expires.
func (d *Dialer) Close() error {
	d.lock.Lock()
	defer d.lock.Unlock()
	for _, i := range d.instances {
		i.Close()
	}
	return nil
}

func (d *Dialer) instance(connName string) (*cloudsql.Instance, error) {
	// Check instance cache
	d.lock.RLock()
	i, ok := d.instances[connName]
	d.lock.RUnlock()
	if !ok {
		d.lock.Lock()
		// Recheck to ensure instance wasn't created between locks
		i, ok = d.instances[connName]
		if !ok {
			// Create a new instance
			var err error
			i, err = cloudsql.NewInstance(connName, d.sqladmin, d.key, d.refreshTimeout, d.iamTokenSource)
			if err != nil {
				d.lock.Unlock()
				return nil, err
			}
			d.instances[connName] = i
		}
		d.lock.Unlock()
	}
	return i, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 y

import (
	"fmt"
	"sync"
	"sync/atomic"

	"os"
	"syscall"
)

var EmptySlice = []byte{}

func OpenSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE
	if sync {
		flags |= syscall.O_DSYNC
	}
	return os.OpenFile(filename, flags, 0666)
}

func Safecopy(a []byte, src []byte) []byte {
	if cap(a) < len(src) {
		a = make([]byte, len(src))
	}
	a = a[:len(src)]
	copy(a, src)
	return a
}

type Slice struct {
	buf []byte
}

func (s *Slice) Resize(sz int) []byte {
	if cap(s.buf) < sz {
		s.buf = make([]byte, sz)
	}
	return s.buf[0:sz]
}

type LevelCloser struct {
	Name    string
	running int32
	nomore  int32
	closed  chan struct{}
	waiting sync.WaitGroup
}

type Closer struct {
	sync.RWMutex
	levels map[string]*LevelCloser
}

func NewCloser() *Closer {
	return &Closer{
		levels: make(map[string]*LevelCloser),
	}
}

func (c *Closer) Register(name string) *LevelCloser {
	c.Lock()
	defer c.Unlock()

	lc, has := c.levels[name]
	if !has {
		lc = &LevelCloser{Name: name, closed: make(chan struct{}, 10)}
		lc.waiting.Add(1)
		c.levels[name] = lc
	}

	AssertTruef(atomic.LoadInt32(&lc.nomore) == 0, "Can't register with closer after signal.")
	atomic.AddInt32(&lc.running, 1)
	return lc
}

func (c *Closer) Get(name string) *LevelCloser {
	c.RLock()
	defer c.RUnlock()

	lc, has := c.levels[name]
	if !has {
		Fatalf("%q not present in Closer", name)
		return nil
	}
	return lc
}

func (c *Closer) SignalAll() {
	c.RLock()
	c.RUnlock()

	for _, l := range c.levels {
		l.Signal()
	}
}

func (c *Closer) WaitForAll() {
	c.RLock()
	c.RUnlock()

	for _, l := range c.levels {
		l.Wait()
	}
}

func (lc *LevelCloser) NumRunning() int {
	return int(atomic.LoadInt32(&lc.running))
}

func (lc *LevelCloser) Signal() {
	if !atomic.CompareAndSwapInt32(&lc.nomore, 0, 1) {
		fmt.Printf("Level %q already got signal\n", lc.Name)
		return
	}
	running := int(atomic.LoadInt32(&lc.running))
	fmt.Printf("Sending signal to %d registered with name %q\n", running, lc.Name)
	for i := 0; i < running; i++ {
		lc.closed <- struct{}{}
	}
}

func (lc *LevelCloser) HasBeenClosed() <-chan struct{} {
	return lc.closed
}

func (lc *LevelCloser) GotSignal() bool {
	return atomic.LoadInt32(&lc.nomore) == 1
}

func (lc *LevelCloser) Done() {
	if atomic.LoadInt32(&lc.running) <= 0 {
		return
	}

	running := atomic.AddInt32(&lc.running, -1)
	if running == 0 {
		lc.waiting.Done()
	}
}

func (lc *LevelCloser) Wait() {
	lc.waiting.Wait()
}

func (lc *LevelCloser) SignalAndWait() {
	lc.Signal()
	lc.Wait()
}
<commit_msg>Fix mutex locking issue<commit_after>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 y

import (
	"fmt"
	"sync"
	"sync/atomic"

	"os"
	"syscall"
)

var EmptySlice = []byte{}

func OpenSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE
	if sync {
		flags |= syscall.O_DSYNC
	}
	return os.OpenFile(filename, flags, 0666)
}

func Safecopy(a []byte, src []byte) []byte {
	if cap(a) < len(src) {
		a = make([]byte, len(src))
	}
	a = a[:len(src)]
	copy(a, src)
	return a
}

type Slice struct {
	buf []byte
}

func (s *Slice) Resize(sz int) []byte {
	if cap(s.buf) < sz {
		s.buf = make([]byte, sz)
	}
	return s.buf[0:sz]
}

type LevelCloser struct {
	Name    string
	running int32
	nomore  int32
	closed  chan struct{}
	waiting sync.WaitGroup
}

type Closer struct {
	sync.RWMutex
	levels map[string]*LevelCloser
}

func NewCloser() *Closer {
	return &Closer{
		levels: make(map[string]*LevelCloser),
	}
}

func (c *Closer) Register(name string) *LevelCloser {
	c.Lock()
	defer c.Unlock()

	lc, has := c.levels[name]
	if !has {
		lc = &LevelCloser{Name: name, closed: make(chan struct{}, 10)}
		lc.waiting.Add(1)
		c.levels[name] = lc
	}

	AssertTruef(atomic.LoadInt32(&lc.nomore) == 0, "Can't register with closer after signal.")
	atomic.AddInt32(&lc.running, 1)
	return lc
}

func (c *Closer) Get(name string) *LevelCloser {
	c.RLock()
	defer c.RUnlock()

	lc, has := c.levels[name]
	if !has {
		Fatalf("%q not present in Closer", name)
		return nil
	}
	return lc
}

func (c *Closer) SignalAll() {
	c.RLock()
	defer c.RUnlock()

	for _, l := range c.levels {
		l.Signal()
	}
}

func (c *Closer) WaitForAll() {
	c.RLock()
	defer c.RUnlock()

	for _, l := range c.levels {
		l.Wait()
	}
}

func (lc *LevelCloser) NumRunning() int {
	return int(atomic.LoadInt32(&lc.running))
}

func (lc *LevelCloser) Signal() {
	if !atomic.CompareAndSwapInt32(&lc.nomore, 0, 1) {
		fmt.Printf("Level %q already got signal\n", lc.Name)
		return
	}
	running := int(atomic.LoadInt32(&lc.running))
	fmt.Printf("Sending signal to %d registered with name %q\n", running, lc.Name)
	for i := 0; i < running; i++ {
		lc.closed <- struct{}{}
	}
}

func (lc *LevelCloser) HasBeenClosed() <-chan struct{} {
	return lc.closed
}

func (lc *LevelCloser) GotSignal() bool {
	return atomic.LoadInt32(&lc.nomore) == 1
}

func (lc *LevelCloser) Done() {
	if atomic.LoadInt32(&lc.running) <= 0 {
		return
	}

	running := atomic.AddInt32(&lc.running, -1)
	if running == 0 {
		lc.waiting.Done()
	}
}

func (lc *LevelCloser) Wait() {
	lc.waiting.Wait()
}

func (lc *LevelCloser) SignalAndWait() {
	lc.Signal()
	lc.Wait()
}
<|endoftext|>


<commit_before>package gcloudzip

import (
	"archive/zip"
	"fmt"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/appengine/log"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"
)

type ZipRequest struct {
	context context.Context
	bucket  *storage.BucketHandle
}

func NewWithBucketNamed(c context.Context, bucketName string) (*ZipRequest, error) {
	client, err := storage.NewClient(c)
	if err != nil {
		return nil, err
	}
	return NewWithBucket(c, client.Bucket(bucketName)), nil
}

func NewWithCredentials(c context.Context, bucketName string, keyFile string) (*ZipRequest, error) {
	// Load the credentials
	jsonKey, err := ioutil.ReadFile(keyFile)
	if err != nil {
		return nil, err
	}
	conf, err := google.JWTConfigFromJSON(
		jsonKey,
		storage.ScopeFullControl,
	)
	if err != nil {
		return nil, err
	}

	// Create the cloud storage client to use
	client, err := storage.NewClient(c, cloud.WithTokenSource(conf.TokenSource(c)))
	if err != nil {
		return nil, err
	}

	return NewWithBucket(c, client.Bucket(bucketName)), nil
}

func NewWithBucket(c context.Context, bucket *storage.BucketHandle) *ZipRequest {
	return &ZipRequest{
		context: c,
		bucket:  bucket,
	}
}

// Pack a folder into zip file
func (r *ZipRequest) Zip(srcFolder string, fileName string, contentType string, metaData *map[string]string) bool {
	c := r.context
	bucket := r.bucket

	log.Infof(c, "Packing bucket %v folder %v to file %v", bucket, srcFolder, fileName)

	srcFolder = fmt.Sprintf("%v/", srcFolder)
	query := &storage.Query{Prefix: srcFolder, Delimiter: "/"}

	// TODO read all pages of the response, see example
	// https://cloud.google.com/appengine/docs/go/googlecloudstorageclient/read-write-to-cloud-storage

	objs, err := bucket.List(c, query)
	if err != nil {
		log.Errorf(c, "Packing failed to list bucket %q: %v", r.bucket, err)
		return false
	}

	totalFiles := len(objs.Results)
	if totalFiles == 0 {
		log.Errorf(c, "Packing failed to find objects found in folder %q: %v", bucket, srcFolder)
		return false
	}

	// create storage file for writing
	log.Infof(c, "Writing new zip file to %v/%v for %v files", bucket, fileName, totalFiles)
	storageWriter := bucket.Object(fileName).NewWriter(c)
	defer storageWriter.Close()

	// add optional content type and meta data
	if len(contentType) > 0 {
		storageWriter.ContentType = contentType
	}
	if metaData != nil {
		storageWriter.Metadata = *metaData
	}

	// Create a new zip archive to memory buffer
	zipWriter := zip.NewWriter(storageWriter)

	// go through each file in the folder
	for _, obj := range objs.Results {

		log.Infof(c, "Packing file %v of size %v to zip file", obj.Name, obj.Size)
		//d.dumpStats(obj)

		// read file in our source folder from storage - io.ReadCloser returned from storage
		storageReader, err := bucket.Object(obj.Name).NewReader(c)
		if err != nil {
			log.Errorf(c, "Packing failed to read from bucket %q file %q: %v", bucket, obj.Name, err)
			return false
		}
		defer storageReader.Close()

		// grab just the filename from directory listing (don't want to store paths in zip)
		_, zipFileName := filepath.Split(obj.Name)
		newFileName := strings.ToLower(zipFileName)

		// add filename to zip
		zipFile, err := zipWriter.Create(newFileName)
		if err != nil {
			log.Errorf(c, "Packing failed to create zip file from bucket %q file %q: %v", bucket, zipFileName, err)
			return false
		}

		// copy from storage reader to zip writer
		_, err = io.Copy(zipFile, storageReader)
		if err != nil {
			log.Errorf(c, "Failed to copy from storage reader to zip file: %v", err)
			return false
		}
	}

	// Make sure to check the error on Close.
	log.Infof(c, "Closing zip writer")
	err = zipWriter.Close()
	if err != nil {
		log.Errorf(c, "Packing failed to close zip file writer from bucket %q file %q : %v", bucket, fileName, err)
		return false
	}

	// success!
	log.Infof(c, "Packed files to new cloud storage file %v successful!", fileName)
	return true
}
<commit_msg>Remove log and report errors<commit_after>package gcloudz

import (
	"archive/zip"
	"errors"
	"fmt"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"
)

var (
	ErrNoFilesFound = errors.New("No files found in folder")
)

type ZipRequest struct {
	context context.Context
	bucket  *storage.BucketHandle
}

func NewWithBucketNamed(c context.Context, bucketName string) (*ZipRequest, error) {
	client, err := storage.NewClient(c)
	if err != nil {
		return nil, err
	}
	return NewWithBucket(c, client.Bucket(bucketName)), nil
}

func NewWithCredentials(c context.Context, bucketName string, keyFile string) (*ZipRequest, error) {
	// Load the credentials
	jsonKey, err := ioutil.ReadFile(keyFile)
	if err != nil {
		return nil, err
	}
	conf, err := google.JWTConfigFromJSON(
		jsonKey,
		storage.ScopeFullControl,
	)
	if err != nil {
		return nil, err
	}

	// Create the cloud storage client to use
	client, err := storage.NewClient(c, cloud.WithTokenSource(conf.TokenSource(c)))
	if err != nil {
		return nil, err
	}

	return NewWithBucket(c, client.Bucket(bucketName)), nil
}

func NewWithBucket(c context.Context, bucket *storage.BucketHandle) *ZipRequest {
	return &ZipRequest{
		context: c,
		bucket:  bucket,
	}
}

// Pack a folder into zip file
func (r *ZipRequest) Zip(srcFolder string, fileName string, contentType string, metaData *map[string]string) error {
	c := r.context
	bucket := r.bucket

	srcFolder = fmt.Sprintf("%v/", srcFolder)
	query := &storage.Query{Prefix: srcFolder, Delimiter: "/"}

	// TODO read all pages of the response, see example
	// https://cloud.google.com/appengine/docs/go/googlecloudstorageclient/read-write-to-cloud-storage

	objs, err := bucket.List(c, query)
	if err != nil {
		return err
	}

	totalFiles := len(objs.Results)
	if totalFiles == 0 {
		return ErrNoFilesFound
	}

	// create storage file for writing
	storageWriter := bucket.Object(fileName).NewWriter(c)
	defer storageWriter.Close()

	// add optional content type and meta data
	if len(contentType) > 0 {
		storageWriter.ContentType = contentType
	}
	if metaData != nil {
		storageWriter.Metadata = *metaData
	}

	// Create a new zip archive to memory buffer
	zipWriter := zip.NewWriter(storageWriter)

	// go through each file in the folder
	for _, obj := range objs.Results {

		// read file in our source folder from storage - io.ReadCloser returned from storage
		storageReader, err := bucket.Object(obj.Name).NewReader(c)
		if err != nil {
			return err
		}
		defer storageReader.Close()

		// grab just the filename from directory listing (don't want to store paths in zip)
		_, zipFileName := filepath.Split(obj.Name)
		newFileName := strings.ToLower(zipFileName)

		// add filename to zip
		zipFile, err := zipWriter.Create(newFileName)
		if err != nil {
			return err
		}

		// copy from storage reader to zip writer
		_, err = io.Copy(zipFile, storageReader)
		if err != nil {
			return err
		}
	}

	// Make sure to check the error on Close.
	err = zipWriter.Close()
	if err != nil {
		return err
	}

	// success!
	return nil
}
<|endoftext|>


<commit_before>//+build ignore

// Makes the docker image xena/xperimental.
package main

import (
	"context"
	"log"
	"path/filepath"

	"within.website/x/internal/yeet"
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tag := "xena/xperimental"
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "build", "-t", tag, ".")

	resTag, err := yeet.DockerTag(ctx, "xena", "xperimental", tag)
	if err != nil {
		log.Fatal(err)
	}
	otherResTag, err := yeet.DockerTag(ctx, "docker.pkg.github.com/Xe/x", "xperimental", tag)
	if err != nil {
		log.Fatal(err)
	}

	gitTag, err := yeet.GitTag(ctx)
	if err != nil {
		log.Fatal(err)
	}

	dnsdTag := "xena/dnsd:" + gitTag

	yeet.ShouldWork(ctx, nil, filepath.Join(yeet.WD, "cmd", "dnsd"), "docker", "build", "-t", dnsdTag, "--build-arg", "X_VERSION="+gitTag, ".")

	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", resTag)
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", dnsdTag)
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", otherResTag)

	log.Printf("use %s", resTag)
}
<commit_msg>docker: also upload dnsd<commit_after>//+build ignore

// Makes the docker image xena/xperimental.
package main

import (
	"context"
	"log"
	"path/filepath"

	"within.website/x/internal/yeet"
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tag := "xena/xperimental"
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "build", "-t", tag, ".")

	resTag, err := yeet.DockerTag(ctx, "xena", "xperimental", tag)
	if err != nil {
		log.Fatal(err)
	}
	otherResTag, err := yeet.DockerTag(ctx, "docker.pkg.github.com/xe/x", "xperimental", tag)
	if err != nil {
		log.Fatal(err)
	}

	gitTag, err := yeet.GitTag(ctx)
	if err != nil {
		log.Fatal(err)
	}

	dnsdTag := "xena/dnsd:" + gitTag

	yeet.ShouldWork(ctx, nil, filepath.Join(yeet.WD, "cmd", "dnsd"), "docker", "build", "-t", dnsdTag, "--build-arg", "X_VERSION="+gitTag, ".")
	dnsdGithubTag, err := yeet.DockerTag(ctx, "docker.pkg.github.com/xe/x", "dnsd", dnsdTag)
	if err != nil {
		log.Fatal(err)
	}

	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", resTag)
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", dnsdTag)
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", otherResTag)
	yeet.ShouldWork(ctx, nil, yeet.WD, "docker", "push", dnsdGithubTag)

	log.Printf("use %s", resTag)
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/fsouza/go-dockerclient"
)

type StackImage struct {
	stack string
	image docker.APIImages
}

type Docker struct {
	c *docker.Client
}

func (d *Docker) Connect() (err error) {
	endpoint := os.Getenv("DOCKER_HOST")
	if endpoint == "" {
		endpoint = "unix:///var/run/docker.sock"
	}
	log.Printf("DOCKER_HOST = %v\n", endpoint)

	certPath := os.Getenv("DOCKER_CERT_PATH")
	log.Printf("DOCKER_CERT_PATH = %v\n", certPath)
	if certPath == "" {
		d.c, err = docker.NewClient(endpoint)
	} else {
		cert := certPath + "/cert.pem"
		key := certPath + "/key.pem"
		ca := certPath + "/ca.pem"
		d.c, err = docker.NewTLSClient(endpoint, cert, key, ca)
	}
	return err
}

func (d *Docker) StackStat(stack string) (*StackImage, error) {
	si := StackImage{
		stack: stack,
	}

	images, err := d.c.ListImages(docker.ListImagesOptions{All: true})
	if err != nil {
		return nil, err
	}

	for _, image := range images {
		for _, tag := range image.RepoTags {
			if tag == stack {
				si.image = image
				return &si, nil
			}
		}
	}

	return nil, nil
}

func (d *Docker) BuildSlugImage(si *StackImage, release *Release) (string, error) {
	slugURL := release.slugURL
	t := time.Now()
	inputBuf, outputBuf := bytes.NewBuffer(nil), bytes.NewBuffer(nil)
	tr := tar.NewWriter(inputBuf)
	dockerContents := fmt.Sprintf(`FROM %s
RUN rm -rf /app
RUN curl '%s' -o /slug.img
RUN (cd / && mkdir /app && tar -xzf /slug.img) && rm -f /app/log /app/tmp && mkdir /app/log /app/tmp &&  chown -R daemon:daemon /app && chmod -R go+r /app && find /app -type d | xargs chmod go+x
WORKDIR /app
`, si.image.ID, slugURL)

	log.Println(dockerContents)
	tr.WriteHeader(&tar.Header{
		Name:    "Dockerfile",
		Size:    int64(len(dockerContents)),
		ModTime: t, AccessTime: t,
		ChangeTime: t})
	tr.Write([]byte(dockerContents))
	tr.Close()

	imageName := release.Name()

	opts := docker.BuildImageOptions{
		Name:           imageName,
		InputStream:    inputBuf,
		OutputStream:   outputBuf,
		SuppressOutput: false,
	}

	if err := d.c.BuildImage(opts); err != nil {
		return "", err
	}

	return imageName, nil
}
<commit_msg>Use a recursive remove<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/fsouza/go-dockerclient"
)

type StackImage struct {
	stack string
	image docker.APIImages
}

type Docker struct {
	c *docker.Client
}

func (d *Docker) Connect() (err error) {
	endpoint := os.Getenv("DOCKER_HOST")
	if endpoint == "" {
		endpoint = "unix:///var/run/docker.sock"
	}
	log.Printf("DOCKER_HOST = %v\n", endpoint)

	certPath := os.Getenv("DOCKER_CERT_PATH")
	log.Printf("DOCKER_CERT_PATH = %v\n", certPath)
	if certPath == "" {
		d.c, err = docker.NewClient(endpoint)
	} else {
		cert := certPath + "/cert.pem"
		key := certPath + "/key.pem"
		ca := certPath + "/ca.pem"
		d.c, err = docker.NewTLSClient(endpoint, cert, key, ca)
	}
	return err
}

func (d *Docker) StackStat(stack string) (*StackImage, error) {
	si := StackImage{
		stack: stack,
	}

	images, err := d.c.ListImages(docker.ListImagesOptions{All: true})
	if err != nil {
		return nil, err
	}

	for _, image := range images {
		for _, tag := range image.RepoTags {
			if tag == stack {
				si.image = image
				return &si, nil
			}
		}
	}

	return nil, nil
}

func (d *Docker) BuildSlugImage(si *StackImage, release *Release) (string, error) {
	slugURL := release.slugURL
	t := time.Now()
	inputBuf, outputBuf := bytes.NewBuffer(nil), bytes.NewBuffer(nil)
	tr := tar.NewWriter(inputBuf)
	dockerContents := fmt.Sprintf(`FROM %s
RUN rm -rf /app
RUN curl '%s' -o /slug.img
RUN (cd / && mkdir /app && tar -xzf /slug.img) && rm -rf /app/log /app/tmp && mkdir /app/log /app/tmp &&  chown -R daemon:daemon /app && chmod -R go+r /app && find /app -type d | xargs chmod go+x
WORKDIR /app
`, si.image.ID, slugURL)

	log.Println(dockerContents)
	tr.WriteHeader(&tar.Header{
		Name:    "Dockerfile",
		Size:    int64(len(dockerContents)),
		ModTime: t, AccessTime: t,
		ChangeTime: t})
	tr.Write([]byte(dockerContents))
	tr.Close()

	imageName := release.Name()

	opts := docker.BuildImageOptions{
		Name:           imageName,
		InputStream:    inputBuf,
		OutputStream:   outputBuf,
		SuppressOutput: false,
	}

	if err := d.c.BuildImage(opts); err != nil {
		return "", err
	}

	return imageName, nil
}
<|endoftext|>


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

// Incomplete source tree on Android.

//go:build !android
// +build !android

package pointer

// This file runs the pointer analysis on all packages and tests beneath
// $GOROOT.  It provides a "smoke test" that the analysis doesn't crash
// on a large input, and a benchmark for performance measurement.
//
// Because it is relatively slow, the --stdlib flag must be enabled for
// this test to run:
//    % go test -v golang.org/x/tools/go/pointer --stdlib

import (
	"flag"
	"go/build"
	"go/token"
	"testing"
	"time"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
)

var runStdlibTest = flag.Bool("stdlib", false, "Run the (slow) stdlib test")

func TestStdlib(t *testing.T) {
	if !*runStdlibTest {
		t.Skip("skipping (slow) stdlib test (use --stdlib)")
	}

	// Load, parse and type-check the program.
	ctxt := build.Default // copy
	ctxt.GOPATH = ""      // disable GOPATH
	conf := loader.Config{Build: &ctxt}
	if _, err := conf.FromArgs(buildutil.AllPackages(conf.Build), true); err != nil {
		t.Errorf("FromArgs failed: %v", err)
		return
	}

	iprog, err := conf.Load()
	if err != nil {
		t.Fatalf("Load failed: %v", err)
	}

	// Create SSA packages.
	prog := ssautil.CreateProgram(iprog, 0)
	prog.Build()

	numPkgs := len(prog.AllPackages())
	if want := 240; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Determine the set of packages/tests to analyze.
	var mains []*ssa.Package
	for _, info := range iprog.InitialPackages() {
		ssapkg := prog.Package(info.Pkg)
		if main := prog.CreateTestMainPackage(ssapkg); main != nil {
			mains = append(mains, main)
		}
	}
	if mains == nil {
		t.Fatal("no tests found in analysis scope")
	}

	// Run the analysis.
	config := &Config{
		Reflection:     false, // TODO(adonovan): fix remaining bug in rVCallConstraint, then enable.
		BuildCallGraph: true,
		Mains:          mains,
	}
	// TODO(adonovan): add some query values (affects track bits).

	t0 := time.Now()

	result, err := Analyze(config)
	if err != nil {
		t.Fatal(err) // internal error in pointer analysis
	}
	_ = result // TODO(adonovan): measure something

	t1 := time.Now()

	// Dump some statistics.
	allFuncs := ssautil.AllFunctions(prog)
	var numInstrs int
	for fn := range allFuncs {
		for _, b := range fn.Blocks {
			numInstrs += len(b.Instrs)
		}
	}

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("#Source lines:          ", lineCount)
	t.Log("#Instructions:          ", numInstrs)
	t.Log("Pointer analysis:       ", t1.Sub(t0))
}
<commit_msg>go/pointer: remove use of deprecated ssa.CreateTestMainPackage<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Incomplete source tree on Android.

//go:build !android
// +build !android

package pointer

// This file runs the pointer analysis on all packages and tests beneath
// $GOROOT.  It provides a "smoke test" that the analysis doesn't crash
// on a large input, and a benchmark for performance measurement.
//
// Because it is relatively slow, the --stdlib flag must be enabled for
// this test to run:
//    % go test -v golang.org/x/tools/go/pointer --stdlib

import (
	"flag"
	"go/token"
	"testing"
	"time"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
)

var runStdlibTest = flag.Bool("stdlib", false, "Run the (slow) stdlib test")

func TestStdlib(t *testing.T) {
	if !*runStdlibTest {
		t.Skip("skipping (slow) stdlib test (use --stdlib)")
	}

	cfg := &packages.Config{
		Mode: packages.LoadAllSyntax,
		// Create test main packages with a main function.
		Tests: true,
	}
	pkgs, err := packages.Load(cfg, "std")
	if err != nil || packages.PrintErrors(pkgs) > 0 {
		t.Fatalf("Load failed: %v", err)
	}

	// Create SSA packages.
	prog, _ := ssautil.AllPackages(pkgs, 0)
	prog.Build()

	numPkgs := len(prog.AllPackages())
	if want := 240; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Determine the set of packages/tests to analyze.
	var mains []*ssa.Package
	for _, ssapkg := range prog.AllPackages() {
		if ssapkg.Pkg.Name() == "main" && ssapkg.Func("main") != nil {
			mains = append(mains, ssapkg)
		}
	}
	if mains == nil {
		t.Fatal("no tests found in analysis scope")
	}

	// Run the analysis.
	config := &Config{
		Reflection:     false, // TODO(adonovan): fix remaining bug in rVCallConstraint, then enable.
		BuildCallGraph: true,
		Mains:          mains,
	}
	// TODO(adonovan): add some query values (affects track bits).

	t0 := time.Now()

	result, err := Analyze(config)
	if err != nil {
		t.Fatal(err) // internal error in pointer analysis
	}
	_ = result // TODO(adonovan): measure something

	t1 := time.Now()

	// Dump some statistics.
	allFuncs := ssautil.AllFunctions(prog)
	var numInstrs int
	for fn := range allFuncs {
		for _, b := range fn.Blocks {
			numInstrs += len(b.Instrs)
		}
	}

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("#Source lines:          ", lineCount)
	t.Log("#Instructions:          ", numInstrs)
	t.Log("Pointer analysis:       ", t1.Sub(t0))
}
<|endoftext|>


<commit_before>package raindrops

import (
	"bytes"
	"fmt"
	"os"
	"strconv"
)

var (
	primes = [...]int{3, 5, 7}
	words  = [...]string{"Pling", "Plang", "Plong"}
)

// Convert returns a string, the contents of which depends on the number's prime factors.
//
// - If the number contains 3 as a prime factor, output 'Pling'.
// - If the number contains 5 as a prime factor, output 'Plang'.
// - If the number contains 7 as a prime factor, output 'Plong'.
// - If the number does not contain 3, 5, or 7 as a prime factor,
//   just pass the number's digits straight through.
func Convert(n int) string {
	var buf bytes.Buffer
	for i, prime := range primes {
		if n%prime == 0 {
			if _, e := buf.WriteString(words[i]); e != nil {
				fmt.Fprintf(os.Stderr, e.Error())
			}
		}
	}

	if buf.Len() == 0 {
		return strconv.Itoa(n)
	}
	return buf.String()
}
<commit_msg>raindrops: use []byte to store an output for speed<commit_after>package raindrops

import "strconv"

var (
	primes = [...]int{3, 5, 7}
	words  = [...][]byte{[]byte("Pling"), []byte("Plang"), []byte("Plong")}
	maxLen = 15
)

// Convert returns a string, the contents of which depends on the number's prime factors.
//
// - If the number contains 3 as a prime factor, output 'Pling'.
// - If the number contains 5 as a prime factor, output 'Plang'.
// - If the number contains 7 as a prime factor, output 'Plong'.
// - If the number does not contain 3, 5, or 7 as a prime factor,
//   just pass the number's digits straight through.
func Convert(n int) string {
	buf := make([]byte, 0, maxLen)
	for i, prime := range primes {
		if n%prime == 0 {
			buf = append(buf, words[i]...)
		}
	}

	if len(buf) == 0 {
		return strconv.Itoa(n)
	}
	return string(buf)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com
package main

// Convert the output of the MIG pkg module OVAL interpreter into JSON events
// for consumption by MozDef

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/ameihm0912/govfeed/src/govfeed"
	"github.com/ameihm0912/gozdef"
	"os"
	"regexp"
)

const sourceName = "mozoval"

var useVFeed string
var vulnEvents []gozdef.VulnEvent

func lineParser(buf string) error {
	var (
		CVE   string
		ID    string
		Title string
	)

	var pTable = []struct {
		expression string
		target     *string
	}{
		{".*(CVE-\\d+-\\d+).*", &CVE},
		{".*id=(\\S+).*", &ID},
		{".*title=\"([^\"]+).*", &Title},
	}

	for i := range pTable {
		r := regexp.MustCompile(pTable[i].expression)
		matches := r.FindStringSubmatch(buf)
		if len(matches) <= 1 {
			continue
		}
		*pTable[i].target = matches[1]
	}

	e, err := gozdef.NewVulnEvent()
	if err != nil {
		return err
	}
	e.SourceName = sourceName

	// XXX These need to be set correctly
	e.Asset.AssetID = 1
	e.Vuln.VulnID = "mozoval-vuln"

	if CVE != "" {
		e.Vuln.CVE = append(e.Vuln.CVE, CVE)
	}

	if err = e.Validate(); err != nil {
		return err
	}

	vulnEvents = append(vulnEvents, e)
	return nil
}

func main() {
	flag.StringVar(&useVFeed, "v", "", "path to vFeed CLI")
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		fmt.Fprint(os.Stderr, "error: specify results output file as argument\n")
		os.Exit(1)
	}

	if useVFeed != "" {
		err := govfeed.GVInit(useVFeed)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
	}

	fd, err := os.Open(args[0])
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		buf := scanner.Text()
		err = lineParser(buf)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
	}
	if err = scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}

	for _, x := range vulnEvents {
		jb, err := json.Marshal(&x)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stdout, "%v\n", string(jb))
	}
}
<commit_msg>incorporate vfeed data for cvss/cve score in mozdef vuln events<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com
package main

// Convert the output of the MIG pkg module OVAL interpreter into JSON events
// for consumption by MozDef

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/ameihm0912/govfeed/src/govfeed"
	"github.com/ameihm0912/gozdef"
	"os"
	"regexp"
)

const sourceName = "mozoval"

var cveCache map[string]govfeed.GVCVE

var useVFeed string
var vulnEvents []gozdef.VulnEvent

func getCVE(cve string) (govfeed.GVCVE, error) {
	x, ok := cveCache[cve]
	if ok {
		return x, nil
	}
	x, err := govfeed.GVQuery(cve)
	if err != nil {
		return x, err
	}
	cveCache[cve] = x
	return x, nil
}

func lineParser(buf string) error {
	var (
		CVE   string
		ID    string
		Title string
	)

	var pTable = []struct {
		expression string
		target     *string
	}{
		{".*(CVE-\\d+-\\d+).*", &CVE},
		{".*id=(\\S+).*", &ID},
		{".*title=\"([^\"]+).*", &Title},
	}

	for i := range pTable {
		r := regexp.MustCompile(pTable[i].expression)
		matches := r.FindStringSubmatch(buf)
		if len(matches) <= 1 {
			continue
		}
		*pTable[i].target = matches[1]
	}

	e, err := gozdef.NewVulnEvent()
	if err != nil {
		return err
	}
	e.SourceName = sourceName

	// XXX These need to be set correctly
	e.Asset.AssetID = 1
	e.Vuln.VulnID = "mozoval-vuln"

	if CVE != "" {
		e.Vuln.Title = CVE
		e.Vuln.CVE = append(e.Vuln.CVE, CVE)
	}

	if useVFeed != "" && CVE != "" {
		cvedata, err := getCVE(CVE)
		if err != nil {
			return err
		}
		if cvedata.Description != "" {
			e.Vuln.Description = cvedata.Description
		}
		if cvedata.CVSS != 0 {
			e.Vuln.CVSS = cvedata.CVSS
		}
	}

	if err = e.Validate(); err != nil {
		return err
	}

	vulnEvents = append(vulnEvents, e)
	return nil
}

func main() {
	flag.StringVar(&useVFeed, "v", "", "path to vFeed CLI")
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		fmt.Fprint(os.Stderr, "error: specify results output file as argument\n")
		os.Exit(1)
	}

	if useVFeed != "" {
		err := govfeed.GVInit(useVFeed)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		cveCache = make(map[string]govfeed.GVCVE)
	}

	fd, err := os.Open(args[0])
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		buf := scanner.Text()
		err = lineParser(buf)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
	}
	if err = scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}

	for _, x := range vulnEvents {
		jb, err := json.Marshal(&x)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stdout, "%v\n", string(jb))
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014,2015,2016 Docker, Inc.
// Copyright (c) 2017 Intel Corporation
//
// 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 (
	"fmt"
	"os"

	vc "github.com/containers/virtcontainers"
	"github.com/containers/virtcontainers/pkg/oci"
	"github.com/urfave/cli"
)

var createCommand = cli.Command{
	Name:  "create",
	Usage: "Create a container",
	ArgsUsage: `<container-id>

   <container-id> is your name for the instance of the container that you
   are starting. The name you provide for the container instance must be unique
   on your host.`,
	Description: `The create command creates an instance of a container for a bundle. The
   bundle is a directory with a specification file named "` + specConfig + `" and a
   root filesystem.
   The specification file includes an args parameter. The args parameter is
   used to specify command(s) that get run when the container is started.
   To change the command(s) that get executed on start, edit the args
   parameter of the spec.`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "bundle, b",
			Value: "",
			Usage: `path to the root of the bundle directory, defaults to the current directory`,
		},
		cli.StringFlag{
			Name:  "console",
			Value: "",
			Usage: "path to a pseudo terminal",
		},
		cli.StringFlag{
			Name:  "pid-file",
			Value: "",
			Usage: "specify the file to write the process id to",
		},
	},
	Action: func(context *cli.Context) error {
		return create(context.Args().First(),
			context.String("bundle"),
			context.String("console"),
			context.String("pid-file"))
	},
}

func create(containerID, bundlePath, console, pidFilePath string) error {
	// Checks the MUST and MUST NOT from OCI runtime specification
	if err := validCreateParams(containerID, bundlePath); err != nil {
		return err
	}

	podConfig, shimConfig, err := getConfigs(bundlePath, containerID, console)
	if err != nil {
		return err
	}

	pod, err := vc.CreatePod(podConfig)
	if err != nil {
		return err
	}

	// Start the shim to retrieve its PID.
	containers := pod.GetContainers()
	if len(containers) != 1 {
		return fmt.Errorf("BUG: Container list from pod is wrong, expecting only one container, found %d containers", len(containers))
	}

	pid, err := startContainerShim(containers[0], shimConfig, pod.URL())
	if err != nil {
		return err
	}

	// Creation of PID file has to be the last thing done in the create
	// because containerd considers the create complete after this file
	// is created.
	if err := createPIDFile(pidFilePath, pid); err != nil {
		return err
	}

	return nil
}

func getConfigs(bundlePath, containerID, console string) (vc.PodConfig, ShimConfig, error) {
	runtimeConfig, shimConfig, err := loadConfiguration("")
	if err != nil {
		return vc.PodConfig{}, ShimConfig{}, err
	}

	podConfig, err := oci.PodConfig(runtimeConfig, bundlePath, containerID, console)
	if err != nil {
		return vc.PodConfig{}, ShimConfig{}, err
	}

	return *podConfig, shimConfig, nil
}

func createPIDFile(pidFilePath string, pid int) error {
	if pidFilePath == "" {
		return fmt.Errorf("Missing PID file path")
	}

	if err := os.RemoveAll(pidFilePath); err != nil {
		return err
	}

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

	pidStr := fmt.Sprintf("%d", pid)

	n, err := f.WriteString(pidStr)
	if err != nil {
		return err
	}

	if n < len(pidStr) {
		return fmt.Errorf("Could not write pid to '%s': only %d bytes written out of %d", pidFilePath, n, len(pidStr))
	}

	return nil
}
<commit_msg>create: Fix broken build after new virtcontainers vendoring<commit_after>// Copyright (c) 2014,2015,2016 Docker, Inc.
// Copyright (c) 2017 Intel Corporation
//
// 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 (
	"fmt"
	"os"

	vc "github.com/containers/virtcontainers"
	"github.com/containers/virtcontainers/pkg/oci"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/urfave/cli"
)

var createCommand = cli.Command{
	Name:  "create",
	Usage: "Create a container",
	ArgsUsage: `<container-id>

   <container-id> is your name for the instance of the container that you
   are starting. The name you provide for the container instance must be unique
   on your host.`,
	Description: `The create command creates an instance of a container for a bundle. The
   bundle is a directory with a specification file named "` + specConfig + `" and a
   root filesystem.
   The specification file includes an args parameter. The args parameter is
   used to specify command(s) that get run when the container is started.
   To change the command(s) that get executed on start, edit the args
   parameter of the spec.`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "bundle, b",
			Value: "",
			Usage: `path to the root of the bundle directory, defaults to the current directory`,
		},
		cli.StringFlag{
			Name:  "console",
			Value: "",
			Usage: "path to a pseudo terminal",
		},
		cli.StringFlag{
			Name:  "pid-file",
			Value: "",
			Usage: "specify the file to write the process id to",
		},
	},
	Action: func(context *cli.Context) error {
		return create(context.Args().First(),
			context.String("bundle"),
			context.String("console"),
			context.String("pid-file"))
	},
}

func create(containerID, bundlePath, console, pidFilePath string) error {
	// Checks the MUST and MUST NOT from OCI runtime specification
	if err := validCreateParams(containerID, bundlePath); err != nil {
		return err
	}

	podConfig, shimConfig, _, err := getConfigs(bundlePath, containerID, console)
	if err != nil {
		return err
	}

	pod, err := vc.CreatePod(podConfig)
	if err != nil {
		return err
	}

	// Start the shim to retrieve its PID.
	containers := pod.GetContainers()
	if len(containers) != 1 {
		return fmt.Errorf("BUG: Container list from pod is wrong, expecting only one container, found %d containers", len(containers))
	}

	pid, err := startContainerShim(containers[0], shimConfig, pod.URL())
	if err != nil {
		return err
	}

	// Creation of PID file has to be the last thing done in the create
	// because containerd considers the create complete after this file
	// is created.
	if err := createPIDFile(pidFilePath, pid); err != nil {
		return err
	}

	return nil
}

func getConfigs(bundlePath, containerID, console string) (vc.PodConfig, ShimConfig, specs.Spec, error) {
	runtimeConfig, shimConfig, err := loadConfiguration("")
	if err != nil {
		return vc.PodConfig{}, ShimConfig{}, specs.Spec{}, err
	}

	podConfig, ociSpec, err := oci.PodConfig(runtimeConfig, bundlePath, containerID, console)
	if err != nil {
		return vc.PodConfig{}, ShimConfig{}, specs.Spec{}, err
	}

	return *podConfig, shimConfig, *ociSpec, nil
}

func createPIDFile(pidFilePath string, pid int) error {
	if pidFilePath == "" {
		return fmt.Errorf("Missing PID file path")
	}

	if err := os.RemoveAll(pidFilePath); err != nil {
		return err
	}

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

	pidStr := fmt.Sprintf("%d", pid)

	n, err := f.WriteString(pidStr)
	if err != nil {
		return err
	}

	if n < len(pidStr) {
		return fmt.Errorf("Could not write pid to '%s': only %d bytes written out of %d", pidFilePath, n, len(pidStr))
	}

	return nil
}
<|endoftext|>


<commit_before>package isolated

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("apply-manifest command", func() {
	var (
		orgName      string
		spaceName    string
		appName      string
		manifestPath string
		appDir       string
	)

	BeforeEach(func() {
		orgName = helpers.NewOrgName()
		spaceName = helpers.NewSpaceName()
		appName = helpers.PrefixedRandomName("app")

		var err error
		appDir, err = ioutil.TempDir("", "simple-app")
		Expect(err).NotTo(HaveOccurred())

		manifestPath = filepath.Join(appDir, "manifest.yml")
		// Ensure the file exists at the minimum
		helpers.WriteManifest(manifestPath, map[string]interface{}{})
	})

	AfterEach(func() {
		Expect(os.RemoveAll(appDir)).ToNot(HaveOccurred())
	})

	Describe("help", func() {
		It("appears in cf help -a", func() {
			session := helpers.CF("help", "-a")
			Eventually(session).Should(Exit(0))
			Expect(session).To(HaveCommandInCategoryWithDescription("apply-manifest", "APPS", "Apply manifest properties to a space"))
		})

		When("--help flag is set", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("apply-manifest", "--help")

				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("apply-manifest - Apply manifest properties to a space"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say("cf apply-manifest -f APP_MANIFEST_PATH"))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("create-app, create-app-manifest, push"))

				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("the environment is not setup correctly", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(true, true, ReadOnlyOrg, "apply-manifest", "-f", manifestPath)
		})
	})

	When("the environment is set up correctly", func() {
		BeforeEach(func() {
			helpers.SetupCF(orgName, spaceName)
		})

		AfterEach(func() {
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app name in the manifest is missing", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"instances": 3,
						},
					},
				})
			})

			It("reports an error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Found an application with no name specified"))
				Eventually(session).Should(Say("FAILED"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("there is a CC error", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name":      appName,
							"instances": -1,
						},
					},
				})
			})

			It("displays the error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Instances must be greater than or equal to 0"))
				Eventually(session).Should(Say("FAILED"))

				Eventually(session).Should(Exit(1))
			})
		})

		When("-f is provided", func() {
			When("the -f flag is not given an arg", func() {
				It("tells the user that the flag requires an arg, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f")

					Eventually(session.Err).Should(Say("Incorrect Usage: expected argument for flag `-f'"))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a directory that does not have a manifest.yml file", func() {
				var (
					emptyDir string
					err      error
				)

				BeforeEach(func() {
					emptyDir, err = ioutil.TempDir("", "empty")
					Expect(err).NotTo(HaveOccurred())
				})

				AfterEach(func() {
					Expect(os.RemoveAll(emptyDir)).ToNot(HaveOccurred())
				})

				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", emptyDir)

					Eventually(session.Err).Should(Say("Incorrect Usage: The specified directory '%s' does not contain a file named 'manifest.yml'.", emptyDir))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a file that does not exist", func() {
				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", "path/that/does/not/exist")

					Eventually(session.Err).Should(Say("Incorrect Usage: The specified path 'path/that/does/not/exist' does not exist."))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the manifest exists where -f points", func() {
				It("applies the manifest successfully", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CF("apply-manifest", "-f", appDir)
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})
		})

		When("-f is not provided", func() {
			When("a properly formatted manifest is present in the pwd", func() {
				It("autodetects and applies the manifest", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: appDir}, "apply-manifest")
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

			When("the current directory does not have a manifest", func() {
				It("fails nicely", func() {
					currentDir, err := os.Getwd()
					Expect(err).NotTo(HaveOccurred())
					session := helpers.CF("apply-manifest")

					Eventually(session.Err).Should(helpers.SayPath(`Could not find 'manifest.yml' file in %s`, currentDir))
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Exit(1))
				})
			})
		})
	})
})
<commit_msg>Use SayPath matcher in another apply-manifest test<commit_after>package isolated

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("apply-manifest command", func() {
	var (
		orgName      string
		spaceName    string
		appName      string
		manifestPath string
		appDir       string
	)

	BeforeEach(func() {
		orgName = helpers.NewOrgName()
		spaceName = helpers.NewSpaceName()
		appName = helpers.PrefixedRandomName("app")

		var err error
		appDir, err = ioutil.TempDir("", "simple-app")
		Expect(err).NotTo(HaveOccurred())

		manifestPath = filepath.Join(appDir, "manifest.yml")
		// Ensure the file exists at the minimum
		helpers.WriteManifest(manifestPath, map[string]interface{}{})
	})

	AfterEach(func() {
		Expect(os.RemoveAll(appDir)).ToNot(HaveOccurred())
	})

	Describe("help", func() {
		It("appears in cf help -a", func() {
			session := helpers.CF("help", "-a")
			Eventually(session).Should(Exit(0))
			Expect(session).To(HaveCommandInCategoryWithDescription("apply-manifest", "APPS", "Apply manifest properties to a space"))
		})

		When("--help flag is set", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("apply-manifest", "--help")

				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("apply-manifest - Apply manifest properties to a space"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say("cf apply-manifest -f APP_MANIFEST_PATH"))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("create-app, create-app-manifest, push"))

				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("the environment is not setup correctly", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(true, true, ReadOnlyOrg, "apply-manifest", "-f", manifestPath)
		})
	})

	When("the environment is set up correctly", func() {
		BeforeEach(func() {
			helpers.SetupCF(orgName, spaceName)
		})

		AfterEach(func() {
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app name in the manifest is missing", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"instances": 3,
						},
					},
				})
			})

			It("reports an error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Found an application with no name specified"))
				Eventually(session).Should(Say("FAILED"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("there is a CC error", func() {
			BeforeEach(func() {
				helpers.WriteManifest(manifestPath, map[string]interface{}{
					"applications": []map[string]interface{}{
						{
							"name":      appName,
							"instances": -1,
						},
					},
				})
			})

			It("displays the error", func() {
				session := helpers.CF("apply-manifest", "-f", manifestPath)
				Eventually(session.Err).Should(Say("Instances must be greater than or equal to 0"))
				Eventually(session).Should(Say("FAILED"))

				Eventually(session).Should(Exit(1))
			})
		})

		When("-f is provided", func() {
			When("the -f flag is not given an arg", func() {
				It("tells the user that the flag requires an arg, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f")

					Eventually(session.Err).Should(Say("Incorrect Usage: expected argument for flag `-f'"))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a directory that does not have a manifest.yml file", func() {
				var (
					emptyDir string
					err      error
				)

				BeforeEach(func() {
					emptyDir, err = ioutil.TempDir("", "empty")
					Expect(err).NotTo(HaveOccurred())
				})

				AfterEach(func() {
					Expect(os.RemoveAll(emptyDir)).ToNot(HaveOccurred())
				})

				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", emptyDir)

					Eventually(session.Err).Should(helpers.SayPath("Incorrect Usage: The specified directory '%s' does not contain a file named 'manifest.yml'.", emptyDir))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the -f flag points to a file that does not exist", func() {
				It("tells the user that the provided path doesn't exist, prints help text, and exits 1", func() {
					session := helpers.CF("apply-manifest", "-f", "path/that/does/not/exist")

					Eventually(session.Err).Should(Say("Incorrect Usage: The specified path 'path/that/does/not/exist' does not exist."))
					Eventually(session).Should(Say("NAME:"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the manifest exists where -f points", func() {
				It("applies the manifest successfully", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CF("apply-manifest", "-f", appDir)
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})
		})

		When("-f is not provided", func() {
			When("a properly formatted manifest is present in the pwd", func() {
				It("autodetects and applies the manifest", func() {
					userName, _ := helpers.GetCredentials()
					helpers.WriteManifest(filepath.Join(appDir, "manifest.yml"), map[string]interface{}{
						"applications": []map[string]interface{}{
							{
								"name":      appName,
								"instances": 3,
							},
						},
					})

					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: appDir}, "apply-manifest")
					Eventually(session).Should(Say("Applying manifest %s in org %s / space %s as %s...", regexp.QuoteMeta(manifestPath), orgName, spaceName, userName))
					Eventually(session).Should(Exit())

					session = helpers.CF("app", appName)
					Eventually(session).Should(Say(`instances:\s+%s`, `\d/3`))
					Eventually(session).Should(Exit())
				})
			})

			When("the current directory does not have a manifest", func() {
				It("fails nicely", func() {
					currentDir, err := os.Getwd()
					Expect(err).NotTo(HaveOccurred())
					session := helpers.CF("apply-manifest")

					Eventually(session.Err).Should(helpers.SayPath(`Could not find 'manifest.yml' file in %s`, currentDir))
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Exit(1))
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package gateway

import (
	"sync"
	"time"

	ttnlog "github.com/TheThingsNetwork/go-utils/log"
	"github.com/TheThingsNetwork/ttn/api/fields"
	pb "github.com/TheThingsNetwork/ttn/api/gateway"
	pb_monitor "github.com/TheThingsNetwork/ttn/api/monitor"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	pb_router "github.com/TheThingsNetwork/ttn/api/router"
)

// NewGateway creates a new in-memory Gateway structure
func NewGateway(ctx ttnlog.Interface, id string) *Gateway {
	ctx = ctx.WithField("GatewayID", id)
	gtw := &Gateway{
		ID:          id,
		Status:      NewStatusStore(),
		Utilization: NewUtilization(),
		Schedule:    NewSchedule(ctx),
		Ctx:         ctx,
	}
	gtw.Schedule.(*schedule).gateway = gtw // FIXME: Issue #420
	return gtw
}

// Gateway contains the state of a gateway
type Gateway struct {
	ID          string
	Status      StatusStore
	Utilization Utilization
	Schedule    Schedule
	LastSeen    time.Time

	mu            sync.RWMutex // Protect token and authenticated
	token         string
	authenticated bool

	Monitor       *pb_monitor.Client
	MonitorStream pb_monitor.GenericStream

	Ctx ttnlog.Interface
}

func (g *Gateway) SetAuth(token string, authenticated bool) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.authenticated = authenticated
	if token == g.token {
		return
	}
	g.token = token
	if g.MonitorStream != nil {
		g.MonitorStream.Close()
	}
	if g.Monitor != nil {
		g.MonitorStream = g.Monitor.NewGatewayStreams(g.ID, token)
	}
}

func (g *Gateway) updateLastSeen() {
	g.LastSeen = time.Now()
}

func (g *Gateway) HandleStatus(status *pb.Status) (err error) {
	g.mu.RLock()
	defer g.mu.RUnlock()
	status.GatewayTrusted = g.authenticated
	if err = g.Status.Update(status); err != nil {
		return err
	}
	g.updateLastSeen()
	if g.MonitorStream != nil {
		clone := *status // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}

func (g *Gateway) HandleUplink(uplink *pb_router.UplinkMessage) (err error) {
	if err = g.Utilization.AddRx(uplink); err != nil {
		return err
	}
	g.Schedule.Sync(uplink.GatewayMetadata.Timestamp)
	g.updateLastSeen()

	status, err := g.Status.Get()
	if err == nil {
		// Inject Gateway location
		if uplink.GatewayMetadata.Gps == nil {
			uplink.GatewayMetadata.Gps = status.GetGps()
		}
		// Inject Gateway region
		if region, ok := pb_lorawan.Region_value[status.Region]; ok {
			if lorawan := uplink.GetProtocolMetadata().GetLorawan(); lorawan != nil {
				lorawan.Region = pb_lorawan.Region(region)
			}
		}
	}

	// Inject authenticated as GatewayTrusted
	g.mu.RLock()
	defer g.mu.RUnlock()
	uplink.GatewayMetadata.GatewayTrusted = g.authenticated
	uplink.GatewayMetadata.GatewayId = g.ID
	if g.MonitorStream != nil {
		clone := *uplink // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}

func (g *Gateway) HandleDownlink(identifier string, downlink *pb_router.DownlinkMessage) (err error) {
	ctx := g.Ctx.WithField("Identifier", identifier).WithFields(fields.Get(downlink))
	if err = g.Schedule.Schedule(identifier, downlink); err != nil {
		ctx.WithError(err).Warn("Could not schedule downlink")
		return err
	}
	ctx.Debug("Scheduled downlink")
	if g.MonitorStream != nil {
		clone := *downlink // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}
<commit_msg>Don't forget to start Gateway monitor streams if there is no token change<commit_after>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package gateway

import (
	"sync"
	"time"

	ttnlog "github.com/TheThingsNetwork/go-utils/log"
	"github.com/TheThingsNetwork/ttn/api/fields"
	pb "github.com/TheThingsNetwork/ttn/api/gateway"
	pb_monitor "github.com/TheThingsNetwork/ttn/api/monitor"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	pb_router "github.com/TheThingsNetwork/ttn/api/router"
)

// NewGateway creates a new in-memory Gateway structure
func NewGateway(ctx ttnlog.Interface, id string) *Gateway {
	ctx = ctx.WithField("GatewayID", id)
	gtw := &Gateway{
		ID:          id,
		Status:      NewStatusStore(),
		Utilization: NewUtilization(),
		Schedule:    NewSchedule(ctx),
		Ctx:         ctx,
	}
	gtw.Schedule.(*schedule).gateway = gtw // FIXME: Issue #420
	return gtw
}

// Gateway contains the state of a gateway
type Gateway struct {
	ID          string
	Status      StatusStore
	Utilization Utilization
	Schedule    Schedule
	LastSeen    time.Time

	mu            sync.RWMutex // Protect token and authenticated
	token         string
	authenticated bool

	Monitor       *pb_monitor.Client
	MonitorStream pb_monitor.GenericStream

	Ctx ttnlog.Interface
}

func (g *Gateway) SetAuth(token string, authenticated bool) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.authenticated = authenticated
	if g.MonitorStream != nil {
		if token == g.token {
			return
		}
		g.Ctx.Debug("Stopping Monitor stream (token changed)")
		g.MonitorStream.Close()
	}
	g.token = token
	if g.Monitor != nil {
		g.Ctx.Debug("Starting Gateway Monitor Stream")
		g.MonitorStream = g.Monitor.NewGatewayStreams(g.ID, g.token)
	}
}

func (g *Gateway) updateLastSeen() {
	g.LastSeen = time.Now()
}

func (g *Gateway) HandleStatus(status *pb.Status) (err error) {
	g.mu.RLock()
	defer g.mu.RUnlock()
	status.GatewayTrusted = g.authenticated
	if err = g.Status.Update(status); err != nil {
		return err
	}
	g.updateLastSeen()
	if g.MonitorStream != nil {
		clone := *status // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}

func (g *Gateway) HandleUplink(uplink *pb_router.UplinkMessage) (err error) {
	if err = g.Utilization.AddRx(uplink); err != nil {
		return err
	}
	g.Schedule.Sync(uplink.GatewayMetadata.Timestamp)
	g.updateLastSeen()

	status, err := g.Status.Get()
	if err == nil {
		// Inject Gateway location
		if uplink.GatewayMetadata.Gps == nil {
			uplink.GatewayMetadata.Gps = status.GetGps()
		}
		// Inject Gateway region
		if region, ok := pb_lorawan.Region_value[status.Region]; ok {
			if lorawan := uplink.GetProtocolMetadata().GetLorawan(); lorawan != nil {
				lorawan.Region = pb_lorawan.Region(region)
			}
		}
	}

	// Inject authenticated as GatewayTrusted
	g.mu.RLock()
	defer g.mu.RUnlock()
	uplink.GatewayMetadata.GatewayTrusted = g.authenticated
	uplink.GatewayMetadata.GatewayId = g.ID
	if g.MonitorStream != nil {
		clone := *uplink // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}

func (g *Gateway) HandleDownlink(identifier string, downlink *pb_router.DownlinkMessage) (err error) {
	ctx := g.Ctx.WithField("Identifier", identifier).WithFields(fields.Get(downlink))
	if err = g.Schedule.Schedule(identifier, downlink); err != nil {
		ctx.WithError(err).Warn("Could not schedule downlink")
		return err
	}
	ctx.Debug("Scheduled downlink")
	if g.MonitorStream != nil {
		clone := *downlink // Avoid race conditions
		g.MonitorStream.Send(&clone)
	}
	return nil
}
<|endoftext|>


<commit_before>// Package nictype is a small package to allow resolving NIC "network" key to "nictype" key.
// It is it's own package to avoid circular dependency issues.
package nictype

import (
	"fmt"

	"github.com/pkg/errors"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/state"
)

// NICType resolves the NIC Type for the supplied NIC device config.
// If the device "type" is "nic" and the "network" property is specified in the device config, then NIC type is
// resolved from the network's type. Otherwise the device's "nictype" property is returned (which may be empty if
// used with non-NIC device configs).
func NICType(s *state.State, d deviceConfig.Device) (string, error) {
	// NIC devices support resolving their "nictype" from their "network" property.
	if d["type"] == "nic" {
		if d["network"] != "" {
			_, netInfo, err := s.Cluster.GetNetworkInAnyState(d["network"])
			if err != nil {
				return "", errors.Wrapf(err, "Failed to load network %q", d["network"])
			}

			var nicType string
			switch netInfo.Type {
			case "bridge":
				nicType = "bridged"
			case "macvlan":
				nicType = "macvlan"
			case "sriov":
				nicType = "sriov"
			default:
				return "", fmt.Errorf("Unrecognised NIC network type for network %q", d["network"])
			}

			return nicType, nil
		}

	}

	// Infiniband devices use "nictype" without supporting "network" property, so just return it directly,
	// which is the same as accessing the property directly from the config.
	return d["nictype"], nil
}
<commit_msg>lxd/device/nictype: Adds ovn support<commit_after>// Package nictype is a small package to allow resolving NIC "network" key to "nictype" key.
// It is it's own package to avoid circular dependency issues.
package nictype

import (
	"fmt"

	"github.com/pkg/errors"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/state"
)

// NICType resolves the NIC Type for the supplied NIC device config.
// If the device "type" is "nic" and the "network" property is specified in the device config, then NIC type is
// resolved from the network's type. Otherwise the device's "nictype" property is returned (which may be empty if
// used with non-NIC device configs).
func NICType(s *state.State, d deviceConfig.Device) (string, error) {
	// NIC devices support resolving their "nictype" from their "network" property.
	if d["type"] == "nic" {
		if d["network"] != "" {
			_, netInfo, err := s.Cluster.GetNetworkInAnyState(d["network"])
			if err != nil {
				return "", errors.Wrapf(err, "Failed to load network %q", d["network"])
			}

			var nicType string
			switch netInfo.Type {
			case "bridge":
				nicType = "bridged"
			case "macvlan":
				nicType = "macvlan"
			case "sriov":
				nicType = "sriov"
			case "ovn":
				nicType = "ovn"
			default:
				return "", fmt.Errorf("Unrecognised NIC network type for network %q", d["network"])
			}

			return nicType, nil
		}

	}

	// Infiniband devices use "nictype" without supporting "network" property, so just return it directly,
	// which is the same as accessing the property directly from the config.
	return d["nictype"], 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 cert

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

// CanReadCertAndKey returns true if the certificate and key files already exists,
// otherwise returns false. If lost one of cert and key, returns error.
func CanReadCertAndKey(certPath, keyPath string) (bool, error) {
	certReadable := canReadFile(certPath)
	keyReadable := canReadFile(keyPath)

	if certReadable == false && keyReadable == false {
		return false, nil
	}

	if certReadable == false {
		return false, fmt.Errorf("error reading %s, certificate and key must be supplied as a pair", certPath)
	}

	if keyReadable == false {
		return false, fmt.Errorf("error reading %s, certificate and key must be supplied as a pair", keyPath)
	}

	return true, nil
}

// If the file represented by path exists and
// readable, returns true otherwise returns false.
func canReadFile(path string) bool {
	f, err := os.Open(path)
	if err != nil {
		return false
	}

	defer f.Close()

	return true
}

// WriteCert writes the pem-encoded certificate data to certPath.
// The certificate file will be created with file mode 0644.
// If the certificate file already exists, it will be overwritten.
// The parent directory of the certPath will be created as needed with file mode 0755.
func WriteCert(certPath string, data []byte) error {
	if err := os.MkdirAll(filepath.Dir(certPath), os.FileMode(0755)); err != nil {
		return err
	}
	return ioutil.WriteFile(certPath, data, os.FileMode(0644))
}

// WriteKey writes the pem-encoded key data to keyPath.
// The key file will be created with file mode 0600.
// If the key file already exists, it will be overwritten.
// The parent directory of the keyPath will be created as needed with file mode 0755.
func WriteKey(keyPath string, data []byte) error {
	if err := os.MkdirAll(filepath.Dir(keyPath), os.FileMode(0755)); err != nil {
		return err
	}
	return ioutil.WriteFile(keyPath, data, os.FileMode(0600))
}

// LoadOrGenerateKeyFile looks for a key in the file at the given path. If it
// can't find one, it will generate a new key and store it there.
func LoadOrGenerateKeyFile(keyPath string) (data []byte, wasGenerated bool, err error) {
	loadedData, err := ioutil.ReadFile(keyPath)
	if err == nil {
		return loadedData, false, err
	}
	if !os.IsNotExist(err) {
		return nil, false, fmt.Errorf("error loading key from %s: %v", keyPath, err)
	}

	generatedData, err := MakeEllipticPrivateKeyPEM()
	if err != nil {
		return nil, false, fmt.Errorf("error generating key: %v", err)
	}
	if err := WriteKey(keyPath, generatedData); err != nil {
		return nil, false, fmt.Errorf("error writing key to %s: %v", keyPath, err)
	}
	return generatedData, true, nil
}

// MarshalPrivateKeyToPEM converts a known private key type of RSA or ECDSA to
// a PEM encoded block or returns an error.
func MarshalPrivateKeyToPEM(privateKey crypto.PrivateKey) ([]byte, error) {
	switch t := privateKey.(type) {
	case *ecdsa.PrivateKey:
		derBytes, err := x509.MarshalECPrivateKey(t)
		if err != nil {
			return nil, err
		}
		privateKeyPemBlock := &pem.Block{
			Type:  ECPrivateKeyBlockType,
			Bytes: derBytes,
		}
		return pem.EncodeToMemory(privateKeyPemBlock), nil
	case *rsa.PrivateKey:
		return EncodePrivateKeyPEM(t), nil
	default:
		return nil, fmt.Errorf("private key is not a recognized type: %T", privateKey)
	}
}

// NewPool returns an x509.CertPool containing the certificates in the given PEM-encoded file.
// Returns an error if the file could not be read, a certificate could not be parsed, or if the file does not contain any certificates
func NewPool(filename string) (*x509.CertPool, error) {
	certs, err := CertsFromFile(filename)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	for _, cert := range certs {
		pool.AddCert(cert)
	}
	return pool, nil
}

// CertsFromFile returns the x509.Certificates contained in the given PEM-encoded file.
// Returns an error if the file could not be read, a certificate could not be parsed, or if the file does not contain any certificates
func CertsFromFile(file string) ([]*x509.Certificate, error) {
	pemBlock, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	certs, err := ParseCertsPEM(pemBlock)
	if err != nil {
		return nil, fmt.Errorf("error reading %s: %s", file, err)
	}
	return certs, nil
}

// PrivateKeyFromFile returns the private key in rsa.PrivateKey or ecdsa.PrivateKey format from a given PEM-encoded file.
// Returns an error if the file could not be read or if the private key could not be parsed.
func PrivateKeyFromFile(file string) (interface{}, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	key, err := ParsePrivateKeyPEM(data)
	if err != nil {
		return nil, fmt.Errorf("error reading private key file %s: %v", file, err)
	}
	return key, nil
}

// PublicKeysFromFile returns the public keys in rsa.PublicKey or ecdsa.PublicKey format from a given PEM-encoded file.
// Reads public keys from both public and private key files.
func PublicKeysFromFile(file string) ([]interface{}, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	keys, err := ParsePublicKeysPEM(data)
	if err != nil {
		return nil, fmt.Errorf("error reading public key file %s: %v", file, err)
	}
	return keys, nil
}
<commit_msg>Re-use private key after failed CSR<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 cert

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

// CanReadCertAndKey returns true if the certificate and key files already exists,
// otherwise returns false. If lost one of cert and key, returns error.
func CanReadCertAndKey(certPath, keyPath string) (bool, error) {
	certReadable := canReadFile(certPath)
	keyReadable := canReadFile(keyPath)

	if certReadable == false && keyReadable == false {
		return false, nil
	}

	if certReadable == false {
		return false, fmt.Errorf("error reading %s, certificate and key must be supplied as a pair", certPath)
	}

	if keyReadable == false {
		return false, fmt.Errorf("error reading %s, certificate and key must be supplied as a pair", keyPath)
	}

	return true, nil
}

// If the file represented by path exists and
// readable, returns true otherwise returns false.
func canReadFile(path string) bool {
	f, err := os.Open(path)
	if err != nil {
		return false
	}

	defer f.Close()

	return true
}

// WriteCert writes the pem-encoded certificate data to certPath.
// The certificate file will be created with file mode 0644.
// If the certificate file already exists, it will be overwritten.
// The parent directory of the certPath will be created as needed with file mode 0755.
func WriteCert(certPath string, data []byte) error {
	if err := os.MkdirAll(filepath.Dir(certPath), os.FileMode(0755)); err != nil {
		return err
	}
	return ioutil.WriteFile(certPath, data, os.FileMode(0644))
}

// WriteKey writes the pem-encoded key data to keyPath.
// The key file will be created with file mode 0600.
// If the key file already exists, it will be overwritten.
// The parent directory of the keyPath will be created as needed with file mode 0755.
func WriteKey(keyPath string, data []byte) error {
	if err := os.MkdirAll(filepath.Dir(keyPath), os.FileMode(0755)); err != nil {
		return err
	}
	return ioutil.WriteFile(keyPath, data, os.FileMode(0600))
}

// LoadOrGenerateKeyFile looks for a key in the file at the given path. If it
// can't find one, it will generate a new key and store it there.
func LoadOrGenerateKeyFile(keyPath string) (data []byte, wasGenerated bool, err error) {
	loadedData, err := ioutil.ReadFile(keyPath)
	// Call verifyKeyData to ensure the file wasn't empty/corrupt.
	if err == nil && verifyKeyData(loadedData) {
		return loadedData, false, err
	}
	if !os.IsNotExist(err) {
		return nil, false, fmt.Errorf("error loading key from %s: %v", keyPath, err)
	}

	generatedData, err := MakeEllipticPrivateKeyPEM()
	if err != nil {
		return nil, false, fmt.Errorf("error generating key: %v", err)
	}
	if err := WriteKey(keyPath, generatedData); err != nil {
		return nil, false, fmt.Errorf("error writing key to %s: %v", keyPath, err)
	}
	return generatedData, true, nil
}

// MarshalPrivateKeyToPEM converts a known private key type of RSA or ECDSA to
// a PEM encoded block or returns an error.
func MarshalPrivateKeyToPEM(privateKey crypto.PrivateKey) ([]byte, error) {
	switch t := privateKey.(type) {
	case *ecdsa.PrivateKey:
		derBytes, err := x509.MarshalECPrivateKey(t)
		if err != nil {
			return nil, err
		}
		privateKeyPemBlock := &pem.Block{
			Type:  ECPrivateKeyBlockType,
			Bytes: derBytes,
		}
		return pem.EncodeToMemory(privateKeyPemBlock), nil
	case *rsa.PrivateKey:
		return EncodePrivateKeyPEM(t), nil
	default:
		return nil, fmt.Errorf("private key is not a recognized type: %T", privateKey)
	}
}

// NewPool returns an x509.CertPool containing the certificates in the given PEM-encoded file.
// Returns an error if the file could not be read, a certificate could not be parsed, or if the file does not contain any certificates
func NewPool(filename string) (*x509.CertPool, error) {
	certs, err := CertsFromFile(filename)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	for _, cert := range certs {
		pool.AddCert(cert)
	}
	return pool, nil
}

// CertsFromFile returns the x509.Certificates contained in the given PEM-encoded file.
// Returns an error if the file could not be read, a certificate could not be parsed, or if the file does not contain any certificates
func CertsFromFile(file string) ([]*x509.Certificate, error) {
	pemBlock, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	certs, err := ParseCertsPEM(pemBlock)
	if err != nil {
		return nil, fmt.Errorf("error reading %s: %s", file, err)
	}
	return certs, nil
}

// PrivateKeyFromFile returns the private key in rsa.PrivateKey or ecdsa.PrivateKey format from a given PEM-encoded file.
// Returns an error if the file could not be read or if the private key could not be parsed.
func PrivateKeyFromFile(file string) (interface{}, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	key, err := ParsePrivateKeyPEM(data)
	if err != nil {
		return nil, fmt.Errorf("error reading private key file %s: %v", file, err)
	}
	return key, nil
}

// PublicKeysFromFile returns the public keys in rsa.PublicKey or ecdsa.PublicKey format from a given PEM-encoded file.
// Reads public keys from both public and private key files.
func PublicKeysFromFile(file string) ([]interface{}, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	keys, err := ParsePublicKeysPEM(data)
	if err != nil {
		return nil, fmt.Errorf("error reading public key file %s: %v", file, err)
	}
	return keys, nil
}

// verifyKeyData returns true if the provided data appears to be a valid private key.
func verifyKeyData(data []byte) bool {
	if len(data) == 0 {
		return false
	}
	_, err := ParsePrivateKeyPEM(data)
	return err == nil
}
<|endoftext|>


<commit_before><commit_msg>naming consistency<commit_after><|endoftext|>