<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/golib/stress"
	"github.com/funkygao/redigo/redis"
)

var (
	addr          string
	mode          string
	topic         string
	loops         int
	suppressError bool
	pubkey        string
	appid         string
	sleep         time.Duration
	ver           string
	sz            int
	async         bool
)

func main() {
	flag.StringVar(&addr, "addr", "http://localhost:9191/", "kateway pub addr, must start with http or https")
	flag.IntVar(&loops, "loops", 1000, "loops in each thread")
	flag.IntVar(&sz, "size", 200, "each pub message size")
	flag.StringVar(&topic, "topic", "foobar", "pub topic")
	flag.StringVar(&mode, "mode", "gw", "<gw|kafka|http|redis>")
	flag.StringVar(&appid, "appid", "app1", "appid of pub")
	flag.StringVar(&pubkey, "pubkey", "mypubkey", "pubkey")
	flag.StringVar(&ver, "ver", "v1", "pub topic version")
	flag.BoolVar(&async, "async", false, "async pub")
	flag.DurationVar(&sleep, "sleep", 0, "sleep between pub")
	flag.BoolVar(&suppressError, "noerr", false, "suppress error output")
	flag.Parse()

	switch mode {
	case "gw":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(pubGatewayLoop)

	case "kafka":
		if async {
			stress.RunStress(pubKafkaAsyncLoop)
		} else {
			stress.RunStress(pubKafkaLoop)
		}

	case "redis":
		stress.RunStress(redisLoop)

	case "http":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(getHttpLoop)
	}

}

func getHttpLoop(seq int) {
	client := createHttpClient()
	req, _ := http.NewRequest("GET", "http://localhost:9090/", nil)
	for i := 0; i < loops; i++ {
		response, err := client.Do(req)
		if err == nil {
			ioutil.ReadAll(response.Body)
			response.Body.Close() // reuse the connection

			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
			//log.Println(err)
		}
	}
}

func pubKafkaLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewSyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, _, err := producer.SendMessage(&sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		})
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
		}
	}

}

func redisLoop(seq int) {
	conn, err := redis.DialTimeout("tcp", ":6379", 0, 1*time.Second, 1*time.Second)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer conn.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, err := conn.Do("SET", "key", msg)
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			if !suppressError {
				log.Println(err)
			}
			stress.IncCounter("fail", 1)
		}
	}

}

func pubKafkaAsyncLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.Flush.Frequency = time.Second * 10
	cf.Producer.Flush.Messages = 1000
	cf.Producer.Flush.MaxMessages = 1000
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewAsyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		producer.Input() <- &sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		}
		stress.IncCounter("ok", 1)
	}

}

func createHttpClient() *http.Client {
	timeout := 3 * time.Second
	httpClient := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   timeout,
				KeepAlive: 60 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: timeout,
		},
	}

	return httpClient
}

func pubGatewayLoop(seq int) {
	httpClient := createHttpClient()
	url := fmt.Sprintf("%s/topics/%s/%s?", addr, topic, ver)
	if async {
		url += "async=1"
	}
	for n := 0; n < loops; n++ {
		req, err := http.NewRequest("POST", url,
			bytes.NewBuffer([]byte(strings.Repeat("X", sz))))
		if err != nil {
			log.Fatalf("Error Occured. %+v", err)
			stress.IncCounter("fail", 1)
			return
		}

		req.Header.Set("Appid", appid)
		req.Header.Set("Pubkey", pubkey)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		// use httpClient to send request
		response, err := httpClient.Do(req)
		if err != nil && response == nil {
			stress.IncCounter("fail", 1)
			if !suppressError {
				log.Printf("Error sending request to API endpoint. %+v", err)
			}
		} else {
			if response.StatusCode != http.StatusOK {
				stress.IncCounter("fail", 1)
				if !suppressError {
					log.Printf("Error sending request to API endpoint. %+v", response.Status)
				}
				return
			}

			// Let's check if the work actually is done
			// We have seen inconsistencies even when we get 200 OK response
			body, err := ioutil.ReadAll(response.Body)
			if err != nil {
				log.Fatalf("Couldn't parse response body. %+v", err)
			}

			// Close the connection to reuse it
			response.Body.Close()

			stress.IncCounter("ok", 1)

			if false {
				log.Println("Response Body:", string(body))
			}
		}

		if sleep > 0 {
			time.Sleep(sleep)
		}
	}
}
<commit_msg>use authority redis pkg<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/golib/stress"
	"github.com/garyburd/redigo/redis"
)

var (
	addr          string
	mode          string
	topic         string
	loops         int
	suppressError bool
	pubkey        string
	appid         string
	sleep         time.Duration
	ver           string
	sz            int
	async         bool
)

func main() {
	flag.StringVar(&addr, "addr", "http://localhost:9191/", "kateway pub addr, must start with http or https")
	flag.IntVar(&loops, "loops", 1000, "loops in each thread")
	flag.IntVar(&sz, "size", 200, "each pub message size")
	flag.StringVar(&topic, "topic", "foobar", "pub topic")
	flag.StringVar(&mode, "mode", "gw", "<gw|kafka|http|redis>")
	flag.StringVar(&appid, "appid", "app1", "appid of pub")
	flag.StringVar(&pubkey, "pubkey", "mypubkey", "pubkey")
	flag.StringVar(&ver, "ver", "v1", "pub topic version")
	flag.BoolVar(&async, "async", false, "async pub")
	flag.DurationVar(&sleep, "sleep", 0, "sleep between pub")
	flag.BoolVar(&suppressError, "noerr", false, "suppress error output")
	flag.Parse()

	switch mode {
	case "gw":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(pubGatewayLoop)

	case "kafka":
		if async {
			stress.RunStress(pubKafkaAsyncLoop)
		} else {
			stress.RunStress(pubKafkaLoop)
		}

	case "redis":
		stress.RunStress(redisLoop)

	case "http":
		http.DefaultClient.Timeout = time.Second * 30
		stress.RunStress(getHttpLoop)
	}

}

func getHttpLoop(seq int) {
	client := createHttpClient()
	req, _ := http.NewRequest("GET", "http://localhost:9090/", nil)
	for i := 0; i < loops; i++ {
		response, err := client.Do(req)
		if err == nil {
			ioutil.ReadAll(response.Body)
			response.Body.Close() // reuse the connection

			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
			//log.Println(err)
		}
	}
}

func pubKafkaLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewSyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, _, err := producer.SendMessage(&sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		})
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			stress.IncCounter("fail", 1)
		}
	}

}

func redisLoop(seq int) {
	conn, err := redis.DialTimeout("tcp", ":6379", 0, 1*time.Second, 1*time.Second)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer conn.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		_, err := conn.Do("SET", "key", msg)
		if err == nil {
			stress.IncCounter("ok", 1)
		} else {
			if !suppressError {
				log.Println(err)
			}
			stress.IncCounter("fail", 1)
		}
	}

}

func pubKafkaAsyncLoop(seq int) {
	cf := sarama.NewConfig()
	cf.Producer.Flush.Frequency = time.Second * 10
	cf.Producer.Flush.Messages = 1000
	cf.Producer.Flush.MaxMessages = 1000
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Partitioner = sarama.NewHashPartitioner
	cf.Producer.Timeout = time.Second
	//cf.Producer.Compression = sarama.CompressionSnappy
	cf.Producer.Retry.Max = 3
	producer, err := sarama.NewAsyncProducer([]string{"localhost:9092"}, cf)
	if err != nil {
		stress.IncCounter("fail", 1)
		log.Println(err)
		return
	}

	defer producer.Close()
	msg := strings.Repeat("X", sz)
	for i := 0; i < loops; i++ {
		producer.Input() <- &sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		}
		stress.IncCounter("ok", 1)
	}

}

func createHttpClient() *http.Client {
	timeout := 3 * time.Second
	httpClient := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   timeout,
				KeepAlive: 60 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: timeout,
		},
	}

	return httpClient
}

func pubGatewayLoop(seq int) {
	httpClient := createHttpClient()
	url := fmt.Sprintf("%s/topics/%s/%s?", addr, topic, ver)
	if async {
		url += "async=1"
	}
	for n := 0; n < loops; n++ {
		req, err := http.NewRequest("POST", url,
			bytes.NewBuffer([]byte(strings.Repeat("X", sz))))
		if err != nil {
			log.Fatalf("Error Occured. %+v", err)
			stress.IncCounter("fail", 1)
			return
		}

		req.Header.Set("Appid", appid)
		req.Header.Set("Pubkey", pubkey)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		// use httpClient to send request
		response, err := httpClient.Do(req)
		if err != nil && response == nil {
			stress.IncCounter("fail", 1)
			if !suppressError {
				log.Printf("Error sending request to API endpoint. %+v", err)
			}
		} else {
			if response.StatusCode != http.StatusOK {
				stress.IncCounter("fail", 1)
				if !suppressError {
					log.Printf("Error sending request to API endpoint. %+v", response.Status)
				}
				return
			}

			// Let's check if the work actually is done
			// We have seen inconsistencies even when we get 200 OK response
			body, err := ioutil.ReadAll(response.Body)
			if err != nil {
				log.Fatalf("Couldn't parse response body. %+v", err)
			}

			// Close the connection to reuse it
			response.Body.Close()

			stress.IncCounter("ok", 1)

			if false {
				log.Println("Response Body:", string(body))
			}
		}

		if sleep > 0 {
			time.Sleep(sleep)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"crypto/x509"
	"encoding/hex"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	cfocsp "github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/ocsp"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/facebookgo/httpdown"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/jmhodges/clock"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/golang.org/x/crypto/ocsp"
	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
	"github.com/letsencrypt/boulder/metrics"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/sa"
)

type cacheCtrlHandler struct {
	http.Handler
	MaxAge time.Duration
}

func (c *cacheCtrlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", c.MaxAge/time.Second))
	c.Handler.ServeHTTP(w, r)
}

/*
DBSource maps a given Database schema to a CA Key Hash, so we can pick
from among them when presented with OCSP requests for different certs.

We assume that OCSP responses are stored in a very simple database table,
with two columns: serialNumber and response

  CREATE TABLE ocsp_responses (serialNumber TEXT, response BLOB);

The serialNumber field may have any type to which Go will match a string,
so you can be more efficient than TEXT if you like.  We use it to store the
serial number in base64.  You probably want to have an index on the
serialNumber field, since we will always query on it.

*/
type DBSource struct {
	dbMap     *gorp.DbMap
	caKeyHash []byte
}

// NewSourceFromDatabase produces a DBSource representing the binding of a
// given DB schema to a CA key.
func NewSourceFromDatabase(dbMap *gorp.DbMap, caKeyHash []byte) (src *DBSource, err error) {
	src = &DBSource{dbMap: dbMap, caKeyHash: caKeyHash}
	return
}

// Response is called by the HTTP server to handle a new OCSP request.
func (src *DBSource) Response(req *ocsp.Request) ([]byte, bool) {
	log := blog.GetAuditLogger()

	// Check that this request is for the proper CA
	if bytes.Compare(req.IssuerKeyHash, src.caKeyHash) != 0 {
		log.Debug(fmt.Sprintf("Request intended for CA Cert ID: %s", hex.EncodeToString(req.IssuerKeyHash)))
		return nil, false
	}

	serialString := core.SerialToString(req.SerialNumber)
	log.Debug(fmt.Sprintf("Searching for OCSP issued by us for serial %s", serialString))

	var response []byte
	defer func() {
		if len(response) != 0 {
			log.Info(fmt.Sprintf("OCSP Response sent for CA=%s, Serial=%s", hex.EncodeToString(src.caKeyHash), serialString))
		}
	}()
	// Note: we first check for an OCSP response in the certificateStatus table (
	// the new method) if we don't find a response there we instead look in the
	// ocspResponses table (the old method) while transitioning between the two
	// tables.
	err := src.dbMap.SelectOne(
		&response,
		"SELECT ocspResponse FROM certificateStatus WHERE serial = :serial",
		map[string]interface{}{"serial": serialString},
	)
	// TODO(#970): Delete this ocspResponses check once the table has been removed
	if len(response) == 0 {
		// Ignoring possible error, if response hasn't been filled, attempt to find
		// response in old table
		err = src.dbMap.SelectOne(
			&response,
			"SELECT response from ocspResponses WHERE serial = :serial ORDER BY id DESC LIMIT 1;",
			map[string]interface{}{"serial": serialString},
		)
	}
	if err != nil {
		return nil, false
	}

	return response, true
}

func makeDBSource(dbConnect, issuerCert string, sqlDebug bool) (*DBSource, error) {
	// Configure DB
	dbMap, err := sa.NewDbMap(dbConnect)
	if err != nil {
		return nil, fmt.Errorf("Could not connect to database: %s", err)
	}
	sa.SetSQLDebug(dbMap, sqlDebug)

	// Load the CA's key so we can store its SubjectKey in the DB
	caCertDER, err := cmd.LoadCert(issuerCert)
	if err != nil {
		return nil, fmt.Errorf("Could not read issuer cert %s: %s", issuerCert, err)
	}
	caCert, err := x509.ParseCertificate(caCertDER)
	if err != nil {
		return nil, fmt.Errorf("Could not parse issuer cert %s: %s", issuerCert, err)
	}
	if len(caCert.SubjectKeyId) == 0 {
		return nil, fmt.Errorf("Empty subjectKeyID")
	}

	// Construct source from DB
	return NewSourceFromDatabase(dbMap, caCert.SubjectKeyId)
}

func main() {
	app := cmd.NewAppShell("boulder-ocsp-responder", "Handles OCSP requests")
	app.Action = func(c cmd.Config) {
		// Set up logging
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")
		auditlogger.Info(app.VersionString())

		// AUDIT[ Error Conditions ] 9cc4d537-8534-4970-8665-4b382abe82f3
		defer auditlogger.AuditPanic()

		blog.SetAuditLogger(auditlogger)

		go cmd.DebugServer(c.OCSPResponder.DebugAddr)

		go cmd.ProfileCmd("OCSP", stats)

		config := c.OCSPResponder
		var source cfocsp.Source
		url, err := url.Parse(config.Source)
		cmd.FailOnError(err, fmt.Sprintf("Source was not a URL: %s", config.Source))

		if url.Scheme == "mysql+tcp" {
			auditlogger.Info(fmt.Sprintf("Loading OCSP Database for CA Cert: %s", c.Common.IssuerCert))
			source, err = makeDBSource(config.Source, c.Common.IssuerCert, c.SQL.SQLDebug)
			cmd.FailOnError(err, "Couldn't load OCSP DB")
		} else if url.Scheme == "file" {
			filename := url.Path
			// Go interprets cwd-relative file urls (file:test/foo.txt) as having the
			// relative part of the path in the 'Opaque' field.
			if filename == "" {
				filename = url.Opaque
			}
			source, err = cfocsp.NewSourceFromFile(filename)
			cmd.FailOnError(err, fmt.Sprintf("Couldn't read file: %s", url.Path))
		} else {
			cmd.FailOnError(errors.New(`"source" parameter not found in JSON config`), "unable to start ocsp-responder")
		}

		stopTimeout, err := time.ParseDuration(c.OCSPResponder.ShutdownStopTimeout)
		cmd.FailOnError(err, "Couldn't parse shutdown stop timeout")
		killTimeout, err := time.ParseDuration(c.OCSPResponder.ShutdownKillTimeout)
		cmd.FailOnError(err, "Couldn't parse shutdown kill timeout")

		m := http.StripPrefix(c.OCSPResponder.Path,
			handler(source, c.OCSPResponder.MaxAge.Duration))

		httpMonitor := metrics.NewHTTPMonitor(stats, m, "OCSP")
		srv := &http.Server{
			Addr:    c.OCSPResponder.ListenAddress,
			Handler: httpMonitor.Handle(),
		}

		hd := &httpdown.HTTP{
			StopTimeout: stopTimeout,
			KillTimeout: killTimeout,
			Stats:       metrics.NewFBAdapter(stats, "OCSP", clock.Default()),
		}
		err = httpdown.ListenAndServe(srv, hd)
		cmd.FailOnError(err, "Error starting HTTP server")
	}

	app.Run()
}

func handler(src cfocsp.Source, maxAge time.Duration) http.Handler {
	return &cacheCtrlHandler{
		Handler: cfocsp.Responder{Source: src},
		MaxAge:  maxAge,
	}
}
<commit_msg>Log OCSP responder SQL errors<commit_after>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"crypto/x509"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	cfocsp "github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/ocsp"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/facebookgo/httpdown"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/jmhodges/clock"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/golang.org/x/crypto/ocsp"
	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
	"github.com/letsencrypt/boulder/metrics"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/sa"
)

type cacheCtrlHandler struct {
	http.Handler
	MaxAge time.Duration
}

func (c *cacheCtrlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", c.MaxAge/time.Second))
	c.Handler.ServeHTTP(w, r)
}

/*
DBSource maps a given Database schema to a CA Key Hash, so we can pick
from among them when presented with OCSP requests for different certs.

We assume that OCSP responses are stored in a very simple database table,
with two columns: serialNumber and response

  CREATE TABLE ocsp_responses (serialNumber TEXT, response BLOB);

The serialNumber field may have any type to which Go will match a string,
so you can be more efficient than TEXT if you like.  We use it to store the
serial number in base64.  You probably want to have an index on the
serialNumber field, since we will always query on it.

*/
type DBSource struct {
	dbMap     *gorp.DbMap
	caKeyHash []byte
}

// NewSourceFromDatabase produces a DBSource representing the binding of a
// given DB schema to a CA key.
func NewSourceFromDatabase(dbMap *gorp.DbMap, caKeyHash []byte) (src *DBSource, err error) {
	src = &DBSource{dbMap: dbMap, caKeyHash: caKeyHash}
	return
}

// Response is called by the HTTP server to handle a new OCSP request.
func (src *DBSource) Response(req *ocsp.Request) ([]byte, bool) {
	log := blog.GetAuditLogger()

	// Check that this request is for the proper CA
	if bytes.Compare(req.IssuerKeyHash, src.caKeyHash) != 0 {
		log.Debug(fmt.Sprintf("Request intended for CA Cert ID: %s", hex.EncodeToString(req.IssuerKeyHash)))
		return nil, false
	}

	serialString := core.SerialToString(req.SerialNumber)
	log.Debug(fmt.Sprintf("Searching for OCSP issued by us for serial %s", serialString))

	var response []byte
	defer func() {
		if len(response) != 0 {
			log.Info(fmt.Sprintf("OCSP Response sent for CA=%s, Serial=%s", hex.EncodeToString(src.caKeyHash), serialString))
		}
	}()
	// Note: we first check for an OCSP response in the certificateStatus table (
	// the new method) if we don't find a response there we instead look in the
	// ocspResponses table (the old method) while transitioning between the two
	// tables.
	err := src.dbMap.SelectOne(
		&response,
		"SELECT ocspResponse FROM certificateStatus WHERE serial = :serial",
		map[string]interface{}{"serial": serialString},
	)
	if err != nil && err != sql.ErrNoRows {
		log.Err(fmt.Sprintf("Failed to retrieve response from certificateStatus table: %s", err))
	}
	// TODO(#970): Delete this ocspResponses check once the table has been removed
	if len(response) == 0 {
		// Ignoring possible error, if response hasn't been filled, attempt to find
		// response in old table
		err = src.dbMap.SelectOne(
			&response,
			"SELECT response from ocspResponses WHERE serial = :serial ORDER BY id DESC LIMIT 1;",
			map[string]interface{}{"serial": serialString},
		)
		if err != nil && err != sql.ErrNoRows {
			log.Err(fmt.Sprintf("Failed to retrieve response from ocspResponses table: %s", err))
		}
	}
	if err != nil {
		return nil, false
	}

	return response, true
}

func makeDBSource(dbConnect, issuerCert string, sqlDebug bool) (*DBSource, error) {
	// Configure DB
	dbMap, err := sa.NewDbMap(dbConnect)
	if err != nil {
		return nil, fmt.Errorf("Could not connect to database: %s", err)
	}
	sa.SetSQLDebug(dbMap, sqlDebug)

	// Load the CA's key so we can store its SubjectKey in the DB
	caCertDER, err := cmd.LoadCert(issuerCert)
	if err != nil {
		return nil, fmt.Errorf("Could not read issuer cert %s: %s", issuerCert, err)
	}
	caCert, err := x509.ParseCertificate(caCertDER)
	if err != nil {
		return nil, fmt.Errorf("Could not parse issuer cert %s: %s", issuerCert, err)
	}
	if len(caCert.SubjectKeyId) == 0 {
		return nil, fmt.Errorf("Empty subjectKeyID")
	}

	// Construct source from DB
	return NewSourceFromDatabase(dbMap, caCert.SubjectKeyId)
}

func main() {
	app := cmd.NewAppShell("boulder-ocsp-responder", "Handles OCSP requests")
	app.Action = func(c cmd.Config) {
		// Set up logging
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")
		auditlogger.Info(app.VersionString())

		// AUDIT[ Error Conditions ] 9cc4d537-8534-4970-8665-4b382abe82f3
		defer auditlogger.AuditPanic()

		blog.SetAuditLogger(auditlogger)

		go cmd.DebugServer(c.OCSPResponder.DebugAddr)

		go cmd.ProfileCmd("OCSP", stats)

		config := c.OCSPResponder
		var source cfocsp.Source
		url, err := url.Parse(config.Source)
		cmd.FailOnError(err, fmt.Sprintf("Source was not a URL: %s", config.Source))

		if url.Scheme == "mysql+tcp" {
			auditlogger.Info(fmt.Sprintf("Loading OCSP Database for CA Cert: %s", c.Common.IssuerCert))
			source, err = makeDBSource(config.Source, c.Common.IssuerCert, c.SQL.SQLDebug)
			cmd.FailOnError(err, "Couldn't load OCSP DB")
		} else if url.Scheme == "file" {
			filename := url.Path
			// Go interprets cwd-relative file urls (file:test/foo.txt) as having the
			// relative part of the path in the 'Opaque' field.
			if filename == "" {
				filename = url.Opaque
			}
			source, err = cfocsp.NewSourceFromFile(filename)
			cmd.FailOnError(err, fmt.Sprintf("Couldn't read file: %s", url.Path))
		} else {
			cmd.FailOnError(errors.New(`"source" parameter not found in JSON config`), "unable to start ocsp-responder")
		}

		stopTimeout, err := time.ParseDuration(c.OCSPResponder.ShutdownStopTimeout)
		cmd.FailOnError(err, "Couldn't parse shutdown stop timeout")
		killTimeout, err := time.ParseDuration(c.OCSPResponder.ShutdownKillTimeout)
		cmd.FailOnError(err, "Couldn't parse shutdown kill timeout")

		m := http.StripPrefix(c.OCSPResponder.Path,
			handler(source, c.OCSPResponder.MaxAge.Duration))

		httpMonitor := metrics.NewHTTPMonitor(stats, m, "OCSP")
		srv := &http.Server{
			Addr:    c.OCSPResponder.ListenAddress,
			Handler: httpMonitor.Handle(),
		}

		hd := &httpdown.HTTP{
			StopTimeout: stopTimeout,
			KillTimeout: killTimeout,
			Stats:       metrics.NewFBAdapter(stats, "OCSP", clock.Default()),
		}
		err = httpdown.ListenAndServe(srv, hd)
		cmd.FailOnError(err, "Error starting HTTP server")
	}

	app.Run()
}

func handler(src cfocsp.Source, maxAge time.Duration) http.Handler {
	return &cacheCtrlHandler{
		Handler: cfocsp.Responder{Source: src},
		MaxAge:  maxAge,
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru 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 tsuru

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/tsuru/tsuru/cmd"
	"io"
	"io/ioutil"
	"launchpad.net/gnuflag"
	"net/http"
	"sort"
	"strings"
)

type ServiceList struct{}

func (s ServiceList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "service-list",
		Usage: "service-list",
		Desc:  "Get all available services, and user's instances for this services",
	}
}

func (s ServiceList) Run(ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/services/instances")
	if err != nil {
		return err
	}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	rslt, err := cmd.ShowServicesInstancesList(b)
	if err != nil {
		return err
	}
	n, err := ctx.Stdout.Write(rslt)
	if n != len(rslt) {
		return errors.New("Failed to write the output of the command")
	}
	return nil
}

type ServiceAdd struct {
	fs        *gnuflag.FlagSet
	teamOwner string
}

func (c *ServiceAdd) Info() *cmd.Info {
	usage := `service-add <servicename> <serviceinstancename> [plan] [-t/--owner-team <team>]
e.g.:

    $ tsuru service-add mongodb tsuru_mongodb small -t myteam

Will add a new instance of the "mongodb" service, named "tsuru_mongodb" with the plan "small".`
	return &cmd.Info{
		Name:    "service-add",
		Usage:   usage,
		Desc:    "Create a service instance to one or more apps make use of.",
		MinArgs: 2,
		MaxArgs: 3,
	}
}

func (c *ServiceAdd) Run(ctx *cmd.Context, client *cmd.Client) error {
	serviceName, instanceName := ctx.Args[0], ctx.Args[1]
	var plan string
	if len(ctx.Args) > 2 {
		plan = ctx.Args[2]
	}
	var b bytes.Buffer
	params := map[string]string{
		"name":         instanceName,
		"service_name": serviceName,
		"plan":         plan,
		"owner":        c.teamOwner,
	}
	err := json.NewEncoder(&b).Encode(params)
	if err != nil {
		return err
	}
	url, err := cmd.GetURL("/services/instances")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("POST", url, &b)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprint(ctx.Stdout, "Service successfully added.\n")
	return nil
}

func (c *ServiceAdd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		flagDesc := "the team that owns te service (mandatory if the user is member of more than one team"
		c.fs = gnuflag.NewFlagSet("service-add", gnuflag.ExitOnError)
		c.fs.StringVar(&c.teamOwner, "team-owner", "", flagDesc)
		c.fs.StringVar(&c.teamOwner, "t", "", flagDesc)
	}
	return c.fs
}

type ServiceBind struct {
	GuessingCommand
}

func (sb *ServiceBind) Run(ctx *cmd.Context, client *cmd.Client) error {
	appName, err := sb.Guess()
	if err != nil {
		return err
	}
	instanceName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instanceName + "/" + appName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("PUT", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var variables []string
	dec := json.NewDecoder(resp.Body)
	msg := fmt.Sprintf("Instance %q is now bound to the app %q.\n", instanceName, appName)
	if err = dec.Decode(&variables); err == nil && len(variables) > 0 {
		msg += fmt.Sprintf(`
The following environment variables are now available for use in your app:

- %s

For more details, please check the documentation for the service, using service-doc command.
`, strings.Join(variables, "\n- "))
	}
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return io.ErrShortWrite
	}
	return nil
}

func (sb *ServiceBind) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "bind",
		Usage: "bind <instancename> [--app appname]",
		Desc: `bind a service instance to an app

If you don't provide the app name, tsuru will try to guess it.`,
		MinArgs: 1,
	}
}

type ServiceUnbind struct {
	GuessingCommand
}

func (su *ServiceUnbind) Run(ctx *cmd.Context, client *cmd.Client) error {
	appName, err := su.Guess()
	if err != nil {
		return err
	}
	instanceName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instanceName + "/" + appName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	msg := fmt.Sprintf("Instance %q is not bound to the app %q anymore.\n", instanceName, appName)
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return errors.New("Failed to write to standard output.\n")
	}
	return nil
}

func (su *ServiceUnbind) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "unbind",
		Usage: "unbind <instancename> [--app appname]",
		Desc: `unbind a service instance from an app

If you don't provide the app name, tsuru will try to guess it.`,
		MinArgs: 1,
	}
}

type ServiceInstanceStatus struct{}

func (c ServiceInstanceStatus) Info() *cmd.Info {
	usg := `service-status <serviceinstancename>
e.g.:

    $ tsuru service-status my_mongodb
`
	return &cmd.Info{
		Name:    "service-status",
		Usage:   usg,
		Desc:    "Check status of a given service instance.",
		MinArgs: 1,
	}
}

func (c ServiceInstanceStatus) Run(ctx *cmd.Context, client *cmd.Client) error {
	instName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instName + "/status")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	bMsg, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	msg := string(bMsg) + "\n"
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return errors.New("Failed to write to standard output.\n")
	}
	return nil
}

type ServiceInfo struct{}

func (c ServiceInfo) Info() *cmd.Info {
	usg := `service-info <service>
e.g.:

    $ tsuru service-info mongodb
`
	return &cmd.Info{
		Name:    "service-info",
		Usage:   usg,
		Desc:    "List all instances of a service",
		MinArgs: 1,
	}
}

type ServiceInstanceModel struct {
	Name string
	Apps []string
	Info map[string]string
}

// in returns true if the list contains the value
func in(value string, list []string) bool {
	for _, item := range list {
		if value == item {
			return true
		}
	}
	return false
}

func (ServiceInfo) ExtraHeaders(instances []ServiceInstanceModel) []string {
	var headers []string
	for _, instance := range instances {
		for key := range instance.Info {
			if !in(key, headers) {
				headers = append(headers, key)
			}
		}
	}
	sort.Sort(sort.StringSlice(headers))
	return headers
}

func (c ServiceInfo) BuildInstancesTable(serviceName string, ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/services/" + serviceName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var instances []ServiceInstanceModel
	err = json.Unmarshal(result, &instances)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte(fmt.Sprintf("Info for \"%s\"\n\n", serviceName)))
	if len(instances) > 0 {
		ctx.Stdout.Write([]byte("Instances\n"))
		table := cmd.NewTable()
		extraHeaders := c.ExtraHeaders(instances)
		for _, instance := range instances {
			apps := strings.Join(instance.Apps, ", ")
			data := []string{instance.Name, apps}
			for _, h := range extraHeaders {
				data = append(data, instance.Info[h])
			}
			table.AddRow(cmd.Row(data))
		}
		headers := []string{"Instances", "Apps"}
		headers = append(headers, extraHeaders...)
		table.Headers = cmd.Row(headers)
		ctx.Stdout.Write(table.Bytes())
	}
	return nil
}

func (c ServiceInfo) BuildPlansTable(serviceName string, ctx *cmd.Context, client *cmd.Client) error {
	ctx.Stdout.Write([]byte("\nPlans\n"))
	url, err := cmd.GetURL(fmt.Sprintf("/services/%s/plans", serviceName))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var plans []map[string]string
	err = json.Unmarshal(result, &plans)
	if err != nil {
		return err
	}
	if len(plans) > 0 {
		table := cmd.NewTable()
		for _, plan := range plans {
			data := []string{plan["Name"], plan["Description"]}
			table.AddRow(cmd.Row(data))
		}
		table.Headers = cmd.Row([]string{"Name", "Description"})
		ctx.Stdout.Write(table.Bytes())
	}
	return nil
}

func (c ServiceInfo) Run(ctx *cmd.Context, client *cmd.Client) error {
	serviceName := ctx.Args[0]
	err := c.BuildInstancesTable(serviceName, ctx, client)
	if err != nil {
		return err
	}
	return c.BuildPlansTable(serviceName, ctx, client)
}

type ServiceDoc struct{}

func (ServiceDoc) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "service-doc",
		Usage:   "service-doc <servicename>",
		Desc:    "Show documentation of a service",
		MinArgs: 1,
	}
}

func (ServiceDoc) Run(ctx *cmd.Context, client *cmd.Client) error {
	sName := ctx.Args[0]
	url := fmt.Sprintf("/services/%s/doc", sName)
	url, err := cmd.GetURL(url)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	ctx.Stdout.Write(result)
	return nil
}

type ServiceRemove struct {
	GuessingCommand
	yes bool
	fs  *gnuflag.FlagSet
}

func (c ServiceRemove) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "service-remove",
		Usage:   "service-remove <serviceinstancename> [--assume-yes]",
		Desc:    "Removes a service instance",
		MinArgs: 1,
	}
}

func (c ServiceRemove) Run(ctx *cmd.Context, client *cmd.Client) error {
	name := ctx.Args[0]
	var answer string
	if !c.yes {
		fmt.Fprintf(ctx.Stdout, `Are you sure you want to remove service "%s"? (y/n) `, name)
		fmt.Fscanf(ctx.Stdin, "%s", &answer)
		if answer != "y" {
			fmt.Fprintln(ctx.Stdout, "Abort.")
			return nil
		}
	}
	url := fmt.Sprintf("/services/instances/%s", name)
	url, err := cmd.GetURL(url)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintf(ctx.Stdout, `Service "%s" successfully removed!`+"\n", name)
	return nil
}

func (c *ServiceRemove) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = c.GuessingCommand.Flags()
		c.fs.BoolVar(&c.yes, "assume-yes", false, "Don't ask for confirmation, just remove the service.")
		c.fs.BoolVar(&c.yes, "y", false, "Don't ask for confirmation, just remove the service.")
	}
	return c.fs
}
<commit_msg>cmd/tsuru-base: use a pointer receiver in ServiceRemove<commit_after>// Copyright 2014 tsuru 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 tsuru

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/tsuru/tsuru/cmd"
	"io"
	"io/ioutil"
	"launchpad.net/gnuflag"
	"net/http"
	"sort"
	"strings"
)

type ServiceList struct{}

func (s ServiceList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "service-list",
		Usage: "service-list",
		Desc:  "Get all available services, and user's instances for this services",
	}
}

func (s ServiceList) Run(ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/services/instances")
	if err != nil {
		return err
	}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	rslt, err := cmd.ShowServicesInstancesList(b)
	if err != nil {
		return err
	}
	n, err := ctx.Stdout.Write(rslt)
	if n != len(rslt) {
		return errors.New("Failed to write the output of the command")
	}
	return nil
}

type ServiceAdd struct {
	fs        *gnuflag.FlagSet
	teamOwner string
}

func (c *ServiceAdd) Info() *cmd.Info {
	usage := `service-add <servicename> <serviceinstancename> [plan] [-t/--owner-team <team>]
e.g.:

    $ tsuru service-add mongodb tsuru_mongodb small -t myteam

Will add a new instance of the "mongodb" service, named "tsuru_mongodb" with the plan "small".`
	return &cmd.Info{
		Name:    "service-add",
		Usage:   usage,
		Desc:    "Create a service instance to one or more apps make use of.",
		MinArgs: 2,
		MaxArgs: 3,
	}
}

func (c *ServiceAdd) Run(ctx *cmd.Context, client *cmd.Client) error {
	serviceName, instanceName := ctx.Args[0], ctx.Args[1]
	var plan string
	if len(ctx.Args) > 2 {
		plan = ctx.Args[2]
	}
	var b bytes.Buffer
	params := map[string]string{
		"name":         instanceName,
		"service_name": serviceName,
		"plan":         plan,
		"owner":        c.teamOwner,
	}
	err := json.NewEncoder(&b).Encode(params)
	if err != nil {
		return err
	}
	url, err := cmd.GetURL("/services/instances")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("POST", url, &b)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprint(ctx.Stdout, "Service successfully added.\n")
	return nil
}

func (c *ServiceAdd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		flagDesc := "the team that owns te service (mandatory if the user is member of more than one team"
		c.fs = gnuflag.NewFlagSet("service-add", gnuflag.ExitOnError)
		c.fs.StringVar(&c.teamOwner, "team-owner", "", flagDesc)
		c.fs.StringVar(&c.teamOwner, "t", "", flagDesc)
	}
	return c.fs
}

type ServiceBind struct {
	GuessingCommand
}

func (sb *ServiceBind) Run(ctx *cmd.Context, client *cmd.Client) error {
	appName, err := sb.Guess()
	if err != nil {
		return err
	}
	instanceName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instanceName + "/" + appName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("PUT", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var variables []string
	dec := json.NewDecoder(resp.Body)
	msg := fmt.Sprintf("Instance %q is now bound to the app %q.\n", instanceName, appName)
	if err = dec.Decode(&variables); err == nil && len(variables) > 0 {
		msg += fmt.Sprintf(`
The following environment variables are now available for use in your app:

- %s

For more details, please check the documentation for the service, using service-doc command.
`, strings.Join(variables, "\n- "))
	}
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return io.ErrShortWrite
	}
	return nil
}

func (sb *ServiceBind) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "bind",
		Usage: "bind <instancename> [--app appname]",
		Desc: `bind a service instance to an app

If you don't provide the app name, tsuru will try to guess it.`,
		MinArgs: 1,
	}
}

type ServiceUnbind struct {
	GuessingCommand
}

func (su *ServiceUnbind) Run(ctx *cmd.Context, client *cmd.Client) error {
	appName, err := su.Guess()
	if err != nil {
		return err
	}
	instanceName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instanceName + "/" + appName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	msg := fmt.Sprintf("Instance %q is not bound to the app %q anymore.\n", instanceName, appName)
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return errors.New("Failed to write to standard output.\n")
	}
	return nil
}

func (su *ServiceUnbind) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "unbind",
		Usage: "unbind <instancename> [--app appname]",
		Desc: `unbind a service instance from an app

If you don't provide the app name, tsuru will try to guess it.`,
		MinArgs: 1,
	}
}

type ServiceInstanceStatus struct{}

func (c ServiceInstanceStatus) Info() *cmd.Info {
	usg := `service-status <serviceinstancename>
e.g.:

    $ tsuru service-status my_mongodb
`
	return &cmd.Info{
		Name:    "service-status",
		Usage:   usg,
		Desc:    "Check status of a given service instance.",
		MinArgs: 1,
	}
}

func (c ServiceInstanceStatus) Run(ctx *cmd.Context, client *cmd.Client) error {
	instName := ctx.Args[0]
	url, err := cmd.GetURL("/services/instances/" + instName + "/status")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	bMsg, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	msg := string(bMsg) + "\n"
	n, err := fmt.Fprint(ctx.Stdout, msg)
	if err != nil {
		return err
	}
	if n != len(msg) {
		return errors.New("Failed to write to standard output.\n")
	}
	return nil
}

type ServiceInfo struct{}

func (c ServiceInfo) Info() *cmd.Info {
	usg := `service-info <service>
e.g.:

    $ tsuru service-info mongodb
`
	return &cmd.Info{
		Name:    "service-info",
		Usage:   usg,
		Desc:    "List all instances of a service",
		MinArgs: 1,
	}
}

type ServiceInstanceModel struct {
	Name string
	Apps []string
	Info map[string]string
}

// in returns true if the list contains the value
func in(value string, list []string) bool {
	for _, item := range list {
		if value == item {
			return true
		}
	}
	return false
}

func (ServiceInfo) ExtraHeaders(instances []ServiceInstanceModel) []string {
	var headers []string
	for _, instance := range instances {
		for key := range instance.Info {
			if !in(key, headers) {
				headers = append(headers, key)
			}
		}
	}
	sort.Sort(sort.StringSlice(headers))
	return headers
}

func (c ServiceInfo) BuildInstancesTable(serviceName string, ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/services/" + serviceName)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var instances []ServiceInstanceModel
	err = json.Unmarshal(result, &instances)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte(fmt.Sprintf("Info for \"%s\"\n\n", serviceName)))
	if len(instances) > 0 {
		ctx.Stdout.Write([]byte("Instances\n"))
		table := cmd.NewTable()
		extraHeaders := c.ExtraHeaders(instances)
		for _, instance := range instances {
			apps := strings.Join(instance.Apps, ", ")
			data := []string{instance.Name, apps}
			for _, h := range extraHeaders {
				data = append(data, instance.Info[h])
			}
			table.AddRow(cmd.Row(data))
		}
		headers := []string{"Instances", "Apps"}
		headers = append(headers, extraHeaders...)
		table.Headers = cmd.Row(headers)
		ctx.Stdout.Write(table.Bytes())
	}
	return nil
}

func (c ServiceInfo) BuildPlansTable(serviceName string, ctx *cmd.Context, client *cmd.Client) error {
	ctx.Stdout.Write([]byte("\nPlans\n"))
	url, err := cmd.GetURL(fmt.Sprintf("/services/%s/plans", serviceName))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var plans []map[string]string
	err = json.Unmarshal(result, &plans)
	if err != nil {
		return err
	}
	if len(plans) > 0 {
		table := cmd.NewTable()
		for _, plan := range plans {
			data := []string{plan["Name"], plan["Description"]}
			table.AddRow(cmd.Row(data))
		}
		table.Headers = cmd.Row([]string{"Name", "Description"})
		ctx.Stdout.Write(table.Bytes())
	}
	return nil
}

func (c ServiceInfo) Run(ctx *cmd.Context, client *cmd.Client) error {
	serviceName := ctx.Args[0]
	err := c.BuildInstancesTable(serviceName, ctx, client)
	if err != nil {
		return err
	}
	return c.BuildPlansTable(serviceName, ctx, client)
}

type ServiceDoc struct{}

func (ServiceDoc) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "service-doc",
		Usage:   "service-doc <servicename>",
		Desc:    "Show documentation of a service",
		MinArgs: 1,
	}
}

func (ServiceDoc) Run(ctx *cmd.Context, client *cmd.Client) error {
	sName := ctx.Args[0]
	url := fmt.Sprintf("/services/%s/doc", sName)
	url, err := cmd.GetURL(url)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(request)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	ctx.Stdout.Write(result)
	return nil
}

type ServiceRemove struct {
	GuessingCommand
	yes bool
	fs  *gnuflag.FlagSet
}

func (c *ServiceRemove) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "service-remove",
		Usage:   "service-remove <serviceinstancename> [--assume-yes]",
		Desc:    "Removes a service instance",
		MinArgs: 1,
	}
}

func (c *ServiceRemove) Run(ctx *cmd.Context, client *cmd.Client) error {
	name := ctx.Args[0]
	var answer string
	if !c.yes {
		fmt.Fprintf(ctx.Stdout, `Are you sure you want to remove service "%s"? (y/n) `, name)
		fmt.Fscanf(ctx.Stdin, "%s", &answer)
		if answer != "y" {
			fmt.Fprintln(ctx.Stdout, "Abort.")
			return nil
		}
	}
	url := fmt.Sprintf("/services/instances/%s", name)
	url, err := cmd.GetURL(url)
	if err != nil {
		return err
	}
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	fmt.Fprintf(ctx.Stdout, `Service "%s" successfully removed!`+"\n", name)
	return nil
}

func (c *ServiceRemove) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = c.GuessingCommand.Flags()
		c.fs.BoolVar(&c.yes, "assume-yes", false, "Don't ask for confirmation, just remove the service.")
		c.fs.BoolVar(&c.yes, "y", false, "Don't ask for confirmation, just remove the service.")
	}
	return c.fs
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/authorizer"
	"github.com/google/fleetspeak/fleetspeak/src/server/comms"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"
	"github.com/google/fleetspeak/fleetspeak/src/server/internal/signatures"
	"github.com/google/fleetspeak/fleetspeak/src/server/stats"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const maxMessagesPerContact = 100
const processingChunkSize = 10

type commsContext struct {
	s *Server
}

// GetClientInfo loads basic information about a client. Returns nil if the client does
// not exist in the datastore.
func (c commsContext) GetClientInfo(ctx context.Context, id common.ClientID) (*comms.ClientInfo, error) {
	cld, err := c.s.dataStore.GetClientData(ctx, id)
	if err != nil {
		if c.s.dataStore.IsNotFound(err) {
			return nil, nil
		}
		return nil, err
	}
	k, err := x509.ParsePKIXPublicKey(cld.Key)
	if err != nil {
		return nil, err
	}
	return &comms.ClientInfo{
		ID:          id,
		Key:         k,
		Labels:      cld.Labels,
		Blacklisted: cld.Blacklisted}, nil
}

// AddClient adds a new client to the system.
func (c commsContext) AddClient(ctx context.Context, id common.ClientID, key crypto.PublicKey) (*comms.ClientInfo, error) {
	k, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return nil, err
	}
	if err := c.s.dataStore.AddClient(ctx, id, &db.ClientData{Key: k}); err != nil {
		return nil, err
	}
	return &comms.ClientInfo{ID: id, Key: key}, nil
}

func (c commsContext) HandleClientContact(ctx context.Context, info *comms.ClientInfo, addr net.Addr, wcd *fspb.WrappedContactData) (*fspb.ContactData, error) {
	sigs, err := signatures.ValidateWrappedContactData(wcd)
	if err != nil {
		return nil, err
	}
	accept, validationInfo := c.s.authorizer.Allow4(
		addr,
		authorizer.ContactInfo{
			ID:           info.ID,
			ContactSize:  len(wcd.ContactData),
			ClientLabels: wcd.ClientLabels,
		},
		authorizer.ClientInfo{
			Labels: info.Labels,
		},
		sigs)
	if !accept {
		return nil, errors.New("contact not authorized")
	}
	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return nil, fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return nil, fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}

	toSend := fspb.ContactData{SequencingNonce: uint64(rand.Int63())}
	ct, err := c.RecordClientContact(ctx, info, toSend.SequencingNonce, cd.SequencingNonce, addr.String())
	if err != nil {
		return nil, err
	}

	err = c.handleMessagesFromClient(ctx, info, ct, &cd, validationInfo)
	if err != nil {
		return nil, err
	}

	toSend.Messages, err = c.FindMessagesForClient(ctx, info, ct, 100)
	if err != nil {
		return nil, err
	}
	return &toSend, nil
}

// FindMessagesForClient finds unprocessed messages for a given client and
// reserves them for processing.
func (c commsContext) FindMessagesForClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, maxMessages int) ([]*fspb.Message, error) {
	if info.Blacklisted {
		m, err := c.MakeBlacklistMessage(ctx, info, contactID)
		return []*fspb.Message{m}, err
	}
	msgs, err := c.s.dataStore.ClientMessagesForProcessing(ctx, info.ID, maxMessages)
	if err != nil {
		if len(msgs) == 0 {
			return nil, err
		}
		log.Warning("Got %v messages along with error, continuing: %v", len(msgs), err)
	}

	bms, err := c.s.broadcastManager.MakeBroadcastMessagesForClient(ctx, info.ID, info.Labels)
	if err != nil {
		return nil, err
	}
	msgs = append(msgs, bms...)

	if len(msgs) == 0 {
		return msgs, nil
	}

	mids := make([]common.MessageID, 0, len(msgs))
	for _, m := range msgs {
		id, err := common.BytesToMessageID(m.MessageId)
		if err != nil {
			return nil, err
		}
		mids = append(mids, id)
	}
	err = c.s.dataStore.LinkMessagesToContact(ctx, contactID, mids)
	if err != nil {
		return nil, err
	}
	return msgs, nil
}

func (c commsContext) MakeBlacklistMessage(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID) (*fspb.Message, error) {
	mid, err := common.RandomMessageID()
	if err != nil {
		return nil, fmt.Errorf("unable to create message id: %v", err)
	}
	msg := &fspb.Message{
		MessageId: mid.Bytes(),
		Source: &fspb.Address{
			ServiceName: "system",
		},
		Destination: &fspb.Address{
			ServiceName: "system",
			ClientId:    info.ID.Bytes(),
		},
		MessageType: "RekeyRequest",
	}
	if err = c.s.dataStore.StoreMessages(ctx, []*fspb.Message{msg}, contactID); err != nil {
		return nil, fmt.Errorf("unable to store RekeyRequest: %v", err)
	}
	return msg, nil
}

func (c commsContext) validateMessageFromClient(id common.ClientID, m *fspb.Message, validationInfo string) error {
	if m.Destination == nil {
		return fmt.Errorf("message must have Destination")
	}
	if m.Destination.ClientId != nil {
		return fmt.Errorf("cannot send a message directly to another client [%v]", m.Destination.ClientId)
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return fmt.Errorf("message must have a source with a ServiceName, got: %v", m.Source)
	}
	if m.SourceMessageId == nil {
		return fmt.Errorf("source message id cannot be empty")
	}

	m.Source.ClientId = id.Bytes()
	m.ValidationInfo = validationInfo
	m.MessageId = common.MakeMessageID(m.Source, m.SourceMessageId).Bytes()
	return nil
}

// RecordClientContact records that a contact occurred. The resulting
// ContactID is guaranteed to be unique.
func (c commsContext) RecordClientContact(ctx context.Context, info *comms.ClientInfo, sentNonce, receivedNonce uint64, addr string) (db.ContactID, error) {
	return c.s.dataStore.RecordClientContact(ctx, info.ID, sentNonce, receivedNonce, addr)
}

// handleMessagesFromClient processes a block of messages from a particular
// client. It saves them to the database, associates them with the contact
// identified by contactTime, and processes them.
func (c commsContext) handleMessagesFromClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, received *fspb.ContactData, validationInfo string) error {
	msgs := make([]*fspb.Message, 0, len(received.Messages))
	for _, m := range received.Messages {
		err := c.validateMessageFromClient(info.ID, m, validationInfo)
		if err != nil {
			log.Errorf("Dropping invalid message from [%v]: %v", info.ID, err)
			continue
		}
		msgs = append(msgs, m)
	}
	if len(msgs) == 0 {
		return nil
	}

	sort.Slice(msgs, func(a, b int) bool {
		return bytes.Compare(msgs[a].MessageId, msgs[b].MessageId) == -1
	})

	for {
		if len(msgs) <= processingChunkSize {
			return c.s.serviceConfig.HandleNewMessages(ctx, msgs, contactID)
		}

		if err := c.s.serviceConfig.HandleNewMessages(ctx, msgs[:processingChunkSize], contactID); err != nil {
			return err
		}
		msgs = msgs[processingChunkSize:]
	}
}

// ReadFile returns the data and modification time of file. Caller is
// responsible for closing data.
//
// Calls to data are permitted to fail if ctx is canceled or expired.
func (c commsContext) ReadFile(ctx context.Context, service, name string) (data db.ReadSeekerCloser, modtime time.Time, err error) {
	return c.s.dataStore.ReadFile(ctx, service, name)
}

// IsNotFound returns whether an error returned by ReadFile indicates that the
// file was not found.
func (c commsContext) IsNotFound(err error) bool {
	return c.s.dataStore.IsNotFound(err)
}

// StatsCollector returns the stats.Collector used by the Fleetspeak
// system. Access is provided to allow collection of stats relating to the
// client communication.
func (c commsContext) StatsCollector() stats.Collector {
	return c.s.statsCollector
}

func (c commsContext) Authorizer() authorizer.Authorizer {
	return c.s.authorizer
}
<commit_msg>Set RekeyRequest creation times.<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/authorizer"
	"github.com/google/fleetspeak/fleetspeak/src/server/comms"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"
	"github.com/google/fleetspeak/fleetspeak/src/server/internal/signatures"
	"github.com/google/fleetspeak/fleetspeak/src/server/stats"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const maxMessagesPerContact = 100
const processingChunkSize = 10

type commsContext struct {
	s *Server
}

// GetClientInfo loads basic information about a client. Returns nil if the client does
// not exist in the datastore.
func (c commsContext) GetClientInfo(ctx context.Context, id common.ClientID) (*comms.ClientInfo, error) {
	cld, err := c.s.dataStore.GetClientData(ctx, id)
	if err != nil {
		if c.s.dataStore.IsNotFound(err) {
			return nil, nil
		}
		return nil, err
	}
	k, err := x509.ParsePKIXPublicKey(cld.Key)
	if err != nil {
		return nil, err
	}
	return &comms.ClientInfo{
		ID:          id,
		Key:         k,
		Labels:      cld.Labels,
		Blacklisted: cld.Blacklisted}, nil
}

// AddClient adds a new client to the system.
func (c commsContext) AddClient(ctx context.Context, id common.ClientID, key crypto.PublicKey) (*comms.ClientInfo, error) {
	k, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return nil, err
	}
	if err := c.s.dataStore.AddClient(ctx, id, &db.ClientData{Key: k}); err != nil {
		return nil, err
	}
	return &comms.ClientInfo{ID: id, Key: key}, nil
}

func (c commsContext) HandleClientContact(ctx context.Context, info *comms.ClientInfo, addr net.Addr, wcd *fspb.WrappedContactData) (*fspb.ContactData, error) {
	sigs, err := signatures.ValidateWrappedContactData(wcd)
	if err != nil {
		return nil, err
	}
	accept, validationInfo := c.s.authorizer.Allow4(
		addr,
		authorizer.ContactInfo{
			ID:           info.ID,
			ContactSize:  len(wcd.ContactData),
			ClientLabels: wcd.ClientLabels,
		},
		authorizer.ClientInfo{
			Labels: info.Labels,
		},
		sigs)
	if !accept {
		return nil, errors.New("contact not authorized")
	}
	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return nil, fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return nil, fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}

	toSend := fspb.ContactData{SequencingNonce: uint64(rand.Int63())}
	ct, err := c.RecordClientContact(ctx, info, toSend.SequencingNonce, cd.SequencingNonce, addr.String())
	if err != nil {
		return nil, err
	}

	err = c.handleMessagesFromClient(ctx, info, ct, &cd, validationInfo)
	if err != nil {
		return nil, err
	}

	toSend.Messages, err = c.FindMessagesForClient(ctx, info, ct, 100)
	if err != nil {
		return nil, err
	}
	return &toSend, nil
}

// FindMessagesForClient finds unprocessed messages for a given client and
// reserves them for processing.
func (c commsContext) FindMessagesForClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, maxMessages int) ([]*fspb.Message, error) {
	if info.Blacklisted {
		m, err := c.MakeBlacklistMessage(ctx, info, contactID)
		return []*fspb.Message{m}, err
	}
	msgs, err := c.s.dataStore.ClientMessagesForProcessing(ctx, info.ID, maxMessages)
	if err != nil {
		if len(msgs) == 0 {
			return nil, err
		}
		log.Warning("Got %v messages along with error, continuing: %v", len(msgs), err)
	}

	bms, err := c.s.broadcastManager.MakeBroadcastMessagesForClient(ctx, info.ID, info.Labels)
	if err != nil {
		return nil, err
	}
	msgs = append(msgs, bms...)

	if len(msgs) == 0 {
		return msgs, nil
	}

	mids := make([]common.MessageID, 0, len(msgs))
	for _, m := range msgs {
		id, err := common.BytesToMessageID(m.MessageId)
		if err != nil {
			return nil, err
		}
		mids = append(mids, id)
	}
	err = c.s.dataStore.LinkMessagesToContact(ctx, contactID, mids)
	if err != nil {
		return nil, err
	}
	return msgs, nil
}

func (c commsContext) MakeBlacklistMessage(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID) (*fspb.Message, error) {
	mid, err := common.RandomMessageID()
	if err != nil {
		return nil, fmt.Errorf("unable to create message id: %v", err)
	}
	msg := &fspb.Message{
		MessageId: mid.Bytes(),
		Source: &fspb.Address{
			ServiceName: "system",
		},
		Destination: &fspb.Address{
			ServiceName: "system",
			ClientId:    info.ID.Bytes(),
		},
		MessageType:  "RekeyRequest",
		CreationTime: db.NowProto(),
	}
	if err = c.s.dataStore.StoreMessages(ctx, []*fspb.Message{msg}, contactID); err != nil {
		return nil, fmt.Errorf("unable to store RekeyRequest: %v", err)
	}
	return msg, nil
}

func (c commsContext) validateMessageFromClient(id common.ClientID, m *fspb.Message, validationInfo string) error {
	if m.Destination == nil {
		return fmt.Errorf("message must have Destination")
	}
	if m.Destination.ClientId != nil {
		return fmt.Errorf("cannot send a message directly to another client [%v]", m.Destination.ClientId)
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return fmt.Errorf("message must have a source with a ServiceName, got: %v", m.Source)
	}
	if m.SourceMessageId == nil {
		return fmt.Errorf("source message id cannot be empty")
	}

	m.Source.ClientId = id.Bytes()
	m.ValidationInfo = validationInfo
	m.MessageId = common.MakeMessageID(m.Source, m.SourceMessageId).Bytes()
	return nil
}

// RecordClientContact records that a contact occurred. The resulting
// ContactID is guaranteed to be unique.
func (c commsContext) RecordClientContact(ctx context.Context, info *comms.ClientInfo, sentNonce, receivedNonce uint64, addr string) (db.ContactID, error) {
	return c.s.dataStore.RecordClientContact(ctx, info.ID, sentNonce, receivedNonce, addr)
}

// handleMessagesFromClient processes a block of messages from a particular
// client. It saves them to the database, associates them with the contact
// identified by contactTime, and processes them.
func (c commsContext) handleMessagesFromClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, received *fspb.ContactData, validationInfo string) error {
	msgs := make([]*fspb.Message, 0, len(received.Messages))
	for _, m := range received.Messages {
		err := c.validateMessageFromClient(info.ID, m, validationInfo)
		if err != nil {
			log.Errorf("Dropping invalid message from [%v]: %v", info.ID, err)
			continue
		}
		msgs = append(msgs, m)
	}
	if len(msgs) == 0 {
		return nil
	}

	sort.Slice(msgs, func(a, b int) bool {
		return bytes.Compare(msgs[a].MessageId, msgs[b].MessageId) == -1
	})

	for {
		if len(msgs) <= processingChunkSize {
			return c.s.serviceConfig.HandleNewMessages(ctx, msgs, contactID)
		}

		if err := c.s.serviceConfig.HandleNewMessages(ctx, msgs[:processingChunkSize], contactID); err != nil {
			return err
		}
		msgs = msgs[processingChunkSize:]
	}
}

// ReadFile returns the data and modification time of file. Caller is
// responsible for closing data.
//
// Calls to data are permitted to fail if ctx is canceled or expired.
func (c commsContext) ReadFile(ctx context.Context, service, name string) (data db.ReadSeekerCloser, modtime time.Time, err error) {
	return c.s.dataStore.ReadFile(ctx, service, name)
}

// IsNotFound returns whether an error returned by ReadFile indicates that the
// file was not found.
func (c commsContext) IsNotFound(err error) bool {
	return c.s.dataStore.IsNotFound(err)
}

// StatsCollector returns the stats.Collector used by the Fleetspeak
// system. Access is provided to allow collection of stats relating to the
// client communication.
func (c commsContext) StatsCollector() stats.Collector {
	return c.s.statsCollector
}

func (c commsContext) Authorizer() authorizer.Authorizer {
	return c.s.authorizer
}
<|endoftext|>


<commit_before>package ttlmap

import (
	"fmt"
	"github.com/mailgun/minheap"
	"github.com/mailgun/timetools"
	"time"
)

type TtlMap struct {
	capacity     int
	elements     map[string]*mapElement
	expiryTimes  *minheap.MinHeap
	timeProvider timetools.TimeProvider
}

type mapElement struct {
	key    string
	value  interface{}
	heapEl *minheap.Element
}

func NewMap(capacity int) (*TtlMap, error) {
	return NewMapWithProvider(capacity, &timetools.RealTime{})
}

func NewMapWithProvider(capacity int, timeProvider timetools.TimeProvider) (*TtlMap, error) {
	if capacity <= 0 {
		return nil, fmt.Errorf("Capacity should be >= 0")
	}
	if timeProvider == nil {
		return nil, fmt.Errorf("Please pass timeProvider")
	}

	return &TtlMap{
		capacity:     capacity,
		elements:     make(map[string]*mapElement),
		expiryTimes:  minheap.NewMinHeap(),
		timeProvider: timeProvider,
	}, nil
}

func (m *TtlMap) Set(key string, value interface{}, ttlSeconds int) error {
	expiryTime, err := m.toEpochSeconds(ttlSeconds)
	if err != nil {
		return err
	}

	mapEl, exists := m.elements[key]
	if !exists {
		if len(m.elements) >= m.capacity {
			m.freeSpace(1)
		}
		heapEl := &minheap.Element{
			Priority: expiryTime,
		}
		mapEl := &mapElement{
			key:    key,
			value:  value,
			heapEl: heapEl,
		}
		heapEl.Value = mapEl
		m.elements[key] = mapEl
		m.expiryTimes.PushEl(heapEl)
	} else {
		mapEl.value = value
		m.expiryTimes.UpdateEl(mapEl.heapEl, expiryTime)
	}
	return nil
}

func (m *TtlMap) toEpochSeconds(ttlSeconds int) (int, error) {
	if ttlSeconds <= 0 {
		return 0, fmt.Errorf("ttlSeconds should be >= 0, got %d", ttlSeconds)
	}
	return int(m.timeProvider.UtcNow().Add(time.Second * time.Duration(ttlSeconds)).Unix()), nil
}

func (m *TtlMap) Len() int {
	return len(m.elements)
}

func (m *TtlMap) Get(key string) (interface{}, bool) {
	mapEl, exists := m.elements[key]
	if !exists {
		return nil, false
	}
	if m.expireElement(mapEl) {
		return nil, false
	}
	return mapEl.value, true
}

func (m *TtlMap) Increment(key string, value int, ttlSeconds int) (int, error) {
	expiryTime, err := m.toEpochSeconds(ttlSeconds)
	if err != nil {
		return 0, err
	}

	mapEl, exists := m.elements[key]
	if !exists {
		m.Set(key, value, ttlSeconds)
		return value, nil
	}
	if m.expireElement(mapEl) {
		m.Set(key, value, ttlSeconds)
		return value, nil
	}
	currentValue, ok := mapEl.value.(int)
	if !ok {
		return 0, fmt.Errorf("Expected existing value to be integer, got %T", mapEl.value)
	}
	currentValue += value
	mapEl.value = currentValue

	m.expiryTimes.UpdateEl(mapEl.heapEl, expiryTime)
	return currentValue, nil
}

func (m *TtlMap) GetInt(key string) (int, bool, error) {
	valueI, exists := m.Get(key)
	if !exists {
		return 0, false, nil
	}
	value, ok := valueI.(int)
	if !ok {
		return 0, false, fmt.Errorf("Expected existing value to be integer, got %T", valueI)
	}
	return value, true, nil
}

func (m *TtlMap) expireElement(mapEl *mapElement) bool {
	now := int(m.timeProvider.UtcNow().Unix())
	if mapEl.heapEl.Priority > now {
		return false
	}
	delete(m.elements, mapEl.key)
	m.expiryTimes.RemoveEl(mapEl.heapEl)
	return true
}

func (m *TtlMap) freeSpace(count int) {
	removed := m.removeExpired(count)
	if removed >= count {
		return
	}
	m.removeLastUsed(count - removed)
}

func (m *TtlMap) removeExpired(iterations int) int {
	removed := 0
	now := int(m.timeProvider.UtcNow().Unix())
	for i := 0; i < iterations; i += 1 {
		if len(m.elements) == 0 {
			break
		}
		heapEl := m.expiryTimes.PeekEl()
		if heapEl.Priority > now {
			break
		}
		m.expiryTimes.PopEl()
		mapEl := heapEl.Value.(*mapElement)
		delete(m.elements, mapEl.key)
		removed += 1
	}
	return removed
}

func (m *TtlMap) removeLastUsed(iterations int) {
	for i := 0; i < iterations; i += 1 {
		if len(m.elements) == 0 {
			return
		}
		heapEl := m.expiryTimes.PopEl()
		mapEl := heapEl.Value.(*mapElement)
		delete(m.elements, mapEl.key)
	}
}
<commit_msg>Expose TimeProvider so other packages can mock time.<commit_after>package ttlmap

import (
	"fmt"
	"github.com/mailgun/minheap"
	"github.com/mailgun/timetools"
	"time"
)

type TtlMap struct {
	capacity     int
	elements     map[string]*mapElement
	expiryTimes  *minheap.MinHeap
	TimeProvider timetools.TimeProvider
}

type mapElement struct {
	key    string
	value  interface{}
	heapEl *minheap.Element
}

func NewMap(capacity int) (*TtlMap, error) {
	return NewMapWithProvider(capacity, &timetools.RealTime{})
}

func NewMapWithProvider(capacity int, timeProvider timetools.TimeProvider) (*TtlMap, error) {
	if capacity <= 0 {
		return nil, fmt.Errorf("Capacity should be >= 0")
	}
	if timeProvider == nil {
		return nil, fmt.Errorf("Please pass timeProvider")
	}

	return &TtlMap{
		capacity:     capacity,
		elements:     make(map[string]*mapElement),
		expiryTimes:  minheap.NewMinHeap(),
		TimeProvider: timeProvider,
	}, nil
}

func (m *TtlMap) Set(key string, value interface{}, ttlSeconds int) error {
	expiryTime, err := m.toEpochSeconds(ttlSeconds)
	if err != nil {
		return err
	}

	mapEl, exists := m.elements[key]
	if !exists {
		if len(m.elements) >= m.capacity {
			m.freeSpace(1)
		}
		heapEl := &minheap.Element{
			Priority: expiryTime,
		}
		mapEl := &mapElement{
			key:    key,
			value:  value,
			heapEl: heapEl,
		}
		heapEl.Value = mapEl
		m.elements[key] = mapEl
		m.expiryTimes.PushEl(heapEl)
	} else {
		mapEl.value = value
		m.expiryTimes.UpdateEl(mapEl.heapEl, expiryTime)
	}
	return nil
}

func (m *TtlMap) toEpochSeconds(ttlSeconds int) (int, error) {
	if ttlSeconds <= 0 {
		return 0, fmt.Errorf("ttlSeconds should be >= 0, got %d", ttlSeconds)
	}
	return int(m.TimeProvider.UtcNow().Add(time.Second * time.Duration(ttlSeconds)).Unix()), nil
}

func (m *TtlMap) Len() int {
	return len(m.elements)
}

func (m *TtlMap) Get(key string) (interface{}, bool) {
	mapEl, exists := m.elements[key]
	if !exists {
		return nil, false
	}
	if m.expireElement(mapEl) {
		return nil, false
	}
	return mapEl.value, true
}

func (m *TtlMap) Increment(key string, value int, ttlSeconds int) (int, error) {
	expiryTime, err := m.toEpochSeconds(ttlSeconds)
	if err != nil {
		return 0, err
	}

	mapEl, exists := m.elements[key]
	if !exists {
		m.Set(key, value, ttlSeconds)
		return value, nil
	}
	if m.expireElement(mapEl) {
		m.Set(key, value, ttlSeconds)
		return value, nil
	}
	currentValue, ok := mapEl.value.(int)
	if !ok {
		return 0, fmt.Errorf("Expected existing value to be integer, got %T", mapEl.value)
	}
	currentValue += value
	mapEl.value = currentValue

	m.expiryTimes.UpdateEl(mapEl.heapEl, expiryTime)
	return currentValue, nil
}

func (m *TtlMap) GetInt(key string) (int, bool, error) {
	valueI, exists := m.Get(key)
	if !exists {
		return 0, false, nil
	}
	value, ok := valueI.(int)
	if !ok {
		return 0, false, fmt.Errorf("Expected existing value to be integer, got %T", valueI)
	}
	return value, true, nil
}

func (m *TtlMap) expireElement(mapEl *mapElement) bool {
	now := int(m.TimeProvider.UtcNow().Unix())
	if mapEl.heapEl.Priority > now {
		return false
	}
	delete(m.elements, mapEl.key)
	m.expiryTimes.RemoveEl(mapEl.heapEl)
	return true
}

func (m *TtlMap) freeSpace(count int) {
	removed := m.removeExpired(count)
	if removed >= count {
		return
	}
	m.removeLastUsed(count - removed)
}

func (m *TtlMap) removeExpired(iterations int) int {
	removed := 0
	now := int(m.TimeProvider.UtcNow().Unix())
	for i := 0; i < iterations; i += 1 {
		if len(m.elements) == 0 {
			break
		}
		heapEl := m.expiryTimes.PeekEl()
		if heapEl.Priority > now {
			break
		}
		m.expiryTimes.PopEl()
		mapEl := heapEl.Value.(*mapElement)
		delete(m.elements, mapEl.key)
		removed += 1
	}
	return removed
}

func (m *TtlMap) removeLastUsed(iterations int) {
	for i := 0; i < iterations; i += 1 {
		if len(m.elements) == 0 {
			return
		}
		heapEl := m.expiryTimes.PopEl()
		mapEl := heapEl.Value.(*mapElement)
		delete(m.elements, mapEl.key)
	}
}
<|endoftext|>


<commit_before>package db

import (
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"testing"
)

type dd struct {
	Name  string
	Value string
}

// Test Basic Query of MongoDB
func TestA(t *testing.T) {
	abc := NewDBConfig("test")
	d, err := abc.Query(func(dbm *mgo.Database) (data interface{}, err error) {
		c := dbm.C("userinfo")
		err = c.Insert(&dd{Name: "Maria", Value: "Hello"})
		if err != nil {
			return
		}
		m := bson.M{"name": "Maria"}
		data = &dd{}
		err = c.Find(m).One(data)
		if err != nil {
			return
		}
		c.RemoveAll(m)
		return
	})
	if err != nil {
		t.Error(err)
	} else {
		t.Log(d)
	}
}

func TestB(t *testing.T) {
	abc := NewDBConfig("test")
	x, err := abc.Read()
	if err != nil {
		t.Error(err)
	} else {
		t.Log(x)
	}
}
<commit_msg>New Wechat<commit_after>package db

import (
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"testing"
)

type dd struct {
	Name  string
	Value string
}

// Test Basic Query of MongoDB
func TestA(t *testing.T) {
	abc := NewDBConfig("test")
	d, err := abc.Query(func(dbm *mgo.Database) (data interface{}, err error) {
		c := dbm.C("userinfo")
		err = c.Insert(&dd{Name: "Maria", Value: "Hello"})
		if err != nil {
			return
		}
		m := bson.M{"name": "Maria"}
		data = &dd{}
		err = c.Find(m).One(data)
		if err != nil {
			return
		}
		c.RemoveAll(m)
		return
	})
	if err != nil {
		t.Error(err)
	} else {
		t.Log(d)
	}
}

func TestB(t *testing.T) {
	abc := NewDBConfig("test")
	x := &wechat.AccessToken{
		Token:      "this is token",
		ExpireTime: time.Now(),
	}
	err = abc.Write(x)
	if err != nil {
		t.Error(err)
	}
	y,err:=abc.Read()
	if err!=nil{
		t.Error(err)
	}
	if x.Token!=y.Token{
		t.Error("Token is not same!"))
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 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 ssa

import "sort"

// cse does common-subexpression elimination on the Function.
// Values are just relinked, nothing is deleted.  A subsequent deadcode
// pass is required to actually remove duplicate expressions.
func cse(f *Func) {
	// Two values are equivalent if they satisfy the following definition:
	// equivalent(v, w):
	//   v.op == w.op
	//   v.type == w.type
	//   v.aux == w.aux
	//   v.auxint == w.auxint
	//   len(v.args) == len(w.args)
	//   equivalent(v.args[i], w.args[i]) for i in 0..len(v.args)-1

	// The algorithm searches for a partition of f's values into
	// equivalence classes using the above definition.
	// It starts with a coarse partition and iteratively refines it
	// until it reaches a fixed point.

	// Make initial partition based on opcode/type-name/aux/auxint/nargs
	type key struct {
		op     Op
		typ    string
		aux    interface{}
		auxint int64
		nargs  int
	}
	m := map[key]eqclass{}
	for _, b := range f.Blocks {
		for _, v := range b.Values {
			k := key{v.Op, v.Type.String(), v.Aux, v.AuxInt, len(v.Args)}
			m[k] = append(m[k], v)
		}
	}

	// A partition is a set of disjoint eqclasses.
	var partition []eqclass
	for _, v := range m {
		partition = append(partition, v)
	}

	// map from value id back to eqclass id
	valueEqClass := make([]int, f.NumValues())
	for i, e := range partition {
		for _, v := range e {
			valueEqClass[v.ID] = i
		}
	}

	// Find an equivalence class where some members of the class have
	// non-equivalent arguments.  Split the equivalence class appropriately.
	// Repeat until we can't find any more splits.
	for {
		changed := false

		for i, e := range partition {
			v := e[0]
			// all values in this equiv class that are not equivalent to v get moved
			// into another equiv class q.
			var q eqclass
		eqloop:
			for j := 1; j < len(e); {
				w := e[j]
				for i := 0; i < len(v.Args); i++ {
					if valueEqClass[v.Args[i].ID] != valueEqClass[w.Args[i].ID] || !v.Type.Equal(w.Type) {
						// w is not equivalent to v.
						// remove w from e
						e, e[j] = e[:len(e)-1], e[len(e)-1]
						// add w to q
						q = append(q, w)
						valueEqClass[w.ID] = len(partition)
						changed = true
						continue eqloop
					}
				}
				// v and w are equivalent.  Keep w in e.
				j++
			}
			partition[i] = e
			if q != nil {
				partition = append(partition, q)
			}
		}

		if !changed {
			break
		}
	}

	// Compute dominator tree
	idom := dominators(f)

	// Compute substitutions we would like to do.  We substitute v for w
	// if v and w are in the same equivalence class and v dominates w.
	rewrite := make([]*Value, f.NumValues())
	for _, e := range partition {
		sort.Sort(e) // ensure deterministic ordering
		for len(e) > 1 {
			// Find a maximal dominant element in e
			v := e[0]
			for _, w := range e[1:] {
				if dom(w.Block, v.Block, idom) {
					v = w
				}
			}

			// Replace all elements of e which v dominates
			for i := 0; i < len(e); {
				w := e[i]
				if w == v {
					e, e[i] = e[:len(e)-1], e[len(e)-1]
				} else if dom(v.Block, w.Block, idom) {
					rewrite[w.ID] = v
					e, e[i] = e[:len(e)-1], e[len(e)-1]
				} else {
					i++
				}
			}
			// TODO(khr): if value is a control value, do we need to keep it block-local?
		}
	}

	// Apply substitutions
	for _, b := range f.Blocks {
		for _, v := range b.Values {
			for i, w := range v.Args {
				if x := rewrite[w.ID]; x != nil {
					v.SetArg(i, x)
				}
			}
		}
	}
}

// returns true if b dominates c.
// TODO(khr): faster
func dom(b, c *Block, idom []*Block) bool {
	// Walk up from c in the dominator tree looking for b.
	for c != nil {
		if c == b {
			return true
		}
		c = idom[c.ID]
	}
	// Reached the entry block, never saw b.
	return false
}

// An eqclass approximates an equivalence class.  During the
// algorithm it may represent the union of several of the
// final equivalence classes.
type eqclass []*Value

// Sort an equivalence class by value ID.
func (e eqclass) Len() int           { return len(e) }
func (e eqclass) Swap(i, j int)      { e[i], e[j] = e[j], e[i] }
func (e eqclass) Less(i, j int) bool { return e[i].ID < e[j].ID }
<commit_msg>[dev.ssa] cmd/compile: don't combine phi vars from different blocks in CSE<commit_after>// Copyright 2015 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 ssa

import "sort"

// cse does common-subexpression elimination on the Function.
// Values are just relinked, nothing is deleted.  A subsequent deadcode
// pass is required to actually remove duplicate expressions.
func cse(f *Func) {
	// Two values are equivalent if they satisfy the following definition:
	// equivalent(v, w):
	//   v.op == w.op
	//   v.type == w.type
	//   v.aux == w.aux
	//   v.auxint == w.auxint
	//   len(v.args) == len(w.args)
	//   v.block == w.block if v.op == OpPhi
	//   equivalent(v.args[i], w.args[i]) for i in 0..len(v.args)-1

	// The algorithm searches for a partition of f's values into
	// equivalence classes using the above definition.
	// It starts with a coarse partition and iteratively refines it
	// until it reaches a fixed point.

	// Make initial partition based on opcode/type-name/aux/auxint/nargs/phi-block
	type key struct {
		op     Op
		typ    string
		aux    interface{}
		auxint int64
		nargs  int
		block  ID // block id for phi vars, -1 otherwise
	}
	m := map[key]eqclass{}
	for _, b := range f.Blocks {
		for _, v := range b.Values {
			bid := ID(-1)
			if v.Op == OpPhi {
				bid = b.ID
			}
			k := key{v.Op, v.Type.String(), v.Aux, v.AuxInt, len(v.Args), bid}
			m[k] = append(m[k], v)
		}
	}

	// A partition is a set of disjoint eqclasses.
	var partition []eqclass
	for _, v := range m {
		partition = append(partition, v)
	}
	// TODO: Sort partition here for perfect reproducibility?
	// Sort by what? Partition size?
	// (Could that improve efficiency by discovering splits earlier?)

	// map from value id back to eqclass id
	valueEqClass := make([]int, f.NumValues())
	for i, e := range partition {
		for _, v := range e {
			valueEqClass[v.ID] = i
		}
	}

	// Find an equivalence class where some members of the class have
	// non-equivalent arguments.  Split the equivalence class appropriately.
	// Repeat until we can't find any more splits.
	for {
		changed := false

		for i, e := range partition {
			v := e[0]
			// all values in this equiv class that are not equivalent to v get moved
			// into another equiv class q.
			var q eqclass
		eqloop:
			for j := 1; j < len(e); {
				w := e[j]
				for i := 0; i < len(v.Args); i++ {
					if valueEqClass[v.Args[i].ID] != valueEqClass[w.Args[i].ID] || !v.Type.Equal(w.Type) {
						// w is not equivalent to v.
						// remove w from e
						e, e[j] = e[:len(e)-1], e[len(e)-1]
						// add w to q
						q = append(q, w)
						valueEqClass[w.ID] = len(partition)
						changed = true
						continue eqloop
					}
				}
				// v and w are equivalent.  Keep w in e.
				j++
			}
			partition[i] = e
			if q != nil {
				partition = append(partition, q)
			}
		}

		if !changed {
			break
		}
	}

	// Compute dominator tree
	idom := dominators(f)

	// Compute substitutions we would like to do.  We substitute v for w
	// if v and w are in the same equivalence class and v dominates w.
	rewrite := make([]*Value, f.NumValues())
	for _, e := range partition {
		sort.Sort(e) // ensure deterministic ordering
		for len(e) > 1 {
			// Find a maximal dominant element in e
			v := e[0]
			for _, w := range e[1:] {
				if dom(w.Block, v.Block, idom) {
					v = w
				}
			}

			// Replace all elements of e which v dominates
			for i := 0; i < len(e); {
				w := e[i]
				if w == v {
					e, e[i] = e[:len(e)-1], e[len(e)-1]
				} else if dom(v.Block, w.Block, idom) {
					rewrite[w.ID] = v
					e, e[i] = e[:len(e)-1], e[len(e)-1]
				} else {
					i++
				}
			}
			// TODO(khr): if value is a control value, do we need to keep it block-local?
		}
	}

	// Apply substitutions
	for _, b := range f.Blocks {
		for _, v := range b.Values {
			for i, w := range v.Args {
				if x := rewrite[w.ID]; x != nil {
					v.SetArg(i, x)
				}
			}
		}
	}
}

// returns true if b dominates c.
// TODO(khr): faster
func dom(b, c *Block, idom []*Block) bool {
	// Walk up from c in the dominator tree looking for b.
	for c != nil {
		if c == b {
			return true
		}
		c = idom[c.ID]
	}
	// Reached the entry block, never saw b.
	return false
}

// An eqclass approximates an equivalence class.  During the
// algorithm it may represent the union of several of the
// final equivalence classes.
type eqclass []*Value

// Sort an equivalence class by value ID.
func (e eqclass) Len() int           { return len(e) }
func (e eqclass) Swap(i, j int)      { e[i], e[j] = e[j], e[i] }
func (e eqclass) Less(i, j int) bool { return e[i].ID < e[j].ID }
<|endoftext|>


<commit_before>package container

import (
	"math/rand"
	"testing"
)

func TestBinaryTree(t *testing.T) {
	const count = 100
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count; j++ {
			k := (<-ch).(int)
			if k != j {
				t.Errorf("%d != %d", k, j)
			}
		}
	}
}

func TestBinaryTreeFind(t *testing.T) {
	const count = 100
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		list = rand.Perm(count)
		for _, j := range list {
			if _, _, n := tree.Find(j); n == nil {
				t.Errorf("Should have found %d, but didn't", j)
			}
		}

	}
}

func TestBinaryTreeDelete(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		for j := 0; j < sub; j++ {
			tree.Delete(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count-sub; j++ {
			k := (<-ch).(int)
			if k != (j + sub) {
				t.Errorf("%d != %d", k, j+sub)
			}
		}
	}
}

func TestBinaryTreeDelete2(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		list = rand.Perm(sub)
		for _, j := range list {
			tree.Delete(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count-sub; j++ {
			k := (<-ch).(int)
			if k != (j + sub) {
				t.Errorf("%d != %d", k, j+sub)
			}
		}
	}
}

func TestBinaryTreeAddDelete(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		for k := 0; k < 10; k++ {
			list = rand.Perm(sub)
			a := rand.Intn(count - sub)
			for _, j := range list {
				tree.Delete(a + j)
			}
			list = rand.Perm(sub)
			for _, j := range list {
				tree.Add(a + j)
			}
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count; j++ {
			k := (<-ch).(int)
			if k != j {
				t.Errorf("%d != %d", k, j)
			}
		}
	}
}
<commit_msg>And some more error checks in TestBinaryTreeFind just to be sure.<commit_after>package container

import (
	"math/rand"
	"testing"
)

func TestBinaryTree(t *testing.T) {
	const count = 100
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count; j++ {
			k := (<-ch).(int)
			if k != j {
				t.Errorf("%d != %d", k, j)
			}
		}
	}
}

func TestBinaryTreeFind(t *testing.T) {
	const count = 100
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		list = rand.Perm(count)
		for _, j := range list {
			if _, _, n := tree.Find(j); n == nil {
				t.Errorf("Should have found %d, but didn't", j)
			} else if v, ok := n.Data.(int); !ok {
				t.Errorf("Unable to cast data to int... %+v", n.Data)
			} else if v != j {
				t.Errorf("Expected to find %d, but got %d", j, v)
			}
		}
	}
}

func TestBinaryTreeDelete(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		for j := 0; j < sub; j++ {
			tree.Delete(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count-sub; j++ {
			k := (<-ch).(int)
			if k != (j + sub) {
				t.Errorf("%d != %d", k, j+sub)
			}
		}
	}
}

func TestBinaryTreeDelete2(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		list = rand.Perm(sub)
		for _, j := range list {
			tree.Delete(j)
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count-sub; j++ {
			k := (<-ch).(int)
			if k != (j + sub) {
				t.Errorf("%d != %d", k, j+sub)
			}
		}
	}
}

func TestBinaryTreeAddDelete(t *testing.T) {
	const count = 100
	const sub = 20
	for i := 0; i < 10; i++ {
		tree := Tree{Compare: func(a, b interface{}) ComparisonResult {
			aa := a.(int)
			bb := b.(int)
			switch {
			case aa < bb:
				return Less
			case aa > bb:
				return Greater
			default:
				return Equal
			}
		}}
		list := rand.Perm(count)
		for _, j := range list {
			tree.Add(j)
		}
		for k := 0; k < 10; k++ {
			list = rand.Perm(sub)
			a := rand.Intn(count - sub)
			for _, j := range list {
				tree.Delete(a + j)
			}
			list = rand.Perm(sub)
			for _, j := range list {
				tree.Add(a + j)
			}
		}

		ch := make(chan interface{})
		go func() {
			tree.Root.Walk(ch)
			close(ch)
		}()
		for j := 0; j < count; j++ {
			k := (<-ch).(int)
			if k != j {
				t.Errorf("%d != %d", k, j)
			}
		}
	}
}
<|endoftext|>