<commit_before>package main

import (
	"context"
	"os"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/grafeas/voucher/v2/cmd/config"
	"github.com/grafeas/voucher/v2/subscriber"
)

var subscriberCmd = &cobra.Command{
	Use:   "subscriber",
	Short: "Runs the subscriber",
	Long: `Run the go subscriber that automatically vouches for images on the specified subscription and project
	use --project=<project> --subscription=<subscription> to specify the project and subscription you want the subscriber to pull from`,
	Run: func(cmd *cobra.Command, args []string) {
		var log = &logrus.Logger{
			Out:       os.Stderr,
			Formatter: new(logrus.JSONFormatter),
			Hooks:     make(logrus.LevelHooks),
			Level:     logrus.DebugLevel,
		}

		secrets, err := config.ReadSecrets()
		if err != nil {
			log.Errorf("error loading EJSON file, no secrets loaded: %s", err)
		}

		metricsClient, err := config.MetricsClient(secrets)
		if err != nil {
			log.Errorf("error configuring metrics client: %s", err)
		}

		config.RegisterDynamicChecks()

		subscriberConfig := subscriber.Config{
			Project:        viper.GetString("pubsub.project"),
			Subscription:   viper.GetString("pubsub.subscription"),
			RequiredChecks: config.GetRequiredChecksFromConfig()["all"],
			DryRun:         viper.GetBool("dryrun"),
			Timeout:        viper.GetInt("pubsub.timeout"),
		}
		voucherSubscriber := subscriber.NewSubscriber(&subscriberConfig, secrets, metricsClient, log)

		err = voucherSubscriber.Subscribe(context.Background())
		if err != nil {
			log.Errorf("couldn't pull pub/sub messages: %s", err)
		}
	},
}

func init() {
	cobra.OnInitialize(config.InitConfig)

	subscriberCmd.Flags().StringP("project", "p", "", "pub/sub project that has the subsciprion (required)")
	subscriberCmd.MarkFlagRequired("project")
	viper.BindPFlag("pubsub.project", subscriberCmd.Flags().Lookup("project"))
	subscriberCmd.Flags().StringP("subscription", "s", "", "pub/sub topic subscription (required)")
	subscriberCmd.MarkFlagRequired("subscription")
	viper.BindPFlag("pubsub.subscription", subscriberCmd.Flags().Lookup("subscription"))
	subscriberCmd.Flags().StringVarP(&config.FileName, "config", "c", "", "path to config")
	subscriberCmd.Flags().IntP("timeout", "", 240, "number of seconds that should be dedicated to a Voucher call")
	viper.BindPFlag("pubsub.timeout", subscriberCmd.Flags().Lookup("timeout"))
}
<commit_msg>flush subscriber on close<commit_after>package main

import (
	"context"
	"io"
	"os"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/grafeas/voucher/v2/cmd/config"
	"github.com/grafeas/voucher/v2/subscriber"
)

var subscriberCmd = &cobra.Command{
	Use:   "subscriber",
	Short: "Runs the subscriber",
	Long: `Run the go subscriber that automatically vouches for images on the specified subscription and project
	use --project=<project> --subscription=<subscription> to specify the project and subscription you want the subscriber to pull from`,
	Run: func(cmd *cobra.Command, args []string) {
		var log = &logrus.Logger{
			Out:       os.Stderr,
			Formatter: new(logrus.JSONFormatter),
			Hooks:     make(logrus.LevelHooks),
			Level:     logrus.DebugLevel,
		}

		secrets, err := config.ReadSecrets()
		if err != nil {
			log.Errorf("error loading EJSON file, no secrets loaded: %s", err)
		}

		metricsClient, err := config.MetricsClient(secrets)
		if err != nil {
			log.Errorf("error configuring metrics client: %s", err)
		} else if closer, ok := metricsClient.(io.Closer); ok {
			defer closer.Close()
		}

		config.RegisterDynamicChecks()

		subscriberConfig := subscriber.Config{
			Project:        viper.GetString("pubsub.project"),
			Subscription:   viper.GetString("pubsub.subscription"),
			RequiredChecks: config.GetRequiredChecksFromConfig()["all"],
			DryRun:         viper.GetBool("dryrun"),
			Timeout:        viper.GetInt("pubsub.timeout"),
		}
		voucherSubscriber := subscriber.NewSubscriber(&subscriberConfig, secrets, metricsClient, log)

		err = voucherSubscriber.Subscribe(context.Background())
		if err != nil {
			log.Errorf("couldn't pull pub/sub messages: %s", err)
		}
	},
}

func init() {
	cobra.OnInitialize(config.InitConfig)

	subscriberCmd.Flags().StringP("project", "p", "", "pub/sub project that has the subsciprion (required)")
	subscriberCmd.MarkFlagRequired("project")
	viper.BindPFlag("pubsub.project", subscriberCmd.Flags().Lookup("project"))
	subscriberCmd.Flags().StringP("subscription", "s", "", "pub/sub topic subscription (required)")
	subscriberCmd.MarkFlagRequired("subscription")
	viper.BindPFlag("pubsub.subscription", subscriberCmd.Flags().Lookup("subscription"))
	subscriberCmd.Flags().StringVarP(&config.FileName, "config", "c", "", "path to config")
	subscriberCmd.Flags().IntP("timeout", "", 240, "number of seconds that should be dedicated to a Voucher call")
	viper.BindPFlag("pubsub.timeout", subscriberCmd.Flags().Lookup("timeout"))
}
<|endoftext|>


<commit_before>package vtls

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	certTime = time.Minute
)

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := map[string]*tls.Certificate{}
	certsM := &sync.RWMutex{}

	addCert := func(name string) (*tls.Certificate, error) {
		log.Println("Generating cert for ", name)
		crt, err := crt.Certify(name, certTime)
		if err != nil {
			return nil, err
		}

		certsM.Lock()
		certs[name] = &crt
		certsM.Unlock()
		return &crt, nil
	}

	getCert := func(name string) (*tls.Certificate, error) {
		lkr := certsM.RLocker()
		lkr.Lock()

		if c, ok := certs[name]; ok {
			log.Println(c.Leaf.NotAfter.String())
			n := time.Now()
			if n.After(c.Leaf.NotBefore) && n.Before(c.Leaf.NotAfter) {
				lkr.Unlock()
				return c, nil
			}
		}
		lkr.Unlock()

		return addCert(name)
	}

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if h.ServerName == "" {
				return nil, fmt.Errorf("Cannot generate certs for IP alone")
			}
			return getCert(h.ServerName)
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<commit_msg>Split certCache into separate type<commit_after>package vtls

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"sync"
	"time"
)

var (
	day  = 24 * time.Hour
	year = 365 * day

	certTime = time.Minute
)

type certCache struct {
	m   map[string]*tls.Certificate
	mut *sync.RWMutex
	crt Certifier
}

func newCertCache(crt Certifier) *certCache {
	return &certCache{
		m:   map[string]*tls.Certificate{},
		mut: &sync.RWMutex{},
		crt: crt,
	}
}

func (cc *certCache) add(name string) (*tls.Certificate, error) {
	crt, err := cc.crt.Certify(name, certTime)
	if err != nil {
		return nil, err
	}

	cc.mut.Lock()
	cc.m[name] = &crt
	cc.mut.Unlock()
	return &crt, nil
}

func (cc *certCache) get(name string) (*tls.Certificate, error) {
	lkr := cc.mut.RLocker()
	lkr.Lock()

	if c, ok := cc.m[name]; ok {
		n := time.Now()
		if n.After(c.Leaf.NotBefore) && n.Before(c.Leaf.NotAfter) {
			lkr.Unlock()
			return c, nil
		}
	}
	lkr.Unlock()

	return cc.add(name)
}

// ListenAndServeTLS mostly mirrors the http.ListenAndServeTLS API, but
// generates the certificates for the server automatically via vault, with a
// short TTL. The function only needs an additional Certifier parameter which
// can generate signed certificates in order to work properly.
func ListenAndServeTLS(addr string, handler http.Handler, crt Certifier) error {
	certs := newCertCache(crt)

	tl, err := tls.Listen("tcp", addr, &tls.Config{
		GetCertificate: func(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if h.ServerName == "" {
				return nil, fmt.Errorf("Cannot generate certs without TLS SNI (no server name was indicated)")
			}
			return certs.get(h.ServerName)
		},
	})

	if err != nil {
		return err
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}

	return http.Serve(tl, handler)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"container/list"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/gogits/gogs/modules/log"

	"github.com/gogits/git"
	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/base"
)

var (
	COMMANDS_READONLY = map[string]int{
		"git-upload-pack":    models.AU_WRITABLE,
		"git upload-pack":    models.AU_WRITABLE,
		"git-upload-archive": models.AU_WRITABLE,
	}

	COMMANDS_WRITE = map[string]int{
		"git-receive-pack": models.AU_READABLE,
		"git receive-pack": models.AU_READABLE,
	}
)

var CmdServ = cli.Command{
	Name:  "serv",
	Usage: "This command just should be called by ssh shell",
	Description: `
gogs serv provide access auth for repositories`,
	Action: runServ,
	Flags:  []cli.Flag{},
}

func parseCmd(cmd string) (string, string) {
	ss := strings.SplitN(cmd, " ", 2)
	if len(ss) != 2 {
		return "", ""
	}

	verb, args := ss[0], ss[1]
	if verb == "git" {
		ss = strings.SplitN(args, " ", 2)
		args = ss[1]
		verb = fmt.Sprintf("%s %s", verb, ss[0])
	}
	return verb, args
}

func In(b string, sl map[string]int) bool {
	_, e := sl[b]
	return e
}

func runServ(k *cli.Context) {
	base.NewConfigContext()
	models.LoadModelsConfig()
	models.NewEngine()

	keys := strings.Split(os.Args[2], "-")
	if len(keys) != 2 {
		fmt.Println("auth file format error")
		return
	}

	keyId, err := strconv.ParseInt(keys[1], 10, 64)
	if err != nil {
		fmt.Println("auth file format error")
		return
	}
	user, err := models.GetUserByKeyId(keyId)
	if err != nil {
		fmt.Println("You have no right to access")
		return
	}

	cmd := os.Getenv("SSH_ORIGINAL_COMMAND")
	if cmd == "" {
		println("Hi", user.Name, "! You've successfully authenticated, but Gogs does not provide shell access.")
		return
	}

	verb, args := parseCmd(cmd)
	rRepo := strings.Trim(args, "'")
	rr := strings.SplitN(rRepo, "/", 2)
	if len(rr) != 2 {
		println("Unavilable repository", args)
		return
	}
	repoName := rr[1]
	if strings.HasSuffix(repoName, ".git") {
		repoName = repoName[:len(repoName)-4]
	}

	repo, err := models.GetRepositoryByName(user.Id, repoName)
	var isExist bool = true
	if err != nil {
		if err == models.ErrRepoNotExist {
			isExist = false
		} else {
			println("Unavilable repository", err)
			return
		}
	}

	isWrite := In(verb, COMMANDS_WRITE)
	isRead := In(verb, COMMANDS_READONLY)

	switch {
	case isWrite:
		has, err := models.HasAccess(user.Name, repoName, models.AU_WRITABLE)
		if err != nil {
			println("Inernel error:", err)
			return
		}
		if !has {
			println("You have no right to write this repository")
			return
		}
	case isRead:
		has, err := models.HasAccess(user.Name, repoName, models.AU_READABLE)
		if err != nil {
			println("Inernel error")
			return
		}
		if !has {
			has, err = models.HasAccess(user.Name, repoName, models.AU_WRITABLE)
			if err != nil {
				println("Inernel error")
				return
			}
		}
		if !has {
			println("You have no right to access this repository")
			return
		}
	default:
		println("Unknown command")
		return
	}

	if !isExist {
		if isRead {
			println("Repository", user.Name+"/"+repoName, "is not exist")
			return
		} else if isWrite {
			_, err := models.CreateRepository(user, repoName, "", "", "", false, true)
			if err != nil {
				println("Create repository failed")
				return
			}
		}
	}

	rep, err := git.OpenRepository(models.RepoPath(user.Name, repoName))
	if err != nil {
		println(err.Error())
		return
	}

	refs, err := rep.AllReferencesMap()
	if err != nil {
		println(err.Error())
		return
	}

	gitcmd := exec.Command(verb, rRepo)
	gitcmd.Dir = base.RepoRootPath

	var s string
	b := bytes.NewBufferString(s)

	gitcmd.Stdout = io.MultiWriter(os.Stdout, b)
	//gitcmd.Stdin = io.MultiReader(os.Stdin, b)
	gitcmd.Stdin = os.Stdin
	gitcmd.Stderr = os.Stderr

	if err = gitcmd.Run(); err != nil {
		println("execute command error:", err.Error())
	}

	if !strings.HasPrefix(cmd, "git-receive-pack") {
		return
	}

	// update
	//w, _ := os.Create("serve.log")
	//defer w.Close()
	//log.SetOutput(w)

	var t = "ok refs/heads/"
	var i int
	var refname string
	for {
		l, err := b.ReadString('\n')
		if err != nil {
			break
		}
		i = i + 1
		l = l[:len(l)-1]
		idx := strings.Index(l, t)
		if idx > 0 {
			refname = l[idx+len(t):]
		}
	}
	var ref *git.Reference
	var ok bool

	var l *list.List
	//log.Info("----", refname, "-----")
	if ref, ok = refs[refname]; !ok {
		refs, err = rep.AllReferencesMap()
		if err != nil {
			println(err.Error())
			return
		}
		if ref, ok = refs[refname]; !ok {
			println("unknow reference name -", refname, "-")
			return
		}
		l, err = ref.AllCommits()
		if err != nil {
			println(err.Error())
			return
		}
	} else {
		//log.Info("----", ref, "-----")
		var last *git.Commit
		//log.Info("00000", ref.Oid.String())
		last, err = ref.LastCommit()
		if err != nil {
			println(err.Error())
			return
		}

		ref2, err := rep.LookupReference(ref.Name)
		if err != nil {
			println(err.Error())
			return
		}

		//log.Info("11111", ref2.Oid.String())
		before, err := ref2.LastCommit()
		if err != nil {
			println(err.Error())
			return
		}
		//log.Info("----", before.Id(), "-----", last.Id())
		l = ref.CommitsBetween(before, last)
	}

	commits := make([][]string, 0)
	var maxCommits = 3
	for e := l.Back(); e != nil; e = e.Prev() {
		commit := e.Value.(*git.Commit)
		commits = append(commits, []string{commit.Id().String(), commit.Message()})
		if len(commits) >= maxCommits {
			break
		}
	}

	if err = models.CommitRepoAction(user.Id, user.Name,
		repo.Id, repoName, refname, &models.PushCommits{l.Len(), commits}); err != nil {
		log.Error("runUpdate.models.CommitRepoAction: %v", err, commits)
	} else {
		//log.Info("refname", refname)
		//log.Info("Listen: %v", cmd)
		//fmt.Println("...", cmd)

		//runUpdate(k)
		c := exec.Command("exec", "git", "update-server-info")
		c.Run()
	}
}
<commit_msg>bug fixed<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"container/list"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/gogits/gogs/modules/log"

	"github.com/gogits/git"
	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/base"
)

var (
	COMMANDS_READONLY = map[string]int{
		"git-upload-pack":    models.AU_WRITABLE,
		"git upload-pack":    models.AU_WRITABLE,
		"git-upload-archive": models.AU_WRITABLE,
	}

	COMMANDS_WRITE = map[string]int{
		"git-receive-pack": models.AU_READABLE,
		"git receive-pack": models.AU_READABLE,
	}
)

var CmdServ = cli.Command{
	Name:  "serv",
	Usage: "This command just should be called by ssh shell",
	Description: `
gogs serv provide access auth for repositories`,
	Action: runServ,
	Flags:  []cli.Flag{},
}

func parseCmd(cmd string) (string, string) {
	ss := strings.SplitN(cmd, " ", 2)
	if len(ss) != 2 {
		return "", ""
	}

	verb, args := ss[0], ss[1]
	if verb == "git" {
		ss = strings.SplitN(args, " ", 2)
		args = ss[1]
		verb = fmt.Sprintf("%s %s", verb, ss[0])
	}
	return verb, args
}

func In(b string, sl map[string]int) bool {
	_, e := sl[b]
	return e
}

func runServ(k *cli.Context) {
	base.NewConfigContext()
	models.LoadModelsConfig()
	models.NewEngine()

	keys := strings.Split(os.Args[2], "-")
	if len(keys) != 2 {
		fmt.Println("auth file format error")
		return
	}

	keyId, err := strconv.ParseInt(keys[1], 10, 64)
	if err != nil {
		fmt.Println("auth file format error")
		return
	}
	user, err := models.GetUserByKeyId(keyId)
	if err != nil {
		fmt.Println("You have no right to access")
		return
	}

	cmd := os.Getenv("SSH_ORIGINAL_COMMAND")
	if cmd == "" {
		println("Hi", user.Name, "! You've successfully authenticated, but Gogs does not provide shell access.")
		return
	}

	verb, args := parseCmd(cmd)
	rRepo := strings.Trim(args, "'")
	rr := strings.SplitN(rRepo, "/", 2)
	if len(rr) != 2 {
		println("Unavilable repository", args)
		return
	}
	repoName := rr[1]
	if strings.HasSuffix(repoName, ".git") {
		repoName = repoName[:len(repoName)-4]
	}

	repo, err := models.GetRepositoryByName(user.Id, repoName)
	var isExist bool = true
	if err != nil {
		if err == models.ErrRepoNotExist {
			isExist = false
		} else {
			println("Unavilable repository", err)
			return
		}
	}

	isWrite := In(verb, COMMANDS_WRITE)
	isRead := In(verb, COMMANDS_READONLY)

	switch {
	case isWrite:
		has, err := models.HasAccess(user.Name, repoName, models.AU_WRITABLE)
		if err != nil {
			println("Inernel error:", err)
			return
		}
		if !has {
			println("You have no right to write this repository")
			return
		}
	case isRead:
		has, err := models.HasAccess(user.Name, repoName, models.AU_READABLE)
		if err != nil {
			println("Inernel error")
			return
		}
		if !has {
			has, err = models.HasAccess(user.Name, repoName, models.AU_WRITABLE)
			if err != nil {
				println("Inernel error")
				return
			}
		}
		if !has {
			println("You have no right to access this repository")
			return
		}
	default:
		println("Unknown command")
		return
	}

	if !isExist {
		if isRead {
			println("Repository", user.Name+"/"+repoName, "is not exist")
			return
		} else if isWrite {
			_, err := models.CreateRepository(user, repoName, "", "", "", false, true)
			if err != nil {
				println("Create repository failed")
				return
			}
		}
	}

	rep, err := git.OpenRepository(models.RepoPath(user.Name, repoName))
	if err != nil {
		println(err.Error())
		return
	}

	refs, err := rep.AllReferencesMap()
	if err != nil {
		println(err.Error())
		return
	}

	gitcmd := exec.Command(verb, rRepo)
	gitcmd.Dir = base.RepoRootPath

	var s string
	b := bytes.NewBufferString(s)

	gitcmd.Stdout = io.MultiWriter(os.Stdout, b)
	//gitcmd.Stdin = io.MultiReader(os.Stdin, b)
	gitcmd.Stdin = os.Stdin
	gitcmd.Stderr = os.Stderr

	if err = gitcmd.Run(); err != nil {
		println("execute command error:", err.Error())
	}

	if !strings.HasPrefix(cmd, "git-receive-pack") {
		return
	}

	// update
	//w, _ := os.Create("serve.log")
	//defer w.Close()
	//log.SetOutput(w)

	var t = "ok refs/heads/"
	var i int
	var refname string
	for {
		l, err := b.ReadString('\n')
		if err != nil {
			break
		}
		i = i + 1
		l = l[:len(l)-1]
		idx := strings.Index(l, t)
		if idx > 0 {
			refname = l[idx+len(t):]
		}
	}
	var ref *git.Reference
	var ok bool

	var l *list.List
	//log.Info("----", refname, "-----")
	if ref, ok = refs[refname]; !ok {
		refs, err = rep.AllReferencesMap()
		if err != nil {
			println(err.Error())
			return
		}
		if ref, ok = refs[refname]; !ok {
			println("unknow reference name -", refname, "-")
			return
		}
		l, err = ref.AllCommits()
		if err != nil {
			println(err.Error())
			return
		}
	} else {
		//log.Info("----", ref, "-----")
		var last *git.Commit
		//log.Info("00000", ref.Oid.String())
		last, err = ref.LastCommit()
		if err != nil {
			println(err.Error())
			return
		}

		ref2, err := rep.LookupReference(ref.Name)
		if err != nil {
			println(err.Error())
			return
		}

		//log.Info("11111", ref2.Oid.String())
		before, err := ref2.LastCommit()
		if err != nil {
			println(err.Error())
			return
		}
		//log.Info("----", before.Id(), "-----", last.Id())
		l = ref.CommitsBetween(before, last)
	}

	commits := make([][]string, 0)
	var maxCommits = 3
	for e := l.Front(); e != nil; e = e.Next() {
		commit := e.Value.(*git.Commit)
		commits = append(commits, []string{commit.Id().String(), commit.Message()})
		if len(commits) >= maxCommits {
			break
		}
	}

	if err = models.CommitRepoAction(user.Id, user.Name,
		repo.Id, repoName, refname, &models.PushCommits{l.Len(), commits}); err != nil {
		log.Error("runUpdate.models.CommitRepoAction: %v", err, commits)
	} else {
		//log.Info("refname", refname)
		//log.Info("Listen: %v", cmd)
		//fmt.Println("...", cmd)

		//runUpdate(k)
		c := exec.Command("exec", "git", "update-server-info")
		c.Run()
	}
}
<|endoftext|>


<commit_before>package dskvs

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

var kvCount int64

const (
	coll       = "games"
	baseKey    = "total annihilation #"
	goroutines = 1
	byteSize   = 100
)

func init() {
	flag.Parse()
	if testing.Short() {
		kvCount = 2048 // To be runable with race detector
	} else {
		kvCount = 10000
	}
	log.Printf("Concurrency Benchmark - Goroutines=%d, unique key-value=%d\n", goroutines, kvCount)
}

type keyValue struct {
	Key   string
	Value []byte
}

type Context struct {
	t      *testing.T
	s      *Store
	kv     keyValue
	wg     *sync.WaitGroup
	sem    chan int
	dur    chan time.Duration
	errors chan error
}

func TestOneOperationWithMultipleConcurrentRequest(t *testing.T) {

	log.Printf("Sequence of %d bytes operations by group, "+
		"%d concurrent request in each groups", byteSize, goroutines)

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount, t)

	log.Println("Put - first time")
	putStats := runTest(doPutRequest, 1, goroutines, store, expectedList, t)
	log.Println(putStats.String())

	log.Println("Put - rewrite")
	rePutStats := runTest(doPutRequest, 1, goroutines, store, expectedList, t)
	log.Println(rePutStats.String())

	log.Printf("Get")
	getStats := runTest(doGetRequest, 1, goroutines, store, expectedList, t)
	log.Println(getStats.String())

	log.Printf("Delete")
	deleteStats := runTest(doDeleteRequest, 1, goroutines, store, expectedList, t)
	log.Println(deleteStats.String())

	log.Printf("by %d cpus, using %d concurrent goroutines\n",
		runtime.NumCPU(), goroutines)

}

func TestManyOperationWithMultipleConcurrentRequest(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount/3, t)

	// Start writing/reading concurrently, in random order
	concurrentFunc := func(ctx Context) {
		switch rand.Intn(5) {
		case 0:
			doFailGetRequest(ctx)
			doPutRequest(ctx)
			doDeleteRequest(ctx)
		case 1:
			doPutRequest(ctx)
			doGetRequest(ctx)
			doDeleteRequest(ctx)
		case 2:
			doPutRequest(ctx)
			doGetRequest(ctx)
			doDeleteRequest(ctx)
		case 3:
			doPutRequest(ctx)
			doDeleteRequest(ctx)
			doFailGetRequest(ctx)
		case 4:
			doFailGetRequest(ctx)
			doPutRequest(ctx)
			doGetRequest(ctx)
		}
	}
	_ = runTest(concurrentFunc, 3, goroutines, store, expectedList, t)

	// Cleanup
	err := store.DeleteAll(coll)
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		checkGetIsEmpty(store, kv.Key, t)
	}

}

func TestConcurrentPutCanBeGetAllAndDeleteAll(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount, t)

	// Don't care about the stats
	_ = runTest(doPutRequest, 1, 100, store, expectedList, t)

	actual, err := store.GetAll(coll)
	if err != nil {
		t.Errorf("Error on GetAll(%s), %v", coll, err)
	} else if len(actual) != len(expectedList) {
		t.Errorf("Expected len(store.GetAll)=<%s> but was <%s>",
			len(expectedList),
			len(actual))
	}

	err = store.DeleteAll(coll)
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		checkGetIsEmpty(store, kv.Key, t)
	}

}

func runTest(
	testedFunc func(Context),
	nGo int,
	goroutines int,
	store *Store,
	expectedList []keyValue,
	t *testing.T,
) stats {

	cErr := make(chan error)
	var wg sync.WaitGroup
	durations := make(chan time.Duration, len(expectedList)*nGo)

	sem := make(chan int, goroutines)
	for _, kv := range expectedList {
		wg.Add(nGo)
		ctx := Context{
			t:      t,
			s:      store,
			kv:     kv,
			wg:     &wg,
			sem:    sem,
			dur:    durations,
			errors: cErr,
		}
		go testedFunc(ctx)

	}

	wg.Wait()
	close(durations)

	if len(cErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(cErr))
	}

	var dTList []time.Duration
	for dT := range durations {
		dTList = append(dTList, dT)
	}

	return newStats(dTList)
}

func doPutRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	err := ctx.s.Put(ctx.kv.Key, ctx.kv.Value)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.t.Fatal("Received an error", err)
		ctx.errors <- err
	}
	ctx.wg.Done()
	<-ctx.sem
}

func doGetRequest(ctx Context) {
	ctx.sem <- 1
	expected := ctx.kv.Value

	t0 := time.Now()
	actual, err := ctx.s.Get(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.errors <- err
	} else if !bytes.Equal(expected, actual) {
		ctx.t.Errorf("Expected <%s> but was <%s>",
			expected,
			actual)
	}
	ctx.wg.Done()
	<-ctx.sem
}

// Get can fail when there's no such key, not true for Put and Delete
func doFailGetRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	_, err := ctx.s.Get(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if _, ok := err.(KeyError); !ok {
		ctx.errors <- errors.New(fmt.Sprintf("Should have failed on Get(%s)", ctx.kv.Key))
	}
	ctx.wg.Done()
	<-ctx.sem
}

func doDeleteRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	err := ctx.s.Delete(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.t.Fatal("Received an error", err)
		ctx.errors <- err
	}
	ctx.wg.Done()
	<-ctx.sem
}

func generateKeyValueList(kvCount int64, t *testing.T) []keyValue {

	coll := "games"
	baseKey := "total annihilation #"

	kvList := make([]keyValue, kvCount)
	for i := int64(0); i < kvCount; i++ {
		key := coll + CollKeySep + baseKey + strconv.FormatInt(i, 10)
		data := make([]byte, byteSize)
		//data := generateData(Data{"It's fun!"}, t)
		kv := keyValue{key, data}
		kvList[i] = kv
	}

	return kvList
}
<commit_msg>Set test to exercise concurrency, not benchmark.<commit_after>package dskvs

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

var kvCount int64

const (
	coll       = "games"
	baseKey    = "total annihilation #"
	goroutines = 10
	byteSize   = 1000
)

func init() {
	flag.Parse()
	if testing.Short() {
		kvCount = 2048 // To be runable with race detector
	} else {
		kvCount = 100000
	}
	log.Printf("Concurrency Benchmark - Goroutines=%d, unique key-value=%d\n", goroutines, kvCount)
}

type keyValue struct {
	Key   string
	Value []byte
}

type Context struct {
	t      *testing.T
	s      *Store
	kv     keyValue
	wg     *sync.WaitGroup
	sem    chan int
	dur    chan time.Duration
	errors chan error
}

func TestOneOperationWithMultipleConcurrentRequest(t *testing.T) {

	log.Printf("Sequence of %d bytes operations by group, "+
		"%d concurrent request in each groups", byteSize, goroutines)

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount, t)

	log.Println("Put - first time")
	putStats := runTest(doPutRequest, 1, goroutines, store, expectedList, t)
	log.Println(putStats.String())

	log.Println("Put - rewrite")
	rePutStats := runTest(doPutRequest, 1, goroutines, store, expectedList, t)
	log.Println(rePutStats.String())

	log.Printf("Get")
	getStats := runTest(doGetRequest, 1, goroutines, store, expectedList, t)
	log.Println(getStats.String())

	log.Printf("Delete")
	deleteStats := runTest(doDeleteRequest, 1, goroutines, store, expectedList, t)
	log.Println(deleteStats.String())

	log.Printf("by %d cpus, using %d concurrent goroutines\n",
		runtime.NumCPU(), goroutines)

}

func TestManyOperationWithMultipleConcurrentRequest(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount/3, t)

	// Start writing/reading concurrently, in random order
	concurrentFunc := func(ctx Context) {
		switch rand.Intn(5) {
		case 0:
			doFailGetRequest(ctx)
			doPutRequest(ctx)
			doDeleteRequest(ctx)
		case 1:
			doPutRequest(ctx)
			doGetRequest(ctx)
			doDeleteRequest(ctx)
		case 2:
			doPutRequest(ctx)
			doGetRequest(ctx)
			doDeleteRequest(ctx)
		case 3:
			doPutRequest(ctx)
			doDeleteRequest(ctx)
			doFailGetRequest(ctx)
		case 4:
			doFailGetRequest(ctx)
			doPutRequest(ctx)
			doGetRequest(ctx)
		}
	}
	_ = runTest(concurrentFunc, 3, goroutines, store, expectedList, t)

	// Cleanup
	err := store.DeleteAll(coll)
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		checkGetIsEmpty(store, kv.Key, t)
	}

}

func TestConcurrentPutCanBeGetAllAndDeleteAll(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	expectedList := generateKeyValueList(kvCount, t)

	// Don't care about the stats
	_ = runTest(doPutRequest, 1, 100, store, expectedList, t)

	actual, err := store.GetAll(coll)
	if err != nil {
		t.Errorf("Error on GetAll(%s), %v", coll, err)
	} else if len(actual) != len(expectedList) {
		t.Errorf("Expected len(store.GetAll)=<%s> but was <%s>",
			len(expectedList),
			len(actual))
	}

	err = store.DeleteAll(coll)
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		checkGetIsEmpty(store, kv.Key, t)
	}

}

func runTest(
	testedFunc func(Context),
	nGo int,
	goroutines int,
	store *Store,
	expectedList []keyValue,
	t *testing.T,
) stats {

	cErr := make(chan error)
	var wg sync.WaitGroup
	durations := make(chan time.Duration, len(expectedList)*nGo)

	sem := make(chan int, goroutines)
	for _, kv := range expectedList {
		wg.Add(nGo)
		ctx := Context{
			t:      t,
			s:      store,
			kv:     kv,
			wg:     &wg,
			sem:    sem,
			dur:    durations,
			errors: cErr,
		}
		go testedFunc(ctx)

	}

	wg.Wait()
	close(durations)

	if len(cErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(cErr))
	}

	var dTList []time.Duration
	for dT := range durations {
		dTList = append(dTList, dT)
	}

	return newStats(dTList)
}

func doPutRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	err := ctx.s.Put(ctx.kv.Key, ctx.kv.Value)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.t.Fatal("Received an error", err)
		ctx.errors <- err
	}
	ctx.wg.Done()
	<-ctx.sem
}

func doGetRequest(ctx Context) {
	ctx.sem <- 1
	expected := ctx.kv.Value

	t0 := time.Now()
	actual, err := ctx.s.Get(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.errors <- err
	} else if !bytes.Equal(expected, actual) {
		ctx.t.Errorf("Expected <%s> but was <%s>",
			expected,
			actual)
	}
	ctx.wg.Done()
	<-ctx.sem
}

// Get can fail when there's no such key, not true for Put and Delete
func doFailGetRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	_, err := ctx.s.Get(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if _, ok := err.(KeyError); !ok {
		ctx.errors <- errors.New(fmt.Sprintf("Should have failed on Get(%s)", ctx.kv.Key))
	}
	ctx.wg.Done()
	<-ctx.sem
}

func doDeleteRequest(ctx Context) {
	ctx.sem <- 1
	t0 := time.Now()
	err := ctx.s.Delete(ctx.kv.Key)
	dT := time.Since(t0)

	ctx.dur <- dT

	if err != nil {
		ctx.t.Fatal("Received an error", err)
		ctx.errors <- err
	}
	ctx.wg.Done()
	<-ctx.sem
}

func generateKeyValueList(kvCount int64, t *testing.T) []keyValue {

	coll := "games"
	baseKey := "total annihilation #"

	kvList := make([]keyValue, kvCount)
	for i := int64(0); i < kvCount; i++ {
		key := coll + CollKeySep + baseKey + strconv.FormatInt(i, 10)
		data := make([]byte, byteSize)
		//data := generateData(Data{"It's fun!"}, t)
		kv := keyValue{key, data}
		kvList[i] = kv
	}

	return kvList
}
<|endoftext|>


<commit_before>package openstack_test

import (
	"crypto/rand"
	"fmt"
	"io"
	. "launchpad.net/gocheck"
	"launchpad.net/goose/client"
	"launchpad.net/goose/identity"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/jujutest"
	"launchpad.net/juju-core/environs/openstack"
	envtesting "launchpad.net/juju-core/environs/testing"
	coretesting "launchpad.net/juju-core/testing"
)

// generate a different bucket name for each config instance, so that
// we are not polluted by previous test state.
func randomName() string {
	buf := make([]byte, 8)
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		panic(fmt.Sprintf("error from crypto rand: %v", err))
	}
	return fmt.Sprintf("%x", buf)
}

func makeTestConfig(cred *identity.Credentials) map[string]interface{} {
	// The following attributes hold the environment configuration
	// for running the OpenStack integration tests.
	//
	// This is missing keys for security reasons; set the following
	// environment variables to make the OpenStack testing work:
	//  access-key: $OS_USERNAME
	//  secret-key: $OS_PASSWORD
	//
	attrs := map[string]interface{}{
		"name":            "sample-" + randomName(),
		"type":            "openstack",
		"auth-mode":       "userpass",
		"control-bucket":  "juju-test-" + randomName(),
		"ca-cert":         coretesting.CACert,
		"ca-private-key":  coretesting.CAKey,
		"authorized-keys": "fakekey",
		"admin-secret":    "secret",
		"username":        cred.User,
		"password":        cred.Secrets,
		"region":          cred.Region,
		"auth-url":        cred.URL,
		"tenant-name":     cred.TenantName,
	}
	return attrs
}

// Register tests to run against a real Openstack instance.
func registerLiveTests(cred *identity.Credentials, testImageDetails openstack.ImageDetails) {
	config := makeTestConfig(cred)
	config["default-image-id"] = testImageDetails.ImageId
	config["default-instance-type"] = testImageDetails.Flavor
	Suite(&LiveTests{
		cred: cred,
		LiveTests: jujutest.LiveTests{
			TestConfig: jujutest.TestConfig{config},
			Attempt:    *openstack.ShortAttempt,
			// TODO: Bug #1133263, once the infrastructure is set up,
			//       enable The state tests on openstack
			CanOpenState: false,
			// TODO: Bug #1133272, enabling this requires mapping from
			//       'series' to an image id, when we have support, set
			//       this flag to True.
			HasProvisioner: false,
		},
	})
}

// LiveTests contains tests that can be run against OpenStack deployments.
// The deployment can be a real live instance or service doubles.
// Each test runs using the same connection.
type LiveTests struct {
	coretesting.LoggingSuite
	jujutest.LiveTests
	cred                   *identity.Credentials
	writeablePublicStorage environs.Storage
}

func (t *LiveTests) SetUpSuite(c *C) {
	t.LoggingSuite.SetUpSuite(c)
	// Update some Config items now that we have services running.
	// This is setting the public-bucket-url and auth-url because that
	// information is set during startup of the localLiveSuite
	cl := client.NewClient(t.cred, identity.AuthUserPass, nil)
	err := cl.Authenticate()
	c.Assert(err, IsNil)
	publicBucketURL, err := cl.MakeServiceURL("object-store", nil)
	c.Assert(err, IsNil)
	t.TestConfig.UpdateConfig(map[string]interface{}{
		"public-bucket-url": publicBucketURL,
		"auth-url":          t.cred.URL,
	})
	t.LiveTests.SetUpSuite(c)
	// Environ.PublicStorage() is read only.
	// For testing, we create a specific storage instance which is authorised to write to
	// the public storage bucket so that we can upload files for testing.
	t.writeablePublicStorage = openstack.WritablePublicStorage(t.Env)
	// Put some fake tools in place so that tests that are simply
	// starting instances without any need to check if those instances
	// are running will find them in the public bucket.
	envtesting.UploadFakeTools(c, t.writeablePublicStorage)
}

func (t *LiveTests) TearDownSuite(c *C) {
	if t.Env == nil {
		// This can happen if SetUpSuite fails.
		return
	}
	if t.writeablePublicStorage != nil {
		envtesting.RemoveFakeTools(c, t.writeablePublicStorage)
	}
	t.LiveTests.TearDownSuite(c)
	t.LoggingSuite.TearDownSuite(c)
}

func (t *LiveTests) SetUpTest(c *C) {
	t.LoggingSuite.SetUpTest(c)
	t.LiveTests.SetUpTest(c)
}

func (t *LiveTests) TearDownTest(c *C) {
	t.LiveTests.TearDownTest(c)
	t.LoggingSuite.TearDownTest(c)
}
<commit_msg>Enable extra live tests for openstack<commit_after>package openstack_test

import (
	"crypto/rand"
	"fmt"
	"io"
	. "launchpad.net/gocheck"
	"launchpad.net/goose/client"
	"launchpad.net/goose/identity"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/jujutest"
	"launchpad.net/juju-core/environs/openstack"
	envtesting "launchpad.net/juju-core/environs/testing"
	coretesting "launchpad.net/juju-core/testing"
)

// generate a different bucket name for each config instance, so that
// we are not polluted by previous test state.
func randomName() string {
	buf := make([]byte, 8)
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		panic(fmt.Sprintf("error from crypto rand: %v", err))
	}
	return fmt.Sprintf("%x", buf)
}

func makeTestConfig(cred *identity.Credentials) map[string]interface{} {
	// The following attributes hold the environment configuration
	// for running the OpenStack integration tests.
	//
	// This is missing keys for security reasons; set the following
	// environment variables to make the OpenStack testing work:
	//  access-key: $OS_USERNAME
	//  secret-key: $OS_PASSWORD
	//
	attrs := map[string]interface{}{
		"name":            "sample-" + randomName(),
		"type":            "openstack",
		"auth-mode":       "userpass",
		"control-bucket":  "juju-test-" + randomName(),
		"ca-cert":         coretesting.CACert,
		"ca-private-key":  coretesting.CAKey,
		"authorized-keys": "fakekey",
		"admin-secret":    "secret",
		"username":        cred.User,
		"password":        cred.Secrets,
		"region":          cred.Region,
		"auth-url":        cred.URL,
		"tenant-name":     cred.TenantName,
	}
	return attrs
}

// Register tests to run against a real Openstack instance.
func registerLiveTests(cred *identity.Credentials, testImageDetails openstack.ImageDetails) {
	config := makeTestConfig(cred)
	config["default-image-id"] = testImageDetails.ImageId
	config["default-instance-type"] = testImageDetails.Flavor
	Suite(&LiveTests{
		cred: cred,
		LiveTests: jujutest.LiveTests{
			TestConfig: jujutest.TestConfig{config},
			Attempt:    *openstack.ShortAttempt,
			CanOpenState: true,
			HasProvisioner: true,
		},
	})
}

// LiveTests contains tests that can be run against OpenStack deployments.
// The deployment can be a real live instance or service doubles.
// Each test runs using the same connection.
type LiveTests struct {
	coretesting.LoggingSuite
	jujutest.LiveTests
	cred                   *identity.Credentials
	writeablePublicStorage environs.Storage
}

func (t *LiveTests) SetUpSuite(c *C) {
	t.LoggingSuite.SetUpSuite(c)
	// Update some Config items now that we have services running.
	// This is setting the public-bucket-url and auth-url because that
	// information is set during startup of the localLiveSuite
	cl := client.NewClient(t.cred, identity.AuthUserPass, nil)
	err := cl.Authenticate()
	c.Assert(err, IsNil)
	publicBucketURL, err := cl.MakeServiceURL("object-store", nil)
	c.Assert(err, IsNil)
	t.TestConfig.UpdateConfig(map[string]interface{}{
		"public-bucket-url": publicBucketURL,
		"auth-url":          t.cred.URL,
	})
	t.LiveTests.SetUpSuite(c)
	// Environ.PublicStorage() is read only.
	// For testing, we create a specific storage instance which is authorised to write to
	// the public storage bucket so that we can upload files for testing.
	t.writeablePublicStorage = openstack.WritablePublicStorage(t.Env)
	// Put some fake tools in place so that tests that are simply
	// starting instances without any need to check if those instances
	// are running will find them in the public bucket.
	envtesting.UploadFakeTools(c, t.writeablePublicStorage)
}

func (t *LiveTests) TearDownSuite(c *C) {
	if t.Env == nil {
		// This can happen if SetUpSuite fails.
		return
	}
	if t.writeablePublicStorage != nil {
		envtesting.RemoveFakeTools(c, t.writeablePublicStorage)
	}
	t.LiveTests.TearDownSuite(c)
	t.LoggingSuite.TearDownSuite(c)
}

func (t *LiveTests) SetUpTest(c *C) {
	t.LoggingSuite.SetUpTest(c)
	t.LiveTests.SetUpTest(c)
}

func (t *LiveTests) TearDownTest(c *C) {
	t.LiveTests.TearDownTest(c)
	t.LoggingSuite.TearDownTest(c)
}
<|endoftext|>


<commit_before>// Copyright 2010 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 tls

import (
	"crypto"
	"crypto/md5"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
)

// a keyAgreement implements the client and server side of a TLS key agreement
// protocol by generating and processing key exchange messages.
type keyAgreement interface {
	// On the server side, the first two methods are called in order.

	// In the case that the key agreement protocol doesn't use a
	// ServerKeyExchange message, generateServerKeyExchange can return nil,
	// nil.
	generateServerKeyExchange(*Config, *Certificate, *clientHelloMsg, *serverHelloMsg) (*serverKeyExchangeMsg, error)
	processClientKeyExchange(*Config, *Certificate, *clientKeyExchangeMsg, uint16) ([]byte, error)

	// On the client side, the next two methods are called in order.

	// This method may not be called if the server doesn't send a
	// ServerKeyExchange message.
	processServerKeyExchange(*Config, *clientHelloMsg, *serverHelloMsg, *x509.Certificate, *serverKeyExchangeMsg) error
	generateClientKeyExchange(*Config, *clientHelloMsg, *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error)
}

var errClientKeyExchange = errors.New("tls: invalid ClientKeyExchange message")
var errServerKeyExchange = errors.New("tls: invalid ServerKeyExchange message")

// rsaKeyAgreement implements the standard TLS key agreement where the client
// encrypts the pre-master secret to the server's public key.
type rsaKeyAgreement struct{}

func (ka rsaKeyAgreement) generateServerKeyExchange(config *Config, cert *Certificate, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
	return nil, nil
}

func (ka rsaKeyAgreement) processClientKeyExchange(config *Config, cert *Certificate, ckx *clientKeyExchangeMsg, version uint16) ([]byte, error) {
	if len(ckx.ciphertext) < 2 {
		return nil, errClientKeyExchange
	}
	ciphertextLen := int(ckx.ciphertext[0])<<8 | int(ckx.ciphertext[1])
	if ciphertextLen != len(ckx.ciphertext)-2 {
		return nil, errClientKeyExchange
	}
	ciphertext := ckx.ciphertext[2:]

	priv, ok := cert.PrivateKey.(crypto.Decrypter)
	if !ok {
		return nil, errors.New("tls: certificate private key does not implement crypto.Decrypter")
	}
	// Perform constant time RSA PKCS #1 v1.5 decryption
	preMasterSecret, err := priv.Decrypt(config.rand(), ciphertext, &rsa.PKCS1v15DecryptOptions{SessionKeyLen: 48})
	if err != nil {
		return nil, err
	}
	// We don't check the version number in the premaster secret. For one,
	// by checking it, we would leak information about the validity of the
	// encrypted pre-master secret. Secondly, it provides only a small
	// benefit against a downgrade attack and some implementations send the
	// wrong version anyway. See the discussion at the end of section
	// 7.4.7.1 of RFC 4346.
	return preMasterSecret, nil
}

func (ka rsaKeyAgreement) processServerKeyExchange(config *Config, clientHello *clientHelloMsg, serverHello *serverHelloMsg, cert *x509.Certificate, skx *serverKeyExchangeMsg) error {
	return errors.New("tls: unexpected ServerKeyExchange")
}

func (ka rsaKeyAgreement) generateClientKeyExchange(config *Config, clientHello *clientHelloMsg, cert *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error) {
	preMasterSecret := make([]byte, 48)
	preMasterSecret[0] = byte(clientHello.vers >> 8)
	preMasterSecret[1] = byte(clientHello.vers)
	_, err := io.ReadFull(config.rand(), preMasterSecret[2:])
	if err != nil {
		return nil, nil, err
	}

	encrypted, err := rsa.EncryptPKCS1v15(config.rand(), cert.PublicKey.(*rsa.PublicKey), preMasterSecret)
	if err != nil {
		return nil, nil, err
	}
	ckx := new(clientKeyExchangeMsg)
	ckx.ciphertext = make([]byte, len(encrypted)+2)
	ckx.ciphertext[0] = byte(len(encrypted) >> 8)
	ckx.ciphertext[1] = byte(len(encrypted))
	copy(ckx.ciphertext[2:], encrypted)
	return preMasterSecret, ckx, nil
}

// sha1Hash calculates a SHA1 hash over the given byte slices.
func sha1Hash(slices [][]byte) []byte {
	hsha1 := sha1.New()
	for _, slice := range slices {
		hsha1.Write(slice)
	}
	return hsha1.Sum(nil)
}

// md5SHA1Hash implements TLS 1.0's hybrid hash function which consists of the
// concatenation of an MD5 and SHA1 hash.
func md5SHA1Hash(slices [][]byte) []byte {
	md5sha1 := make([]byte, md5.Size+sha1.Size)
	hmd5 := md5.New()
	for _, slice := range slices {
		hmd5.Write(slice)
	}
	copy(md5sha1, hmd5.Sum(nil))
	copy(md5sha1[md5.Size:], sha1Hash(slices))
	return md5sha1
}

// hashForServerKeyExchange hashes the given slices and returns their digest
// using the given hash function (for >= TLS 1.2) or using a default based on
// the sigType (for earlier TLS versions). For Ed25519 signatures, which don't
// do pre-hashing, it returns the concatenation of the slices.
func hashForServerKeyExchange(sigType uint8, hashFunc crypto.Hash, version uint16, slices ...[]byte) []byte {
	if sigType == signatureEd25519 {
		var signed []byte
		for _, slice := range slices {
			signed = append(signed, slice...)
		}
		return signed
	}
	if version >= VersionTLS12 {
		h := hashFunc.New()
		for _, slice := range slices {
			h.Write(slice)
		}
		digest := h.Sum(nil)
		return digest
	}
	if sigType == signatureECDSA {
		return sha1Hash(slices)
	}
	return md5SHA1Hash(slices)
}

// ecdheKeyAgreement implements a TLS key agreement where the server
// generates an ephemeral EC public/private key pair and signs it. The
// pre-master secret is then calculated using ECDH. The signature may
// be ECDSA, Ed25519 or RSA.
type ecdheKeyAgreement struct {
	version uint16
	isRSA   bool
	params  ecdheParameters

	// ckx and preMasterSecret are generated in processServerKeyExchange
	// and returned in generateClientKeyExchange.
	ckx             *clientKeyExchangeMsg
	preMasterSecret []byte
}

func (ka *ecdheKeyAgreement) generateServerKeyExchange(config *Config, cert *Certificate, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
	var curveID CurveID
	for _, c := range clientHello.supportedCurves {
		if config.supportsCurve(c) {
			curveID = c
			break
		}
	}

	if curveID == 0 {
		return nil, errors.New("tls: no supported elliptic curves offered")
	}
	if _, ok := curveForCurveID(curveID); curveID != X25519 && !ok {
		return nil, errors.New("tls: CurvePreferences includes unsupported curve")
	}

	params, err := generateECDHEParameters(config.rand(), curveID)
	if err != nil {
		return nil, err
	}
	ka.params = params

	// See RFC 4492, Section 5.4.
	ecdhePublic := params.PublicKey()
	serverECDHEParams := make([]byte, 1+2+1+len(ecdhePublic))
	serverECDHEParams[0] = 3 // named curve
	serverECDHEParams[1] = byte(curveID >> 8)
	serverECDHEParams[2] = byte(curveID)
	serverECDHEParams[3] = byte(len(ecdhePublic))
	copy(serverECDHEParams[4:], ecdhePublic)

	priv, ok := cert.PrivateKey.(crypto.Signer)
	if !ok {
		return nil, fmt.Errorf("tls: certificate private key of type %T does not implement crypto.Signer", cert.PrivateKey)
	}

	var signatureAlgorithm SignatureScheme
	var sigType uint8
	var sigHash crypto.Hash
	if ka.version >= VersionTLS12 {
		signatureAlgorithm, err = selectSignatureScheme(ka.version, cert, clientHello.supportedSignatureAlgorithms)
		if err != nil {
			return nil, err
		}
		sigType, sigHash, err = typeAndHashFromSignatureScheme(signatureAlgorithm)
		if err != nil {
			return nil, err
		}
	} else {
		sigType, sigHash, err = legacyTypeAndHashFromPublicKey(priv.Public())
		if err != nil {
			return nil, err
		}
	}
	if (sigType == signaturePKCS1v15 || sigType == signatureRSAPSS) != ka.isRSA {
		return nil, errors.New("tls: certificate cannot be used with the selected cipher suite")
	}

	signed := hashForServerKeyExchange(sigType, sigHash, ka.version, clientHello.random, hello.random, serverECDHEParams)

	signOpts := crypto.SignerOpts(sigHash)
	if sigType == signatureRSAPSS {
		signOpts = &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash, Hash: sigHash}
	}
	sig, err := priv.Sign(config.rand(), signed, signOpts)
	if err != nil {
		return nil, errors.New("tls: failed to sign ECDHE parameters: " + err.Error())
	}

	skx := new(serverKeyExchangeMsg)
	sigAndHashLen := 0
	if ka.version >= VersionTLS12 {
		sigAndHashLen = 2
	}
	skx.key = make([]byte, len(serverECDHEParams)+sigAndHashLen+2+len(sig))
	copy(skx.key, serverECDHEParams)
	k := skx.key[len(serverECDHEParams):]
	if ka.version >= VersionTLS12 {
		k[0] = byte(signatureAlgorithm >> 8)
		k[1] = byte(signatureAlgorithm)
		k = k[2:]
	}
	k[0] = byte(len(sig) >> 8)
	k[1] = byte(len(sig))
	copy(k[2:], sig)

	return skx, nil
}

func (ka *ecdheKeyAgreement) processClientKeyExchange(config *Config, cert *Certificate, ckx *clientKeyExchangeMsg, version uint16) ([]byte, error) {
	if len(ckx.ciphertext) == 0 || int(ckx.ciphertext[0]) != len(ckx.ciphertext)-1 {
		return nil, errClientKeyExchange
	}

	preMasterSecret := ka.params.SharedKey(ckx.ciphertext[1:])
	if preMasterSecret == nil {
		return nil, errClientKeyExchange
	}

	return preMasterSecret, nil
}

func (ka *ecdheKeyAgreement) processServerKeyExchange(config *Config, clientHello *clientHelloMsg, serverHello *serverHelloMsg, cert *x509.Certificate, skx *serverKeyExchangeMsg) error {
	if len(skx.key) < 4 {
		return errServerKeyExchange
	}
	if skx.key[0] != 3 { // named curve
		return errors.New("tls: server selected unsupported curve")
	}
	curveID := CurveID(skx.key[1])<<8 | CurveID(skx.key[2])

	publicLen := int(skx.key[3])
	if publicLen+4 > len(skx.key) {
		return errServerKeyExchange
	}
	serverECDHEParams := skx.key[:4+publicLen]
	publicKey := serverECDHEParams[4:]

	sig := skx.key[4+publicLen:]
	if len(sig) < 2 {
		return errServerKeyExchange
	}

	if _, ok := curveForCurveID(curveID); curveID != X25519 && !ok {
		return errors.New("tls: server selected unsupported curve")
	}

	params, err := generateECDHEParameters(config.rand(), curveID)
	if err != nil {
		return err
	}
	ka.params = params

	ka.preMasterSecret = params.SharedKey(publicKey)
	if ka.preMasterSecret == nil {
		return errServerKeyExchange
	}

	ourPublicKey := params.PublicKey()
	ka.ckx = new(clientKeyExchangeMsg)
	ka.ckx.ciphertext = make([]byte, 1+len(ourPublicKey))
	ka.ckx.ciphertext[0] = byte(len(ourPublicKey))
	copy(ka.ckx.ciphertext[1:], ourPublicKey)

	var sigType uint8
	var sigHash crypto.Hash
	if ka.version >= VersionTLS12 {
		signatureAlgorithm := SignatureScheme(sig[0])<<8 | SignatureScheme(sig[1])
		sig = sig[2:]
		if len(sig) < 2 {
			return errServerKeyExchange
		}

		if !isSupportedSignatureAlgorithm(signatureAlgorithm, clientHello.supportedSignatureAlgorithms) {
			return errors.New("tls: certificate used with invalid signature algorithm")
		}
		sigType, sigHash, err = typeAndHashFromSignatureScheme(signatureAlgorithm)
		if err != nil {
			return err
		}
	} else {
		sigType, sigHash, err = legacyTypeAndHashFromPublicKey(cert.PublicKey)
		if err != nil {
			return err
		}
	}
	if (sigType == signaturePKCS1v15 || sigType == signatureRSAPSS) != ka.isRSA {
		return errServerKeyExchange
	}

	sigLen := int(sig[0])<<8 | int(sig[1])
	if sigLen+2 != len(sig) {
		return errServerKeyExchange
	}
	sig = sig[2:]

	signed := hashForServerKeyExchange(sigType, sigHash, ka.version, clientHello.random, serverHello.random, serverECDHEParams)
	if err := verifyHandshakeSignature(sigType, cert.PublicKey, sigHash, signed, sig); err != nil {
		return errors.New("tls: invalid signature by the server certificate: " + err.Error())
	}
	return nil
}

func (ka *ecdheKeyAgreement) generateClientKeyExchange(config *Config, clientHello *clientHelloMsg, cert *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error) {
	if ka.ckx == nil {
		return nil, nil, errors.New("tls: missing ServerKeyExchange message")
	}

	return ka.preMasterSecret, ka.ckx, nil
}
<commit_msg>crypto/tls: test key type when casting<commit_after>// Copyright 2010 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 tls

import (
	"crypto"
	"crypto/md5"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
)

// a keyAgreement implements the client and server side of a TLS key agreement
// protocol by generating and processing key exchange messages.
type keyAgreement interface {
	// On the server side, the first two methods are called in order.

	// In the case that the key agreement protocol doesn't use a
	// ServerKeyExchange message, generateServerKeyExchange can return nil,
	// nil.
	generateServerKeyExchange(*Config, *Certificate, *clientHelloMsg, *serverHelloMsg) (*serverKeyExchangeMsg, error)
	processClientKeyExchange(*Config, *Certificate, *clientKeyExchangeMsg, uint16) ([]byte, error)

	// On the client side, the next two methods are called in order.

	// This method may not be called if the server doesn't send a
	// ServerKeyExchange message.
	processServerKeyExchange(*Config, *clientHelloMsg, *serverHelloMsg, *x509.Certificate, *serverKeyExchangeMsg) error
	generateClientKeyExchange(*Config, *clientHelloMsg, *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error)
}

var errClientKeyExchange = errors.New("tls: invalid ClientKeyExchange message")
var errServerKeyExchange = errors.New("tls: invalid ServerKeyExchange message")

// rsaKeyAgreement implements the standard TLS key agreement where the client
// encrypts the pre-master secret to the server's public key.
type rsaKeyAgreement struct{}

func (ka rsaKeyAgreement) generateServerKeyExchange(config *Config, cert *Certificate, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
	return nil, nil
}

func (ka rsaKeyAgreement) processClientKeyExchange(config *Config, cert *Certificate, ckx *clientKeyExchangeMsg, version uint16) ([]byte, error) {
	if len(ckx.ciphertext) < 2 {
		return nil, errClientKeyExchange
	}
	ciphertextLen := int(ckx.ciphertext[0])<<8 | int(ckx.ciphertext[1])
	if ciphertextLen != len(ckx.ciphertext)-2 {
		return nil, errClientKeyExchange
	}
	ciphertext := ckx.ciphertext[2:]

	priv, ok := cert.PrivateKey.(crypto.Decrypter)
	if !ok {
		return nil, errors.New("tls: certificate private key does not implement crypto.Decrypter")
	}
	// Perform constant time RSA PKCS #1 v1.5 decryption
	preMasterSecret, err := priv.Decrypt(config.rand(), ciphertext, &rsa.PKCS1v15DecryptOptions{SessionKeyLen: 48})
	if err != nil {
		return nil, err
	}
	// We don't check the version number in the premaster secret. For one,
	// by checking it, we would leak information about the validity of the
	// encrypted pre-master secret. Secondly, it provides only a small
	// benefit against a downgrade attack and some implementations send the
	// wrong version anyway. See the discussion at the end of section
	// 7.4.7.1 of RFC 4346.
	return preMasterSecret, nil
}

func (ka rsaKeyAgreement) processServerKeyExchange(config *Config, clientHello *clientHelloMsg, serverHello *serverHelloMsg, cert *x509.Certificate, skx *serverKeyExchangeMsg) error {
	return errors.New("tls: unexpected ServerKeyExchange")
}

func (ka rsaKeyAgreement) generateClientKeyExchange(config *Config, clientHello *clientHelloMsg, cert *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error) {
	preMasterSecret := make([]byte, 48)
	preMasterSecret[0] = byte(clientHello.vers >> 8)
	preMasterSecret[1] = byte(clientHello.vers)
	_, err := io.ReadFull(config.rand(), preMasterSecret[2:])
	if err != nil {
		return nil, nil, err
	}

	rsaKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return nil, nil, errors.New("tls: server certificate contains incorrect key type for selected ciphersuite")
	}
	encrypted, err := rsa.EncryptPKCS1v15(config.rand(), rsaKey, preMasterSecret)
	if err != nil {
		return nil, nil, err
	}
	ckx := new(clientKeyExchangeMsg)
	ckx.ciphertext = make([]byte, len(encrypted)+2)
	ckx.ciphertext[0] = byte(len(encrypted) >> 8)
	ckx.ciphertext[1] = byte(len(encrypted))
	copy(ckx.ciphertext[2:], encrypted)
	return preMasterSecret, ckx, nil
}

// sha1Hash calculates a SHA1 hash over the given byte slices.
func sha1Hash(slices [][]byte) []byte {
	hsha1 := sha1.New()
	for _, slice := range slices {
		hsha1.Write(slice)
	}
	return hsha1.Sum(nil)
}

// md5SHA1Hash implements TLS 1.0's hybrid hash function which consists of the
// concatenation of an MD5 and SHA1 hash.
func md5SHA1Hash(slices [][]byte) []byte {
	md5sha1 := make([]byte, md5.Size+sha1.Size)
	hmd5 := md5.New()
	for _, slice := range slices {
		hmd5.Write(slice)
	}
	copy(md5sha1, hmd5.Sum(nil))
	copy(md5sha1[md5.Size:], sha1Hash(slices))
	return md5sha1
}

// hashForServerKeyExchange hashes the given slices and returns their digest
// using the given hash function (for >= TLS 1.2) or using a default based on
// the sigType (for earlier TLS versions). For Ed25519 signatures, which don't
// do pre-hashing, it returns the concatenation of the slices.
func hashForServerKeyExchange(sigType uint8, hashFunc crypto.Hash, version uint16, slices ...[]byte) []byte {
	if sigType == signatureEd25519 {
		var signed []byte
		for _, slice := range slices {
			signed = append(signed, slice...)
		}
		return signed
	}
	if version >= VersionTLS12 {
		h := hashFunc.New()
		for _, slice := range slices {
			h.Write(slice)
		}
		digest := h.Sum(nil)
		return digest
	}
	if sigType == signatureECDSA {
		return sha1Hash(slices)
	}
	return md5SHA1Hash(slices)
}

// ecdheKeyAgreement implements a TLS key agreement where the server
// generates an ephemeral EC public/private key pair and signs it. The
// pre-master secret is then calculated using ECDH. The signature may
// be ECDSA, Ed25519 or RSA.
type ecdheKeyAgreement struct {
	version uint16
	isRSA   bool
	params  ecdheParameters

	// ckx and preMasterSecret are generated in processServerKeyExchange
	// and returned in generateClientKeyExchange.
	ckx             *clientKeyExchangeMsg
	preMasterSecret []byte
}

func (ka *ecdheKeyAgreement) generateServerKeyExchange(config *Config, cert *Certificate, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
	var curveID CurveID
	for _, c := range clientHello.supportedCurves {
		if config.supportsCurve(c) {
			curveID = c
			break
		}
	}

	if curveID == 0 {
		return nil, errors.New("tls: no supported elliptic curves offered")
	}
	if _, ok := curveForCurveID(curveID); curveID != X25519 && !ok {
		return nil, errors.New("tls: CurvePreferences includes unsupported curve")
	}

	params, err := generateECDHEParameters(config.rand(), curveID)
	if err != nil {
		return nil, err
	}
	ka.params = params

	// See RFC 4492, Section 5.4.
	ecdhePublic := params.PublicKey()
	serverECDHEParams := make([]byte, 1+2+1+len(ecdhePublic))
	serverECDHEParams[0] = 3 // named curve
	serverECDHEParams[1] = byte(curveID >> 8)
	serverECDHEParams[2] = byte(curveID)
	serverECDHEParams[3] = byte(len(ecdhePublic))
	copy(serverECDHEParams[4:], ecdhePublic)

	priv, ok := cert.PrivateKey.(crypto.Signer)
	if !ok {
		return nil, fmt.Errorf("tls: certificate private key of type %T does not implement crypto.Signer", cert.PrivateKey)
	}

	var signatureAlgorithm SignatureScheme
	var sigType uint8
	var sigHash crypto.Hash
	if ka.version >= VersionTLS12 {
		signatureAlgorithm, err = selectSignatureScheme(ka.version, cert, clientHello.supportedSignatureAlgorithms)
		if err != nil {
			return nil, err
		}
		sigType, sigHash, err = typeAndHashFromSignatureScheme(signatureAlgorithm)
		if err != nil {
			return nil, err
		}
	} else {
		sigType, sigHash, err = legacyTypeAndHashFromPublicKey(priv.Public())
		if err != nil {
			return nil, err
		}
	}
	if (sigType == signaturePKCS1v15 || sigType == signatureRSAPSS) != ka.isRSA {
		return nil, errors.New("tls: certificate cannot be used with the selected cipher suite")
	}

	signed := hashForServerKeyExchange(sigType, sigHash, ka.version, clientHello.random, hello.random, serverECDHEParams)

	signOpts := crypto.SignerOpts(sigHash)
	if sigType == signatureRSAPSS {
		signOpts = &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash, Hash: sigHash}
	}
	sig, err := priv.Sign(config.rand(), signed, signOpts)
	if err != nil {
		return nil, errors.New("tls: failed to sign ECDHE parameters: " + err.Error())
	}

	skx := new(serverKeyExchangeMsg)
	sigAndHashLen := 0
	if ka.version >= VersionTLS12 {
		sigAndHashLen = 2
	}
	skx.key = make([]byte, len(serverECDHEParams)+sigAndHashLen+2+len(sig))
	copy(skx.key, serverECDHEParams)
	k := skx.key[len(serverECDHEParams):]
	if ka.version >= VersionTLS12 {
		k[0] = byte(signatureAlgorithm >> 8)
		k[1] = byte(signatureAlgorithm)
		k = k[2:]
	}
	k[0] = byte(len(sig) >> 8)
	k[1] = byte(len(sig))
	copy(k[2:], sig)

	return skx, nil
}

func (ka *ecdheKeyAgreement) processClientKeyExchange(config *Config, cert *Certificate, ckx *clientKeyExchangeMsg, version uint16) ([]byte, error) {
	if len(ckx.ciphertext) == 0 || int(ckx.ciphertext[0]) != len(ckx.ciphertext)-1 {
		return nil, errClientKeyExchange
	}

	preMasterSecret := ka.params.SharedKey(ckx.ciphertext[1:])
	if preMasterSecret == nil {
		return nil, errClientKeyExchange
	}

	return preMasterSecret, nil
}

func (ka *ecdheKeyAgreement) processServerKeyExchange(config *Config, clientHello *clientHelloMsg, serverHello *serverHelloMsg, cert *x509.Certificate, skx *serverKeyExchangeMsg) error {
	if len(skx.key) < 4 {
		return errServerKeyExchange
	}
	if skx.key[0] != 3 { // named curve
		return errors.New("tls: server selected unsupported curve")
	}
	curveID := CurveID(skx.key[1])<<8 | CurveID(skx.key[2])

	publicLen := int(skx.key[3])
	if publicLen+4 > len(skx.key) {
		return errServerKeyExchange
	}
	serverECDHEParams := skx.key[:4+publicLen]
	publicKey := serverECDHEParams[4:]

	sig := skx.key[4+publicLen:]
	if len(sig) < 2 {
		return errServerKeyExchange
	}

	if _, ok := curveForCurveID(curveID); curveID != X25519 && !ok {
		return errors.New("tls: server selected unsupported curve")
	}

	params, err := generateECDHEParameters(config.rand(), curveID)
	if err != nil {
		return err
	}
	ka.params = params

	ka.preMasterSecret = params.SharedKey(publicKey)
	if ka.preMasterSecret == nil {
		return errServerKeyExchange
	}

	ourPublicKey := params.PublicKey()
	ka.ckx = new(clientKeyExchangeMsg)
	ka.ckx.ciphertext = make([]byte, 1+len(ourPublicKey))
	ka.ckx.ciphertext[0] = byte(len(ourPublicKey))
	copy(ka.ckx.ciphertext[1:], ourPublicKey)

	var sigType uint8
	var sigHash crypto.Hash
	if ka.version >= VersionTLS12 {
		signatureAlgorithm := SignatureScheme(sig[0])<<8 | SignatureScheme(sig[1])
		sig = sig[2:]
		if len(sig) < 2 {
			return errServerKeyExchange
		}

		if !isSupportedSignatureAlgorithm(signatureAlgorithm, clientHello.supportedSignatureAlgorithms) {
			return errors.New("tls: certificate used with invalid signature algorithm")
		}
		sigType, sigHash, err = typeAndHashFromSignatureScheme(signatureAlgorithm)
		if err != nil {
			return err
		}
	} else {
		sigType, sigHash, err = legacyTypeAndHashFromPublicKey(cert.PublicKey)
		if err != nil {
			return err
		}
	}
	if (sigType == signaturePKCS1v15 || sigType == signatureRSAPSS) != ka.isRSA {
		return errServerKeyExchange
	}

	sigLen := int(sig[0])<<8 | int(sig[1])
	if sigLen+2 != len(sig) {
		return errServerKeyExchange
	}
	sig = sig[2:]

	signed := hashForServerKeyExchange(sigType, sigHash, ka.version, clientHello.random, serverHello.random, serverECDHEParams)
	if err := verifyHandshakeSignature(sigType, cert.PublicKey, sigHash, signed, sig); err != nil {
		return errors.New("tls: invalid signature by the server certificate: " + err.Error())
	}
	return nil
}

func (ka *ecdheKeyAgreement) generateClientKeyExchange(config *Config, clientHello *clientHelloMsg, cert *x509.Certificate) ([]byte, *clientKeyExchangeMsg, error) {
	if ka.ckx == nil {
		return nil, nil, errors.New("tls: missing ServerKeyExchange message")
	}

	return ka.preMasterSecret, ka.ckx, nil
}
<|endoftext|>


<commit_before>// Tor websocket server transport plugin.
//
// Usage:
// ServerTransportPlugin websocket exec ./websocket-server --port 9901

package main

import (
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
)

var ptInfo PtServerInfo

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

// An abstraction that makes an underlying WebSocket connection look like an
// io.ReadWriteCloser. It internally takes care of things like base64 encoding and
// decoding.
type websocketConn struct {
	Ws         *Websocket
	Base64     bool
	messageBuf []byte
}

// Implements io.Reader.
func (conn *websocketConn) Read(b []byte) (n int, err error) {
	for len(conn.messageBuf) == 0 {
		var m WebsocketMessage
		m, err = conn.Ws.ReadMessage()
		if err != nil {
			return
		}
		if m.Opcode == 8 {
			err = io.EOF
			return
		}
		if conn.Base64 {
			if m.Opcode != 1 {
				err = errors.New(fmt.Sprintf("got non-text opcode %d with the base64 subprotocol", m.Opcode))
				return
			}
			conn.messageBuf = make([]byte, base64.StdEncoding.DecodedLen(len(m.Payload)))
			var num int
			num, err = base64.StdEncoding.Decode(conn.messageBuf, m.Payload)
			if err != nil {
				return
			}
			conn.messageBuf = conn.messageBuf[:num]
		} else {
			if m.Opcode != 2 {
				err = errors.New(fmt.Sprintf("got non-binary opcode %d with no subprotocol", m.Opcode))
				return
			}
			conn.messageBuf = m.Payload
		}
	}

	n = copy(b, conn.messageBuf)
	conn.messageBuf = conn.messageBuf[n:]

	return
}

// Implements io.Writer.
func (conn *websocketConn) Write(b []byte) (n int, err error) {
	if conn.Base64 {
		buf := make([]byte, base64.StdEncoding.EncodedLen(len(b)))
		base64.StdEncoding.Encode(buf, b)
		err = conn.Ws.WriteMessage(1, buf)
		if err != nil {
			return
		}
		n = len(b)
	} else {
		err = conn.Ws.WriteMessage(2, b)
		n = len(b)
	}
	return
}

// Implements io.Closer.
func (conn *websocketConn) Close() (err error) {
	err = conn.Ws.WriteFrame(8, nil)
	if err != nil {
		conn.Ws.Conn.Close()
		return
	}
	err = conn.Ws.Conn.Close()
	return
}

// Create a new websocketConn.
func NewWebsocketConn(ws *Websocket) websocketConn {
	var conn websocketConn
	conn.Ws = ws
	conn.Base64 = (ws.Subprotocol == "base64")
	return conn
}

// Copy from WebSocket to socket and vice versa.
func proxy(local *net.TCPConn, conn *websocketConn) {
	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		_, err := io.Copy(conn, local)
		if err != nil {
			logDebug("error copying ORPort to WebSocket: " + err.Error())
		}
		local.CloseRead()
		conn.Close()
		wg.Done()
	}()

	go func() {
		_, err := io.Copy(local, conn)
		if err != nil {
			logDebug("error copying WebSocket to ORPort: " + err.Error())
		}
		local.CloseWrite()
		conn.Close()
		wg.Done()
	}()

	wg.Wait()
}

func websocketHandler(ws *Websocket) {
	conn := NewWebsocketConn(ws)

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	s, err := net.DialTCP("tcp", nil, ptInfo.OrAddr)
	if err != nil {
		logDebug("Failed to connect to ORPort: " + err.Error())
		return
	}

	proxy(s, &conn)
}

func startListener(addr *net.TCPAddr) (*net.TCPListener, error) {
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		var config WebsocketConfig
		config.Subprotocols = []string{"base64"}
		config.MaxMessageSize = 2500
		http.Handle("/", config.Handler(websocketHandler))
		err = http.Serve(ln, nil)
		if err != nil {
			logDebug("http.Serve: " + err.Error())
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var defaultPort int

	flag.IntVar(&defaultPort, "port", 0, "port to listen on if unspecified by Tor")
	flag.Parse()

	ptInfo = PtServerSetup([]string{ptMethodName})

	listeners := make([]*net.TCPListener, 0)
	for _, bindAddr := range ptInfo.BindAddrs {
		// Override tor's requested port (which is 0 if this transport
		// has not been run before) with the one requested by the --port
		// option.
		if defaultPort != 0 {
			bindAddr.Addr.Port = defaultPort
		}

		ln, err := startListener(bindAddr.Addr)
		if err != nil {
			PtSmethodError(bindAddr.MethodName, err.Error())
		}
		PtSmethod(bindAddr.MethodName, ln.Addr())
		listeners = append(listeners, ln)
	}
	PtSmethodsDone()

	var numHandlers int = 0

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	var sigint bool = false
	for !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}

	for _, ln := range listeners {
		ln.Close()
	}

	sigint = false
	for numHandlers != 0 && !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}
}
<commit_msg>Add --log option to websocket-server.<commit_after>// Tor websocket server transport plugin.
//
// Usage:
// ServerTransportPlugin websocket exec ./websocket-server --port 9901

package main

import (
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
)

var logFile = os.Stderr

var ptInfo PtServerInfo

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(logFile, format+"\n", v...)
}

// An abstraction that makes an underlying WebSocket connection look like an
// io.ReadWriteCloser. It internally takes care of things like base64 encoding and
// decoding.
type websocketConn struct {
	Ws         *Websocket
	Base64     bool
	messageBuf []byte
}

// Implements io.Reader.
func (conn *websocketConn) Read(b []byte) (n int, err error) {
	for len(conn.messageBuf) == 0 {
		var m WebsocketMessage
		m, err = conn.Ws.ReadMessage()
		if err != nil {
			return
		}
		if m.Opcode == 8 {
			err = io.EOF
			return
		}
		if conn.Base64 {
			if m.Opcode != 1 {
				err = errors.New(fmt.Sprintf("got non-text opcode %d with the base64 subprotocol", m.Opcode))
				return
			}
			conn.messageBuf = make([]byte, base64.StdEncoding.DecodedLen(len(m.Payload)))
			var num int
			num, err = base64.StdEncoding.Decode(conn.messageBuf, m.Payload)
			if err != nil {
				return
			}
			conn.messageBuf = conn.messageBuf[:num]
		} else {
			if m.Opcode != 2 {
				err = errors.New(fmt.Sprintf("got non-binary opcode %d with no subprotocol", m.Opcode))
				return
			}
			conn.messageBuf = m.Payload
		}
	}

	n = copy(b, conn.messageBuf)
	conn.messageBuf = conn.messageBuf[n:]

	return
}

// Implements io.Writer.
func (conn *websocketConn) Write(b []byte) (n int, err error) {
	if conn.Base64 {
		buf := make([]byte, base64.StdEncoding.EncodedLen(len(b)))
		base64.StdEncoding.Encode(buf, b)
		err = conn.Ws.WriteMessage(1, buf)
		if err != nil {
			return
		}
		n = len(b)
	} else {
		err = conn.Ws.WriteMessage(2, b)
		n = len(b)
	}
	return
}

// Implements io.Closer.
func (conn *websocketConn) Close() (err error) {
	err = conn.Ws.WriteFrame(8, nil)
	if err != nil {
		conn.Ws.Conn.Close()
		return
	}
	err = conn.Ws.Conn.Close()
	return
}

// Create a new websocketConn.
func NewWebsocketConn(ws *Websocket) websocketConn {
	var conn websocketConn
	conn.Ws = ws
	conn.Base64 = (ws.Subprotocol == "base64")
	return conn
}

// Copy from WebSocket to socket and vice versa.
func proxy(local *net.TCPConn, conn *websocketConn) {
	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		_, err := io.Copy(conn, local)
		if err != nil {
			logDebug("error copying ORPort to WebSocket: " + err.Error())
		}
		local.CloseRead()
		conn.Close()
		wg.Done()
	}()

	go func() {
		_, err := io.Copy(local, conn)
		if err != nil {
			logDebug("error copying WebSocket to ORPort: " + err.Error())
		}
		local.CloseWrite()
		conn.Close()
		wg.Done()
	}()

	wg.Wait()
}

func websocketHandler(ws *Websocket) {
	conn := NewWebsocketConn(ws)

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	s, err := net.DialTCP("tcp", nil, ptInfo.OrAddr)
	if err != nil {
		logDebug("Failed to connect to ORPort: " + err.Error())
		return
	}

	proxy(s, &conn)
}

func startListener(addr *net.TCPAddr) (*net.TCPListener, error) {
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		var config WebsocketConfig
		config.Subprotocols = []string{"base64"}
		config.MaxMessageSize = 2500
		http.Handle("/", config.Handler(websocketHandler))
		err = http.Serve(ln, nil)
		if err != nil {
			logDebug("http.Serve: " + err.Error())
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var defaultPort int
	var logFilename string

	flag.IntVar(&defaultPort, "port", 0, "port to listen on if unspecified by Tor")
	flag.StringVar(&logFilename, "log", "", "log file to write to")
	flag.Parse()

	if logFilename != "" {
		f, err := os.OpenFile(logFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Can't open log file %q: %s.\n", logFilename, err.Error())
			os.Exit(1)
		}
		logFile = f
	}

	ptInfo = PtServerSetup([]string{ptMethodName})

	listeners := make([]*net.TCPListener, 0)
	for _, bindAddr := range ptInfo.BindAddrs {
		// Override tor's requested port (which is 0 if this transport
		// has not been run before) with the one requested by the --port
		// option.
		if defaultPort != 0 {
			bindAddr.Addr.Port = defaultPort
		}

		ln, err := startListener(bindAddr.Addr)
		if err != nil {
			PtSmethodError(bindAddr.MethodName, err.Error())
		}
		PtSmethod(bindAddr.MethodName, ln.Addr())
		listeners = append(listeners, ln)
	}
	PtSmethodsDone()

	var numHandlers int = 0

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	var sigint bool = false
	for !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}

	for _, ln := range listeners {
		ln.Close()
	}

	sigint = false
	for numHandlers != 0 && !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}
}
<|endoftext|>


<commit_before>package container

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
)

// Runner runs commands specified by a Command. The default implementation uses
// Docker, but MockExecutor also implements this interface and may be used
// during testing (by temporarily changing the DefaultRunner variable). The sole
// purpose of this interface is to allow mocking.
type Runner interface {
	Run(*Command) ([]byte, error)
}

var DefaultRunner Runner = dockerRunner{}

var (
	RunRetries   = 3
	BuildRetries = 3
)

type dockerRunner struct{}

func (_ dockerRunner) Run(c *Command) ([]byte, error) {
	tmpDir, err := ioutil.TempDir("", "sg-docker")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpDir)

	image := "sg-temp:" + filepath.Base(tmpDir)
	log.Printf("Running image %s", image)

	cmdJSON, err := json.Marshal(c.Cmd)
	if err != nil {
		return nil, err
	}

	// Build Docker container.
	dockerfile := c.Dockerfile

	// AddFiles
	for i, f := range c.AddFiles {
		if i == 0 {
			dockerfile = append(dockerfile, '\n')
		}

		// Copy to the Docker context dir so we can ADD it.
		name := filepath.Base(f[0])
		err = cp(f[0], filepath.Join(tmpDir, name))
		if err != nil {
			return nil, err
		}
		dockerfile = append(dockerfile, []byte(fmt.Sprintf("ADD %s %s\n", "./"+name, f[1]))...)
	}

	// AddDirs
	for i, d := range c.AddDirs {
		if i == 0 {
			dockerfile = append(dockerfile, '\n')
		}

		// Move repository to Docker context dir and ADD it.
		dirName := filepath.Base(d[0])
		err := cp(d[0], filepath.Join(tmpDir, dirName))
		if err != nil {
			return nil, err
		}
		dockerfile = append(dockerfile, []byte(fmt.Sprintf("ADD %s %s\n", "./"+dirName, d[1]))...)
	}

	dockerfile = append(dockerfile, c.PreCmdDockerfile...)

	dockerfile = append(dockerfile, []byte(fmt.Sprintf("\nCMD %s", cmdJSON))...)
	err = ioutil.WriteFile(filepath.Join(tmpDir, "Dockerfile"), []byte(dockerfile), 0700)
	if err != nil {
		return nil, err
	}
	buildCmd := exec.Command("docker", "build", "--rm=false", "-t", image, ".")
	buildCmd.Dir = tmpDir
	buildCmd.Stdout, buildCmd.Stderr = c.Stderr, c.Stderr

	for i := 0; i < BuildRetries; i++ {
		remainingAttempts := RunRetries - i - 1
		err = buildCmd.Run()
		if err != nil {
			if remainingAttempts == 0 {
				return nil, err
			} else {
				log.Printf("Command failed: %v: %s (retrying %d more times)", buildCmd.Args, err, remainingAttempts)
				continue
			}
		}
		break
	}

	for i := 0; i < RunRetries; i++ {
		remainingAttempts := RunRetries - i - 1
		runOptions := append([]string{}, c.RunOptions...)
		if c.Dir != "" {
			runOptions = append(runOptions, "--workdir="+c.Dir)
		}
		args := append([]string{"run"}, runOptions...)
		args = append(args, image)
		runCmd := exec.Command("docker", args...)
		runCmd.Stderr = os.Stderr

		// Get original output.
		data, err := runCmd.Output()
		if err != nil {
			if remainingAttempts == 0 {
				return nil, err
			} else {
				log.Printf("Command failed: %v: %s (retrying %d more times)", runCmd.Args, err, remainingAttempts)
				continue
			}
		}

		// Transform.
		if c.Transform != nil {
			data, err = c.Transform(data)
			if err != nil {
				return nil, err
			}
		}

		return data, nil
	}

	panic("unreachable")
}

type MockRunner struct {
	Cmds []*Command
}

func (r *MockRunner) Run(c *Command) ([]byte, error) {
	r.Cmds = append(r.Cmds, c)
	return nil, nil
}

var _ Runner = &MockRunner{}

func cp(src, dst string) error {
	cmd := exec.Command("cp", "-R", "--no-dereference", "--preserve=mode,ownership,timestamps", src, dst)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cp %s %s failed: %s: %q", src, dst, err, out)
	}
	return nil
}
<commit_msg>fix build cmd retry (reinitialize command each time; otherwise fails)<commit_after>package container

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
)

// Runner runs commands specified by a Command. The default implementation uses
// Docker, but MockExecutor also implements this interface and may be used
// during testing (by temporarily changing the DefaultRunner variable). The sole
// purpose of this interface is to allow mocking.
type Runner interface {
	Run(*Command) ([]byte, error)
}

var DefaultRunner Runner = dockerRunner{}

var (
	RunRetries   = 3
	BuildRetries = 3
)

type dockerRunner struct{}

func (_ dockerRunner) Run(c *Command) ([]byte, error) {
	tmpDir, err := ioutil.TempDir("", "sg-docker")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpDir)

	image := "sg-temp:" + filepath.Base(tmpDir)
	log.Printf("Running image %s", image)

	cmdJSON, err := json.Marshal(c.Cmd)
	if err != nil {
		return nil, err
	}

	// Build Docker container.
	dockerfile := c.Dockerfile

	// AddFiles
	for i, f := range c.AddFiles {
		if i == 0 {
			dockerfile = append(dockerfile, '\n')
		}

		// Copy to the Docker context dir so we can ADD it.
		name := filepath.Base(f[0])
		err = cp(f[0], filepath.Join(tmpDir, name))
		if err != nil {
			return nil, err
		}
		dockerfile = append(dockerfile, []byte(fmt.Sprintf("ADD %s %s\n", "./"+name, f[1]))...)
	}

	// AddDirs
	for i, d := range c.AddDirs {
		if i == 0 {
			dockerfile = append(dockerfile, '\n')
		}

		// Move repository to Docker context dir and ADD it.
		dirName := filepath.Base(d[0])
		err := cp(d[0], filepath.Join(tmpDir, dirName))
		if err != nil {
			return nil, err
		}
		dockerfile = append(dockerfile, []byte(fmt.Sprintf("ADD %s %s\n", "./"+dirName, d[1]))...)
	}

	dockerfile = append(dockerfile, c.PreCmdDockerfile...)

	dockerfile = append(dockerfile, []byte(fmt.Sprintf("\nCMD %s", cmdJSON))...)
	err = ioutil.WriteFile(filepath.Join(tmpDir, "Dockerfile"), []byte(dockerfile), 0700)
	if err != nil {
		return nil, err
	}

	for i := 0; i < BuildRetries; i++ {
		remainingAttempts := RunRetries - i - 1
		buildCmd := exec.Command("docker", "build", "--rm=false", "-t", image, ".")
		buildCmd.Dir = tmpDir
		buildCmd.Stdout, buildCmd.Stderr = c.Stderr, c.Stderr
		err = buildCmd.Run()
		if err != nil {
			if remainingAttempts == 0 {
				return nil, err
			} else {
				log.Printf("Command failed: %v: %s (retrying %d more times)", buildCmd.Args, err, remainingAttempts)
				continue
			}
		}
		break
	}

	for i := 0; i < RunRetries; i++ {
		remainingAttempts := RunRetries - i - 1
		runOptions := append([]string{}, c.RunOptions...)
		if c.Dir != "" {
			runOptions = append(runOptions, "--workdir="+c.Dir)
		}
		args := append([]string{"run"}, runOptions...)
		args = append(args, image)
		runCmd := exec.Command("docker", args...)
		runCmd.Stderr = os.Stderr

		// Get original output.
		data, err := runCmd.Output()
		if err != nil {
			if remainingAttempts == 0 {
				return nil, err
			} else {
				log.Printf("Command failed: %v: %s (retrying %d more times)", runCmd.Args, err, remainingAttempts)
				continue
			}
		}

		// Transform.
		if c.Transform != nil {
			data, err = c.Transform(data)
			if err != nil {
				return nil, err
			}
		}

		return data, nil
	}

	panic("unreachable")
}

type MockRunner struct {
	Cmds []*Command
}

func (r *MockRunner) Run(c *Command) ([]byte, error) {
	r.Cmds = append(r.Cmds, c)
	return nil, nil
}

var _ Runner = &MockRunner{}

func cp(src, dst string) error {
	cmd := exec.Command("cp", "-R", "--no-dereference", "--preserve=mode,ownership,timestamps", src, dst)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cp %s %s failed: %s: %q", src, dst, err, out)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/testing/testbase"
)

type aggregateSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	addresses []instance.Address
	status    string
}

var _ instance.Instance = (*testInstance)(nil)

func (t testInstance) Addresses() ([]instance.Address, error) {
	return t.addresses, nil
}

func (t testInstance) Status() string {
	return t.status
}

type testInstanceGetter struct {
	ids     []instance.Id
	results []testInstance
	err     bool
}

func (i *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	i.ids = ids
	if i.err {
		return nil, fmt.Errorf("Some error")
	}
	for _, inst := range i.results {
		result = append(result, inst)
	}
	return
}

func newTestInstance(status string, addresses []string) *testInstance {
	thisInstance := &testInstance{status: status}
	for _, address := range addresses {
		thisInstance.addresses = append(thisInstance.addresses, instance.NewAddress(address))
	}
	return thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []testInstance{*instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)
	c.Assert(reply.info, gc.DeepEquals, instanceInfo{
		status:    "foobar",
		addresses: instance1.addresses,
	})
	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo")})
}

func (s *aggregateSuite) TestRequestBatching(c *gc.C) {
	s.PatchValue(&GatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []testInstance{*instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	instance2 := newTestInstance("not foobar", []string{"192.168.1.2"})
	instance3 := newTestInstance("ok-ish", []string{"192.168.1.3"})

	replyChan2 := make(chan instanceInfoReply)
	replyChan3 := make(chan instanceInfoReply)

	aggregator.reqc <- instanceInfoReq{reply: replyChan2, instId: instance.Id("foo2")}
	aggregator.reqc <- instanceInfoReq{reply: replyChan3, instId: instance.Id("foo3")}

	testGetter.results = []testInstance{*instance2, *instance3}
	reply2 := <-replyChan2
	reply3 := <-replyChan3
	c.Assert(reply2.err, gc.IsNil)
	c.Assert(reply3.err, gc.IsNil)
	c.Assert(reply2.info.status, gc.Equals, "not foobar")
	c.Assert(reply3.info.status, gc.Equals, "ok-ish")

	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo2"), instance.Id("foo3")})
}
<commit_msg>Skeleton test for error<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/testing/testbase"
)

type aggregateSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	addresses []instance.Address
	status    string
}

var _ instance.Instance = (*testInstance)(nil)

func (t testInstance) Addresses() ([]instance.Address, error) {
	return t.addresses, nil
}

func (t testInstance) Status() string {
	return t.status
}

type testInstanceGetter struct {
	ids     []instance.Id
	results []testInstance
	err     bool
}

func (i *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	i.ids = ids
	if i.err {
		return nil, fmt.Errorf("Some error")
	}
	for _, inst := range i.results {
		result = append(result, inst)
	}
	return
}

func newTestInstance(status string, addresses []string) *testInstance {
	thisInstance := &testInstance{status: status}
	for _, address := range addresses {
		thisInstance.addresses = append(thisInstance.addresses, instance.NewAddress(address))
	}
	return thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []testInstance{*instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)
	c.Assert(reply.info, gc.DeepEquals, instanceInfo{
		status:    "foobar",
		addresses: instance1.addresses,
	})
	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo")})
}

func (s *aggregateSuite) TestRequestBatching(c *gc.C) {
	s.PatchValue(&GatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []testInstance{*instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	instance2 := newTestInstance("not foobar", []string{"192.168.1.2"})
	instance3 := newTestInstance("ok-ish", []string{"192.168.1.3"})

	replyChan2 := make(chan instanceInfoReply)
	replyChan3 := make(chan instanceInfoReply)

	aggregator.reqc <- instanceInfoReq{reply: replyChan2, instId: instance.Id("foo2")}
	aggregator.reqc <- instanceInfoReq{reply: replyChan3, instId: instance.Id("foo3")}

	testGetter.results = []testInstance{*instance2, *instance3}
	reply2 := <-replyChan2
	reply3 := <-replyChan3
	c.Assert(reply2.err, gc.IsNil)
	c.Assert(reply3.err, gc.IsNil)
	c.Assert(reply2.info.status, gc.Equals, "not foobar")
	c.Assert(reply3.info.status, gc.Equals, "ok-ish")

	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo2"), instance.Id("foo3")})
}

func (s *aggregateSuite) TestError(c *gc.C) {
	testGetter := new(testInstanceGetter)

	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package terminationworker_test

import (
	"os"
	"os/signal"
	"runtime"
	stdtesting "testing"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/testing"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/terminationworker"
)

func TestPackage(t *stdtesting.T) {
	gc.TestingT(t)
}

var _ = gc.Suite(&TerminationWorkerSuite{})

type TerminationWorkerSuite struct {
	testing.BaseSuite
	// c is a channel that will wait for the termination
	// signal, to prevent signals terminating the process.
	c chan os.Signal
}

func (s *TerminationWorkerSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.c = make(chan os.Signal, 1)
	signal.Notify(s.c, terminationworker.TerminationSignal)
}

func (s *TerminationWorkerSuite) TearDownTest(c *gc.C) {
	close(s.c)
	signal.Stop(s.c)
	s.BaseSuite.TearDownTest(c)
}

func (s *TerminationWorkerSuite) TestStartStop(c *gc.C) {
	w := terminationworker.NewWorker()
	w.Kill()
	err := w.Wait()
	c.Assert(err, jc.ErrorIsNil)
}

func (s *TerminationWorkerSuite) TestSignal(c *gc.C) {
	//TODO(bogdanteleaga): Inspect this further on windows
	if runtime.GOOS == "windows" {
		c.Skip("bug 1403084: sending this signal is not supported on windows")
	}
	w := terminationworker.NewWorker()
	proc, err := os.FindProcess(os.Getpid())
	c.Assert(err, jc.ErrorIsNil)
	defer proc.Release()
	err = proc.Signal(terminationworker.TerminationSignal)
	c.Assert(err, jc.ErrorIsNil)
	err = w.Wait()
	c.Assert(err, gc.Equals, worker.ErrTerminateAgent)
}
<commit_msg>worker/terminationworker: stop listening for signal before closing signal channel<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package terminationworker_test

import (
	"os"
	"os/signal"
	"runtime"
	stdtesting "testing"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/testing"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/terminationworker"
)

func TestPackage(t *stdtesting.T) {
	gc.TestingT(t)
}

var _ = gc.Suite(&TerminationWorkerSuite{})

type TerminationWorkerSuite struct {
	testing.BaseSuite
	// c is a channel that will wait for the termination
	// signal, to prevent signals terminating the process.
	c chan os.Signal
}

func (s *TerminationWorkerSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.c = make(chan os.Signal, 1)
	signal.Notify(s.c, terminationworker.TerminationSignal)
}

func (s *TerminationWorkerSuite) TearDownTest(c *gc.C) {
	signal.Stop(s.c)
	close(s.c)
	s.BaseSuite.TearDownTest(c)
}

func (s *TerminationWorkerSuite) TestStartStop(c *gc.C) {
	w := terminationworker.NewWorker()
	w.Kill()
	err := w.Wait()
	c.Assert(err, jc.ErrorIsNil)
}

func (s *TerminationWorkerSuite) TestSignal(c *gc.C) {
	//TODO(bogdanteleaga): Inspect this further on windows
	if runtime.GOOS == "windows" {
		c.Skip("bug 1403084: sending this signal is not supported on windows")
	}
	w := terminationworker.NewWorker()
	proc, err := os.FindProcess(os.Getpid())
	c.Assert(err, jc.ErrorIsNil)
	defer proc.Release()
	err = proc.Signal(terminationworker.TerminationSignal)
	c.Assert(err, jc.ErrorIsNil)
	err = w.Wait()
	c.Assert(err, gc.Equals, worker.ErrTerminateAgent)
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package operation_test

import (
	"github.com/juju/errors"
	"github.com/juju/names"
	"github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/worker/uniter/operation"
)

type UpdateStorageSuite struct {
	testing.IsolationSuite
}

var _ = gc.Suite(&UpdateStorageSuite{})

func (s *UpdateStorageSuite) TestPrepare(c *gc.C) {
	factory := operation.NewFactory(operation.FactoryParams{})
	op, err := factory.NewUpdateStorage(nil)
	c.Assert(err, jc.ErrorIsNil)
	state, err := op.Prepare(operation.State{})
	c.Check(err, jc.ErrorIsNil)
	c.Check(state, gc.IsNil)
}

func (s *UpdateStorageSuite) TestExecuteError(c *gc.C) {
	updater := &mockStorageUpdater{err: errors.New("meep")}
	factory := operation.NewFactory(operation.FactoryParams{StorageUpdater: updater})

	tag0 := names.NewStorageTag("data/0")
	tag1 := names.NewStorageTag("data/1")
	tags := []names.StorageTag{tag0, tag1}
	op, err := factory.NewUpdateStorage(tags)
	c.Assert(err, jc.ErrorIsNil)
	state, err := op.Prepare(operation.State{})
	c.Check(err, jc.ErrorIsNil)
	c.Check(state, gc.IsNil)

	state, err = op.Execute(operation.State{})
	c.Check(err, gc.ErrorMatches, "meep")
	c.Check(state, gc.IsNil)
	c.Check(updater.tags, jc.DeepEquals, [][]names.StorageTag{tags})
}

func (s *UpdateStorageSuite) TestExecuteSuccess(c *gc.C) {
	updater := &mockStorageUpdater{}
	factory := operation.NewFactory(operation.FactoryParams{StorageUpdater: updater})

	tag0 := names.NewStorageTag("data/0")
	tag1 := names.NewStorageTag("data/1")
	tags := []names.StorageTag{tag0, tag1}
	op, err := factory.NewUpdateStorage(tags)
	c.Assert(err, jc.ErrorIsNil)
	state, err := op.Prepare(operation.State{})
	c.Check(err, jc.ErrorIsNil)
	c.Check(state, gc.IsNil)

	state, err = op.Execute(operation.State{})
	c.Check(err, jc.ErrorIsNil)
	c.Check(state, gc.IsNil)
	c.Check(updater.tags, jc.DeepEquals, [][]names.StorageTag{tags})
}

func (s *UpdateStorageSuite) TestCommit(c *gc.C) {
	factory := operation.NewFactory(operation.FactoryParams{})
	op, err := factory.NewUpdateStorage(nil)
	c.Assert(err, jc.ErrorIsNil)
	state, err := op.Commit(operation.State{})
	c.Check(err, jc.ErrorIsNil)
	c.Check(state, gc.IsNil)
}

func (s *UpdateStorageSuite) TestDoesNotNeedGlobalMachineLock(c *gc.C) {
	factory := operation.NewFactory(operation.FactoryParams{})
	op, err := factory.NewUpdateStorage(nil)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(op.NeedsGlobalMachineLock(), jc.IsFalse)
}

type mockStorageUpdater struct {
	tags [][]names.StorageTag
	err  error
}

func (u *mockStorageUpdater) UpdateStorage(tags []names.StorageTag) error {
	u.tags = append(u.tags, tags)
	return u.err
}
<commit_msg>Delete some more code<commit_after><|endoftext|>


<commit_before>// Copyright (c) 2015 Andrea Masi. All rights reserved.
// Use of this source code is governed by MIT license
// which that can be found in the LICENSE.txt file.

package learn

import (
	"math"
	"testing"
)

func loadDataset(t *testing.T) Table {
	// Load all data in memory.
	trainData, err := ReadAllCSV("datasets/linear_test.csv")
	if err != nil {
		t.Fatal(err)
	}
	return trainData
}

func TestLinearRegression_Predict(t *testing.T) {
	trainData := loadDataset(t)
	var tab MemoryTable = make([][]interface{}, 1)
	// Hand crafted case.
	tab[0] = []interface{}{1650.0, 3.0}
	lr, err := NewLinearRegression(trainData)
	if err != nil {
		t.Fatal(err)
	}
	y, err := lr.Predict(tab)
	if err != nil {
		t.Fatal(err)
	}
	if e := math.Abs(y[0] - 293081.464335); e > tolerance {
		t.Fatal("prediction over tolerance:", e)
	}
}
<commit_msg>Refactor linear regression test<commit_after>// Copyright (c) 2015 Andrea Masi. All rights reserved.
// Use of this source code is governed by MIT license
// which that can be found in the LICENSE.txt file.

package learn

import (
	"math"
	"testing"
)

func loadDataset(t *testing.T) Table {
	// Load all data in memory.
	trainData, err := ReadAllCSV("datasets/linear_test.csv")
	if err != nil {
		t.Fatal(err)
	}
	return trainData
}

const tolerance = 1e-7

func TestLinearRegression_Predict(t *testing.T) {
	trainData := loadDataset(t)
	var tab MemoryTable = make([][]interface{}, 1)
	// Hand crafted case.
	tab[0] = []interface{}{1650.0, 3.0}
	lr, err := NewLinearRegression(trainData)
	if err != nil {
		t.Fatal(err)
	}
	y, err := lr.Predict(tab)
	if err != nil {
		t.Fatal(err)
	}
	if e := math.Abs(y[0] - 293081.464335); e > tolerance {
		t.Fatal("prediction over tolerance:", e)
	}
}
<|endoftext|>