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

// Daemon agentd holds a private key in memory and makes it available to a
// subprocess via the agent protocol.
package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"

	"v.io/v23"
	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/x/lib/vlog"
	"v.io/x/ref/envvar"
	vsecurity "v.io/x/ref/lib/security"
	vsignals "v.io/x/ref/lib/signals"
	"v.io/x/ref/services/agent/internal/server"

	_ "v.io/x/ref/profiles"
)

const childAgentFd = 3
const pkgPath = "v.io/x/ref/services/agent/agentd"

var (
	errCantReadPassphrase       = verror.Register(pkgPath+".errCantReadPassphrase", verror.NoRetry, "{1:}{2:} failed to read passphrase{:_}")
	errNeedPassphrase           = verror.Register(pkgPath+".errNeedPassphrase", verror.NoRetry, "{1:}{2:} Passphrase required for decrypting principal{:_}")
	errCantParseRestartExitCode = verror.Register(pkgPath+".errCantParseRestartExitCode", verror.NoRetry, "{1:}{2:} Failed to parse restart exit code{:_}")
)

var (
	keypath      = flag.String("additional-principals", "", "If non-empty, allow for the creation of new principals and save them in this directory.")
	noPassphrase = flag.Bool("no-passphrase", false, "If true, user will not be prompted for principal encryption passphrase.")

	// TODO(caprita): We use the exit code of the child to determine if the
	// agent should restart it.  Consider changing this to use the unix
	// socket for this purpose.
	restartExitCode = flag.String("restart-exit-code", "", "If non-empty, will restart the command when it exits, provided that the command's exit code matches the value of this flag.  The value must be an integer, or an integer preceded by '!' (in which case all exit codes except the flag will trigger a restart.")
)

func main() {
	os.Exit(Main())
}

func Main() int {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage: %s [agent options] command command_args...

Loads the private key specified in privatekey.pem in %v into memory, then
starts the specified command with access to the private key via the
agent protocol instead of directly reading from disk.

`, os.Args[0], envvar.Credentials)
		flag.PrintDefaults()
	}
	flag.Parse()
	if len(flag.Args()) < 1 {
		fmt.Fprintln(os.Stderr, "Need at least one argument.")
		flag.Usage()
		return 1
	}
	var restartOpts restartOptions
	if err := restartOpts.parse(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		flag.Usage()
		return 1
	}

	// This is a bit tricky. We're trying to share the runtime's
	// v23.credentials flag.  However we need to parse it before
	// creating the runtime.  We depend on the profile's init() function
	// calling flags.CreateAndRegister(flag.CommandLine, flags.Runtime)
	// This will read the envvar.Credentials env var, then our call to
	// flag.Parse() will take any override passed on the command line.
	var dir string
	if f := flag.Lookup("v23.credentials").Value; true {
		dir = f.String()
		// Clear out the flag value to prevent v23.Init from
		// trying to load this password protected principal.
		f.Set("")
	}
	if len(dir) == 0 {
		vlog.Fatalf("The %v environment variable must be set to a directory: %q", envvar.Credentials, os.Getenv(envvar.Credentials))
	}

	p, passphrase, err := newPrincipalFromDir(dir)
	if err != nil {
		vlog.Fatalf("failed to create new principal from dir(%s): %v", dir, err)
	}

	// Clear out the environment variable before v23.Init.
	if err = envvar.ClearCredentials(); err != nil {
		vlog.Fatalf("envvar.ClearCredentials: %v", err)
	}
	ctx, shutdown := v23.Init()
	defer shutdown()

	if ctx, err = v23.WithPrincipal(ctx, p); err != nil {
		vlog.Panic("failed to set principal for ctx: %v", err)
	}

	if *keypath == "" && passphrase != nil {
		// If we're done with the passphrase, zero it out so it doesn't stay in memory
		for i := range passphrase {
			passphrase[i] = 0
		}
		passphrase = nil
	}

	// Start running our server.
	var sock, mgrSock *os.File
	var endpoint string
	if sock, endpoint, err = server.RunAnonymousAgent(ctx, p, childAgentFd); err != nil {
		vlog.Fatalf("RunAnonymousAgent: %v", err)
	}
	if err = os.Setenv(envvar.AgentEndpoint, endpoint); err != nil {
		vlog.Fatalf("setenv: %v", err)
	}

	if *keypath != "" {
		if mgrSock, err = server.RunKeyManager(ctx, *keypath, passphrase); err != nil {
			vlog.Fatalf("RunKeyManager: %v", err)
		}
	}

	exitCode := 0
	for {
		// Run the client and wait for it to finish.
		cmd := exec.Command(flag.Args()[0], flag.Args()[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.ExtraFiles = []*os.File{sock}

		if mgrSock != nil {
			cmd.ExtraFiles = append(cmd.ExtraFiles, mgrSock)
		}

		err = cmd.Start()
		if err != nil {
			vlog.Fatalf("Error starting child: %v", err)
		}
		shutdown := make(chan struct{})
		go func() {
			select {
			case sig := <-vsignals.ShutdownOnSignals(ctx):
				// TODO(caprita): Should we also relay double
				// signal to the child?  That currently just
				// force exits the current process.
				if sig == vsignals.STOP {
					sig = syscall.SIGTERM
				}
				cmd.Process.Signal(sig)
			case <-shutdown:
			}
		}()
		cmd.Wait()
		close(shutdown)
		exitCode = cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
		if !restartOpts.restart(exitCode) {
			break
		}
	}
	// TODO(caprita): If restartOpts.enabled is false, we could close these
	// right after cmd.Start().
	sock.Close()
	mgrSock.Close()
	return exitCode
}

func newPrincipalFromDir(dir string) (security.Principal, []byte, error) {
	p, err := vsecurity.LoadPersistentPrincipal(dir, nil)
	if os.IsNotExist(err) {
		return handleDoesNotExist(dir)
	}
	if verror.ErrorID(err) == vsecurity.ErrBadPassphrase.ID {
		return handlePassphrase(dir)
	}
	return p, nil, err
}

func handleDoesNotExist(dir string) (security.Principal, []byte, error) {
	fmt.Println("Private key file does not exist. Creating new private key...")
	var pass []byte
	if !*noPassphrase {
		var err error
		if pass, err = getPassword("Enter passphrase (entering nothing will store unencrypted): "); err != nil {
			return nil, nil, verror.New(errCantReadPassphrase, nil, err)
		}
	}
	p, err := vsecurity.CreatePersistentPrincipal(dir, pass)
	if err != nil {
		return nil, pass, err
	}
	vsecurity.InitDefaultBlessings(p, "agent_principal")
	return p, pass, nil
}

func handlePassphrase(dir string) (security.Principal, []byte, error) {
	if *noPassphrase {
		return nil, nil, verror.New(errNeedPassphrase, nil)
	}
	pass, err := getPassword("Private key file is encrypted. Please enter passphrase.\nEnter passphrase: ")
	if err != nil {
		return nil, nil, verror.New(errCantReadPassphrase, nil, err)
	}
	p, err := vsecurity.LoadPersistentPrincipal(dir, pass)
	return p, pass, err
}

func getPassword(prompt string) ([]byte, error) {
	if !terminal.IsTerminal(int(os.Stdin.Fd())) {
		// If the standard input is not a terminal, the password is obtained by reading a line from it.
		return readPassword()
	}
	fmt.Printf(prompt)
	stop := make(chan bool)
	defer close(stop)
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	if err != nil {
		return nil, err
	}
	go catchTerminationSignals(stop, state)
	defer fmt.Printf("\n")
	return terminal.ReadPassword(int(os.Stdin.Fd()))
}

// readPassword reads form Stdin until it sees '\n' or EOF.
func readPassword() ([]byte, error) {
	var pass []byte
	var total int
	for {
		b := make([]byte, 1)
		count, err := os.Stdin.Read(b)
		if err != nil && err != io.EOF {
			return nil, err
		}
		if err == io.EOF || b[0] == '\n' {
			return pass[:total], nil
		}
		total += count
		pass = secureAppend(pass, b)
	}
}

func secureAppend(s, t []byte) []byte {
	res := append(s, t...)
	if len(res) > cap(s) {
		// When append needs to allocate a new array, clear out the old one.
		for i := range s {
			s[i] = '0'
		}
	}
	// Clear out the second array.
	for i := range t {
		t[i] = '0'
	}
	return res
}

// catchTerminationSignals catches signals to allow us to turn terminal echo back on.
func catchTerminationSignals(stop <-chan bool, state *terminal.State) {
	var successErrno syscall.Errno
	sig := make(chan os.Signal, 4)
	// Catch the blockable termination signals.
	signal.Notify(sig, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGHUP)
	select {
	case <-sig:
		// Start on new line in terminal.
		fmt.Printf("\n")
		if err := terminal.Restore(int(os.Stdin.Fd()), state); err != successErrno {
			vlog.Errorf("Failed to restore terminal state (%v), you words may not show up when you type, enter 'stty echo' to fix this.", err)
		}
		os.Exit(-1)
	case <-stop:
		signal.Stop(sig)
	}
}

type restartOptions struct {
	enabled, unless bool
	code            int
}

func (opts *restartOptions) parse() error {
	code := *restartExitCode
	if code == "" {
		return nil
	}
	opts.enabled = true
	if code[0] == '!' {
		opts.unless = true
		code = code[1:]
	}
	var err error
	if opts.code, err = strconv.Atoi(code); err != nil {
		return verror.New(errCantParseRestartExitCode, nil, err)
	}
	return nil
}

func (opts *restartOptions) restart(exitCode int) bool {
	return opts.enabled && opts.unless != (exitCode == opts.code)
}
<commit_msg>services/agent/agentd: Be able to choose a name other than "agent_principal" when creating a new principal.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Daemon agentd holds a private key in memory and makes it available to a
// subprocess via the agent protocol.
package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"

	"v.io/v23"
	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/x/lib/vlog"
	"v.io/x/ref/envvar"
	vsecurity "v.io/x/ref/lib/security"
	vsignals "v.io/x/ref/lib/signals"
	"v.io/x/ref/services/agent/internal/server"

	_ "v.io/x/ref/profiles"
)

const childAgentFd = 3
const pkgPath = "v.io/x/ref/services/agent/agentd"

var (
	errCantReadPassphrase       = verror.Register(pkgPath+".errCantReadPassphrase", verror.NoRetry, "{1:}{2:} failed to read passphrase{:_}")
	errNeedPassphrase           = verror.Register(pkgPath+".errNeedPassphrase", verror.NoRetry, "{1:}{2:} Passphrase required for decrypting principal{:_}")
	errCantParseRestartExitCode = verror.Register(pkgPath+".errCantParseRestartExitCode", verror.NoRetry, "{1:}{2:} Failed to parse restart exit code{:_}")
)

var (
	keypath      = flag.String("additional-principals", "", "If non-empty, allow for the creation of new principals and save them in this directory.")
	noPassphrase = flag.Bool("no-passphrase", false, "If true, user will not be prompted for principal encryption passphrase.")

	// TODO(caprita): We use the exit code of the child to determine if the
	// agent should restart it.  Consider changing this to use the unix
	// socket for this purpose.
	restartExitCode = flag.String("restart-exit-code", "", "If non-empty, will restart the command when it exits, provided that the command's exit code matches the value of this flag.  The value must be an integer, or an integer preceded by '!' (in which case all exit codes except the flag will trigger a restart.")

	newname = flag.String("new-principal-blessing-name", "", "If creating a new principal (--v23.credentials does not exist), then have it blessed with this name.")
)

func main() {
	os.Exit(Main())
}

func Main() int {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage: %s [agent options] command command_args...

Loads the private key specified in privatekey.pem in %v into memory, then
starts the specified command with access to the private key via the
agent protocol instead of directly reading from disk.

`, os.Args[0], envvar.Credentials)
		flag.PrintDefaults()
	}
	flag.Parse()
	if len(flag.Args()) < 1 {
		fmt.Fprintln(os.Stderr, "Need at least one argument.")
		flag.Usage()
		return 1
	}
	var restartOpts restartOptions
	if err := restartOpts.parse(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		flag.Usage()
		return 1
	}

	// This is a bit tricky. We're trying to share the runtime's
	// v23.credentials flag.  However we need to parse it before
	// creating the runtime.  We depend on the profile's init() function
	// calling flags.CreateAndRegister(flag.CommandLine, flags.Runtime)
	// This will read the envvar.Credentials env var, then our call to
	// flag.Parse() will take any override passed on the command line.
	var dir string
	if f := flag.Lookup("v23.credentials").Value; true {
		dir = f.String()
		// Clear out the flag value to prevent v23.Init from
		// trying to load this password protected principal.
		f.Set("")
	}
	if len(dir) == 0 {
		vlog.Fatalf("The %v environment variable must be set to a directory: %q", envvar.Credentials, os.Getenv(envvar.Credentials))
	}

	p, passphrase, err := newPrincipalFromDir(dir)
	if err != nil {
		vlog.Fatalf("failed to create new principal from dir(%s): %v", dir, err)
	}

	// Clear out the environment variable before v23.Init.
	if err = envvar.ClearCredentials(); err != nil {
		vlog.Fatalf("envvar.ClearCredentials: %v", err)
	}
	ctx, shutdown := v23.Init()
	defer shutdown()

	if ctx, err = v23.WithPrincipal(ctx, p); err != nil {
		vlog.Panic("failed to set principal for ctx: %v", err)
	}

	if *keypath == "" && passphrase != nil {
		// If we're done with the passphrase, zero it out so it doesn't stay in memory
		for i := range passphrase {
			passphrase[i] = 0
		}
		passphrase = nil
	}

	// Start running our server.
	var sock, mgrSock *os.File
	var endpoint string
	if sock, endpoint, err = server.RunAnonymousAgent(ctx, p, childAgentFd); err != nil {
		vlog.Fatalf("RunAnonymousAgent: %v", err)
	}
	if err = os.Setenv(envvar.AgentEndpoint, endpoint); err != nil {
		vlog.Fatalf("setenv: %v", err)
	}

	if *keypath != "" {
		if mgrSock, err = server.RunKeyManager(ctx, *keypath, passphrase); err != nil {
			vlog.Fatalf("RunKeyManager: %v", err)
		}
	}

	exitCode := 0
	for {
		// Run the client and wait for it to finish.
		cmd := exec.Command(flag.Args()[0], flag.Args()[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.ExtraFiles = []*os.File{sock}

		if mgrSock != nil {
			cmd.ExtraFiles = append(cmd.ExtraFiles, mgrSock)
		}

		err = cmd.Start()
		if err != nil {
			vlog.Fatalf("Error starting child: %v", err)
		}
		shutdown := make(chan struct{})
		go func() {
			select {
			case sig := <-vsignals.ShutdownOnSignals(ctx):
				// TODO(caprita): Should we also relay double
				// signal to the child?  That currently just
				// force exits the current process.
				if sig == vsignals.STOP {
					sig = syscall.SIGTERM
				}
				cmd.Process.Signal(sig)
			case <-shutdown:
			}
		}()
		cmd.Wait()
		close(shutdown)
		exitCode = cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
		if !restartOpts.restart(exitCode) {
			break
		}
	}
	// TODO(caprita): If restartOpts.enabled is false, we could close these
	// right after cmd.Start().
	sock.Close()
	mgrSock.Close()
	return exitCode
}

func newPrincipalFromDir(dir string) (security.Principal, []byte, error) {
	p, err := vsecurity.LoadPersistentPrincipal(dir, nil)
	if os.IsNotExist(err) {
		return handleDoesNotExist(dir)
	}
	if verror.ErrorID(err) == vsecurity.ErrBadPassphrase.ID {
		return handlePassphrase(dir)
	}
	return p, nil, err
}

func handleDoesNotExist(dir string) (security.Principal, []byte, error) {
	fmt.Println("Private key file does not exist. Creating new private key...")
	var pass []byte
	if !*noPassphrase {
		var err error
		if pass, err = getPassword("Enter passphrase (entering nothing will store unencrypted): "); err != nil {
			return nil, nil, verror.New(errCantReadPassphrase, nil, err)
		}
	}
	p, err := vsecurity.CreatePersistentPrincipal(dir, pass)
	if err != nil {
		return nil, pass, err
	}
	name := *newname
	if len(name) == 0 {
		name = "agent_principal"
	}
	vsecurity.InitDefaultBlessings(p, name)
	return p, pass, nil
}

func handlePassphrase(dir string) (security.Principal, []byte, error) {
	if *noPassphrase {
		return nil, nil, verror.New(errNeedPassphrase, nil)
	}
	pass, err := getPassword("Private key file is encrypted. Please enter passphrase.\nEnter passphrase: ")
	if err != nil {
		return nil, nil, verror.New(errCantReadPassphrase, nil, err)
	}
	p, err := vsecurity.LoadPersistentPrincipal(dir, pass)
	return p, pass, err
}

func getPassword(prompt string) ([]byte, error) {
	if !terminal.IsTerminal(int(os.Stdin.Fd())) {
		// If the standard input is not a terminal, the password is obtained by reading a line from it.
		return readPassword()
	}
	fmt.Printf(prompt)
	stop := make(chan bool)
	defer close(stop)
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	if err != nil {
		return nil, err
	}
	go catchTerminationSignals(stop, state)
	defer fmt.Printf("\n")
	return terminal.ReadPassword(int(os.Stdin.Fd()))
}

// readPassword reads form Stdin until it sees '\n' or EOF.
func readPassword() ([]byte, error) {
	var pass []byte
	var total int
	for {
		b := make([]byte, 1)
		count, err := os.Stdin.Read(b)
		if err != nil && err != io.EOF {
			return nil, err
		}
		if err == io.EOF || b[0] == '\n' {
			return pass[:total], nil
		}
		total += count
		pass = secureAppend(pass, b)
	}
}

func secureAppend(s, t []byte) []byte {
	res := append(s, t...)
	if len(res) > cap(s) {
		// When append needs to allocate a new array, clear out the old one.
		for i := range s {
			s[i] = '0'
		}
	}
	// Clear out the second array.
	for i := range t {
		t[i] = '0'
	}
	return res
}

// catchTerminationSignals catches signals to allow us to turn terminal echo back on.
func catchTerminationSignals(stop <-chan bool, state *terminal.State) {
	var successErrno syscall.Errno
	sig := make(chan os.Signal, 4)
	// Catch the blockable termination signals.
	signal.Notify(sig, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGHUP)
	select {
	case <-sig:
		// Start on new line in terminal.
		fmt.Printf("\n")
		if err := terminal.Restore(int(os.Stdin.Fd()), state); err != successErrno {
			vlog.Errorf("Failed to restore terminal state (%v), you words may not show up when you type, enter 'stty echo' to fix this.", err)
		}
		os.Exit(-1)
	case <-stop:
		signal.Stop(sig)
	}
}

type restartOptions struct {
	enabled, unless bool
	code            int
}

func (opts *restartOptions) parse() error {
	code := *restartExitCode
	if code == "" {
		return nil
	}
	opts.enabled = true
	if code[0] == '!' {
		opts.unless = true
		code = code[1:]
	}
	var err error
	if opts.code, err = strconv.Atoi(code); err != nil {
		return verror.New(errCantParseRestartExitCode, nil, err)
	}
	return nil
}

func (opts *restartOptions) restart(exitCode int) bool {
	return opts.enabled && opts.unless != (exitCode == opts.code)
}
<|endoftext|>


<commit_before>// Package artifice provides seqencing of timed triggers for pulling information.
package artifice

import (
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/antihax/evedata/internal/apicache"
	"github.com/antihax/evedata/internal/redisqueue"
	"github.com/antihax/goesi"
	"github.com/antihax/goesi/esi"
	"github.com/garyburd/redigo/redis"
	"github.com/jmoiron/sqlx"
	"golang.org/x/oauth2"
)

// Artifice handles the scheduling of routine tasks.
type Artifice struct {
	stop     chan bool
	hammerWG *sync.WaitGroup
	inQueue  *redisqueue.RedisQueue
	esi      *goesi.APIClient
	redis    *redis.Pool
	db       *sqlx.DB
	mail     chan esi.PostCharactersCharacterIdMailMail

	// authentication
	token       *oauth2.TokenSource
	tokenCharID int32
	auth        *goesi.SSOAuthenticator
}

// NewArtifice Service.
func NewArtifice(redis *redis.Pool, db *sqlx.DB, clientID string, secret string, refresh string, refreshCharID string) *Artifice {

	if clientID == "" {
		log.Fatalln("Missing clientID")
	}
	if secret == "" {
		log.Fatalln("Missing secret")
	}
	if refresh == "" {
		log.Fatalln("Missing refresh token")
	}
	if refreshCharID == "" {
		log.Fatalln("Missing refresh CharID")
	}
	// Get a caching http client
	cache := apicache.CreateHTTPClientCache()

	// Create our ESI API Client
	esiClient := goesi.NewAPIClient(cache, "EVEData-API-Artifice")

	// Setup an authenticator
	auth := goesi.NewSSOAuthenticator(cache, clientID, secret, "", []string{})

	tok := &oauth2.Token{
		Expiry:       time.Now(),
		AccessToken:  "",
		RefreshToken: refresh,
		TokenType:    "Bearer",
	}

	charID, err := strconv.Atoi(refreshCharID)
	if err != nil {
		log.Fatalln(err)
	}

	// Build our token
	token, err := auth.TokenSource(tok)
	if err != nil {
		log.Fatalln(err)
	}

	// Setup a new artifice
	s := &Artifice{
		stop:     make(chan bool),
		hammerWG: &sync.WaitGroup{},
		inQueue: redisqueue.NewRedisQueue(
			redis,
			"evedata-hammer",
		),
		db:    db,
		auth:  auth,
		mail:  make(chan esi.PostCharactersCharacterIdMailMail),
		esi:   esiClient,
		redis: redis,

		tokenCharID: int32(charID),
		token:       &token,
	}

	return s
}

// Close the hammer service
func (s *Artifice) Close() {
	close(s.stop)
	s.hammerWG.Wait()
}

// ChangeBasePath for ESI (sisi/mock/tranquility)
func (s *Artifice) ChangeBasePath(path string) {
	s.esi.ChangeBasePath(path)
}

// ChangeTokenPath for ESI (sisi/mock/tranquility)
func (s *Artifice) ChangeTokenPath(path string) {
	s.auth.ChangeTokenURL(path)
	s.auth.ChangeAuthURL(path)
}

// QueueWork directly
func (s *Artifice) QueueWork(work []redisqueue.Work, priority int) error {
	return s.inQueue.QueueWork(work, priority)
}

// QueueSize returns the size of the queue
func (s *Artifice) QueueSize() (int, error) {
	return s.inQueue.Size()
}

// Run the hammer service
func (s *Artifice) Run() {
	go s.startup()
	go s.zkillboardPost()
	go s.warKillmails()
	go s.runMetrics()
	go s.mailRunner()
	s.runTriggers()
}

// RetryTransaction on deadlocks
func retryTransaction(tx *sqlx.Tx) error {
	for {
		err := tx.Commit()
		if err != nil {
			if !strings.Contains(err.Error(), "1213") {
				return err
			}
			time.Sleep(250 * time.Millisecond)
			continue
		} else {
			return err
		}
	}
}

// DoSQL executes a sql statement
func (s *Artifice) doSQL(stmt string, args ...interface{}) error {
	for {
		err := s.doSQLTranq(stmt, args...)
		if err != nil {
			if !strings.Contains(err.Error(), "1213") {

				return err
			}
			time.Sleep(250 * time.Millisecond)
			continue
		} else {
			return err
		}
	}
}

// DoSQL executes a sql statement
func (s *Artifice) doSQLTranq(stmt string, args ...interface{}) error {
	tx, err := s.db.Beginx()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	_, err = tx.Exec(stmt, args...)
	if err != nil {
		return err
	}

	err = retryTransaction(tx)
	if err != nil {
		return err
	}
	return nil
}

// RetryExecTillNoRows retries the exec until we get no error (deadlocks) and no results are returned
func (s *Artifice) RetryExecTillNoRows(sql string, args ...interface{}) error {
	for {
		rows, err := s.RetryExec(sql, args...)
		if err != nil {
			return err
		}
		if rows == 0 {
			break
		}
	}
	return nil
}

// RetryExec retries the exec until we get no error (deadlocks)
func (s *Artifice) RetryExec(sql string, args ...interface{}) (int64, error) {
	var rows int64
	for {
		res, err := s.db.Exec(sql, args...)
		if err == nil {
			rows, err = res.RowsAffected()
			return rows, err
		} else if strings.Contains(err.Error(), "1213") == false {
			return rows, err
		}
	}
}

type CharacterPairs struct {
	CharacterID      int32 `db:"characterID"`
	TokenCharacterID int32 `db:"tokenCharacterID"`
	AllianceID       int32 `db:"allianceID"`
	CorporationID    int32 `db:"corporationID"`
}

func (s *Artifice) GetCharactersForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ?`, "%"+scope+"%")
	return pairs, err
}

func (s *Artifice) GetAllianceForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID, allianceID, corporationID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ? AND allianceID > 0
			GROUP BY allianceID
			`, "%"+scope+"%")
	return pairs, err
}

func (s *Artifice) GetCorporationForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID, allianceID, corporationID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ?
			GROUP BY corporationID
			`, "%"+scope+"%")
	return pairs, err
}
<commit_msg>cleanup<commit_after>// Package artifice provides seqencing of timed triggers for pulling information.
package artifice

import (
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/antihax/evedata/internal/apicache"
	"github.com/antihax/evedata/internal/redisqueue"
	"github.com/antihax/goesi"
	"github.com/antihax/goesi/esi"
	"github.com/garyburd/redigo/redis"
	"github.com/jmoiron/sqlx"
	"golang.org/x/oauth2"
)

// Artifice handles the scheduling of routine tasks.
type Artifice struct {
	stop    chan bool
	inQueue *redisqueue.RedisQueue
	esi     *goesi.APIClient
	redis   *redis.Pool
	db      *sqlx.DB
	mail    chan esi.PostCharactersCharacterIdMailMail

	// authentication
	token       *oauth2.TokenSource
	tokenCharID int32
	auth        *goesi.SSOAuthenticator
}

// NewArtifice Service.
func NewArtifice(redis *redis.Pool, db *sqlx.DB, clientID string, secret string, refresh string, refreshCharID string) *Artifice {

	if clientID == "" {
		log.Fatalln("Missing clientID")
	}
	if secret == "" {
		log.Fatalln("Missing secret")
	}
	if refresh == "" {
		log.Fatalln("Missing refresh token")
	}
	if refreshCharID == "" {
		log.Fatalln("Missing refresh CharID")
	}
	// Get a caching http client
	cache := apicache.CreateHTTPClientCache()

	// Create our ESI API Client
	esiClient := goesi.NewAPIClient(cache, "EVEData-API-Artifice")

	// Setup an authenticator
	auth := goesi.NewSSOAuthenticator(cache, clientID, secret, "", []string{})

	tok := &oauth2.Token{
		Expiry:       time.Now(),
		AccessToken:  "",
		RefreshToken: refresh,
		TokenType:    "Bearer",
	}

	charID, err := strconv.Atoi(refreshCharID)
	if err != nil {
		log.Fatalln(err)
	}

	// Build our token
	token, err := auth.TokenSource(tok)
	if err != nil {
		log.Fatalln(err)
	}

	// Setup a new artifice
	s := &Artifice{
		stop: make(chan bool),
		inQueue: redisqueue.NewRedisQueue(
			redis,
			"evedata-hammer",
		),
		db:    db,
		auth:  auth,
		mail:  make(chan esi.PostCharactersCharacterIdMailMail),
		esi:   esiClient,
		redis: redis,

		tokenCharID: int32(charID),
		token:       &token,
	}

	return s
}

// Close the hammer service
func (s *Artifice) Close() {
	close(s.stop)
}

// ChangeBasePath for ESI (sisi/mock/tranquility)
func (s *Artifice) ChangeBasePath(path string) {
	s.esi.ChangeBasePath(path)
}

// ChangeTokenPath for ESI (sisi/mock/tranquility)
func (s *Artifice) ChangeTokenPath(path string) {
	s.auth.ChangeTokenURL(path)
	s.auth.ChangeAuthURL(path)
}

// QueueWork directly
func (s *Artifice) QueueWork(work []redisqueue.Work, priority int) error {
	return s.inQueue.QueueWork(work, priority)
}

// QueueSize returns the size of the queue
func (s *Artifice) QueueSize() (int, error) {
	return s.inQueue.Size()
}

// Run the hammer service
func (s *Artifice) Run() {
	go s.startup()
	go s.zkillboardPost()
	go s.warKillmails()
	go s.runMetrics()
	go s.mailRunner()
	s.runTriggers()
}

// RetryTransaction on deadlocks
func retryTransaction(tx *sqlx.Tx) error {
	for {
		err := tx.Commit()
		if err != nil {
			if !strings.Contains(err.Error(), "1213") {
				return err
			}
			time.Sleep(250 * time.Millisecond)
			continue
		} else {
			return err
		}
	}
}

// DoSQL executes a sql statement
func (s *Artifice) doSQL(stmt string, args ...interface{}) error {
	for {
		err := s.doSQLTranq(stmt, args...)
		if err != nil {
			if !strings.Contains(err.Error(), "1213") {

				return err
			}
			time.Sleep(250 * time.Millisecond)
			continue
		} else {
			return err
		}
	}
}

// DoSQL executes a sql statement
func (s *Artifice) doSQLTranq(stmt string, args ...interface{}) error {
	tx, err := s.db.Beginx()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	_, err = tx.Exec(stmt, args...)
	if err != nil {
		return err
	}

	err = retryTransaction(tx)
	if err != nil {
		return err
	}
	return nil
}

// RetryExecTillNoRows retries the exec until we get no error (deadlocks) and no results are returned
func (s *Artifice) RetryExecTillNoRows(sql string, args ...interface{}) error {
	for {
		rows, err := s.RetryExec(sql, args...)
		if err != nil {
			return err
		}
		if rows == 0 {
			break
		}
	}
	return nil
}

// RetryExec retries the exec until we get no error (deadlocks)
func (s *Artifice) RetryExec(sql string, args ...interface{}) (int64, error) {
	var rows int64
	for {
		res, err := s.db.Exec(sql, args...)
		if err == nil {
			rows, err = res.RowsAffected()
			return rows, err
		} else if strings.Contains(err.Error(), "1213") == false {
			return rows, err
		}
	}
}

type CharacterPairs struct {
	CharacterID      int32 `db:"characterID"`
	TokenCharacterID int32 `db:"tokenCharacterID"`
	AllianceID       int32 `db:"allianceID"`
	CorporationID    int32 `db:"corporationID"`
}

func (s *Artifice) GetCharactersForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ?`, "%"+scope+"%")
	return pairs, err
}

func (s *Artifice) GetAllianceForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID, allianceID, corporationID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ? AND allianceID > 0
			GROUP BY allianceID
			`, "%"+scope+"%")
	return pairs, err
}

func (s *Artifice) GetCorporationForScope(scope string) ([]CharacterPairs, error) {
	pairs := []CharacterPairs{}
	err := s.db.Select(&pairs,
		`SELECT characterID, tokenCharacterID, allianceID, corporationID FROM evedata.crestTokens T
			WHERE lastStatus != "invalid_token" AND scopes LIKE ?
			GROUP BY corporationID
			`, "%"+scope+"%")
	return pairs, err
}
<|endoftext|>


<commit_before>package rest

import (
	"connectordb/streamdb/operator"
	"io"
	"net/http"
	"time"

	"github.com/apcera/nats"
	"github.com/gorilla/websocket"

	log "github.com/Sirupsen/logrus"
)

const (
	//The max size of a websocket message
	messageSizeLimit = 1 * Mb

	//The time allowed to write a message
	writeWait = 2 * time.Second

	//Ping pong stuff - making sure that the connection still exists
	pongWait   = 60 * time.Second
	pingPeriod = (pongWait * 9) / 10

	//The number of messages to buffer
	messageBuffer = 3
)

//The websocket upgrader
var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}
)

//WebsocketConnection is the general connection with a websocket that is run.
//Loosely based on github.com/gorilla/websocket/blob/master/examples/chat/conn.go
//No need for mutex because only reader reads and implements commands
type WebsocketConnection struct {
	ws *websocket.Conn

	subscriptions map[string]*nats.Subscription

	c chan operator.Message

	logger *log.Entry //logrus uses a mutex internally
	o      operator.Operator
}

//NewWebsocketConnection creates a new websocket connection based on the operators and stuff
func NewWebsocketConnection(o operator.Operator, writer http.ResponseWriter, request *http.Request, logger *log.Entry) (*WebsocketConnection, error) {
	logger = logger.WithField("op", "ws")

	ws, err := upgrader.Upgrade(writer, request, nil)
	if err != nil {
		logger.Errorln(err)
		return nil, err
	}

	ws.SetReadLimit(messageSizeLimit)

	return &WebsocketConnection{ws, make(map[string]*nats.Subscription), make(chan operator.Message, messageBuffer), logger, o}, nil
}

func (c *WebsocketConnection) write(obj interface{}) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteJSON(obj)
}

//Close the websocket connection
func (c *WebsocketConnection) Close() {
	c.UnsubscribeAll()
	close(c.c)
	c.ws.Close()
	c.logger.WithField("cmd", "close").Debugln()
}

//Insert a datapoint using the websocket
func (c *WebsocketConnection) Insert(ws *websocketCommand) {
	logger := c.logger.WithFields(log.Fields{"cmd": "insert", "arg": ws.Arg})
	logger.Infoln("Inserting", len(ws.D), "dp")
	err := c.o.InsertStream(ws.Arg, ws.D)
	if err != nil {
		//TODO: Notify user of insert failure
		logger.Warn(err.Error())
	}
}

//Subscribe to the given data stream
func (c *WebsocketConnection) Subscribe(s string) {
	logger := c.logger.WithFields(log.Fields{"cmd": "subscribe", "arg": s})
	if _, ok := c.subscriptions[s]; !ok {
		subs, err := c.o.Subscribe(s, c.c)
		if err != nil {
			logger.Warningln(err)
		} else {
			logger.Debugln()
			c.subscriptions[s] = subs
		}
	} else {
		logger.Warningln("Already subscribed")
	}
}

//Unsubscribe from the given data stream
func (c *WebsocketConnection) Unsubscribe(s string) {
	logger := c.logger.WithFields(log.Fields{"cmd": "unsubscribe", "arg": s})
	if val, ok := c.subscriptions[s]; ok {
		logger.Debugln()
		val.Unsubscribe()
		delete(c.subscriptions, s)
	} else {
		logger.Warningln("subscription DNE")
	}
}

//UnsubscribeAll from all streams of data
func (c *WebsocketConnection) UnsubscribeAll() {
	c.logger.WithField("cmd", "unsubscribeALL").Debugln()
	for _, val := range c.subscriptions {
		val.Unsubscribe()
	}
	c.subscriptions = make(map[string]*nats.Subscription)
}

//A command is a cmd and the arg operation
type websocketCommand struct {
	Cmd string
	Arg string
	D   []operator.Datapoint //If the command is "insert", it needs an additional datapoint
}

//RunReader runs the reading routine. It also maps the commands to actual subscriptions
func (c *WebsocketConnection) RunReader(readmessenger chan string) {

	//Set up the heartbeat reader(makes sure that sockets are alive)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error {
		c.logger.WithField("cmd", "PingPong").Debugln()
		c.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	var cmd websocketCommand
	for {
		err := c.ws.ReadJSON(&cmd)
		if err != nil {
			if err == io.EOF {
				readmessenger <- "EXIT"
				return //On EOF, do nothing - it is just a close
			}
			c.logger.Warningln(err)
			break
		}
		switch cmd.Cmd {
		default:
			c.logger.Warningln("Command not recognized:", cmd.Cmd)
			//Do nothing - the command is not recognized
		case "insert":
			c.Insert(&cmd)
		case "subscribe":
			c.Subscribe(cmd.Arg)
		case "unsubscribe":
			c.Unsubscribe(cmd.Arg)
		case "unsubscribe_all":
			c.UnsubscribeAll()
		}
	}
	//Since the reader is exiting, notify the writer to send close message
	readmessenger <- "@EXIT"
}

//RunWriter writes the subscription data as well as the heartbeat pings.
func (c *WebsocketConnection) RunWriter(readmessenger chan string, exitchan chan bool) {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
loop:
	for {
		select {
		case dp, ok := <-c.c:
			if !ok {
				c.ws.SetWriteDeadline(time.Now().Add(writeWait))
				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
				break loop
			}
			c.logger.WithFields(log.Fields{"cmd": "MSG", "arg": dp.Stream}).Debugln()
			if err := c.write(dp); err != nil {
				break loop
			}
		case <-ticker.C:
			//c.logger.WithField("cmd", "PING").Debugln()
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				break loop
			}
		case msg := <-readmessenger:
			if msg == "EXIT" {
				break loop
			} else if msg == "@EXIT" {
				c.ws.SetWriteDeadline(time.Now().Add(writeWait))
				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
				break loop
			}
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			c.ws.WriteMessage(websocket.TextMessage, []byte(msg))
		}
	}
	exitchan <- true
}

//Run the websocket operations
func (c *WebsocketConnection) Run() error {
	c.logger.Debugln("Running websocket...")

	//The reader can communicate with the writer through the channel
	msgchn := make(chan string, 1)
	exitchan := make(chan bool, 1)
	go c.RunWriter(msgchn, exitchan)
	c.RunReader(msgchn)
	//Wait for writer to exit, or for the exit timeout to happen
	go func() {
		time.Sleep(writeWait)
		exitchan <- false
	}()

	if !<-exitchan {
		c.logger.Error("writer exit timeout")
	}
	return nil
}

//RunWebsocket runs the websocket handler
func RunWebsocket(o operator.Operator, writer http.ResponseWriter, request *http.Request, logger *log.Entry) error {
	conn, err := NewWebsocketConnection(o, writer, request, logger)
	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		return err
	}
	defer conn.Close()

	return conn.Run()
}
<commit_msg>REST: Logging websocket ping/pong off<commit_after>package rest

import (
	"connectordb/streamdb/operator"
	"io"
	"net/http"
	"time"

	"github.com/apcera/nats"
	"github.com/gorilla/websocket"

	log "github.com/Sirupsen/logrus"
)

const (
	//The max size of a websocket message
	messageSizeLimit = 1 * Mb

	//The time allowed to write a message
	writeWait = 2 * time.Second

	//Ping pong stuff - making sure that the connection still exists
	pongWait   = 60 * time.Second
	pingPeriod = (pongWait * 9) / 10

	//The number of messages to buffer
	messageBuffer = 3
)

//The websocket upgrader
var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}
)

//WebsocketConnection is the general connection with a websocket that is run.
//Loosely based on github.com/gorilla/websocket/blob/master/examples/chat/conn.go
//No need for mutex because only reader reads and implements commands
type WebsocketConnection struct {
	ws *websocket.Conn

	subscriptions map[string]*nats.Subscription

	c chan operator.Message

	logger *log.Entry //logrus uses a mutex internally
	o      operator.Operator
}

//NewWebsocketConnection creates a new websocket connection based on the operators and stuff
func NewWebsocketConnection(o operator.Operator, writer http.ResponseWriter, request *http.Request, logger *log.Entry) (*WebsocketConnection, error) {
	logger = logger.WithField("op", "ws")

	ws, err := upgrader.Upgrade(writer, request, nil)
	if err != nil {
		logger.Errorln(err)
		return nil, err
	}

	ws.SetReadLimit(messageSizeLimit)

	return &WebsocketConnection{ws, make(map[string]*nats.Subscription), make(chan operator.Message, messageBuffer), logger, o}, nil
}

func (c *WebsocketConnection) write(obj interface{}) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteJSON(obj)
}

//Close the websocket connection
func (c *WebsocketConnection) Close() {
	c.UnsubscribeAll()
	close(c.c)
	c.ws.Close()
	c.logger.WithField("cmd", "close").Debugln()
}

//Insert a datapoint using the websocket
func (c *WebsocketConnection) Insert(ws *websocketCommand) {
	logger := c.logger.WithFields(log.Fields{"cmd": "insert", "arg": ws.Arg})
	logger.Infoln("Inserting", len(ws.D), "dp")
	err := c.o.InsertStream(ws.Arg, ws.D)
	if err != nil {
		//TODO: Notify user of insert failure
		logger.Warn(err.Error())
	}
}

//Subscribe to the given data stream
func (c *WebsocketConnection) Subscribe(s string) {
	logger := c.logger.WithFields(log.Fields{"cmd": "subscribe", "arg": s})
	if _, ok := c.subscriptions[s]; !ok {
		subs, err := c.o.Subscribe(s, c.c)
		if err != nil {
			logger.Warningln(err)
		} else {
			logger.Debugln()
			c.subscriptions[s] = subs
		}
	} else {
		logger.Warningln("Already subscribed")
	}
}

//Unsubscribe from the given data stream
func (c *WebsocketConnection) Unsubscribe(s string) {
	logger := c.logger.WithFields(log.Fields{"cmd": "unsubscribe", "arg": s})
	if val, ok := c.subscriptions[s]; ok {
		logger.Debugln()
		val.Unsubscribe()
		delete(c.subscriptions, s)
	} else {
		logger.Warningln("subscription DNE")
	}
}

//UnsubscribeAll from all streams of data
func (c *WebsocketConnection) UnsubscribeAll() {
	c.logger.WithField("cmd", "unsubscribeALL").Debugln()
	for _, val := range c.subscriptions {
		val.Unsubscribe()
	}
	c.subscriptions = make(map[string]*nats.Subscription)
}

//A command is a cmd and the arg operation
type websocketCommand struct {
	Cmd string
	Arg string
	D   []operator.Datapoint //If the command is "insert", it needs an additional datapoint
}

//RunReader runs the reading routine. It also maps the commands to actual subscriptions
func (c *WebsocketConnection) RunReader(readmessenger chan string) {

	//Set up the heartbeat reader(makes sure that sockets are alive)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error {
		//c.logger.WithField("cmd", "PingPong").Debugln()
		c.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	var cmd websocketCommand
	for {
		err := c.ws.ReadJSON(&cmd)
		if err != nil {
			if err == io.EOF {
				readmessenger <- "EXIT"
				return //On EOF, do nothing - it is just a close
			}
			c.logger.Warningln(err)
			break
		}
		switch cmd.Cmd {
		default:
			c.logger.Warningln("Command not recognized:", cmd.Cmd)
			//Do nothing - the command is not recognized
		case "insert":
			c.Insert(&cmd)
		case "subscribe":
			c.Subscribe(cmd.Arg)
		case "unsubscribe":
			c.Unsubscribe(cmd.Arg)
		case "unsubscribe_all":
			c.UnsubscribeAll()
		}
	}
	//Since the reader is exiting, notify the writer to send close message
	readmessenger <- "@EXIT"
}

//RunWriter writes the subscription data as well as the heartbeat pings.
func (c *WebsocketConnection) RunWriter(readmessenger chan string, exitchan chan bool) {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
loop:
	for {
		select {
		case dp, ok := <-c.c:
			if !ok {
				c.ws.SetWriteDeadline(time.Now().Add(writeWait))
				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
				break loop
			}
			c.logger.WithFields(log.Fields{"cmd": "MSG", "arg": dp.Stream}).Debugln()
			if err := c.write(dp); err != nil {
				break loop
			}
		case <-ticker.C:
			//c.logger.WithField("cmd", "PING").Debugln()
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				break loop
			}
		case msg := <-readmessenger:
			if msg == "EXIT" {
				break loop
			} else if msg == "@EXIT" {
				c.ws.SetWriteDeadline(time.Now().Add(writeWait))
				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
				break loop
			}
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			c.ws.WriteMessage(websocket.TextMessage, []byte(msg))
		}
	}
	exitchan <- true
}

//Run the websocket operations
func (c *WebsocketConnection) Run() error {
	c.logger.Debugln("Running websocket...")

	//The reader can communicate with the writer through the channel
	msgchn := make(chan string, 1)
	exitchan := make(chan bool, 1)
	go c.RunWriter(msgchn, exitchan)
	c.RunReader(msgchn)
	//Wait for writer to exit, or for the exit timeout to happen
	go func() {
		time.Sleep(writeWait)
		exitchan <- false
	}()

	if !<-exitchan {
		c.logger.Error("writer exit timeout")
	}
	return nil
}

//RunWebsocket runs the websocket handler
func RunWebsocket(o operator.Operator, writer http.ResponseWriter, request *http.Request, logger *log.Entry) error {
	conn, err := NewWebsocketConnection(o, writer, request, logger)
	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		return err
	}
	defer conn.Close()

	return conn.Run()
}
<|endoftext|>


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

package main

// Test that SIGPROF received in C code does not crash the process
// looking for the C code's func pointer.

// The test fails when the function is the first C function.
// The exported functions are the first C functions, so we use that.

// extern void GoNop();
import "C"

import (
	"bytes"
	"fmt"
	"runtime/pprof"
)

func init() {
	register("CgoCCodeSIGPROF", CgoCCodeSIGPROF)
}

//export GoNop
func GoNop() {}

func CgoCCodeSIGPROF() {
	c := make(chan bool)
	go func() {
		for {
			<-c
			for i := 0; i < 1e7; i++ {
				C.GoNop()
			}
			c <- true
		}
	}()

	var buf bytes.Buffer
	pprof.StartCPUProfile(&buf)
	c <- true
	<-c
	pprof.StopCPUProfile()

	fmt.Println("OK")
}
<commit_msg>runtime: limit TestCgoCCodeSIGPROF test to 1 second<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

// Test that SIGPROF received in C code does not crash the process
// looking for the C code's func pointer.

// The test fails when the function is the first C function.
// The exported functions are the first C functions, so we use that.

// extern void GoNop();
import "C"

import (
	"bytes"
	"fmt"
	"runtime/pprof"
	"time"
)

func init() {
	register("CgoCCodeSIGPROF", CgoCCodeSIGPROF)
}

//export GoNop
func GoNop() {}

func CgoCCodeSIGPROF() {
	c := make(chan bool)
	go func() {
		for {
			<-c
			start := time.Now()
			for i := 0; i < 1e7; i++ {
				if i%1000 == 0 {
					if time.Since(start) > time.Second {
						break
					}
				}
				C.GoNop()
			}
			c <- true
		}
	}()

	var buf bytes.Buffer
	pprof.StartCPUProfile(&buf)
	c <- true
	<-c
	pprof.StopCPUProfile()

	fmt.Println("OK")
}
<|endoftext|>


<commit_before>package assets

import (
	"fmt"
	"path/filepath"

	"github.com/pachyderm/pachyderm/src/server/pkg/dbutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TODO: Refactor the stateful set setup to better capture the shared functionality between the etcd / postgres setup.
// New / existing features that apply to both should be captured in one place.
// TODO: Move off of kubernetes Deployment object entirely since it is not well suited for stateful applications.
// The primary motivation for this would be to avoid the deadlock that can occur when using a ReadWriteOnce volume mount
// with a kubernetes Deployment.

var (
	postgresImage = "postgres:13.0-alpine"

	postgresHeadlessServiceName     = "postgres-headless"
	postgresName                    = "postgres"
	postgresVolumeName              = "postgres-volume"
	postgresInitVolumeName          = "postgres-init"
	postgresInitConfigMapName       = "postgres-init-cm"
	postgresVolumeClaimName         = "postgres-storage"
	defaultPostgresStorageClassName = "postgres-storage-class"
)

// PostgresOpts are options that are applicable to postgres.
type PostgresOpts struct {
	Nodes  int
	Volume string

	// CPURequest is the amount of CPU (in cores) we request for each
	// postgres node. If empty, assets.go will choose a default size.
	CPURequest string

	// MemRequest is the amount of memory we request for each postgres
	// node. If empty, assets.go will choose a default size.
	MemRequest string

	// StorageClassName is the name of an existing StorageClass to use when
	// creating a StatefulSet for dynamic postgres storage. If unset, a new
	// StorageClass will be created for the StatefulSet.
	StorageClassName string
}

// PostgresDeployment generates a Deployment for the pachyderm postgres instance.
func PostgresDeployment(opts *AssetOpts, hostPath string) *apps.Deployment {
	cpu := resource.MustParse(opts.PostgresOpts.CPURequest)
	mem := resource.MustParse(opts.PostgresOpts.MemRequest)
	var volumes []v1.Volume
	if hostPath == "" {
		volumes = []v1.Volume{
			{
				Name: "postgres-storage",
				VolumeSource: v1.VolumeSource{
					PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
						ClaimName: postgresVolumeClaimName,
					},
				},
			},
		}
	} else {
		volumes = []v1.Volume{
			{
				Name: "postgres-storage",
				VolumeSource: v1.VolumeSource{
					HostPath: &v1.HostPathVolumeSource{
						Path: filepath.Join(hostPath, "postgres"),
					},
				},
			},
		}
	}
	volumes = append(volumes, v1.Volume{
		Name: postgresInitVolumeName,
		VolumeSource: v1.VolumeSource{
			ConfigMap: &v1.ConfigMapVolumeSource{
				LocalObjectReference: v1.LocalObjectReference{Name: postgresInitConfigMapName},
			},
		},
	})
	resourceRequirements := v1.ResourceRequirements{
		Requests: v1.ResourceList{
			v1.ResourceCPU:    cpu,
			v1.ResourceMemory: mem,
		},
	}
	if !opts.NoGuaranteed {
		resourceRequirements.Limits = v1.ResourceList{
			v1.ResourceCPU:    cpu,
			v1.ResourceMemory: mem,
		}
	}
	image := postgresImage
	if opts.Registry != "" {
		image = AddRegistry(opts.Registry, postgresImage)
	}
	return &apps.Deployment{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Deployment",
			APIVersion: "apps/v1",
		},
		ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
		Spec: apps.DeploymentSpec{
			Replicas: replicas(1),
			Selector: &metav1.LabelSelector{
				MatchLabels: labels(postgresName),
			},
			Template: v1.PodTemplateSpec{
				ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  postgresName,
							Image: image,
							//TODO figure out how to get a cluster of these to talk to each other
							Ports: []v1.ContainerPort{
								{
									ContainerPort: 5432,
									Name:          "client-port",
								},
							},
							VolumeMounts: []v1.VolumeMount{
								{
									Name:      "postgres-storage",
									MountPath: "/var/lib/postgresql/data",
								},
								{
									Name:      postgresInitVolumeName,
									MountPath: "/docker-entrypoint-initdb.d",
								},
							},
							ImagePullPolicy: "IfNotPresent",
							Resources:       resourceRequirements,
							Env: []v1.EnvVar{
								// TODO: Figure out how we want to handle auth in real deployments.
								// The auth has been removed for now to allow PFS tests to run against
								// a deployed Postgres instance.
								{Name: "POSTGRES_DB", Value: dbutil.DefaultDBName},
								{Name: "POSTGRES_HOST_AUTH_METHOD", Value: "trust"},
							},
						},
					},
					Volumes:          volumes,
					ImagePullSecrets: imagePullSecrets(opts),
				},
			},
		},
	}
}

// PostgresStorageClass creates a storage class used for dynamic volume
// provisioning.  Currently dynamic volume provisioning only works
// on AWS and GCE.
func PostgresStorageClass(opts *AssetOpts, backend Backend) (interface{}, error) {
	return makeStorageClass(opts, backend, defaultPostgresStorageClassName, labels(postgresName))
}

// PostgresHeadlessService returns a headless postgres service, which is only for DNS
// resolution.
func PostgresHeadlessService(opts *AssetOpts) *v1.Service {
	ports := []v1.ServicePort{
		{
			Name: "client-port",
			Port: 5432,
		},
	}
	return makeHeadlessService(opts, postgresName, postgresHeadlessServiceName, ports)
}

// PostgresStatefulSet returns a stateful set that manages an etcd cluster
func PostgresStatefulSet(opts *AssetOpts, backend Backend, diskSpace int) interface{} {
	mem := resource.MustParse(opts.PostgresOpts.MemRequest)
	cpu := resource.MustParse(opts.PostgresOpts.CPURequest)
	var pvcTemplates []interface{}
	switch backend {
	case GoogleBackend, AmazonBackend:
		storageClassName := opts.PostgresOpts.StorageClassName
		if storageClassName == "" {
			storageClassName = defaultPostgresStorageClassName
		}
		pvcTemplates = []interface{}{
			map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":   postgresVolumeClaimName,
					"labels": labels(postgresName),
					"annotations": map[string]string{
						"volume.beta.kubernetes.io/storage-class": storageClassName,
					},
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"resources": map[string]interface{}{
						"requests": map[string]interface{}{
							"storage": resource.MustParse(fmt.Sprintf("%vGi", diskSpace)),
						},
					},
					"accessModes": []string{"ReadWriteOnce"},
				},
			},
		}
	default:
		pvcTemplates = []interface{}{
			map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":      postgresVolumeClaimName,
					"labels":    labels(postgresName),
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"resources": map[string]interface{}{
						"requests": map[string]interface{}{
							"storage": resource.MustParse(fmt.Sprintf("%vGi", diskSpace)),
						},
					},
					"accessModes": []string{"ReadWriteOnce"},
				},
			},
		}
	}
	var imagePullSecrets []map[string]string
	if opts.ImagePullSecret != "" {
		imagePullSecrets = append(imagePullSecrets, map[string]string{"name": opts.ImagePullSecret})
	}
	// As of March 17, 2017, the Kubernetes client does not include structs for
	// Stateful Set, so we generate the kubernetes manifest using raw json.
	// TODO(msteffen): we're now upgrading our kubernetes client, so we should be
	// abe to rewrite this spec using k8s client structs
	image := postgresImage
	if opts.Registry != "" {
		image = AddRegistry(opts.Registry, postgresImage)
	}
	return map[string]interface{}{
		"apiVersion": "apps/v1",
		"kind":       "StatefulSet",
		"metadata": map[string]interface{}{
			"name":      postgresName,
			"labels":    labels(postgresName),
			"namespace": opts.Namespace,
		},
		"spec": map[string]interface{}{
			// Effectively configures a RC
			"serviceName": postgresHeadlessServiceName,
			"replicas":    int(opts.PostgresOpts.Nodes),
			"selector": map[string]interface{}{
				"matchLabels": labels(postgresName),
			},

			// pod template
			"template": map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":      postgresName,
					"labels":    labels(postgresName),
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"imagePullSecrets": imagePullSecrets,
					"containers": []interface{}{
						map[string]interface{}{
							"name":  postgresName,
							"image": image,
							// TODO: Figure out how we want to handle auth in real deployments.
							// The auth has been removed for now to allow PFS tests to run against
							// a deployed Postgres instance.
							"env": []map[string]interface{}{{
								"name":  "POSTGRES_DB",
								"value": dbutil.DefaultDBName,
							}, {
								"name":  "POSTGRES_HOST_AUTH_METHOD",
								"value": "trust",
							}},
							"ports": []interface{}{
								map[string]interface{}{
									"containerPort": 5432,
									"name":          "client-port",
								},
							},
							"volumeMounts": []interface{}{
								map[string]interface{}{
									"name":      postgresVolumeClaimName,
									"mountPath": "/var/lib/postgresql/data",
								},
								map[string]interface{}{
									"name":      postgresInitVolumeName,
									"mountPath": "/docker-entrypoint-initdb.d",
								},
							},
							"imagePullPolicy": "IfNotPresent",
							"resources": map[string]interface{}{
								"requests": map[string]interface{}{
									string(v1.ResourceCPU):    cpu.String(),
									string(v1.ResourceMemory): mem.String(),
								},
							},
						},
					},
				},
			},
			"volumeClaimTemplates": pvcTemplates,
		},
	}
}

// PostgresVolume creates a persistent volume backed by a volume with name "name"
func PostgresVolume(persistentDiskBackend Backend, opts *AssetOpts,
	hostPath string, name string, size int) (*v1.PersistentVolume, error) {
	return makePersistentVolume(opts, persistentDiskBackend, hostPath, name, size, postgresVolumeName, labels(postgresName))
}

// PostgresVolumeClaim creates a persistent volume claim of 'size' GB.
//
// Note that if you're controlling Postgres with a Stateful Set, this is
// unnecessary (the stateful set controller will create PVCs automatically).
func PostgresVolumeClaim(size int, opts *AssetOpts) *v1.PersistentVolumeClaim {
	return makeVolumeClaim(opts, size, postgresVolumeName, postgresVolumeClaimName, labels(postgresName))
}

// PostgresService generates a Service for the pachyderm postgres instance.
func PostgresService(local bool, opts *AssetOpts) *v1.Service {
	var clientNodePort int32
	if local {
		clientNodePort = 32228
	}
	return &v1.Service{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
		Spec: v1.ServiceSpec{
			Type: v1.ServiceTypeNodePort,
			Selector: map[string]string{
				"app": postgresName,
			},
			Ports: []v1.ServicePort{
				{
					Port:     5432,
					Name:     "client-port",
					NodePort: clientNodePort,
				},
			},
		},
	}
}

// PostgresInitConfigMap generates a configmap which can be mounted into
// the postgres container to initialize the database.
func PostgresInitConfigMap(opts *AssetOpts) *v1.ConfigMap {
	return &v1.ConfigMap{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ConfigMap",
			APIVersion: "v1",
		},
		ObjectMeta: objectMeta(postgresInitConfigMapName, labels(postgresName), nil, opts.Namespace),
		Data: map[string]string{
			"init-db.sh": `
#!/bin/bash
set -e

psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
    CREATE DATABASE dex;
    GRANT ALL PRIVILEGES ON DATABASE dex TO postgres;
EOSQL
`,
		},
	}
}
<commit_msg>fix postgres deployment manifest on windows using incorrect path separators (#5611)<commit_after>package assets

import (
	"fmt"
	"path"

	"github.com/pachyderm/pachyderm/src/server/pkg/dbutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TODO: Refactor the stateful set setup to better capture the shared functionality between the etcd / postgres setup.
// New / existing features that apply to both should be captured in one place.
// TODO: Move off of kubernetes Deployment object entirely since it is not well suited for stateful applications.
// The primary motivation for this would be to avoid the deadlock that can occur when using a ReadWriteOnce volume mount
// with a kubernetes Deployment.

var (
	postgresImage = "postgres:13.0-alpine"

	postgresHeadlessServiceName     = "postgres-headless"
	postgresName                    = "postgres"
	postgresVolumeName              = "postgres-volume"
	postgresInitVolumeName          = "postgres-init"
	postgresInitConfigMapName       = "postgres-init-cm"
	postgresVolumeClaimName         = "postgres-storage"
	defaultPostgresStorageClassName = "postgres-storage-class"
)

// PostgresOpts are options that are applicable to postgres.
type PostgresOpts struct {
	Nodes  int
	Volume string

	// CPURequest is the amount of CPU (in cores) we request for each
	// postgres node. If empty, assets.go will choose a default size.
	CPURequest string

	// MemRequest is the amount of memory we request for each postgres
	// node. If empty, assets.go will choose a default size.
	MemRequest string

	// StorageClassName is the name of an existing StorageClass to use when
	// creating a StatefulSet for dynamic postgres storage. If unset, a new
	// StorageClass will be created for the StatefulSet.
	StorageClassName string
}

// PostgresDeployment generates a Deployment for the pachyderm postgres instance.
func PostgresDeployment(opts *AssetOpts, hostPath string) *apps.Deployment {
	cpu := resource.MustParse(opts.PostgresOpts.CPURequest)
	mem := resource.MustParse(opts.PostgresOpts.MemRequest)
	var volumes []v1.Volume
	if hostPath == "" {
		volumes = []v1.Volume{
			{
				Name: "postgres-storage",
				VolumeSource: v1.VolumeSource{
					PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
						ClaimName: postgresVolumeClaimName,
					},
				},
			},
		}
	} else {
		volumes = []v1.Volume{
			{
				Name: "postgres-storage",
				VolumeSource: v1.VolumeSource{
					HostPath: &v1.HostPathVolumeSource{
						Path: path.Join(hostPath, "postgres"),
					},
				},
			},
		}
	}
	volumes = append(volumes, v1.Volume{
		Name: postgresInitVolumeName,
		VolumeSource: v1.VolumeSource{
			ConfigMap: &v1.ConfigMapVolumeSource{
				LocalObjectReference: v1.LocalObjectReference{Name: postgresInitConfigMapName},
			},
		},
	})
	resourceRequirements := v1.ResourceRequirements{
		Requests: v1.ResourceList{
			v1.ResourceCPU:    cpu,
			v1.ResourceMemory: mem,
		},
	}
	if !opts.NoGuaranteed {
		resourceRequirements.Limits = v1.ResourceList{
			v1.ResourceCPU:    cpu,
			v1.ResourceMemory: mem,
		}
	}
	image := postgresImage
	if opts.Registry != "" {
		image = AddRegistry(opts.Registry, postgresImage)
	}
	return &apps.Deployment{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Deployment",
			APIVersion: "apps/v1",
		},
		ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
		Spec: apps.DeploymentSpec{
			Replicas: replicas(1),
			Selector: &metav1.LabelSelector{
				MatchLabels: labels(postgresName),
			},
			Template: v1.PodTemplateSpec{
				ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  postgresName,
							Image: image,
							//TODO figure out how to get a cluster of these to talk to each other
							Ports: []v1.ContainerPort{
								{
									ContainerPort: 5432,
									Name:          "client-port",
								},
							},
							VolumeMounts: []v1.VolumeMount{
								{
									Name:      "postgres-storage",
									MountPath: "/var/lib/postgresql/data",
								},
								{
									Name:      postgresInitVolumeName,
									MountPath: "/docker-entrypoint-initdb.d",
								},
							},
							ImagePullPolicy: "IfNotPresent",
							Resources:       resourceRequirements,
							Env: []v1.EnvVar{
								// TODO: Figure out how we want to handle auth in real deployments.
								// The auth has been removed for now to allow PFS tests to run against
								// a deployed Postgres instance.
								{Name: "POSTGRES_DB", Value: dbutil.DefaultDBName},
								{Name: "POSTGRES_HOST_AUTH_METHOD", Value: "trust"},
							},
						},
					},
					Volumes:          volumes,
					ImagePullSecrets: imagePullSecrets(opts),
				},
			},
		},
	}
}

// PostgresStorageClass creates a storage class used for dynamic volume
// provisioning.  Currently dynamic volume provisioning only works
// on AWS and GCE.
func PostgresStorageClass(opts *AssetOpts, backend Backend) (interface{}, error) {
	return makeStorageClass(opts, backend, defaultPostgresStorageClassName, labels(postgresName))
}

// PostgresHeadlessService returns a headless postgres service, which is only for DNS
// resolution.
func PostgresHeadlessService(opts *AssetOpts) *v1.Service {
	ports := []v1.ServicePort{
		{
			Name: "client-port",
			Port: 5432,
		},
	}
	return makeHeadlessService(opts, postgresName, postgresHeadlessServiceName, ports)
}

// PostgresStatefulSet returns a stateful set that manages an etcd cluster
func PostgresStatefulSet(opts *AssetOpts, backend Backend, diskSpace int) interface{} {
	mem := resource.MustParse(opts.PostgresOpts.MemRequest)
	cpu := resource.MustParse(opts.PostgresOpts.CPURequest)
	var pvcTemplates []interface{}
	switch backend {
	case GoogleBackend, AmazonBackend:
		storageClassName := opts.PostgresOpts.StorageClassName
		if storageClassName == "" {
			storageClassName = defaultPostgresStorageClassName
		}
		pvcTemplates = []interface{}{
			map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":   postgresVolumeClaimName,
					"labels": labels(postgresName),
					"annotations": map[string]string{
						"volume.beta.kubernetes.io/storage-class": storageClassName,
					},
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"resources": map[string]interface{}{
						"requests": map[string]interface{}{
							"storage": resource.MustParse(fmt.Sprintf("%vGi", diskSpace)),
						},
					},
					"accessModes": []string{"ReadWriteOnce"},
				},
			},
		}
	default:
		pvcTemplates = []interface{}{
			map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":      postgresVolumeClaimName,
					"labels":    labels(postgresName),
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"resources": map[string]interface{}{
						"requests": map[string]interface{}{
							"storage": resource.MustParse(fmt.Sprintf("%vGi", diskSpace)),
						},
					},
					"accessModes": []string{"ReadWriteOnce"},
				},
			},
		}
	}
	var imagePullSecrets []map[string]string
	if opts.ImagePullSecret != "" {
		imagePullSecrets = append(imagePullSecrets, map[string]string{"name": opts.ImagePullSecret})
	}
	// As of March 17, 2017, the Kubernetes client does not include structs for
	// Stateful Set, so we generate the kubernetes manifest using raw json.
	// TODO(msteffen): we're now upgrading our kubernetes client, so we should be
	// abe to rewrite this spec using k8s client structs
	image := postgresImage
	if opts.Registry != "" {
		image = AddRegistry(opts.Registry, postgresImage)
	}
	return map[string]interface{}{
		"apiVersion": "apps/v1",
		"kind":       "StatefulSet",
		"metadata": map[string]interface{}{
			"name":      postgresName,
			"labels":    labels(postgresName),
			"namespace": opts.Namespace,
		},
		"spec": map[string]interface{}{
			// Effectively configures a RC
			"serviceName": postgresHeadlessServiceName,
			"replicas":    int(opts.PostgresOpts.Nodes),
			"selector": map[string]interface{}{
				"matchLabels": labels(postgresName),
			},

			// pod template
			"template": map[string]interface{}{
				"metadata": map[string]interface{}{
					"name":      postgresName,
					"labels":    labels(postgresName),
					"namespace": opts.Namespace,
				},
				"spec": map[string]interface{}{
					"imagePullSecrets": imagePullSecrets,
					"containers": []interface{}{
						map[string]interface{}{
							"name":  postgresName,
							"image": image,
							// TODO: Figure out how we want to handle auth in real deployments.
							// The auth has been removed for now to allow PFS tests to run against
							// a deployed Postgres instance.
							"env": []map[string]interface{}{{
								"name":  "POSTGRES_DB",
								"value": dbutil.DefaultDBName,
							}, {
								"name":  "POSTGRES_HOST_AUTH_METHOD",
								"value": "trust",
							}},
							"ports": []interface{}{
								map[string]interface{}{
									"containerPort": 5432,
									"name":          "client-port",
								},
							},
							"volumeMounts": []interface{}{
								map[string]interface{}{
									"name":      postgresVolumeClaimName,
									"mountPath": "/var/lib/postgresql/data",
								},
								map[string]interface{}{
									"name":      postgresInitVolumeName,
									"mountPath": "/docker-entrypoint-initdb.d",
								},
							},
							"imagePullPolicy": "IfNotPresent",
							"resources": map[string]interface{}{
								"requests": map[string]interface{}{
									string(v1.ResourceCPU):    cpu.String(),
									string(v1.ResourceMemory): mem.String(),
								},
							},
						},
					},
				},
			},
			"volumeClaimTemplates": pvcTemplates,
		},
	}
}

// PostgresVolume creates a persistent volume backed by a volume with name "name"
func PostgresVolume(persistentDiskBackend Backend, opts *AssetOpts,
	hostPath string, name string, size int) (*v1.PersistentVolume, error) {
	return makePersistentVolume(opts, persistentDiskBackend, hostPath, name, size, postgresVolumeName, labels(postgresName))
}

// PostgresVolumeClaim creates a persistent volume claim of 'size' GB.
//
// Note that if you're controlling Postgres with a Stateful Set, this is
// unnecessary (the stateful set controller will create PVCs automatically).
func PostgresVolumeClaim(size int, opts *AssetOpts) *v1.PersistentVolumeClaim {
	return makeVolumeClaim(opts, size, postgresVolumeName, postgresVolumeClaimName, labels(postgresName))
}

// PostgresService generates a Service for the pachyderm postgres instance.
func PostgresService(local bool, opts *AssetOpts) *v1.Service {
	var clientNodePort int32
	if local {
		clientNodePort = 32228
	}
	return &v1.Service{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: objectMeta(postgresName, labels(postgresName), nil, opts.Namespace),
		Spec: v1.ServiceSpec{
			Type: v1.ServiceTypeNodePort,
			Selector: map[string]string{
				"app": postgresName,
			},
			Ports: []v1.ServicePort{
				{
					Port:     5432,
					Name:     "client-port",
					NodePort: clientNodePort,
				},
			},
		},
	}
}

// PostgresInitConfigMap generates a configmap which can be mounted into
// the postgres container to initialize the database.
func PostgresInitConfigMap(opts *AssetOpts) *v1.ConfigMap {
	return &v1.ConfigMap{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ConfigMap",
			APIVersion: "v1",
		},
		ObjectMeta: objectMeta(postgresInitConfigMapName, labels(postgresName), nil, opts.Namespace),
		Data: map[string]string{
			"init-db.sh": `
#!/bin/bash
set -e

psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
    CREATE DATABASE dex;
    GRANT ALL PRIVILEGES ON DATABASE dex TO postgres;
EOSQL
`,
		},
	}
}
<|endoftext|>


<commit_before>package service

import (
	"github.com/ReneGa/tweetcount-microservices/stopwordfilter/domain"
	"github.com/ReneGa/tweetcount-microservices/stopwordfilter/repository"
	"golang.org/x/text/language"
)

// StopWordFilter is a service that filters stopwords from a stream of tweets
type StopWordFilter interface {
	TweetsWords(tweets domain.Tweets) domain.TweetsWords
}

func NewStopWordFilter(repository repository.WordSet) StopWordFilter {
	return &stopWordFilter{repository}
}

type stopWordFilter struct {
	repository repository.WordSet
}

func (s *stopWordFilter) TweetsWords(tweets domain.Tweets) domain.TweetsWords {

	stopWordSetIDs := s.repository.List()
	languages := make([]language.Tag, len(stopWordSetIDs))
	for i, ID := range stopWordSetIDs {
		languages[i] = language.Make(ID)
	}
	matcher := language.NewMatcher(languages)

	data := make(chan domain.TweetWords)
	stop := make(chan bool)

	go func() {
		tweets.Stop <- <-stop
	}()

	go func() {
		defer close(data)
		for tweet := range tweets.Data {
			_, i, _ := matcher.Match(language.Make(tweet.Language))
			stopWords := s.repository.Get(stopWordSetIDs[i])
			data <- domain.FilterStopWords(stopWords, tweet)
		}
	}()

	return domain.TweetsWords{
		Data: data,
		Stop: stop,
	}
}
<commit_msg>Add comment<commit_after>package service

import (
	"github.com/ReneGa/tweetcount-microservices/stopwordfilter/domain"
	"github.com/ReneGa/tweetcount-microservices/stopwordfilter/repository"
	"golang.org/x/text/language"
)

// StopWordFilter is a service that filters stopwords from a stream of tweets
type StopWordFilter interface {
	TweetsWords(tweets domain.Tweets) domain.TweetsWords
}

// NewStopWordFilter creates a new stop word filter service
func NewStopWordFilter(repository repository.WordSet) StopWordFilter {
	return &stopWordFilter{repository}
}

type stopWordFilter struct {
	repository repository.WordSet
}

func (s *stopWordFilter) TweetsWords(tweets domain.Tweets) domain.TweetsWords {

	stopWordSetIDs := s.repository.List()
	languages := make([]language.Tag, len(stopWordSetIDs))
	for i, ID := range stopWordSetIDs {
		languages[i] = language.Make(ID)
	}
	matcher := language.NewMatcher(languages)

	data := make(chan domain.TweetWords)
	stop := make(chan bool)

	go func() {
		tweets.Stop <- <-stop
	}()

	go func() {
		defer close(data)
		for tweet := range tweets.Data {
			_, i, _ := matcher.Match(language.Make(tweet.Language))
			stopWords := s.repository.Get(stopWordSetIDs[i])
			data <- domain.FilterStopWords(stopWords, tweet)
		}
	}()

	return domain.TweetsWords{
		Data: data,
		Stop: stop,
	}
}
<|endoftext|>


<commit_before>package ethutil

import (
	"reflect"
	"testing"
)

const LONG_WORD = "1234567890abcdefghijklmnopqrstuvwxxzABCEFGHIJKLMNOPQRSTUVWXYZ"

type MemDatabase struct {
	db map[string][]byte
}

func NewMemDatabase() (*MemDatabase, error) {
	db := &MemDatabase{db: make(map[string][]byte)}
	return db, nil
}
func (db *MemDatabase) Put(key []byte, value []byte) {
	db.db[string(key)] = value
}
func (db *MemDatabase) Get(key []byte) ([]byte, error) {
	return db.db[string(key)], nil
}
func (db *MemDatabase) Delete(key []byte) error {
	delete(db.db, string(key))
	return nil
}
func (db *MemDatabase) Print()              {}
func (db *MemDatabase) Close()              {}
func (db *MemDatabase) LastKnownTD() []byte { return nil }

func New() (*MemDatabase, *Trie) {
	db, _ := NewMemDatabase()
	return db, NewTrie(db, "")
}

func TestTrieSync(t *testing.T) {
	db, trie := New()

	trie.Update("dog", LONG_WORD)
	if len(db.db) != 0 {
		t.Error("Expected no data in database")
	}

	trie.Sync()
	if len(db.db) == 0 {
		t.Error("Expected data to be persisted")
	}
}

func TestTrieDirtyTracking(t *testing.T) {
	_, trie := New()
	trie.Update("dog", LONG_WORD)
	if !trie.cache.IsDirty {
		t.Error("Expected trie to be dirty")
	}

	trie.Sync()
	if trie.cache.IsDirty {
		t.Error("Expected trie not to be dirty")
	}

	trie.Update("test", LONG_WORD)
	trie.cache.Undo()
	if trie.cache.IsDirty {
		t.Error("Expected trie not to be dirty")
	}

}

func TestTrieReset(t *testing.T) {
	_, trie := New()

	trie.Update("cat", LONG_WORD)
	if len(trie.cache.nodes) == 0 {
		t.Error("Expected cached nodes")
	}

	trie.cache.Undo()

	if len(trie.cache.nodes) != 0 {
		t.Error("Expected no nodes after undo")
	}
}

func TestTrieGet(t *testing.T) {
	_, trie := New()

	trie.Update("cat", LONG_WORD)
	x := trie.Get("cat")
	if x != LONG_WORD {
		t.Error("expected %s, got %s", LONG_WORD, x)
	}
}

func TestTrieUpdating(t *testing.T) {
	_, trie := New()
	trie.Update("cat", LONG_WORD)
	trie.Update("cat", LONG_WORD+"1")
	x := trie.Get("cat")
	if x != LONG_WORD+"1" {
		t.Error("expected %S, got %s", LONG_WORD+"1", x)
	}
}

func TestTrieCmp(t *testing.T) {
	_, trie1 := New()
	_, trie2 := New()

	trie1.Update("doge", LONG_WORD)
	trie2.Update("doge", LONG_WORD)
	if !trie1.Cmp(trie2) {
		t.Error("Expected tries to be equal")
	}

	trie1.Update("dog", LONG_WORD)
	trie2.Update("cat", LONG_WORD)
	if trie1.Cmp(trie2) {
		t.Errorf("Expected tries not to be equal %x %x", trie1.Root, trie2.Root)
	}
}

func TestTrieDelete(t *testing.T) {
	_, trie := New()
	trie.Update("cat", LONG_WORD)
	exp := trie.Root
	trie.Update("dog", LONG_WORD)
	trie.Delete("dog")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}

	trie.Update("dog", LONG_WORD)
	exp = trie.Root
	trie.Update("dude", LONG_WORD)
	trie.Delete("dude")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}
}

func TestTrieDeleteWithValue(t *testing.T) {
	_, trie := New()
	trie.Update("c", LONG_WORD)
	exp := trie.Root
	trie.Update("ca", LONG_WORD)
	trie.Update("cat", LONG_WORD)
	trie.Delete("ca")
	trie.Delete("cat")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}

}

func TestTrieIterator(t *testing.T) {
	_, trie := New()
	trie.Update("c", LONG_WORD)
	trie.Update("ca", LONG_WORD)
	trie.Update("cat", LONG_WORD)

	lenBefore := len(trie.cache.nodes)
	it := trie.NewIterator()
	if num := it.Purge(); num != 3 {
		t.Errorf("Expected purge to return 3, got %d", num)
	}

	if lenBefore == len(trie.cache.nodes) {
		t.Errorf("Expected cached nodes to be deleted")
	}
}
<commit_msg>Conform to Db interface<commit_after>package ethutil

import (
	"reflect"
	"testing"
)

const LONG_WORD = "1234567890abcdefghijklmnopqrstuvwxxzABCEFGHIJKLMNOPQRSTUVWXYZ"

type MemDatabase struct {
	db map[string][]byte
}

func NewMemDatabase() (*MemDatabase, error) {
	db := &MemDatabase{db: make(map[string][]byte)}
	return db, nil
}
func (db *MemDatabase) Put(key []byte, value []byte) {
	db.db[string(key)] = value
}
func (db *MemDatabase) Get(key []byte) ([]byte, error) {
	return db.db[string(key)], nil
}
func (db *MemDatabase) Delete(key []byte) error {
	delete(db.db, string(key))
	return nil
}
func (db *MemDatabase) GetKeys() []*Key     { return nil }
func (db *MemDatabase) Print()              {}
func (db *MemDatabase) Close()              {}
func (db *MemDatabase) LastKnownTD() []byte { return nil }

func New() (*MemDatabase, *Trie) {
	db, _ := NewMemDatabase()
	return db, NewTrie(db, "")
}

func TestTrieSync(t *testing.T) {
	db, trie := New()

	trie.Update("dog", LONG_WORD)
	if len(db.db) != 0 {
		t.Error("Expected no data in database")
	}

	trie.Sync()
	if len(db.db) == 0 {
		t.Error("Expected data to be persisted")
	}
}

func TestTrieDirtyTracking(t *testing.T) {
	_, trie := New()
	trie.Update("dog", LONG_WORD)
	if !trie.cache.IsDirty {
		t.Error("Expected trie to be dirty")
	}

	trie.Sync()
	if trie.cache.IsDirty {
		t.Error("Expected trie not to be dirty")
	}

	trie.Update("test", LONG_WORD)
	trie.cache.Undo()
	if trie.cache.IsDirty {
		t.Error("Expected trie not to be dirty")
	}

}

func TestTrieReset(t *testing.T) {
	_, trie := New()

	trie.Update("cat", LONG_WORD)
	if len(trie.cache.nodes) == 0 {
		t.Error("Expected cached nodes")
	}

	trie.cache.Undo()

	if len(trie.cache.nodes) != 0 {
		t.Error("Expected no nodes after undo")
	}
}

func TestTrieGet(t *testing.T) {
	_, trie := New()

	trie.Update("cat", LONG_WORD)
	x := trie.Get("cat")
	if x != LONG_WORD {
		t.Error("expected %s, got %s", LONG_WORD, x)
	}
}

func TestTrieUpdating(t *testing.T) {
	_, trie := New()
	trie.Update("cat", LONG_WORD)
	trie.Update("cat", LONG_WORD+"1")
	x := trie.Get("cat")
	if x != LONG_WORD+"1" {
		t.Error("expected %S, got %s", LONG_WORD+"1", x)
	}
}

func TestTrieCmp(t *testing.T) {
	_, trie1 := New()
	_, trie2 := New()

	trie1.Update("doge", LONG_WORD)
	trie2.Update("doge", LONG_WORD)
	if !trie1.Cmp(trie2) {
		t.Error("Expected tries to be equal")
	}

	trie1.Update("dog", LONG_WORD)
	trie2.Update("cat", LONG_WORD)
	if trie1.Cmp(trie2) {
		t.Errorf("Expected tries not to be equal %x %x", trie1.Root, trie2.Root)
	}
}

func TestTrieDelete(t *testing.T) {
	_, trie := New()
	trie.Update("cat", LONG_WORD)
	exp := trie.Root
	trie.Update("dog", LONG_WORD)
	trie.Delete("dog")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}

	trie.Update("dog", LONG_WORD)
	exp = trie.Root
	trie.Update("dude", LONG_WORD)
	trie.Delete("dude")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}
}

func TestTrieDeleteWithValue(t *testing.T) {
	_, trie := New()
	trie.Update("c", LONG_WORD)
	exp := trie.Root
	trie.Update("ca", LONG_WORD)
	trie.Update("cat", LONG_WORD)
	trie.Delete("ca")
	trie.Delete("cat")
	if !reflect.DeepEqual(exp, trie.Root) {
		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
	}

}

func TestTrieIterator(t *testing.T) {
	_, trie := New()
	trie.Update("c", LONG_WORD)
	trie.Update("ca", LONG_WORD)
	trie.Update("cat", LONG_WORD)

	lenBefore := len(trie.cache.nodes)
	it := trie.NewIterator()
	if num := it.Purge(); num != 3 {
		t.Errorf("Expected purge to return 3, got %d", num)
	}

	if lenBefore == len(trie.cache.nodes) {
		t.Errorf("Expected cached nodes to be deleted")
	}
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 The Dename 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 verifier

import (
	"bytes"
	"crypto/sha256"
	"crypto/tls"
	"log"
	"math"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/agl/ed25519"
	"github.com/yahoo/coname/common"
	"github.com/yahoo/coname/proto"
	"github.com/yahoo/coname/server/kv"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// Config encapsulates everything that needs to be specified about a verifer.
// TODO: make this a protobuf, Unmarshal from JSON
type Config struct {
	Realm          string
	KeyserverVerif *proto.AuthorizationPolicy
	KeyserverAddr  string

	ID              uint64
	RatificationKey *[ed25519.PrivateKeySize]byte // [32]byte: secret; [32]byte: public
	TLS             *tls.Config                   // FIXME: tls.Config is not serializable, replicate relevant fields
}

// Verifier verifies that the Keyserver of the realm is not cheating.
// The veirifier is not replicated because one can just run several.
type Verifier struct {
	realm          string
	keyserverVerif *proto.AuthorizationPolicy
	keyserverAddr  string
	auth           credentials.TransportAuthenticator

	id              uint64
	ratificationKey *[ed25519.PrivateKeySize]byte

	db kv.DB
	vs proto.VerifierState

	keyserver proto.E2EKSVerificationClient

	stopOnce sync.Once
	stop     chan struct{}
	waitStop sync.WaitGroup
}

// Start initializes a new verifier based on config and db, or returns an error
// if initialization fails. It then starts the worker goroutine(s).
func Start(cfg *Config, db kv.DB) (*Verifier, error) {
	vr := &Verifier{
		realm:          cfg.Realm,
		keyserverVerif: cfg.KeyserverVerif,
		keyserverAddr:  cfg.KeyserverAddr,
		auth:           credentials.NewTLS(cfg.TLS),

		id:              cfg.ID,
		ratificationKey: cfg.RatificationKey,

		db: db,

		stop: make(chan struct{}),
	}

	switch verifierStateBytes, err := db.Get(tableVerifierState); err {
	case vr.db.ErrNotFound():
		vr.vs.NextEpoch = 1
	case nil:
		if err := vr.vs.Unmarshal(verifierStateBytes); err != nil {
			return nil, err
		}
	default:
		return nil, err
	}

	vr.waitStop.Add(1)
	go func() { vr.run(); vr.waitStop.Done() }()
	return vr, nil
}

// Stop cleanly shuts down the verifier and then returns.
func (vr *Verifier) Stop() {
	vr.stopOnce.Do(func() {
		close(vr.stop)
		vr.waitStop.Wait()
	})
}

func (vr *Verifier) shuttingDown() bool {
	select {
	case <-vr.stop:
		return true
	default:
		return false
	}
}

// run is the CSP-style main loop of the verifier. All code critical for safe
// persistence should be directly in run. All functions called from run should
// either interpret data and modify their mutable arguments OR interact with the
// network and disk, but not both.
func (vr *Verifier) run() {
	keyserverConnection, err := grpc.Dial(vr.keyserverAddr, grpc.WithTimeout(time.Second), grpc.WithTransportCredentials(vr.auth))
	if err != nil {
		log.Fatalf("dial %s: %s", vr.keyserverAddr, err)
	}
	vr.keyserver = proto.NewE2EKSVerificationClient(keyserverConnection)
	stream, err := vr.keyserver.VerifierStream(context.TODO(), &proto.VerifierStreamRequest{
		Start:    vr.vs.NextIndex,
		PageSize: math.MaxUint64,
	})
	if err != nil {
		keyserverConnection.Close()
		log.Fatalf("VerifierStream: %s", err)
	}

	wb := vr.db.NewBatch()
	for !vr.shuttingDown() {
		var step *proto.VerifierStep
		step, err = stream.Recv()
		if err != nil {
			log.Printf("VerifierStream.Recv: %s", err)
			break
		}
		wb.Put(tableVerifierLog(vr.vs.NextIndex), proto.MustMarshal(step))
		deferredIO := vr.step(step, &vr.vs, wb)
		vr.vs.NextIndex++
		if deferredIO != nil {
			wb.Put(tableVerifierState, proto.MustMarshal(&vr.vs))
			if err := vr.db.Write(wb); err != nil {
				log.Panicf("sync step to db: %s", err)
			}
			wb.Reset()
			deferredIO()
		}
	}
}

// step is called by run and changes the in-memory state. No i/o allowed.
func (vr *Verifier) step(step *proto.VerifierStep, vs *proto.VerifierState, wb kv.Batch) (deferredIO func()) {
	// vr: &const
	// step, vs, wb: &mut
	switch {
	case step.Update != nil:
		if err := common.VerifyUpdate( /*TODO(dmz): tree lookup*/ nil, step.Update); err != nil {
			return // TODO: decide whether keyserver should guarantee that verifiers don't see bad updates
		}
		// TODO(dmz): set entry in tree
	case step.Epoch != nil:
		ok := common.VerifyPolicy(
			vr.keyserverVerif,
			step.Epoch.Head.PreservedEncoding,
			step.Epoch.Signatures)
		// the bad steps here will not get persisted to disk right now. do we want them to?
		if !ok {
			log.Fatalf("%d: keyserver signature verification failed: %#v", vs.NextIndex, *step)
		}
		r := step.Epoch.Head
		if r.Head.Realm != vr.realm {
			log.Fatalf("%d: seh for realm %q, expected %q: %#v", vr.vs.NextEpoch, r.Head.Realm, vr.realm, *step)
		}
		if r.Head.Epoch != vr.vs.NextEpoch {
			log.Fatalf("%d: got epoch %d instead: %#v", vr.vs.NextEpoch, r.Head.Epoch, *step)
		}
		s := &r.Head
		if !bytes.Equal(s.PreviousSummaryHash, vr.vs.PreviousSummaryHash) {
			log.Fatalf("%d: seh with previous summary hash %q, expected %q: %#v", vr.vs.NextEpoch, s.PreviousSummaryHash, vr.vs.PreviousSummaryHash, *step)
		}
		if s.RootHash != nil /*TODO(dmz): merkletree.GetRootHash()*/ {
			log.Fatalf("%d: seh with root hash %q, expected nil: %#v", vr.vs.NextEpoch, s.RootHash /*, TODO*/, *step)
		}
		seh := &proto.SignedEpochHead{
			Head: proto.TimestampedEpochHead_PreserveEncoding{proto.TimestampedEpochHead{
				Head: proto.EpochHead_PreserveEncoding{proto.EpochHead{
					RootHash:            nil, // TODO(dmz): merklemap.GetRootHash()
					PreviousSummaryHash: vr.vs.PreviousSummaryHash,
					Realm:               vr.realm,
					Epoch:               vr.vs.NextEpoch,
				}, nil},
				Timestamp: proto.Time(time.Now()),
			}, nil},
			Signatures: make(map[uint64][]byte, 1),
		}
		seh.Head.Head.UpdateEncoding()
		h := sha256.Sum256(seh.Head.Head.PreservedEncoding)
		vr.vs.PreviousSummaryHash = h[:]
		seh.Head.UpdateEncoding()
		seh.Signatures[vr.id] = ed25519.Sign(vr.ratificationKey, proto.MustMarshal(&seh.Head))[:]
		wb.Put(tableRatifications(vr.vs.NextEpoch, vr.id), proto.MustMarshal(seh))
		vs.NextEpoch++
		return func() {
			_, err := vr.keyserver.PushRatification(context.TODO(), seh)
			if err != nil { // TODO: how should this error be handled (grpc issue #238 may be relevant)
				log.Printf("PushRatification: %s", err)
			}
		}
	default:
		log.Fatalf("%d: unknown step: %#v", vs.NextIndex, *step)
	}
	return
}
<commit_msg>remove timeout from verifier dial (oops)<commit_after>// Copyright 2014-2015 The Dename 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 verifier

import (
	"bytes"
	"crypto/sha256"
	"crypto/tls"
	"log"
	"math"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/agl/ed25519"
	"github.com/yahoo/coname/common"
	"github.com/yahoo/coname/proto"
	"github.com/yahoo/coname/server/kv"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// Config encapsulates everything that needs to be specified about a verifer.
// TODO: make this a protobuf, Unmarshal from JSON
type Config struct {
	Realm          string
	KeyserverVerif *proto.AuthorizationPolicy
	KeyserverAddr  string

	ID              uint64
	RatificationKey *[ed25519.PrivateKeySize]byte // [32]byte: secret; [32]byte: public
	TLS             *tls.Config                   // FIXME: tls.Config is not serializable, replicate relevant fields
}

// Verifier verifies that the Keyserver of the realm is not cheating.
// The veirifier is not replicated because one can just run several.
type Verifier struct {
	realm          string
	keyserverVerif *proto.AuthorizationPolicy
	keyserverAddr  string
	auth           credentials.TransportAuthenticator

	id              uint64
	ratificationKey *[ed25519.PrivateKeySize]byte

	db kv.DB
	vs proto.VerifierState

	keyserver proto.E2EKSVerificationClient

	stopOnce sync.Once
	stop     chan struct{}
	waitStop sync.WaitGroup
}

// Start initializes a new verifier based on config and db, or returns an error
// if initialization fails. It then starts the worker goroutine(s).
func Start(cfg *Config, db kv.DB) (*Verifier, error) {
	vr := &Verifier{
		realm:          cfg.Realm,
		keyserverVerif: cfg.KeyserverVerif,
		keyserverAddr:  cfg.KeyserverAddr,
		auth:           credentials.NewTLS(cfg.TLS),

		id:              cfg.ID,
		ratificationKey: cfg.RatificationKey,

		db: db,

		stop: make(chan struct{}),
	}

	switch verifierStateBytes, err := db.Get(tableVerifierState); err {
	case vr.db.ErrNotFound():
		vr.vs.NextEpoch = 1
	case nil:
		if err := vr.vs.Unmarshal(verifierStateBytes); err != nil {
			return nil, err
		}
	default:
		return nil, err
	}

	vr.waitStop.Add(1)
	go func() { vr.run(); vr.waitStop.Done() }()
	return vr, nil
}

// Stop cleanly shuts down the verifier and then returns.
func (vr *Verifier) Stop() {
	vr.stopOnce.Do(func() {
		close(vr.stop)
		vr.waitStop.Wait()
	})
}

func (vr *Verifier) shuttingDown() bool {
	select {
	case <-vr.stop:
		return true
	default:
		return false
	}
}

// run is the CSP-style main loop of the verifier. All code critical for safe
// persistence should be directly in run. All functions called from run should
// either interpret data and modify their mutable arguments OR interact with the
// network and disk, but not both.
func (vr *Verifier) run() {
	keyserverConnection, err := grpc.Dial(vr.keyserverAddr, grpc.WithTransportCredentials(vr.auth))
	if err != nil {
		log.Fatalf("dial %s: %s", vr.keyserverAddr, err)
	}
	vr.keyserver = proto.NewE2EKSVerificationClient(keyserverConnection)
	stream, err := vr.keyserver.VerifierStream(context.TODO(), &proto.VerifierStreamRequest{
		Start:    vr.vs.NextIndex,
		PageSize: math.MaxUint64,
	})
	if err != nil {
		keyserverConnection.Close()
		log.Fatalf("VerifierStream: %s", err)
	}

	wb := vr.db.NewBatch()
	for !vr.shuttingDown() {
		var step *proto.VerifierStep
		step, err = stream.Recv()
		if err != nil {
			log.Printf("VerifierStream.Recv: %s", err)
			break
		}
		wb.Put(tableVerifierLog(vr.vs.NextIndex), proto.MustMarshal(step))
		deferredIO := vr.step(step, &vr.vs, wb)
		vr.vs.NextIndex++
		if deferredIO != nil {
			wb.Put(tableVerifierState, proto.MustMarshal(&vr.vs))
			if err := vr.db.Write(wb); err != nil {
				log.Panicf("sync step to db: %s", err)
			}
			wb.Reset()
			deferredIO()
		}
	}
}

// step is called by run and changes the in-memory state. No i/o allowed.
func (vr *Verifier) step(step *proto.VerifierStep, vs *proto.VerifierState, wb kv.Batch) (deferredIO func()) {
	// vr: &const
	// step, vs, wb: &mut
	switch {
	case step.Update != nil:
		if err := common.VerifyUpdate( /*TODO(dmz): tree lookup*/ nil, step.Update); err != nil {
			return // TODO: decide whether keyserver should guarantee that verifiers don't see bad updates
		}
		// TODO(dmz): set entry in tree
	case step.Epoch != nil:
		ok := common.VerifyPolicy(
			vr.keyserverVerif,
			step.Epoch.Head.PreservedEncoding,
			step.Epoch.Signatures)
		// the bad steps here will not get persisted to disk right now. do we want them to?
		if !ok {
			log.Fatalf("%d: keyserver signature verification failed: %#v", vs.NextIndex, *step)
		}
		r := step.Epoch.Head
		if r.Head.Realm != vr.realm {
			log.Fatalf("%d: seh for realm %q, expected %q: %#v", vr.vs.NextEpoch, r.Head.Realm, vr.realm, *step)
		}
		if r.Head.Epoch != vr.vs.NextEpoch {
			log.Fatalf("%d: got epoch %d instead: %#v", vr.vs.NextEpoch, r.Head.Epoch, *step)
		}
		s := &r.Head
		if !bytes.Equal(s.PreviousSummaryHash, vr.vs.PreviousSummaryHash) {
			log.Fatalf("%d: seh with previous summary hash %q, expected %q: %#v", vr.vs.NextEpoch, s.PreviousSummaryHash, vr.vs.PreviousSummaryHash, *step)
		}
		if s.RootHash != nil /*TODO(dmz): merkletree.GetRootHash()*/ {
			log.Fatalf("%d: seh with root hash %q, expected nil: %#v", vr.vs.NextEpoch, s.RootHash /*, TODO*/, *step)
		}
		seh := &proto.SignedEpochHead{
			Head: proto.TimestampedEpochHead_PreserveEncoding{proto.TimestampedEpochHead{
				Head: proto.EpochHead_PreserveEncoding{proto.EpochHead{
					RootHash:            nil, // TODO(dmz): merklemap.GetRootHash()
					PreviousSummaryHash: vr.vs.PreviousSummaryHash,
					Realm:               vr.realm,
					Epoch:               vr.vs.NextEpoch,
				}, nil},
				Timestamp: proto.Time(time.Now()),
			}, nil},
			Signatures: make(map[uint64][]byte, 1),
		}
		seh.Head.Head.UpdateEncoding()
		h := sha256.Sum256(seh.Head.Head.PreservedEncoding)
		vr.vs.PreviousSummaryHash = h[:]
		seh.Head.UpdateEncoding()
		seh.Signatures[vr.id] = ed25519.Sign(vr.ratificationKey, proto.MustMarshal(&seh.Head))[:]
		wb.Put(tableRatifications(vr.vs.NextEpoch, vr.id), proto.MustMarshal(seh))
		vs.NextEpoch++
		return func() {
			_, err := vr.keyserver.PushRatification(context.TODO(), seh)
			if err != nil { // TODO: how should this error be handled (grpc issue #238 may be relevant)
				log.Printf("PushRatification: %s", err)
			}
		}
	default:
		log.Fatalf("%d: unknown step: %#v", vs.NextIndex, *step)
	}
	return
}
<|endoftext|>