<commit_before>package charond

import (
	"database/sql"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"sync"
	"testing"

	"github.com/go-kit/kit/log"
	libldap "github.com/go-ldap/ldap"
	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/charon/charonrpc"
	"github.com/piotrkowalczuk/charon/internal/ldap"
	"github.com/piotrkowalczuk/charon/internal/password"
	"github.com/piotrkowalczuk/mnemosyne/mnemosynerpc"
	"github.com/piotrkowalczuk/promgrpc"
	"github.com/piotrkowalczuk/sklog"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
)

// DaemonOpts ...
type DaemonOpts struct {
	Test               bool
	Monitoring         bool
	TLS                bool
	TLSCertFile        string
	TLSKeyFile         string
	PostgresAddress    string
	PostgresDebug      bool
	PasswordBCryptCost int
	MnemosyneAddress   string
	LDAP               bool
	LDAPAddress        string
	LDAPBaseDN         string
	LDAPSearchDN       string
	LDAPBasePassword   string
	LDAPMappings       *ldap.Mappings
	Logger             log.Logger
	RPCListener        net.Listener
	DebugListener      net.Listener
}

// TestDaemonOpts represent set of options that can be passed to the TestDaemon constructor.
type TestDaemonOpts struct {
	MnemosyneAddress string
	PostgresAddress  string
}

// Daemon ...
type Daemon struct {
	opts          DaemonOpts
	monitor       *monitoring
	ldap          *libldap.Conn
	logger        log.Logger
	rpcListener   net.Listener
	debugListener net.Listener
	mnemosyneConn *grpc.ClientConn
	mnemosyne     mnemosynerpc.SessionManagerClient
}

// NewDaemon ...
func NewDaemon(opts DaemonOpts) *Daemon {
	d := &Daemon{
		opts:          opts,
		logger:        opts.Logger,
		rpcListener:   opts.RPCListener,
		debugListener: opts.DebugListener,
	}

	return d
}

// TestDaemon returns address of fully started in-memory daemon and closer to close it.
func TestDaemon(t *testing.T, opts TestDaemonOpts) (net.Addr, io.Closer) {
	l, err := net.Listen("tcp", "127.0.0.1:0") // any available address
	if err != nil {
		t.Fatalf("charon daemon tcp listener setup error: %s", err.Error())
	}

	logger := sklog.NewTestLogger(t)
	grpclog.SetLogger(sklog.NewGRPCLogger(logger))

	d := NewDaemon(DaemonOpts{
		Test:               true,
		Monitoring:         false,
		MnemosyneAddress:   opts.MnemosyneAddress,
		Logger:             logger,
		PostgresAddress:    opts.PostgresAddress,
		RPCListener:        l,
		PasswordBCryptCost: bcrypt.MinCost,
	})
	if err := d.Run(); err != nil {
		t.Fatalf("charon daemon start error: %s", err.Error())
	}

	return d.Addr(), d
}

// Run ...
func (d *Daemon) Run() (err error) {
	interceptor := promgrpc.NewInterceptor()
	if err = d.initMonitoring(); err != nil {
		return
	}

	var db *sql.DB
	db, err = initPostgres(d.opts.PostgresAddress, d.opts.Test, d.logger)
	if err != nil {
		return err
	}
	repos := newRepositories(db)

	d.mnemosyne, d.mnemosyneConn = initMnemosyne(d.opts.MnemosyneAddress, interceptor, d.logger)

	var passwordHasher password.Hasher
	if d.opts.LDAP {
		// open connection to check if it is reachable
		if d.ldap, err = initLDAP(d.opts.LDAPAddress, d.opts.LDAPBaseDN, d.opts.LDAPBasePassword, d.logger); err != nil {
			return
		}
		d.ldap.Close()
	}

	passwordHasher = initHasher(d.opts.PasswordBCryptCost, d.logger)
	if d.opts.Test {
		if _, err = createDummyTestUser(repos.user, passwordHasher); err != nil {
			return
		}
		sklog.Info(d.logger, "test super user has been created")
	}

	permissionReg := initPermissionRegistry(repos.permission, charon.AllPermissions, d.logger)

	opts := []grpc.ServerOption{
		// No stream endpoint available at the moment.
		grpc.UnaryInterceptor(unaryServerInterceptors(
			func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
				res, err := handler(ctx, req)

				if err != nil && grpc.Code(err) != codes.OK {
					switch grpc.Code(err) {
					case codes.OK:
						sklog.Debug(d.logger, "request handled successfully", "handler", info.FullMethod)
					default:
						sklog.Error(d.logger, errors.New(grpc.ErrorDesc(err)), "handler", info.FullMethod, "code", grpc.Code(err).String())
					}
					return nil, handleError(err)
				}
				return res, err
			},
			interceptor.UnaryServer(),
		)),
	}
	if d.opts.TLS {
		creds, err := credentials.NewServerTLSFromFile(d.opts.TLSCertFile, d.opts.TLSKeyFile)
		if err != nil {
			return err
		}
		opts = append(opts, grpc.Creds(creds))
	}

	grpclog.SetLogger(sklog.NewGRPCLogger(d.logger))
	gRPCServer := grpc.NewServer(opts...)
	server := &rpcServer{
		opts:               d.opts,
		logger:             d.logger,
		session:            d.mnemosyne,
		passwordHasher:     passwordHasher,
		permissionRegistry: permissionReg,
		repository:         repos,
		ldap: &sync.Pool{
			New: func() interface{} {
				conn, err := initLDAP(d.opts.LDAPAddress, d.opts.LDAPBaseDN, d.opts.LDAPBasePassword, d.logger)
				if err != nil {
					return err
				}
				return conn
			},
		},
	}
	charonrpc.RegisterAuthServer(gRPCServer, newAuth(server))
	charonrpc.RegisterUserManagerServer(gRPCServer, newUserManager(server))
	charonrpc.RegisterGroupManagerServer(gRPCServer, newGroupManager(server))
	charonrpc.RegisterPermissionManagerServer(gRPCServer, newPermissionManager(server))

	go func() {
		sklog.Info(d.logger, "rpc server is running", "address", d.rpcListener.Addr().String())

		if err := gRPCServer.Serve(d.rpcListener); err != nil {
			if err == grpc.ErrServerStopped {
				sklog.Info(d.logger, "grpc server has been stoped")
				return
			}

			sklog.Error(d.logger, err)
		}
	}()

	if d.debugListener != nil {
		go func() {
			sklog.Info(d.logger, "debug server is running", "address", d.debugListener.Addr().String())
			// TODO: implement keep alive

			mux := http.NewServeMux()
			mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
			mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
			mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
			mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
			mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
			mux.Handle("/metrics", prometheus.Handler())
			sklog.Error(d.logger, http.Serve(d.debugListener, mux))
		}()
	}

	return
}

// Close implements io.Closer interface.
func (d *Daemon) Close() (err error) {
	if d.ldap != nil {
		d.ldap.Close()
	}
	if err = d.mnemosyneConn.Close(); err != nil {
		return
	}
	if err = d.rpcListener.Close(); err != nil {
		return
	}
	if d.debugListener != nil {
		err = d.debugListener.Close()
	}
	return
}

// Addr returns net.Addr that rpc service is listening on.
func (d *Daemon) Addr() net.Addr {
	return d.rpcListener.Addr()
}

func (d *Daemon) initMonitoring() (err error) {
	hostname, err := os.Hostname()
	if err != nil {
		return errors.New("getting hostname failed")
	}
	d.monitor = initPrometheus("charon", d.opts.Monitoring, prometheus.Labels{"server": hostname})
	return nil
}
<commit_msg>tweak grpc interceptor to log properly debug msg<commit_after>package charond

import (
	"database/sql"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/go-kit/kit/log"
	libldap "github.com/go-ldap/ldap"
	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/charon/charonrpc"
	"github.com/piotrkowalczuk/charon/internal/ldap"
	"github.com/piotrkowalczuk/charon/internal/password"
	"github.com/piotrkowalczuk/mnemosyne/mnemosynerpc"
	"github.com/piotrkowalczuk/promgrpc"
	"github.com/piotrkowalczuk/sklog"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
)

// DaemonOpts ...
type DaemonOpts struct {
	Test               bool
	Monitoring         bool
	TLS                bool
	TLSCertFile        string
	TLSKeyFile         string
	PostgresAddress    string
	PostgresDebug      bool
	PasswordBCryptCost int
	MnemosyneAddress   string
	LDAP               bool
	LDAPAddress        string
	LDAPBaseDN         string
	LDAPSearchDN       string
	LDAPBasePassword   string
	LDAPMappings       *ldap.Mappings
	Logger             log.Logger
	RPCListener        net.Listener
	DebugListener      net.Listener
}

// TestDaemonOpts represent set of options that can be passed to the TestDaemon constructor.
type TestDaemonOpts struct {
	MnemosyneAddress string
	PostgresAddress  string
}

// Daemon ...
type Daemon struct {
	opts          DaemonOpts
	monitor       *monitoring
	ldap          *libldap.Conn
	logger        log.Logger
	rpcListener   net.Listener
	debugListener net.Listener
	mnemosyneConn *grpc.ClientConn
	mnemosyne     mnemosynerpc.SessionManagerClient
}

// NewDaemon ...
func NewDaemon(opts DaemonOpts) *Daemon {
	d := &Daemon{
		opts:          opts,
		logger:        opts.Logger,
		rpcListener:   opts.RPCListener,
		debugListener: opts.DebugListener,
	}

	return d
}

// TestDaemon returns address of fully started in-memory daemon and closer to close it.
func TestDaemon(t *testing.T, opts TestDaemonOpts) (net.Addr, io.Closer) {
	l, err := net.Listen("tcp", "127.0.0.1:0") // any available address
	if err != nil {
		t.Fatalf("charon daemon tcp listener setup error: %s", err.Error())
	}

	logger := sklog.NewTestLogger(t)
	grpclog.SetLogger(sklog.NewGRPCLogger(logger))

	d := NewDaemon(DaemonOpts{
		Test:               true,
		Monitoring:         false,
		MnemosyneAddress:   opts.MnemosyneAddress,
		Logger:             logger,
		PostgresAddress:    opts.PostgresAddress,
		RPCListener:        l,
		PasswordBCryptCost: bcrypt.MinCost,
	})
	if err := d.Run(); err != nil {
		t.Fatalf("charon daemon start error: %s", err.Error())
	}

	return d.Addr(), d
}

// Run ...
func (d *Daemon) Run() (err error) {
	interceptor := promgrpc.NewInterceptor()
	if err = d.initMonitoring(); err != nil {
		return
	}

	var db *sql.DB
	db, err = initPostgres(d.opts.PostgresAddress, d.opts.Test, d.logger)
	if err != nil {
		return err
	}
	repos := newRepositories(db)

	d.mnemosyne, d.mnemosyneConn = initMnemosyne(d.opts.MnemosyneAddress, interceptor, d.logger)

	var passwordHasher password.Hasher
	if d.opts.LDAP {
		// open connection to check if it is reachable
		if d.ldap, err = initLDAP(d.opts.LDAPAddress, d.opts.LDAPBaseDN, d.opts.LDAPBasePassword, d.logger); err != nil {
			return
		}
		d.ldap.Close()
	}

	passwordHasher = initHasher(d.opts.PasswordBCryptCost, d.logger)
	if d.opts.Test {
		if _, err = createDummyTestUser(repos.user, passwordHasher); err != nil {
			return
		}
		sklog.Info(d.logger, "test super user has been created")
	}

	permissionReg := initPermissionRegistry(repos.permission, charon.AllPermissions, d.logger)

	opts := []grpc.ServerOption{
		// No stream endpoint available at the moment.
		grpc.UnaryInterceptor(unaryServerInterceptors(
			func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
				start := time.Now()
				res, err := handler(ctx, req)

				if err != nil && grpc.Code(err) != codes.OK {
					sklog.Error(d.logger, errors.New(grpc.ErrorDesc(err)), "handler", info.FullMethod, "code", grpc.Code(err).String(), "elapsed", time.Since(start))
					return nil, handleError(err)
				}
				sklog.Debug(d.logger, "request handled successfully", "handler", info.FullMethod, "elapsed", time.Since(start))
				return res, err
			},
			interceptor.UnaryServer(),
		)),
	}
	if d.opts.TLS {
		creds, err := credentials.NewServerTLSFromFile(d.opts.TLSCertFile, d.opts.TLSKeyFile)
		if err != nil {
			return err
		}
		opts = append(opts, grpc.Creds(creds))
	}

	grpclog.SetLogger(sklog.NewGRPCLogger(d.logger))
	gRPCServer := grpc.NewServer(opts...)
	server := &rpcServer{
		opts:               d.opts,
		logger:             d.logger,
		session:            d.mnemosyne,
		passwordHasher:     passwordHasher,
		permissionRegistry: permissionReg,
		repository:         repos,
		ldap: &sync.Pool{
			New: func() interface{} {
				conn, err := initLDAP(d.opts.LDAPAddress, d.opts.LDAPBaseDN, d.opts.LDAPBasePassword, d.logger)
				if err != nil {
					return err
				}
				return conn
			},
		},
	}
	charonrpc.RegisterAuthServer(gRPCServer, newAuth(server))
	charonrpc.RegisterUserManagerServer(gRPCServer, newUserManager(server))
	charonrpc.RegisterGroupManagerServer(gRPCServer, newGroupManager(server))
	charonrpc.RegisterPermissionManagerServer(gRPCServer, newPermissionManager(server))

	go func() {
		sklog.Info(d.logger, "rpc server is running", "address", d.rpcListener.Addr().String())

		if err := gRPCServer.Serve(d.rpcListener); err != nil {
			if err == grpc.ErrServerStopped {
				sklog.Info(d.logger, "grpc server has been stoped")
				return
			}

			sklog.Error(d.logger, err)
		}
	}()

	if d.debugListener != nil {
		go func() {
			sklog.Info(d.logger, "debug server is running", "address", d.debugListener.Addr().String())
			// TODO: implement keep alive

			mux := http.NewServeMux()
			mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
			mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
			mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
			mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
			mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
			mux.Handle("/metrics", prometheus.Handler())
			sklog.Error(d.logger, http.Serve(d.debugListener, mux))
		}()
	}

	return
}

// Close implements io.Closer interface.
func (d *Daemon) Close() (err error) {
	if d.ldap != nil {
		d.ldap.Close()
	}
	if err = d.mnemosyneConn.Close(); err != nil {
		return
	}
	if err = d.rpcListener.Close(); err != nil {
		return
	}
	if d.debugListener != nil {
		err = d.debugListener.Close()
	}
	return
}

// Addr returns net.Addr that rpc service is listening on.
func (d *Daemon) Addr() net.Addr {
	return d.rpcListener.Addr()
}

func (d *Daemon) initMonitoring() (err error) {
	hostname, err := os.Hostname()
	if err != nil {
		return errors.New("getting hostname failed")
	}
	d.monitor = initPrometheus("charon", d.opts.Monitoring, prometheus.Labels{"server": hostname})
	return nil
}
<|endoftext|>


<commit_before>/*
 * Minio Cloud Storage, (C) 2016 Minio, Inc.
 *
 * 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"
	"net"
	"runtime"
	"testing"
)

func TestCheckPortAvailability(t *testing.T) {
	tests := []struct {
		port int
	}{
		{9000},
		{10000},
	}
	for _, test := range tests {
		// This test should pass if the ports are available
		err := checkPortAvailability(test.port)
		if err != nil {
			t.Fatalf("checkPortAvailability test failed for port: %d. Error: %v", test.port, err)
		}

		// Now use the ports and check again
		ln, err := net.Listen("tcp", fmt.Sprintf(":%d", test.port))
		if err != nil {
			t.Fail()
		}
		defer ln.Close()

		err = checkPortAvailability(test.port)

		// Skip if the os is windows due to https://github.com/golang/go/issues/7598
		if err == nil && runtime.GOOS != "windows" {
			t.Fatalf("checkPortAvailability should fail for port: %d. Error: %v", test.port, err)
		}
	}
}
<commit_msg>tests: Make sure we try tests on free ports. (#2402)<commit_after>/*
 * Minio Cloud Storage, (C) 2016 Minio, Inc.
 *
 * 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"
	"net"
	"runtime"
	"testing"
)

// Tests for port availability logic written for server startup sequence.
func TestCheckPortAvailability(t *testing.T) {
	tests := []struct {
		port int
	}{
		{getFreePort()},
		{getFreePort()},
	}
	for _, test := range tests {
		// This test should pass if the ports are available
		err := checkPortAvailability(test.port)
		if err != nil {
			t.Fatalf("checkPortAvailability test failed for port: %d. Error: %v", test.port, err)
		}

		// Now use the ports and check again
		ln, err := net.Listen("tcp", fmt.Sprintf(":%d", test.port))
		if err != nil {
			t.Fail()
		}
		defer ln.Close()

		err = checkPortAvailability(test.port)

		// Skip if the os is windows due to https://github.com/golang/go/issues/7598
		if err == nil && runtime.GOOS != "windows" {
			t.Fatalf("checkPortAvailability should fail for port: %d. Error: %v", test.port, err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Google Inc.
//
// 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 resolve

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/context/keys"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/api"
	"github.com/google/gapid/gapis/api/sync"
	"github.com/google/gapid/gapis/capture"
	"github.com/google/gapid/gapis/database"
	"github.com/google/gapid/gapis/service"
	"github.com/google/gapid/gapis/service/path"
)

// CommandTree resolves the specified command tree path.
func CommandTree(ctx context.Context, c *path.CommandTree) (*service.CommandTree, error) {
	id, err := database.Store(ctx, &CommandTreeResolvable{c})
	if err != nil {
		return nil, err
	}
	return &service.CommandTree{
		Root: &path.CommandTreeNode{Tree: path.NewID(id)},
	}, nil
}

type commandTree struct {
	path *path.CommandTree
	root api.CmdIDGroup
}

func (t *commandTree) index(indices []uint64) api.SpanItem {
	group := api.CmdGroupOrRoot(t.root)
	for _, idx := range indices {
		switch item := group.Index(idx).(type) {
		case api.CmdIDGroup:
			group = item
		case api.SubCmdRoot:
			group = item
		default:
			return item
		}
	}
	return group
}

func (t *commandTree) indices(id api.CmdID) []uint64 {
	out := []uint64{}
	group := t.root
	for {
		i := group.IndexOf(id)
		out = append(out, i)
		switch item := group.Index(i).(type) {
		case api.CmdIDGroup:
			group = item
		default:
			return out
		}
	}
}

// CommandTreeNode resolves the specified command tree node path.
func CommandTreeNode(ctx context.Context, c *path.CommandTreeNode) (*service.CommandTreeNode, error) {
	boxed, err := database.Resolve(ctx, c.Tree.ID())
	if err != nil {
		return nil, err
	}

	cmdTree := boxed.(*commandTree)

	switch item := cmdTree.index(c.Indices).(type) {
	case api.SubCmdIdx:
		return &service.CommandTreeNode{
			NumChildren: 0, // TODO: Subcommands
			Commands:    cmdTree.path.Capture.SubCommandRange(item, item),
		}, nil
	case api.CmdIDGroup:
		return &service.CommandTreeNode{
			NumChildren: item.Count(),
			Commands:    cmdTree.path.Capture.CommandRange(uint64(item.Range.First()), uint64(item.Range.Last())),
			Group:       item.Name,
			NumCommands: item.DeepCount(func(g api.CmdIDGroup) bool { return true /* TODO: Subcommands */ }),
		}, nil
	case api.SubCmdRoot:
		count := uint64(1)
		g := ""
		commandStart := append([]uint64{}, item.Id...)
		commandEnd := append([]uint64{}, item.Id...)
		if len(item.Id) > 1 {
			g = fmt.Sprintf("%v", item.Id)
			commandStart = append(commandStart, uint64(0))
			commandEnd = append(commandEnd, uint64(item.SubGroup.Count()-1))
			count = uint64(item.SubGroup.Count())
		}
		return &service.CommandTreeNode{
			NumChildren: item.SubGroup.Bounds().Length(),
			Commands:    cmdTree.path.Capture.SubCommandRange(commandStart, commandEnd),
			Group:       g,
			NumCommands: count,
		}, nil
	default:
		panic(fmt.Errorf("Unexpected type: %T", item))
	}
}

// CommandTreeNodeForCommand returns the path to the CommandTreeNode that
// represents the specified command.
func CommandTreeNodeForCommand(ctx context.Context, p *path.CommandTreeNodeForCommand) (*path.CommandTreeNode, error) {
	boxed, err := database.Resolve(ctx, p.Tree.ID())
	if err != nil {
		return nil, err
	}

	cmdTree := boxed.(*commandTree)

	atomIdx := p.Command.Indices[0]
	if len(p.Command.Indices) > 1 {
		return nil, fmt.Errorf("Subcommands currently not supported for Command Tree") // TODO: Subcommands
	}

	return &path.CommandTreeNode{
		Tree:    p.Tree,
		Indices: cmdTree.indices(api.CmdID(atomIdx)),
	}, nil
}

type group struct {
	start api.CmdID
	end   api.CmdID
	name  string
}

type grouper interface {
	process(context.Context, api.CmdID, api.Cmd, *api.State)
	flush(count uint64)
	groups() []group
}

type runGrouper struct {
	f       func(cmd api.Cmd, s *api.State) (value interface{}, name string)
	start   api.CmdID
	current interface{}
	name    string
	out     []group
}

func (g *runGrouper) process(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	val, name := g.f(cmd, s)
	if val != g.current {
		if g.current != nil {
			g.out = append(g.out, group{g.start, id, g.name})
		}
		g.start = id
	}
	g.current, g.name = val, name
}

func (g *runGrouper) flush(count uint64) {
	end := api.CmdID(count)
	if g.current != nil && g.start != end {
		g.out = append(g.out, group{g.start, end, g.name})
	}
}

func (g *runGrouper) groups() []group { return g.out }

type markerGrouper struct {
	stack []group
	count int
	out   []group
}

func (g *markerGrouper) push(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	var name string
	if l, ok := cmd.(api.Labeled); ok {
		name = l.Label(ctx, s)
	}
	if len(name) > 0 {
		g.stack = append(g.stack, group{start: id, name: fmt.Sprintf("\"%s\"", name)})
	} else {
		g.stack = append(g.stack, group{start: id, name: fmt.Sprintf("Marker %d", g.count)})
		g.count++
	}
}

func (g *markerGrouper) pop(id api.CmdID) {
	m := g.stack[len(g.stack)-1]
	m.end = id + 1 // +1 to include pop marker
	g.out = append(g.out, m)
	g.stack = g.stack[:len(g.stack)-1]
}

func (g *markerGrouper) process(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	if cmd.CmdFlags().IsPushUserMarker() {
		g.push(ctx, id, cmd, s)
	}
	if cmd.CmdFlags().IsPopUserMarker() && len(g.stack) > 0 {
		g.pop(id)
	}
}

func (g *markerGrouper) flush(count uint64) {
	for len(g.stack) > 0 {
		g.pop(api.CmdID(count) - 1)
	}
}

func (g *markerGrouper) groups() []group { return g.out }

// Resolve builds and returns a *commandTree for the path.CommandTreeNode.
// Resolve implements the database.Resolver interface.
func (r *CommandTreeResolvable) Resolve(ctx context.Context) (interface{}, error) {
	p := r.Path
	ctx = capture.Put(ctx, p.Capture)

	c, err := capture.Resolve(ctx)
	if err != nil {
		return nil, err
	}

	syncData, err := database.Build(ctx, &SynchronizationResolvable{p.Capture})
	if err != nil {
		return nil, log.Errf(ctx, nil, "Error building sync data")
	}
	snc, ok := syncData.(*sync.Data)
	if !ok {
		return nil, log.Errf(ctx, nil, "Could not find valid Synchronization Data")
	}

	filter, err := buildFilter(ctx, p.Capture, p.Filter)
	if err != nil {
		return nil, err
	}

	groupers := []grouper{}

	if p.GroupByApi {
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			if api := cmd.API(); api != nil {
				return api.ID(), api.Name()
			}
			return nil, "No context"
		}})
	}

	if p.GroupByContext {
		var noContextID interface{}
		if p.IncludeNoContextGroups {
			noContextID = api.ContextID{}
		}
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			if api := cmd.API(); api != nil {
				if context := api.Context(s, cmd.Thread()); context != nil {
					return context.ID(), context.Name()
				}
			}
			return noContextID, "No context"
		}})
	}

	if p.GroupByThread {
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			thread := cmd.Thread()
			return thread, fmt.Sprintf("Thread: 0x%x", thread)
		}})
	}

	if p.GroupByUserMarkers {
		groupers = append(groupers, &markerGrouper{})
	}

	// Walk the list of unfiltered atoms to build the groups.
	s := c.NewState()
	api.ForeachCmd(ctx, c.Commands, func(ctx context.Context, id api.CmdID, cmd api.Cmd) error {
		cmd.Mutate(ctx, s, nil)
		if filter(cmd, s) {
			for _, g := range groupers {
				g.process(ctx, id, cmd, s)
			}
		}
		return nil
	})
	for _, g := range groupers {
		g.flush(uint64(len(c.Commands)))
	}

	// Build the command tree
	out := &commandTree{
		path: p,
		root: api.CmdIDGroup{
			Name:  "root",
			Range: api.CmdIDRange{End: api.CmdID(len(c.Commands))},
		},
	}
	for _, g := range groupers {
		for _, l := range g.groups() {
			out.root.AddGroup(l.start, l.end, l.name)
		}
	}

	if p.GroupByDrawCall || p.GroupByFrame {
		events, err := Events(ctx, &path.Events{
			Capture:      p.Capture,
			Filter:       p.Filter,
			DrawCalls:    true,
			FirstInFrame: true,
			LastInFrame:  true,
		})
		if err != nil {
			return nil, log.Errf(ctx, err, "Couldn't get events")
		}
		if p.GroupByFrame {
			addFrameEvents(ctx, events, p, out, api.CmdID(len(c.Commands)))
		}
		if p.GroupByDrawCall {
			addDrawEvents(ctx, events, p, out, api.CmdID(len(c.Commands)))
		}
	}

	for k, v := range snc.SubcommandGroups {
		r := out.root.AddRoot([]uint64{uint64(k)})
		for _, x := range v {
			r.Insert([]uint64{uint64(k)}, append([]uint64{}, x...))
		}
	}

	// Now we have all the groups, we finally need to add the filtered atoms.
	{
		// Clone the context to prevent cancellation occuring in Mutate() calls.
		ctx := keys.Clone(context.Background(), ctx)

		s = c.NewState()
		out.root.AddAtoms(func(i api.CmdID) bool {
			cmd := c.Commands[i]
			cmd.Mutate(ctx, s, nil)
			return filter(cmd, s)
		}, uint64(p.MaxChildren), uint64(p.MaxNeighbours))
	}

	return out, nil
}

func addDrawEvents(ctx context.Context, events *service.Events, p *path.CommandTree, t *commandTree, last api.CmdID) {
	drawCount, drawStart := 0, api.CmdID(0)
	for _, e := range events.List {
		i := api.CmdID(e.Command.Indices[0])
		switch e.Kind {
		case service.EventKind_DrawCall:
			t.root.AddGroup(drawStart, i+1, fmt.Sprintf("Draw %v", drawCount+1))
			drawCount++
			drawStart = i + 1

		case service.EventKind_FirstInFrame:
			drawCount, drawStart = 0, i
		}
	}
}

func addFrameEvents(ctx context.Context, events *service.Events, p *path.CommandTree, t *commandTree, last api.CmdID) {
	frameCount, frameStart, frameEnd := 0, api.CmdID(0), api.CmdID(0)
	for _, e := range events.List {
		i := api.CmdID(e.Command.Indices[0])
		switch e.Kind {
		case service.EventKind_FirstInFrame:
			frameStart = i

		case service.EventKind_LastInFrame:
			t.root.AddGroup(frameStart, i+1, fmt.Sprintf("Frame %v", frameCount+1))
			frameCount++
			frameEnd = i
		}
	}
	if p.AllowIncompleteFrame && frameCount > 0 && frameStart > frameEnd {
		t.root.AddGroup(frameStart, last, "Incomplete Frame")
	}
}
<commit_msg>Extend the end of frame group to past user-group<commit_after>// Copyright (C) 2017 Google Inc.
//
// 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 resolve

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/context/keys"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/api"
	"github.com/google/gapid/gapis/api/sync"
	"github.com/google/gapid/gapis/capture"
	"github.com/google/gapid/gapis/database"
	"github.com/google/gapid/gapis/service"
	"github.com/google/gapid/gapis/service/path"
)

// CommandTree resolves the specified command tree path.
func CommandTree(ctx context.Context, c *path.CommandTree) (*service.CommandTree, error) {
	id, err := database.Store(ctx, &CommandTreeResolvable{c})
	if err != nil {
		return nil, err
	}
	return &service.CommandTree{
		Root: &path.CommandTreeNode{Tree: path.NewID(id)},
	}, nil
}

type commandTree struct {
	path *path.CommandTree
	root api.CmdIDGroup
}

func (t *commandTree) index(indices []uint64) api.SpanItem {
	group := api.CmdGroupOrRoot(t.root)
	for _, idx := range indices {
		switch item := group.Index(idx).(type) {
		case api.CmdIDGroup:
			group = item
		case api.SubCmdRoot:
			group = item
		default:
			return item
		}
	}
	return group
}

func (t *commandTree) indices(id api.CmdID) []uint64 {
	out := []uint64{}
	group := t.root
	for {
		i := group.IndexOf(id)
		out = append(out, i)
		switch item := group.Index(i).(type) {
		case api.CmdIDGroup:
			group = item
		default:
			return out
		}
	}
}

// CommandTreeNode resolves the specified command tree node path.
func CommandTreeNode(ctx context.Context, c *path.CommandTreeNode) (*service.CommandTreeNode, error) {
	boxed, err := database.Resolve(ctx, c.Tree.ID())
	if err != nil {
		return nil, err
	}

	cmdTree := boxed.(*commandTree)

	switch item := cmdTree.index(c.Indices).(type) {
	case api.SubCmdIdx:
		return &service.CommandTreeNode{
			NumChildren: 0, // TODO: Subcommands
			Commands:    cmdTree.path.Capture.SubCommandRange(item, item),
		}, nil
	case api.CmdIDGroup:
		return &service.CommandTreeNode{
			NumChildren: item.Count(),
			Commands:    cmdTree.path.Capture.CommandRange(uint64(item.Range.First()), uint64(item.Range.Last())),
			Group:       item.Name,
			NumCommands: item.DeepCount(func(g api.CmdIDGroup) bool { return true /* TODO: Subcommands */ }),
		}, nil
	case api.SubCmdRoot:
		count := uint64(1)
		g := ""
		commandStart := append([]uint64{}, item.Id...)
		commandEnd := append([]uint64{}, item.Id...)
		if len(item.Id) > 1 {
			g = fmt.Sprintf("%v", item.Id)
			commandStart = append(commandStart, uint64(0))
			commandEnd = append(commandEnd, uint64(item.SubGroup.Count()-1))
			count = uint64(item.SubGroup.Count())
		}
		return &service.CommandTreeNode{
			NumChildren: item.SubGroup.Bounds().Length(),
			Commands:    cmdTree.path.Capture.SubCommandRange(commandStart, commandEnd),
			Group:       g,
			NumCommands: count,
		}, nil
	default:
		panic(fmt.Errorf("Unexpected type: %T", item))
	}
}

// CommandTreeNodeForCommand returns the path to the CommandTreeNode that
// represents the specified command.
func CommandTreeNodeForCommand(ctx context.Context, p *path.CommandTreeNodeForCommand) (*path.CommandTreeNode, error) {
	boxed, err := database.Resolve(ctx, p.Tree.ID())
	if err != nil {
		return nil, err
	}

	cmdTree := boxed.(*commandTree)

	atomIdx := p.Command.Indices[0]
	if len(p.Command.Indices) > 1 {
		return nil, fmt.Errorf("Subcommands currently not supported for Command Tree") // TODO: Subcommands
	}

	return &path.CommandTreeNode{
		Tree:    p.Tree,
		Indices: cmdTree.indices(api.CmdID(atomIdx)),
	}, nil
}

type group struct {
	start api.CmdID
	end   api.CmdID
	name  string
}

type grouper interface {
	process(context.Context, api.CmdID, api.Cmd, *api.State)
	flush(count uint64)
	groups() []group
}

type runGrouper struct {
	f       func(cmd api.Cmd, s *api.State) (value interface{}, name string)
	start   api.CmdID
	current interface{}
	name    string
	out     []group
}

func (g *runGrouper) process(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	val, name := g.f(cmd, s)
	if val != g.current {
		if g.current != nil {
			g.out = append(g.out, group{g.start, id, g.name})
		}
		g.start = id
	}
	g.current, g.name = val, name
}

func (g *runGrouper) flush(count uint64) {
	end := api.CmdID(count)
	if g.current != nil && g.start != end {
		g.out = append(g.out, group{g.start, end, g.name})
	}
}

func (g *runGrouper) groups() []group { return g.out }

type markerGrouper struct {
	stack []group
	count int
	out   []group
}

func (g *markerGrouper) push(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	var name string
	if l, ok := cmd.(api.Labeled); ok {
		name = l.Label(ctx, s)
	}
	if len(name) > 0 {
		g.stack = append(g.stack, group{start: id, name: fmt.Sprintf("\"%s\"", name)})
	} else {
		g.stack = append(g.stack, group{start: id, name: fmt.Sprintf("Marker %d", g.count)})
		g.count++
	}
}

func (g *markerGrouper) pop(id api.CmdID) {
	m := g.stack[len(g.stack)-1]
	m.end = id + 1 // +1 to include pop marker
	g.out = append(g.out, m)
	g.stack = g.stack[:len(g.stack)-1]
}

func (g *markerGrouper) process(ctx context.Context, id api.CmdID, cmd api.Cmd, s *api.State) {
	if cmd.CmdFlags().IsPushUserMarker() {
		g.push(ctx, id, cmd, s)
	}
	if cmd.CmdFlags().IsPopUserMarker() && len(g.stack) > 0 {
		g.pop(id)
	}
}

func (g *markerGrouper) flush(count uint64) {
	for len(g.stack) > 0 {
		g.pop(api.CmdID(count) - 1)
	}
}

func (g *markerGrouper) groups() []group { return g.out }

// Resolve builds and returns a *commandTree for the path.CommandTreeNode.
// Resolve implements the database.Resolver interface.
func (r *CommandTreeResolvable) Resolve(ctx context.Context) (interface{}, error) {
	p := r.Path
	ctx = capture.Put(ctx, p.Capture)

	c, err := capture.Resolve(ctx)
	if err != nil {
		return nil, err
	}

	syncData, err := database.Build(ctx, &SynchronizationResolvable{p.Capture})
	if err != nil {
		return nil, log.Errf(ctx, nil, "Error building sync data")
	}
	snc, ok := syncData.(*sync.Data)
	if !ok {
		return nil, log.Errf(ctx, nil, "Could not find valid Synchronization Data")
	}

	filter, err := buildFilter(ctx, p.Capture, p.Filter)
	if err != nil {
		return nil, err
	}

	groupers := []grouper{}

	if p.GroupByApi {
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			if api := cmd.API(); api != nil {
				return api.ID(), api.Name()
			}
			return nil, "No context"
		}})
	}

	if p.GroupByContext {
		var noContextID interface{}
		if p.IncludeNoContextGroups {
			noContextID = api.ContextID{}
		}
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			if api := cmd.API(); api != nil {
				if context := api.Context(s, cmd.Thread()); context != nil {
					return context.ID(), context.Name()
				}
			}
			return noContextID, "No context"
		}})
	}

	if p.GroupByThread {
		groupers = append(groupers, &runGrouper{f: func(cmd api.Cmd, s *api.State) (interface{}, string) {
			thread := cmd.Thread()
			return thread, fmt.Sprintf("Thread: 0x%x", thread)
		}})
	}

	if p.GroupByUserMarkers {
		groupers = append(groupers, &markerGrouper{})
	}

	// Walk the list of unfiltered atoms to build the groups.
	s := c.NewState()
	api.ForeachCmd(ctx, c.Commands, func(ctx context.Context, id api.CmdID, cmd api.Cmd) error {
		cmd.Mutate(ctx, s, nil)
		if filter(cmd, s) {
			for _, g := range groupers {
				g.process(ctx, id, cmd, s)
			}
		}
		return nil
	})
	for _, g := range groupers {
		g.flush(uint64(len(c.Commands)))
	}

	// Build the command tree
	out := &commandTree{
		path: p,
		root: api.CmdIDGroup{
			Name:  "root",
			Range: api.CmdIDRange{End: api.CmdID(len(c.Commands))},
		},
	}
	for _, g := range groupers {
		for _, l := range g.groups() {
			out.root.AddGroup(l.start, l.end, l.name)
		}
	}

	if p.GroupByDrawCall || p.GroupByFrame {
		events, err := Events(ctx, &path.Events{
			Capture:      p.Capture,
			Filter:       p.Filter,
			DrawCalls:    true,
			FirstInFrame: true,
			LastInFrame:  true,
		})
		if err != nil {
			return nil, log.Errf(ctx, err, "Couldn't get events")
		}
		if p.GroupByFrame {
			addFrameEvents(ctx, events, p, out, api.CmdID(len(c.Commands)))
		}
		if p.GroupByDrawCall {
			addDrawEvents(ctx, events, p, out, api.CmdID(len(c.Commands)))
		}
	}

	for k, v := range snc.SubcommandGroups {
		r := out.root.AddRoot([]uint64{uint64(k)})
		for _, x := range v {
			r.Insert([]uint64{uint64(k)}, append([]uint64{}, x...))
		}
	}

	// Now we have all the groups, we finally need to add the filtered atoms.
	{
		// Clone the context to prevent cancellation occuring in Mutate() calls.
		ctx := keys.Clone(context.Background(), ctx)

		s = c.NewState()
		out.root.AddAtoms(func(i api.CmdID) bool {
			cmd := c.Commands[i]
			cmd.Mutate(ctx, s, nil)
			return filter(cmd, s)
		}, uint64(p.MaxChildren), uint64(p.MaxNeighbours))
	}

	return out, nil
}

func addDrawEvents(ctx context.Context, events *service.Events, p *path.CommandTree, t *commandTree, last api.CmdID) {
	drawCount, drawStart := 0, api.CmdID(0)
	for _, e := range events.List {
		i := api.CmdID(e.Command.Indices[0])
		switch e.Kind {
		case service.EventKind_DrawCall:
			t.root.AddGroup(drawStart, i+1, fmt.Sprintf("Draw %v", drawCount+1))
			drawCount++
			drawStart = i + 1

		case service.EventKind_FirstInFrame:
			drawCount, drawStart = 0, i
		}
	}
}

func addFrameEvents(ctx context.Context, events *service.Events, p *path.CommandTree, t *commandTree, last api.CmdID) {
	frameCount, frameStart, frameEnd := 0, api.CmdID(0), api.CmdID(0)
	for _, e := range events.List {
		i := api.CmdID(e.Command.Indices[0])
		switch e.Kind {
		case service.EventKind_FirstInFrame:
			frameStart = i

			// If the start is within existing group, move it past the end of the group
			if idx := t.root.Spans.IndexOf(uint64(frameStart)); idx != -1 {
				if subgroup, ok := t.root.Spans[idx].(*api.CmdIDGroup); ok {
					frameStart = subgroup.Range.End
				}
			}

		case service.EventKind_LastInFrame:
			frameCount++
			frameEnd = i

			// If the end is within existing group, move it to the end of the group
			if idx := t.root.Spans.IndexOf(uint64(frameEnd)); idx != -1 {
				if subgroup, ok := t.root.Spans[idx].(*api.CmdIDGroup); ok {
					frameEnd = subgroup.Range.Last()
				}
			}

			// If the app properly annotates frames as well, we will end up with
			// both groupings, where one is the only child of the other.
			// However, we can not reliably detect this situation as the user
			// group might be surrounded by (potentially filtered) commands.

			t.root.AddGroup(frameStart, frameEnd+1, fmt.Sprintf("Frame %v", frameCount))
		}
	}
	if p.AllowIncompleteFrame && frameCount > 0 && frameStart > frameEnd {
		t.root.AddGroup(frameStart, last, "Incomplete Frame")
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI 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 backend

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/golang/protobuf/proto"

	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/memcache"
	"go.chromium.org/luci/appengine/tq"
	"go.chromium.org/luci/common/api/swarming/swarming/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/gce/api/tasks/v1"
	"go.chromium.org/luci/gce/appengine/model"
)

// utcRFC3339 is the timestamp format used by Swarming.
// Similar to RFC3339 but with an implicit UTC time zone.
const utcRFC3339 = "2006-01-02T15:04:05"

// setConnected sets the Swarming bot as connected in the datastore if it isn't already.
func setConnected(c context.Context, id, hostname string, at time.Time) error {
	// Check dscache in case vm.Connected is set now, to avoid a costly transaction.
	vm := &model.VM{
		ID: id,
	}
	switch err := datastore.Get(c, vm); {
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != hostname:
		return errors.Reason("bot %q does not exist", hostname).Err()
	case vm.Connected > 0:
		return nil
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Hostname != hostname:
			return errors.Reason("bot %q does not exist", hostname).Err()
		case vm.Connected > 0:
			return nil
		}
		vm.Connected = at.Unix()
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// manageMissingBot manages a missing Swarming bot.
func manageMissingBot(c context.Context, vm *model.VM) error {
	// Set that the bot has not yet connected to Swarming.
	switch {
	case vm.Lifetime > 0 && vm.Created+vm.Lifetime < time.Now().Unix():
		logging.Debugf(c, "deadline %d exceeded", vm.Created+vm.Lifetime)
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case vm.Drained:
		logging.Debugf(c, "VM drained")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case vm.Timeout > 0 && vm.Created+vm.Timeout < time.Now().Unix():
		logging.Debugf(c, "timeout %d exceeded", vm.Created+vm.Timeout)
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	default:
		return nil
	}
}

// manageExistingBot manages an existing Swarming bot.
func manageExistingBot(c context.Context, bot *swarming.SwarmingRpcsBotInfo, vm *model.VM) error {
	// This value of vm.Connected may be several seconds old, because the VM was fetched
	// prior to sending an RPC to Swarming. Still, check it here to save a costly operation
	// in setConnected, since manageExistingBot may be called thousands of times per minute.
	if vm.Connected == 0 {
		t, err := time.Parse(utcRFC3339, bot.FirstSeenTs)
		if err != nil {
			return errors.Annotate(err, "failed to parse bot connection time").Err()
		}
		if err := setConnected(c, vm.ID, vm.Hostname, t); err != nil {
			return err
		}
	}
	// A bot connected to Swarming may be executing workload.
	// To destroy the instance, terminate the bot first to avoid interruptions.
	// Termination can be skipped if the bot is deleted, dead, or already terminated.
	switch {
	case bot.Deleted:
		logging.Debugf(c, "bot deleted")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case bot.IsDead:
		logging.Debugf(c, "bot dead")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	}
	srv := getSwarming(c, vm.Swarming).Bot
	// bot_terminate occurs when the bot starts the termination task and is normally followed
	// by task_completed and bot_shutdown. A terminated bot has no further events. Responses
	// also include the full set of dimensions when the event was recorded. Limit response size
	// by fetching only recent events, and only the type of each.
	events, err := srv.Events(vm.Hostname).Context(c).Fields("items/event_type").Limit(5).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to fetch bot events").Err()
	}
	for _, e := range events.Items {
		if e.EventType == "bot_terminate" {
			logging.Debugf(c, "bot terminated")
			return destroyInstanceAsync(c, vm.ID, vm.URL)
		}
	}
	switch {
	case vm.Lifetime > 0 && vm.Created+vm.Lifetime < time.Now().Unix():
		logging.Debugf(c, "deadline %d exceeded", vm.Created+vm.Lifetime)
		return terminateBotAsync(c, vm.ID, vm.Hostname)
	case vm.Drained:
		logging.Debugf(c, "VM drained")
		return terminateBotAsync(c, vm.ID, vm.Hostname)
	}
	return nil
}

// manageBotQueue is the name of the manage bot task handler queue.
const manageBotQueue = "manage-bot"

// manageBot manages a Swarming bot.
func manageBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ManageBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.URL == "":
		logging.Debugf(c, "instance %q does not exist", vm.Hostname)
		return nil
	}
	// Drain the VM if necessary. No-op if unnecessary or already drained.
	if err := drainVM(c, vm); err != nil {
		return errors.Annotate(err, "failed to drain VM").Err()
	}
	logging.Debugf(c, "fetching bot %q: %s", vm.Hostname, vm.Swarming)
	bot, err := getSwarming(c, vm.Swarming).Bot.Get(vm.Hostname).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				logging.Debugf(c, "bot not found")
				return manageMissingBot(c, vm)
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to fetch bot").Err()
	}
	logging.Debugf(c, "found bot")
	return manageExistingBot(c, bot, vm)
}

// terminateBotAsync schedules a task queue task to terminate a Swarming bot.
func terminateBotAsync(c context.Context, id, hostname string) error {
	t := &tq.Task{
		Payload: &tasks.TerminateBot{
			Id:       id,
			Hostname: hostname,
		},
	}
	if err := getDispatcher(c).AddTask(c, t); err != nil {
		return errors.Annotate(err, "failed to schedule terminate task").Err()
	}
	return nil
}

// terminateBotQueue is the name of the terminate bot task handler queue.
const terminateBotQueue = "terminate-bot"

// terminateBot terminates an existing Swarming bot.
func terminateBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.TerminateBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetHostname() == "":
		return errors.Reason("hostname is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != task.Hostname:
		// Instance is already destroyed and replaced. Don't terminate the new bot.
		logging.Debugf(c, "bot %q does not exist", task.Hostname)
		return nil
	}
	// TODO(crbug/982840): swarming should perform de-duplication itself.
	mi := memcache.NewItem(c, fmt.Sprintf("terminate-%s/%s", vm.Swarming, vm.Hostname))
	if err := memcache.Get(c, mi); err == nil {
		logging.Debugf(c, "bot %q already has terminate task from us", task.Hostname)
		return nil
	}
	logging.Debugf(c, "terminating bot %q: %s", vm.Hostname, vm.Swarming)
	srv := getSwarming(c, vm.Swarming)
	if _, err := srv.Bot.Terminate(vm.Hostname).Context(c).Do(); err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				// Bot is already deleted.
				logging.Debugf(c, "bot not found")
				return nil
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to terminate bot").Err()
	}
	if err := memcache.Set(c, mi.SetExpiration(time.Hour)); err != nil {
		logging.Warningf(c, "failed to record terminate task in memcache: %s", err)
	}
	return nil
}

// deleteBotAsync schedules a task queue task to delete a Swarming bot.
func deleteBotAsync(c context.Context, id, hostname string) error {
	t := &tq.Task{
		Payload: &tasks.DeleteBot{
			Id:       id,
			Hostname: hostname,
		},
	}
	if err := getDispatcher(c).AddTask(c, t); err != nil {
		return errors.Annotate(err, "failed to schedule delete task").Err()
	}
	return nil
}

// deleteBotQueue is the name of the delete bot task handler queue.
const deleteBotQueue = "delete-bot"

// deleteBot deletes an existing Swarming bot.
func deleteBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.DeleteBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetHostname() == "":
		return errors.Reason("hostname is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != task.Hostname:
		// Instance is already destroyed and replaced. Don't delete the new bot.
		return errors.Reason("bot %q does not exist", task.Hostname).Err()
	}
	logging.Debugf(c, "deleting bot %q: %s", vm.Hostname, vm.Swarming)
	srv := getSwarming(c, vm.Swarming).Bot
	_, err := srv.Delete(vm.Hostname).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				// Bot is already deleted.
				logging.Debugf(c, "bot not found")
				return deleteVM(c, task.Id, vm.Hostname)
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to delete bot").Err()
	}
	return deleteVM(c, task.Id, vm.Hostname)
}

// deleteVM deletes the given VM from the datastore if it exists.
func deleteVM(c context.Context, id, hostname string) error {
	vm := &model.VM{
		ID: id,
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
			return nil
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Hostname != hostname:
			logging.Debugf(c, "VM %q does not exist", hostname)
			return nil
		}
		if err := datastore.Delete(c, vm); err != nil {
			return errors.Annotate(err, "failed to delete VM").Err()
		}
		return nil
	}, nil)
}
<commit_msg>gce: change TODO to explanation.<commit_after>// Copyright 2019 The LUCI 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 backend

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/golang/protobuf/proto"

	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/memcache"
	"go.chromium.org/luci/appengine/tq"
	"go.chromium.org/luci/common/api/swarming/swarming/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/gce/api/tasks/v1"
	"go.chromium.org/luci/gce/appengine/model"
)

// utcRFC3339 is the timestamp format used by Swarming.
// Similar to RFC3339 but with an implicit UTC time zone.
const utcRFC3339 = "2006-01-02T15:04:05"

// setConnected sets the Swarming bot as connected in the datastore if it isn't already.
func setConnected(c context.Context, id, hostname string, at time.Time) error {
	// Check dscache in case vm.Connected is set now, to avoid a costly transaction.
	vm := &model.VM{
		ID: id,
	}
	switch err := datastore.Get(c, vm); {
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != hostname:
		return errors.Reason("bot %q does not exist", hostname).Err()
	case vm.Connected > 0:
		return nil
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Hostname != hostname:
			return errors.Reason("bot %q does not exist", hostname).Err()
		case vm.Connected > 0:
			return nil
		}
		vm.Connected = at.Unix()
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// manageMissingBot manages a missing Swarming bot.
func manageMissingBot(c context.Context, vm *model.VM) error {
	// Set that the bot has not yet connected to Swarming.
	switch {
	case vm.Lifetime > 0 && vm.Created+vm.Lifetime < time.Now().Unix():
		logging.Debugf(c, "deadline %d exceeded", vm.Created+vm.Lifetime)
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case vm.Drained:
		logging.Debugf(c, "VM drained")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case vm.Timeout > 0 && vm.Created+vm.Timeout < time.Now().Unix():
		logging.Debugf(c, "timeout %d exceeded", vm.Created+vm.Timeout)
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	default:
		return nil
	}
}

// manageExistingBot manages an existing Swarming bot.
func manageExistingBot(c context.Context, bot *swarming.SwarmingRpcsBotInfo, vm *model.VM) error {
	// This value of vm.Connected may be several seconds old, because the VM was fetched
	// prior to sending an RPC to Swarming. Still, check it here to save a costly operation
	// in setConnected, since manageExistingBot may be called thousands of times per minute.
	if vm.Connected == 0 {
		t, err := time.Parse(utcRFC3339, bot.FirstSeenTs)
		if err != nil {
			return errors.Annotate(err, "failed to parse bot connection time").Err()
		}
		if err := setConnected(c, vm.ID, vm.Hostname, t); err != nil {
			return err
		}
	}
	// A bot connected to Swarming may be executing workload.
	// To destroy the instance, terminate the bot first to avoid interruptions.
	// Termination can be skipped if the bot is deleted, dead, or already terminated.
	switch {
	case bot.Deleted:
		logging.Debugf(c, "bot deleted")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	case bot.IsDead:
		logging.Debugf(c, "bot dead")
		return destroyInstanceAsync(c, vm.ID, vm.URL)
	}
	srv := getSwarming(c, vm.Swarming).Bot
	// bot_terminate occurs when the bot starts the termination task and is normally followed
	// by task_completed and bot_shutdown. A terminated bot has no further events. Responses
	// also include the full set of dimensions when the event was recorded. Limit response size
	// by fetching only recent events, and only the type of each.
	events, err := srv.Events(vm.Hostname).Context(c).Fields("items/event_type").Limit(5).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to fetch bot events").Err()
	}
	for _, e := range events.Items {
		if e.EventType == "bot_terminate" {
			logging.Debugf(c, "bot terminated")
			return destroyInstanceAsync(c, vm.ID, vm.URL)
		}
	}
	switch {
	case vm.Lifetime > 0 && vm.Created+vm.Lifetime < time.Now().Unix():
		logging.Debugf(c, "deadline %d exceeded", vm.Created+vm.Lifetime)
		return terminateBotAsync(c, vm.ID, vm.Hostname)
	case vm.Drained:
		logging.Debugf(c, "VM drained")
		return terminateBotAsync(c, vm.ID, vm.Hostname)
	}
	return nil
}

// manageBotQueue is the name of the manage bot task handler queue.
const manageBotQueue = "manage-bot"

// manageBot manages a Swarming bot.
func manageBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ManageBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.URL == "":
		logging.Debugf(c, "instance %q does not exist", vm.Hostname)
		return nil
	}
	// Drain the VM if necessary. No-op if unnecessary or already drained.
	if err := drainVM(c, vm); err != nil {
		return errors.Annotate(err, "failed to drain VM").Err()
	}
	logging.Debugf(c, "fetching bot %q: %s", vm.Hostname, vm.Swarming)
	bot, err := getSwarming(c, vm.Swarming).Bot.Get(vm.Hostname).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				logging.Debugf(c, "bot not found")
				return manageMissingBot(c, vm)
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to fetch bot").Err()
	}
	logging.Debugf(c, "found bot")
	return manageExistingBot(c, bot, vm)
}

// terminateBotAsync schedules a task queue task to terminate a Swarming bot.
func terminateBotAsync(c context.Context, id, hostname string) error {
	t := &tq.Task{
		Payload: &tasks.TerminateBot{
			Id:       id,
			Hostname: hostname,
		},
	}
	if err := getDispatcher(c).AddTask(c, t); err != nil {
		return errors.Annotate(err, "failed to schedule terminate task").Err()
	}
	return nil
}

// terminateBotQueue is the name of the terminate bot task handler queue.
const terminateBotQueue = "terminate-bot"

// terminateBot terminates an existing Swarming bot.
func terminateBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.TerminateBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetHostname() == "":
		return errors.Reason("hostname is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != task.Hostname:
		// Instance is already destroyed and replaced. Don't terminate the new bot.
		logging.Debugf(c, "bot %q does not exist", task.Hostname)
		return nil
	}
	// 1 terminate task from GCE-P should normally suffice to terminate a bot,
	// but in case users or bugs interfere, resend terminate task after 1 hour.
	// Even if memcache content is cleared, more frequent terminate requests can
	// now be handled by the swarming server and after crbug/982840 will be
	// auto-deduped.
	mi := memcache.NewItem(c, fmt.Sprintf("terminate-%s/%s", vm.Swarming, vm.Hostname))
	if err := memcache.Get(c, mi); err == nil {
		logging.Debugf(c, "bot %q already has terminate task from us", task.Hostname)
		return nil
	}
	logging.Debugf(c, "terminating bot %q: %s", vm.Hostname, vm.Swarming)
	srv := getSwarming(c, vm.Swarming)
	if _, err := srv.Bot.Terminate(vm.Hostname).Context(c).Do(); err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				// Bot is already deleted.
				logging.Debugf(c, "bot not found")
				return nil
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to terminate bot").Err()
	}
	if err := memcache.Set(c, mi.SetExpiration(time.Hour)); err != nil {
		logging.Warningf(c, "failed to record terminate task in memcache: %s", err)
	}
	return nil
}

// deleteBotAsync schedules a task queue task to delete a Swarming bot.
func deleteBotAsync(c context.Context, id, hostname string) error {
	t := &tq.Task{
		Payload: &tasks.DeleteBot{
			Id:       id,
			Hostname: hostname,
		},
	}
	if err := getDispatcher(c).AddTask(c, t); err != nil {
		return errors.Annotate(err, "failed to schedule delete task").Err()
	}
	return nil
}

// deleteBotQueue is the name of the delete bot task handler queue.
const deleteBotQueue = "delete-bot"

// deleteBot deletes an existing Swarming bot.
func deleteBot(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.DeleteBot)
	switch {
	case !ok:
		return errors.Reason("unexpected payload %q", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetHostname() == "":
		return errors.Reason("hostname is required").Err()
	}
	vm := &model.VM{
		ID: task.Id,
	}
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
		return nil
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	case vm.Hostname != task.Hostname:
		// Instance is already destroyed and replaced. Don't delete the new bot.
		return errors.Reason("bot %q does not exist", task.Hostname).Err()
	}
	logging.Debugf(c, "deleting bot %q: %s", vm.Hostname, vm.Swarming)
	srv := getSwarming(c, vm.Swarming).Bot
	_, err := srv.Delete(vm.Hostname).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			if gerr.Code == http.StatusNotFound {
				// Bot is already deleted.
				logging.Debugf(c, "bot not found")
				return deleteVM(c, task.Id, vm.Hostname)
			}
			logErrors(c, gerr)
		}
		return errors.Annotate(err, "failed to delete bot").Err()
	}
	return deleteVM(c, task.Id, vm.Hostname)
}

// deleteVM deletes the given VM from the datastore if it exists.
func deleteVM(c context.Context, id, hostname string) error {
	vm := &model.VM{
		ID: id,
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
			return nil
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Hostname != hostname:
			logging.Debugf(c, "VM %q does not exist", hostname)
			return nil
		}
		if err := datastore.Delete(c, vm); err != nil {
			return errors.Annotate(err, "failed to delete VM").Err()
		}
		return nil
	}, nil)
}
<|endoftext|>


<commit_before>// Copyright 2017 Orinoco Payments
//
// 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 cli

import (
	"errors"
	"fmt"
	"math/big"
	"path/filepath"
	"runtime"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/accounts/keystore"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/node"
	"github.com/ethereum/go-ethereum/params"
	homedir "github.com/mitchellh/go-homedir"
)

// ObtainWallet fetches the wallet for a given address
func ObtainWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	wallet, err := obtainGethWallet(chainID, address)
	if err == nil {
		return wallet, nil
	}

	wallet, err = obtainParityWallet(chainID, address)
	if err == nil {
		return wallet, err
	}

	return wallet, fmt.Errorf("Failed to obtain wallet")
}

func obtainGethWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	keydir := node.DefaultDataDir()
	if chainID.Cmp(params.MainnetChainConfig.ChainId) == 0 {
		// Nothing to add for mainnet
	} else if chainID.Cmp(params.TestnetChainConfig.ChainId) == 0 {
		keydir = filepath.Join(keydir, "testnet")
	} else if chainID.Cmp(params.RinkebyChainConfig.ChainId) == 0 {
		keydir = filepath.Join(keydir, "rinkeby")
	}
	keydir = filepath.Join(keydir, "keystore")
	backends := []accounts.Backend{keystore.NewKeyStore(keydir, keystore.StandardScryptN, keystore.StandardScryptP)}
	accountManager := accounts.NewManager(backends...)
	defer accountManager.Close()
	account := accounts.Account{Address: address}
	wallet, err := accountManager.Find(account)
	return wallet, err
}

func obtainParityWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	keydir, err := homedir.Dir()
	if err != nil {
		return nil, fmt.Errorf("Failed to find home directory")
	}
	if runtime.GOOS == "windows" {
		keydir = filepath.Join(keydir, "AppData\\Roaming\\Parity\\Ethereum\\keys")
	} else if runtime.GOOS == "darwin" {
		keydir = filepath.Join(keydir, "Library/Application Support/io.parity.ethereum/keys")
	} else if runtime.GOOS == "linux" {
		keydir = filepath.Join(keydir, ".local/share/io.parity.ethereum/keys")
	} else {
		return nil, fmt.Errorf("Unsupported operating system")
	}

	if chainID.Cmp(params.MainnetChainConfig.ChainId) == 0 {
		keydir = filepath.Join(keydir, "ethereum")
	} else if chainID.Cmp(params.TestnetChainConfig.ChainId) == 0 {
		keydir = filepath.Join(keydir, "test")
	}

	backends := []accounts.Backend{keystore.NewKeyStore(keydir, keystore.StandardScryptN, keystore.StandardScryptP)}
	accountManager := accounts.NewManager(backends...)
	defer accountManager.Close()
	account := accounts.Account{Address: address}
	wallet, err := accountManager.Find(account)
	return wallet, err
}

// ObtainAccount fetches the account for a given address
func ObtainAccount(wallet *accounts.Wallet, address *common.Address, passphrase string) (*accounts.Account, error) {
	for _, account := range (*wallet).Accounts() {
		if *address == account.Address {
			if !VerifyPassphrase(*wallet, account, passphrase) {
				return nil, errors.New("invalid passphrase")
			}
			return &account, nil
		}
	}
	return nil, errors.New("account not found")
}

// VerifyPassphrase confirms that a passphrase is correct for an account
func VerifyPassphrase(wallet accounts.Wallet, account accounts.Account, passphrase string) bool {

	_, err := wallet.SignHashWithPassphrase(account, passphrase, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
	return err == nil
}
<commit_msg>also fix here<commit_after>// Copyright 2017 Orinoco Payments
//
// 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 cli

import (
	"errors"
	"fmt"
	"math/big"
	"path/filepath"
	"runtime"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/accounts/keystore"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/node"
	"github.com/ethereum/go-ethereum/params"
	homedir "github.com/mitchellh/go-homedir"
)

// ObtainWallet fetches the wallet for a given address
func ObtainWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	wallet, err := obtainGethWallet(chainID, address)
	if err == nil {
		return wallet, nil
	}

	wallet, err = obtainParityWallet(chainID, address)
	if err == nil {
		return wallet, err
	}

	return wallet, fmt.Errorf("Failed to obtain wallet")
}

func obtainGethWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	keydir := node.DefaultDataDir()
	if chainID.Cmp(params.MainnetChainConfig.ChainID) == 0 {
		// Nothing to add for mainnet
	} else if chainID.Cmp(params.TestnetChainConfig.ChainID) == 0 {
		keydir = filepath.Join(keydir, "testnet")
	} else if chainID.Cmp(params.RinkebyChainConfig.ChainID) == 0 {
		keydir = filepath.Join(keydir, "rinkeby")
	}
	keydir = filepath.Join(keydir, "keystore")
	backends := []accounts.Backend{keystore.NewKeyStore(keydir, keystore.StandardScryptN, keystore.StandardScryptP)}
	accountManager := accounts.NewManager(backends...)
	defer accountManager.Close()
	account := accounts.Account{Address: address}
	wallet, err := accountManager.Find(account)
	return wallet, err
}

func obtainParityWallet(chainID *big.Int, address common.Address) (accounts.Wallet, error) {
	keydir, err := homedir.Dir()
	if err != nil {
		return nil, fmt.Errorf("Failed to find home directory")
	}
	if runtime.GOOS == "windows" {
		keydir = filepath.Join(keydir, "AppData\\Roaming\\Parity\\Ethereum\\keys")
	} else if runtime.GOOS == "darwin" {
		keydir = filepath.Join(keydir, "Library/Application Support/io.parity.ethereum/keys")
	} else if runtime.GOOS == "linux" {
		keydir = filepath.Join(keydir, ".local/share/io.parity.ethereum/keys")
	} else {
		return nil, fmt.Errorf("Unsupported operating system")
	}

	if chainID.Cmp(params.MainnetChainConfig.ChainID) == 0 {
		keydir = filepath.Join(keydir, "ethereum")
	} else if chainID.Cmp(params.TestnetChainConfig.ChainID) == 0 {
		keydir = filepath.Join(keydir, "test")
	}

	backends := []accounts.Backend{keystore.NewKeyStore(keydir, keystore.StandardScryptN, keystore.StandardScryptP)}
	accountManager := accounts.NewManager(backends...)
	defer accountManager.Close()
	account := accounts.Account{Address: address}
	wallet, err := accountManager.Find(account)
	return wallet, err
}

// ObtainAccount fetches the account for a given address
func ObtainAccount(wallet *accounts.Wallet, address *common.Address, passphrase string) (*accounts.Account, error) {
	for _, account := range (*wallet).Accounts() {
		if *address == account.Address {
			if !VerifyPassphrase(*wallet, account, passphrase) {
				return nil, errors.New("invalid passphrase")
			}
			return &account, nil
		}
	}
	return nil, errors.New("account not found")
}

// VerifyPassphrase confirms that a passphrase is correct for an account
func VerifyPassphrase(wallet accounts.Wallet, account accounts.Account, passphrase string) bool {

	_, err := wallet.SignHashWithPassphrase(account, passphrase, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
	return err == nil
}
<|endoftext|>


<commit_before>package clickhouse

import (
	"bufio"
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"net"
	"reflect"
	"regexp"
	"sync"
	"time"

	"github.com/kshvakov/clickhouse/lib/binary"
	"github.com/kshvakov/clickhouse/lib/data"
	"github.com/kshvakov/clickhouse/lib/protocol"
	"github.com/kshvakov/clickhouse/lib/types"
)

var (
	ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
	ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
)

var (
	splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
)

type logger func(format string, v ...interface{})

type clickhouse struct {
	sync.Mutex
	data.ServerInfo
	data.ClientInfo
	logf          logger
	conn          *connect
	block         *data.Block
	buffer        *bufio.Writer
	decoder       *binary.Decoder
	encoder       *binary.Encoder
	compress      bool
	blockSize     int
	inTransaction bool
	readTimeout   time.Duration
	writeTimeout  time.Duration
}

func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
	return ch.prepareContext(context.Background(), query)
}

func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return ch.prepareContext(ctx, query)
}

func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	ch.logf("[prepare] %s", query)
	switch {
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	case ch.block != nil:
		return nil, ErrLimitDataRequestInTx
	case isInsert(query):
		if !ch.inTransaction {
			return nil, ErrInsertInNotBatchMode
		}
		return ch.insert(query)
	}
	return &stmt{
		ch:       ch,
		query:    query,
		numInput: numInput(query),
	}, nil
}

func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
	if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
		return nil, err
	}
	if ch.block, err = ch.readMeta(); err != nil {
		return nil, err
	}
	return &stmt{
		ch:       ch,
		isInsert: true,
	}, nil
}

func (ch *clickhouse) Begin() (driver.Tx, error) {
	return ch.beginTx(context.Background(), txOptions{})
}

func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	return ch.beginTx(ctx, txOptions{
		Isolation: int(opts.Isolation),
		ReadOnly:  opts.ReadOnly,
	})
}

type txOptions struct {
	Isolation int
	ReadOnly  bool
}

func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
	ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	switch {
	case ch.inTransaction:
		return nil, sql.ErrTxDone
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	}
	if finish := ch.watchCancel(ctx); finish != nil {
		defer finish()
	}
	ch.block = nil
	ch.inTransaction = true
	return ch, nil
}

func (ch *clickhouse) Commit() error {
	ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	defer func() {
		if ch.block != nil {
			ch.block.Reset()
			ch.block = nil
		}
		ch.inTransaction = false
	}()
	switch {
	case !ch.inTransaction:
		return sql.ErrTxDone
	case ch.conn.closed:
		return driver.ErrBadConn
	}
	if ch.block != nil {
		if err := ch.writeBlock(ch.block); err != nil {
			return err
		}
		// Send empty block as marker of end of data.
		if err := ch.writeBlock(&data.Block{}); err != nil {
			return err
		}
		if err := ch.buffer.Flush(); err != nil {
			return err
		}
		return ch.process()
	}
	return nil
}

func (ch *clickhouse) Rollback() error {
	ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	if !ch.inTransaction {
		return sql.ErrTxDone
	}
	ch.block = nil
	ch.inTransaction = false
	return nil
}

func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
	switch v := nv.Value.(type) {
	case Date:
		nv.Value = v.convert()
	case DateTime:
		nv.Value = v.convert()
	case IP, *types.Array, UUID:
		nv.Value = v
	case net.IP:
		nv.Value = IP(v)
	default:
		switch value := reflect.ValueOf(nv.Value); value.Kind() {
		case reflect.Bool:
			nv.Value = uint8(0)
			if value.Bool() {
				nv.Value = uint8(1)
			}
		case reflect.Int8:
			nv.Value = int8(value.Int())
		case reflect.Int16:
			nv.Value = int16(value.Int())
		case reflect.Int32:
			nv.Value = int32(value.Int())
		case reflect.Int64:
			nv.Value = value.Int()
		case reflect.Uint8:
			nv.Value = uint8(value.Uint())
		case reflect.Uint16:
			nv.Value = uint16(value.Uint())
		case reflect.Uint32:
			nv.Value = uint32(value.Uint())
		case reflect.Uint64:
			nv.Value = uint64(value.Uint())
		case reflect.Float32:
			nv.Value = float32(value.Float())
		case reflect.Float64:
			nv.Value = float64(value.Float())
		case reflect.String:
			nv.Value = value.String()
		case reflect.Slice:
			nv.Value = types.NewArray(value.Interface())
		}
	}
	return nil
}

func (ch *clickhouse) Close() error {
	ch.block = nil
	return ch.conn.Close()
}

func (ch *clickhouse) process() error {
	packet, err := ch.decoder.Uvarint()
	if err != nil {
		return err
	}
	for {
		switch packet {
		case protocol.ServerPong:
			ch.logf("[process] <- pong")
			return nil
		case protocol.ServerException:
			ch.logf("[process] <- exception")
			return ch.exception()
		case protocol.ServerProgress:
			progress, err := ch.progress()
			if err != nil {
				return err
			}
			ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
				progress.rows,
				progress.bytes,
				progress.totalRows,
			)
		case protocol.ServerProfileInfo:
			profileInfo, err := ch.profileInfo()
			if err != nil {
				return err
			}
			ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
		case protocol.ServerData:
			block, err := ch.readBlock()
			if err != nil {
				return err
			}
			ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
		case protocol.ServerEndOfStream:
			ch.logf("[process] <- end of stream")
			return nil
		default:
			ch.conn.Close()
			return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
		}
		if packet, err = ch.decoder.Uvarint(); err != nil {
			return err
		}
	}
}

func (ch *clickhouse) cancel() error {
	ch.logf("[cancel request]")
	if err := ch.encoder.Uvarint(protocol.ClientCancel); err != nil {
		return err
	}
	return ch.conn.Close()
}

func (ch *clickhouse) watchCancel(ctx context.Context) func() {
	if done := ctx.Done(); done != nil {
		finished := make(chan struct{})
		go func() {
			select {
			case <-done:
				ch.cancel()
				finished <- struct{}{}
				ch.logf("[cancel] <- done")
			case <-finished:
				ch.logf("[cancel] <- finished")
			}
		}()
		return func() {
			select {
			case <-finished:
			case finished <- struct{}{}:
			}
		}
	}
	return nil
}
<commit_msg>CheckNamedValue: minor improvements<commit_after>package clickhouse

import (
	"bufio"
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"net"
	"reflect"
	"regexp"
	"sync"
	"time"

	"github.com/kshvakov/clickhouse/lib/binary"
	"github.com/kshvakov/clickhouse/lib/data"
	"github.com/kshvakov/clickhouse/lib/protocol"
	"github.com/kshvakov/clickhouse/lib/types"
)

var (
	ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
	ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
)

var (
	splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
)

type logger func(format string, v ...interface{})

type clickhouse struct {
	sync.Mutex
	data.ServerInfo
	data.ClientInfo
	logf          logger
	conn          *connect
	block         *data.Block
	buffer        *bufio.Writer
	decoder       *binary.Decoder
	encoder       *binary.Encoder
	compress      bool
	blockSize     int
	inTransaction bool
	readTimeout   time.Duration
	writeTimeout  time.Duration
}

func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
	return ch.prepareContext(context.Background(), query)
}

func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return ch.prepareContext(ctx, query)
}

func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	ch.logf("[prepare] %s", query)
	switch {
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	case ch.block != nil:
		return nil, ErrLimitDataRequestInTx
	case isInsert(query):
		if !ch.inTransaction {
			return nil, ErrInsertInNotBatchMode
		}
		return ch.insert(query)
	}
	return &stmt{
		ch:       ch,
		query:    query,
		numInput: numInput(query),
	}, nil
}

func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
	if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
		return nil, err
	}
	if ch.block, err = ch.readMeta(); err != nil {
		return nil, err
	}
	return &stmt{
		ch:       ch,
		isInsert: true,
	}, nil
}

func (ch *clickhouse) Begin() (driver.Tx, error) {
	return ch.beginTx(context.Background(), txOptions{})
}

func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	return ch.beginTx(ctx, txOptions{
		Isolation: int(opts.Isolation),
		ReadOnly:  opts.ReadOnly,
	})
}

type txOptions struct {
	Isolation int
	ReadOnly  bool
}

func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
	ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	switch {
	case ch.inTransaction:
		return nil, sql.ErrTxDone
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	}
	if finish := ch.watchCancel(ctx); finish != nil {
		defer finish()
	}
	ch.block = nil
	ch.inTransaction = true
	return ch, nil
}

func (ch *clickhouse) Commit() error {
	ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	defer func() {
		if ch.block != nil {
			ch.block.Reset()
			ch.block = nil
		}
		ch.inTransaction = false
	}()
	switch {
	case !ch.inTransaction:
		return sql.ErrTxDone
	case ch.conn.closed:
		return driver.ErrBadConn
	}
	if ch.block != nil {
		if err := ch.writeBlock(ch.block); err != nil {
			return err
		}
		// Send empty block as marker of end of data.
		if err := ch.writeBlock(&data.Block{}); err != nil {
			return err
		}
		if err := ch.buffer.Flush(); err != nil {
			return err
		}
		return ch.process()
	}
	return nil
}

func (ch *clickhouse) Rollback() error {
	ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	if !ch.inTransaction {
		return sql.ErrTxDone
	}
	ch.block = nil
	ch.inTransaction = false
	return nil
}

func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
	switch v := nv.Value.(type) {
	case IP, *types.Array, UUID:
		return nil
	case int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, string, time.Time:
		return nil
	case
		[]int, []int8, []int16, []int32, []int64,
		[]uint, []uint8, []uint16, []uint32, []uint64,
		[]float32, []float64,
		[]string:
		nv.Value = types.NewArray(v)
	case Date:
		nv.Value = v.convert()
	case DateTime:
		nv.Value = v.convert()
	case net.IP:
		nv.Value = IP(v)
	default:
		switch value := reflect.ValueOf(nv.Value); value.Kind() {
		case reflect.Bool:
			nv.Value = uint8(0)
			if value.Bool() {
				nv.Value = uint8(1)
			}
		case reflect.Int8:
			nv.Value = int8(value.Int())
		case reflect.Int16:
			nv.Value = int16(value.Int())
		case reflect.Int32:
			nv.Value = int32(value.Int())
		case reflect.Int64:
			nv.Value = value.Int()
		case reflect.Uint8:
			nv.Value = uint8(value.Uint())
		case reflect.Uint16:
			nv.Value = uint16(value.Uint())
		case reflect.Uint32:
			nv.Value = uint32(value.Uint())
		case reflect.Uint64:
			nv.Value = uint64(value.Uint())
		case reflect.Float32:
			nv.Value = float32(value.Float())
		case reflect.Float64:
			nv.Value = float64(value.Float())
		case reflect.String:
			nv.Value = value.String()
		}
	}
	return nil
}

func (ch *clickhouse) Close() error {
	ch.block = nil
	return ch.conn.Close()
}

func (ch *clickhouse) process() error {
	packet, err := ch.decoder.Uvarint()
	if err != nil {
		return err
	}
	for {
		switch packet {
		case protocol.ServerPong:
			ch.logf("[process] <- pong")
			return nil
		case protocol.ServerException:
			ch.logf("[process] <- exception")
			return ch.exception()
		case protocol.ServerProgress:
			progress, err := ch.progress()
			if err != nil {
				return err
			}
			ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
				progress.rows,
				progress.bytes,
				progress.totalRows,
			)
		case protocol.ServerProfileInfo:
			profileInfo, err := ch.profileInfo()
			if err != nil {
				return err
			}
			ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
		case protocol.ServerData:
			block, err := ch.readBlock()
			if err != nil {
				return err
			}
			ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
		case protocol.ServerEndOfStream:
			ch.logf("[process] <- end of stream")
			return nil
		default:
			ch.conn.Close()
			return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
		}
		if packet, err = ch.decoder.Uvarint(); err != nil {
			return err
		}
	}
}

func (ch *clickhouse) cancel() error {
	ch.logf("[cancel request]")
	if err := ch.encoder.Uvarint(protocol.ClientCancel); err != nil {
		return err
	}
	return ch.conn.Close()
}

func (ch *clickhouse) watchCancel(ctx context.Context) func() {
	if done := ctx.Done(); done != nil {
		finished := make(chan struct{})
		go func() {
			select {
			case <-done:
				ch.cancel()
				finished <- struct{}{}
				ch.logf("[cancel] <- done")
			case <-finished:
				ch.logf("[cancel] <- finished")
			}
		}()
		return func() {
			select {
			case <-finished:
			case finished <- struct{}{}:
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package coprhd

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"

	log "github.com/Sirupsen/logrus"

	"github.com/portworx/kvdb"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"

	napping "gopkg.in/jmcvetta/napping.v3"
)

const (
	Name = "coprhd"
	Type = api.Block

	// URI_LOGIN path to create a authentication token
	URI_LOGIN = "login.json"
	// URI_LOGIN path to create volume
	URI_CREATE_VOL = "block/volumes.json"
	// URI_EXPORT_VOL path to export a volume
	URI_EXPORT_VOL = "block/export.json"
	// URI_TPL_DEL_VOL template path to delete/deactivate a volume
	URI_TPL_DEL_VOL = "block/volumes/%s/deactivate.json"
	// URL_TPL_NEW_SNAP path to create a volume snapshot
	URL_TPL_NEW_SNAP = "block/volumes/%s/protections/snapshots.json"
	// URI_TPL_UNEXP_VOL path template to remove a volume export
	URI_TPL_UNEXP_VOL = "block/export/%s/deactivate.json"
)

type (
	driver struct {
		*volume.IoNotSupported
		*volume.DefaultEnumerator
		consistency_group string
		project           string
		varray            string
		vpool             string
		url               string
		httpClient        *http.Client
		creds             *url.Userinfo
	}

	// ApiError represents the default api error code
	ApiError struct {
		Code        string `json:"code"`
		Retryable   string `json:"retryable"`
		Description string `json:"description"`
		Details     string `json:"details"`
	}

	// CreateVolumeArgs represents the json parameters for the create volume REST call
	CreateVolumeArgs struct {
		ConsistencyGroup string `json:"consistency_group"`
		Count            int    `json:"count"`
		Name             string `json:"name"`
		Project          string `json:"project"`
		Size             string `json:"size"`
		VArray           string `json:"varray"`
		VPool            string `json:"vpool"`
	}

	// CreateVolumeReply is the reply from the create volume REST call
	CreateVolumeReply struct {
		Task []struct {
			Resource struct {
				Name string       `json:"name"`
				Id   api.VolumeID `json:"id"`
			} `json:"resource"`
		} `json:"task"`
	}
)

func Init(params volume.DriverParams) (volume.VolumeDriver, error) {

	restUrl, ok := params["restUrl"]
	if !ok {
		return nil, fmt.Errorf("rest api 'url' configuration parameter must be set")
	}

	user, ok := params["user"]
	if !ok {
		return nil, fmt.Errorf("rest auth 'user' must be set")
	}

	pass, ok := params["password"]
	if !ok {
		return nil, fmt.Errorf("rest auth 'password' must be set")
	}

	consistency_group, ok := params["consistency_group"]
	if !ok {
		return nil, fmt.Errorf("'consistency_group' configuration parameter must be set")
	}

	project, ok := params["project"]
	if !ok {
		return nil, fmt.Errorf("'project' configuration parameter must be set")
	}

	varray, ok := params["varray"]
	if !ok {
		return nil, fmt.Errorf("'varray' configuration parameter must be set")
	}

	vpool, ok := params["vpool"]
	if !ok {
		return nil, fmt.Errorf("'vpool' configuration parameter must be set")
	}

	d := &driver{
		DefaultEnumerator: volume.NewDefaultEnumerator(Name, kvdb.Instance()),
		consistency_group: consistency_group,
		project:           project,
		varray:            varray,
		vpool:             vpool,
		url:               restUrl,
		creds:             url.UserPassword(user, pass),
		httpClient: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}

	return d, nil
}

func (d *driver) String() string {
	return Name
}

func (d *driver) Type() api.DriverType {
	return Type
}

func init() {
	volume.Register(Name, Init)
}

func (d *driver) Create(
	locator api.VolumeLocator,
	source *api.Source,
	spec *api.VolumeSpec) (api.VolumeID, error) {

	s, err := d.getAuthSession()

	if err != nil {
		log.Errorf("Failed to create session: %s", err.Error())
		return api.BadVolumeID, err
	}

	e := ApiError{}

	res := &CreateVolumeReply{}

	sz := int64(spec.Size / (1024 * 1024 * 1000))

	payload := CreateVolumeArgs{
		"Default",                 // ConsistencyGroup
		1,                         // Count
		locator.Name,              // Name
		d.project,                 // Project
		fmt.Sprintf("%.6fGB", sz), // Volume Size
		d.varray,                  // Virtual Block Array
		d.vpool,                   // Virtual Block Pool
	}

	url := d.url + URI_CREATE_VOL

	resp, err := s.Post(url, &payload, res, &e)

	if resp.Status() != http.StatusAccepted {

		return api.BadVolumeID, fmt.Errorf("Failed to create volume: %s", resp.Status())
	}

	return res.Task[0].Resource.Id, err
}

func (d *driver) Delete(volumeID api.VolumeID) error {
	return nil
}

func (d *driver) Stats(volumeID api.VolumeID) (api.Stats, error) {
	return api.Stats{}, volume.ErrNotSupported
}

func (d *driver) Alerts(volumeID api.VolumeID) (api.Alerts, error) {
	return api.Alerts{}, volume.ErrNotSupported
}

func (d *driver) Attach(volumeID api.VolumeID) (path string, err error) {
	return "", nil
}

func (d *driver) Detach(volumeID api.VolumeID) error {
	return nil
}

func (d *driver) Mount(volumeID api.VolumeID, mountpath string) error {
	return nil
}

func (d *driver) Unmount(volumeID api.VolumeID, mountpath string) error {

	return nil
}

func (d *driver) Set(volumeID api.VolumeID, locator *api.VolumeLocator, spec *api.VolumeSpec) error {
	return volume.ErrNotSupported
}

func (d *driver) Shutdown() {
	log.Infof("%s Shutting down", Name)
}

func (d *driver) Snapshot(volumeID api.VolumeID, readonly bool, locator api.VolumeLocator) (api.VolumeID, error) {
	return "", nil
}

func (v *driver) Status() [][2]string {
	return [][2]string{}
}

// getAuthSession returns an authenticated API Session
func (d *driver) getAuthSession() (session *napping.Session, err error) {

	e := ApiError{}

	s := napping.Session{
		Userinfo: d.creds,
		Client:   d.httpClient,
	}

	url := d.url + URI_LOGIN

	resp, err := s.Get(url, nil, nil, &e)

	if err != nil {
		return
	}

	token := resp.HttpResponse().Header.Get("X-SDS-AUTH-TOKEN")

	h := http.Header{}

	h.Set("X-SDS-AUTH-TOKEN", token)

	session = &napping.Session{
		Client: d.httpClient,
		Header: &h,
	}

	return
}
<commit_msg>Moving init to top of source<commit_after>package coprhd

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"

	log "github.com/Sirupsen/logrus"

	"github.com/portworx/kvdb"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"

	napping "gopkg.in/jmcvetta/napping.v3"
)

const (
	Name = "coprhd"
	Type = api.Block

	// URI_LOGIN path to create a authentication token
	URI_LOGIN = "login.json"
	// URI_LOGIN path to create volume
	URI_CREATE_VOL = "block/volumes.json"
	// URI_EXPORT_VOL path to export a volume
	URI_EXPORT_VOL = "block/export.json"
	// URI_TPL_DEL_VOL template path to delete/deactivate a volume
	URI_TPL_DEL_VOL = "block/volumes/%s/deactivate.json"
	// URL_TPL_NEW_SNAP path to create a volume snapshot
	URL_TPL_NEW_SNAP = "block/volumes/%s/protections/snapshots.json"
	// URI_TPL_UNEXP_VOL path template to remove a volume export
	URI_TPL_UNEXP_VOL = "block/export/%s/deactivate.json"
)

type (
	driver struct {
		*volume.IoNotSupported
		*volume.DefaultEnumerator
		consistency_group string
		project           string
		varray            string
		vpool             string
		url               string
		httpClient        *http.Client
		creds             *url.Userinfo
	}

	// ApiError represents the default api error code
	ApiError struct {
		Code        string `json:"code"`
		Retryable   string `json:"retryable"`
		Description string `json:"description"`
		Details     string `json:"details"`
	}

	// CreateVolumeArgs represents the json parameters for the create volume REST call
	CreateVolumeArgs struct {
		ConsistencyGroup string `json:"consistency_group"`
		Count            int    `json:"count"`
		Name             string `json:"name"`
		Project          string `json:"project"`
		Size             string `json:"size"`
		VArray           string `json:"varray"`
		VPool            string `json:"vpool"`
	}

	// CreateVolumeReply is the reply from the create volume REST call
	CreateVolumeReply struct {
		Task []struct {
			Resource struct {
				Name string       `json:"name"`
				Id   api.VolumeID `json:"id"`
			} `json:"resource"`
		} `json:"task"`
	}
)

func init() {
	volume.Register(Name, Init)
}

func Init(params volume.DriverParams) (volume.VolumeDriver, error) {

	restUrl, ok := params["restUrl"]
	if !ok {
		return nil, fmt.Errorf("rest api 'url' configuration parameter must be set")
	}

	user, ok := params["user"]
	if !ok {
		return nil, fmt.Errorf("rest auth 'user' must be set")
	}

	pass, ok := params["password"]
	if !ok {
		return nil, fmt.Errorf("rest auth 'password' must be set")
	}

	consistency_group, ok := params["consistency_group"]
	if !ok {
		return nil, fmt.Errorf("'consistency_group' configuration parameter must be set")
	}

	project, ok := params["project"]
	if !ok {
		return nil, fmt.Errorf("'project' configuration parameter must be set")
	}

	varray, ok := params["varray"]
	if !ok {
		return nil, fmt.Errorf("'varray' configuration parameter must be set")
	}

	vpool, ok := params["vpool"]
	if !ok {
		return nil, fmt.Errorf("'vpool' configuration parameter must be set")
	}

	d := &driver{
		DefaultEnumerator: volume.NewDefaultEnumerator(Name, kvdb.Instance()),
		consistency_group: consistency_group,
		project:           project,
		varray:            varray,
		vpool:             vpool,
		url:               restUrl,
		creds:             url.UserPassword(user, pass),
		httpClient: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}

	return d, nil
}

func (d *driver) String() string {
	return Name
}

func (d *driver) Type() api.DriverType {
	return Type
}

func (d *driver) Create(
	locator api.VolumeLocator,
	source *api.Source,
	spec *api.VolumeSpec) (api.VolumeID, error) {

	s, err := d.getAuthSession()

	if err != nil {
		log.Errorf("Failed to create session: %s", err.Error())
		return api.BadVolumeID, err
	}

	e := ApiError{}

	res := &CreateVolumeReply{}

	sz := int64(spec.Size / (1024 * 1024 * 1000))

	payload := CreateVolumeArgs{
		"Default",                 // ConsistencyGroup
		1,                         // Count
		locator.Name,              // Name
		d.project,                 // Project
		fmt.Sprintf("%.6fGB", sz), // Volume Size
		d.varray,                  // Virtual Block Array
		d.vpool,                   // Virtual Block Pool
	}

	url := d.url + URI_CREATE_VOL

	resp, err := s.Post(url, &payload, res, &e)

	if resp.Status() != http.StatusAccepted {

		return api.BadVolumeID, fmt.Errorf("Failed to create volume: %s", resp.Status())
	}

	return res.Task[0].Resource.Id, err
}

func (d *driver) Delete(volumeID api.VolumeID) error {
	return nil
}

func (d *driver) Stats(volumeID api.VolumeID) (api.Stats, error) {
	return api.Stats{}, volume.ErrNotSupported
}

func (d *driver) Alerts(volumeID api.VolumeID) (api.Alerts, error) {
	return api.Alerts{}, volume.ErrNotSupported
}

func (d *driver) Attach(volumeID api.VolumeID) (path string, err error) {
	return "", nil
}

func (d *driver) Detach(volumeID api.VolumeID) error {
	return nil
}

func (d *driver) Mount(volumeID api.VolumeID, mountpath string) error {
	return nil
}

func (d *driver) Unmount(volumeID api.VolumeID, mountpath string) error {

	return nil
}

func (d *driver) Set(volumeID api.VolumeID, locator *api.VolumeLocator, spec *api.VolumeSpec) error {
	return volume.ErrNotSupported
}

func (d *driver) Shutdown() {
	log.Infof("%s Shutting down", Name)
}

func (d *driver) Snapshot(volumeID api.VolumeID, readonly bool, locator api.VolumeLocator) (api.VolumeID, error) {
	return "", nil
}

func (v *driver) Status() [][2]string {
	return [][2]string{}
}

// getAuthSession returns an authenticated API Session
func (d *driver) getAuthSession() (session *napping.Session, err error) {

	e := ApiError{}

	s := napping.Session{
		Userinfo: d.creds,
		Client:   d.httpClient,
	}

	url := d.url + URI_LOGIN

	resp, err := s.Get(url, nil, nil, &e)

	if err != nil {
		return
	}

	token := resp.HttpResponse().Header.Get("X-SDS-AUTH-TOKEN")

	h := http.Header{}

	h.Set("X-SDS-AUTH-TOKEN", token)

	session = &napping.Session{
		Client: d.httpClient,
		Header: &h,
	}

	return
}
<|endoftext|>