<commit_before>package engine

import (
	"bytes"
	"fmt"

	"github.com/agl/ed25519"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
	"golang.org/x/crypto/nacl/box"
)

// SelfProof = true: runs the detkey engine and uses the eddsa key as
// the signing key.  This is currently only used for testing to
// generate a fake users who only has a detkey, but perhaps it
// will be useful for something else...
type DetKeyArgs struct {
	PPStream    *libkb.PassphraseStream
	SelfProof   bool
	Me          *libkb.User
	SigningKey  libkb.GenericKey
	EldestKeyID keybase1.KID
	SkipPush    bool
}

type DetKeyEngine struct {
	arg         *DetKeyArgs
	newEddsaKey libkb.GenericKey
	dhKey       libkb.GenericKey
	dev         *libkb.Device
	libkb.Contextified
}

func NewDetKeyEngine(arg *DetKeyArgs, g *libkb.GlobalContext) *DetKeyEngine {
	return &DetKeyEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (d *DetKeyEngine) Name() string {
	return "DetKey"
}

func (d *DetKeyEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (d *DetKeyEngine) SubConsumers() []libkb.UIConsumer {
	return nil
}

func (d *DetKeyEngine) Prereqs() Prereqs { return Prereqs{} }

func (d *DetKeyEngine) SigKey() libkb.GenericKey {
	return d.newEddsaKey
}

func (d *DetKeyEngine) EncKey() libkb.GenericKey {
	return d.dhKey
}

// Run runs the detkey engine.
func (d *DetKeyEngine) Run(ctx *Context) error {
	d.dev = libkb.NewWebDevice()

	if err := d.eddsa(ctx, d.arg.PPStream); err != nil {
		return fmt.Errorf("eddsa error: %s", err)
	}

	// turn off self proof
	d.arg.SelfProof = false

	if err := d.dh(ctx, d.arg.PPStream.DHSeed()); err != nil {
		return fmt.Errorf("dh error: %s", err)
	}
	return nil
}

func (d *DetKeyEngine) eddsa(ctx *Context, tpk *libkb.PassphraseStream) error {
	serverHalf, err := libkb.RandBytes(len(tpk.EdDSASeed()))
	if err != nil {
		return err
	}
	key, err := GenSigningDetKey(tpk, serverHalf)
	if err != nil {
		return err
	}

	var signingKey libkb.GenericKey

	if !d.arg.SelfProof {
		signingKey = d.arg.SigningKey
	}
	d.newEddsaKey = key

	return d.push(ctx, newPusher(key, signingKey, serverHalf).EdDSA())
}

func GenSigningDetKey(tpk *libkb.PassphraseStream, serverHalf []byte) (gkey libkb.GenericKey, err error) {
	xseed, err := serverSeed(tpk.EdDSASeed(), serverHalf)
	if err != nil {
		return nil, err
	}
	pub, priv, err := ed25519.GenerateKey(bytes.NewBuffer(xseed))
	if err != nil {
		return nil, err
	}

	var key libkb.NaclSigningKeyPair
	copy(key.Public[:], (*pub)[:])
	key.Private = &libkb.NaclSigningKeyPrivate{}
	copy(key.Private[:], (*priv)[:])

	return key, nil
}

func (d *DetKeyEngine) dh(ctx *Context, seed []byte) error {
	serverHalf, err := libkb.RandBytes(len(seed))
	if err != nil {
		return err
	}
	xseed, err := serverSeed(seed, serverHalf)
	if err != nil {
		return err
	}
	pub, priv, err := box.GenerateKey(bytes.NewBuffer(xseed))
	if err != nil {
		return err
	}

	var key libkb.NaclDHKeyPair
	copy(key.Public[:], (*pub)[:])
	key.Private = &libkb.NaclDHKeyPrivate{}
	copy(key.Private[:], (*priv)[:])

	d.dhKey = key

	return d.push(ctx, newPusher(key, d.newEddsaKey, serverHalf).DH())
}

func (d *DetKeyEngine) push(ctx *Context, p *pusher) error {
	if d.arg.SkipPush {
		return nil
	}
	return p.push(ctx, d.arg.Me, d.dev)
}

type pusher struct {
	key        libkb.GenericKey
	signing    libkb.GenericKey
	serverHalf []byte
	expire     int
	sibkey     bool
}

func newPusher(key, signing libkb.GenericKey, serverHalf []byte) *pusher {
	return &pusher{
		key:        key,
		signing:    signing,
		serverHalf: serverHalf,
	}
}

func (p *pusher) EdDSA() *pusher {
	p.expire = libkb.NaclEdDSAExpireIn
	p.sibkey = true
	return p
}

func (p *pusher) DH() *pusher {
	p.expire = libkb.NaclDHExpireIn
	p.sibkey = false
	return p
}

func (p *pusher) push(ctx *Context, me *libkb.User, device *libkb.Device) error {
	if device == nil {
		return libkb.ErrCannotGenerateDevice
	}

	g := libkb.Delegator{
		NewKey:      p.key,
		Sibkey:      p.sibkey,
		Expire:      p.expire,
		ExistingKey: p.signing,
		ServerHalf:  p.serverHalf,
		Me:          me,
		Device:      device,
	}

	return g.Run(ctx.LoginContext)
}

func serverSeed(seed, serverHalf []byte) (newseed []byte, err error) {
	newseed = make([]byte, len(seed))
	libkb.XORBytes(newseed, seed, serverHalf)
	return newseed, nil
}
<commit_msg>multi support for detkey flow<commit_after>package engine

import (
	"bytes"
	"fmt"

	"github.com/agl/ed25519"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
	"golang.org/x/crypto/nacl/box"
)

// SelfProof = true: runs the detkey engine and uses the eddsa key as
// the signing key.  This is currently only used for testing to
// generate a fake users who only has a detkey, but perhaps it
// will be useful for something else...
type DetKeyArgs struct {
	PPStream    *libkb.PassphraseStream
	SelfProof   bool
	Me          *libkb.User
	SigningKey  libkb.GenericKey
	EldestKeyID keybase1.KID
	SkipPush    bool
}

type DetKeyEngine struct {
	arg         *DetKeyArgs
	newEddsaKey libkb.GenericKey
	dhKey       libkb.GenericKey
	dev         *libkb.Device
	libkb.Contextified
}

func NewDetKeyEngine(arg *DetKeyArgs, g *libkb.GlobalContext) *DetKeyEngine {
	return &DetKeyEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (d *DetKeyEngine) Name() string {
	return "DetKey"
}

func (d *DetKeyEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (d *DetKeyEngine) SubConsumers() []libkb.UIConsumer {
	return nil
}

func (d *DetKeyEngine) Prereqs() Prereqs { return Prereqs{} }

func (d *DetKeyEngine) SigKey() libkb.GenericKey {
	return d.newEddsaKey
}

func (d *DetKeyEngine) EncKey() libkb.GenericKey {
	return d.dhKey
}

// Run runs the detkey engine.
func (d *DetKeyEngine) Run(ctx *Context) (err error) {
	d.dev = libkb.NewWebDevice()

	delegators := []libkb.Delegator{}

	var delegator libkb.Delegator
	if delegator, err = d.eddsa(ctx, d.arg.PPStream); err != nil {
		err = fmt.Errorf("eddsa error: %s", err)
		return
	}

	delegators = append(delegators, delegator)

	// turn off self proof
	d.arg.SelfProof = false

	if delegator, err = d.dh(ctx, d.arg.PPStream.DHSeed()); err != nil {
		err = fmt.Errorf("dh error: %s", err)
		return
	}

	delegators = append(delegators, delegator)

	err = libkb.DelegatorAggregator(ctx.LoginContext, delegators)
	return
}

func (d *DetKeyEngine) eddsa(ctx *Context, tpk *libkb.PassphraseStream) (delegator libkb.Delegator, err error) {
	var serverHalf []byte
	serverHalf, err = libkb.RandBytes(len(tpk.EdDSASeed()))
	if err != nil {
		return
	}

	var key libkb.GenericKey
	key, err = GenSigningDetKey(tpk, serverHalf)
	if err != nil {
		return
	}

	var signingKey libkb.GenericKey

	if !d.arg.SelfProof {
		signingKey = d.arg.SigningKey
	}
	d.newEddsaKey = key

	delegator, err = d.push(ctx, newPusher(key, signingKey, serverHalf).EdDSA())
	return
}

func GenSigningDetKey(tpk *libkb.PassphraseStream, serverHalf []byte) (gkey libkb.GenericKey, err error) {
	xseed, err := serverSeed(tpk.EdDSASeed(), serverHalf)
	if err != nil {
		return nil, err
	}
	pub, priv, err := ed25519.GenerateKey(bytes.NewBuffer(xseed))
	if err != nil {
		return nil, err
	}

	var key libkb.NaclSigningKeyPair
	copy(key.Public[:], (*pub)[:])
	key.Private = &libkb.NaclSigningKeyPrivate{}
	copy(key.Private[:], (*priv)[:])

	return key, nil
}

func (d *DetKeyEngine) dh(ctx *Context, seed []byte) (delegator libkb.Delegator, err error) {
	var serverHalf []byte
	serverHalf, err = libkb.RandBytes(len(seed))
	if err != nil {
		return
	}

	var xseed []byte
	xseed, err = serverSeed(seed, serverHalf)
	if err != nil {
		return
	}

	var pub *[32]byte
	var priv *[32]byte
	pub, priv, err = box.GenerateKey(bytes.NewBuffer(xseed))
	if err != nil {
		return
	}

	var key libkb.NaclDHKeyPair
	copy(key.Public[:], (*pub)[:])
	key.Private = &libkb.NaclDHKeyPrivate{}
	copy(key.Private[:], (*priv)[:])

	d.dhKey = key

	delegator, err = d.push(ctx, newPusher(key, d.newEddsaKey, serverHalf).DH())
	return
}

func (d *DetKeyEngine) push(ctx *Context, p *pusher) (delegator libkb.Delegator, err error) {
	if d.arg.SkipPush {
		return
	}

	delegator, err = p.push(ctx, d.arg.Me, d.dev)
	return
}

type pusher struct {
	key        libkb.GenericKey
	signing    libkb.GenericKey
	serverHalf []byte
	expire     int
	sibkey     bool
}

func newPusher(key, signing libkb.GenericKey, serverHalf []byte) *pusher {
	return &pusher{
		key:        key,
		signing:    signing,
		serverHalf: serverHalf,
	}
}

func (p *pusher) EdDSA() *pusher {
	p.expire = libkb.NaclEdDSAExpireIn
	p.sibkey = true
	return p
}

func (p *pusher) DH() *pusher {
	p.expire = libkb.NaclDHExpireIn
	p.sibkey = false
	return p
}

func (p *pusher) push(ctx *Context, me *libkb.User, device *libkb.Device) (delegator libkb.Delegator, err error) {
	if device == nil {
		err = libkb.ErrCannotGenerateDevice
		return
	}

	delegator = libkb.Delegator{
		NewKey:      p.key,
		Sibkey:      p.sibkey,
		Expire:      p.expire,
		ExistingKey: p.signing,
		ServerHalf:  p.serverHalf,
		Me:          me,
		Device:      device,
	}

	return
}

func serverSeed(seed, serverHalf []byte) (newseed []byte, err error) {
	newseed = make([]byte, len(seed))
	libkb.XORBytes(newseed, seed, serverHalf)
	return newseed, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"crypto/rc4"
	"strconv"
	"net"
	// "os"
)

//---------------------------------------------------------------------
// config
//---------------------------------------------------------------------
type Config struct {
	run_mode string
	server_host string
	client_host string
	server_port int
	client_port int
	key string
	method string
}

var config = Config {}

//---------------------------------------------------------------------
// Protocol
//---------------------------------------------------------------------
type Protocol struct {
	crypt_send *rc4.Cipher
	crypt_recv *rc4.Cipher
}

//---------------------------------------------------------------------
// arguments_init()
//---------------------------------------------------------------------
func arguments_init() bool {
	server := flag.Bool("server", false, "Running mode")
	server_host := flag.String("s", "", "Server address")
	client_host := flag.String("b", "", "Local address")
	server_port := flag.Int("p", 1030, "Server port")
	client_port := flag.Int("l", 1030, "Client port")
	method := flag.String("m", "rc4", "Encryption method")
	key := flag.String("k", "", "Key string")
	flag.Parse()
	config.run_mode = "client"
	if *server {
		config.run_mode = "server"
	}
	config.server_host = *server_host
	config.client_host = *client_host
	config.server_port = *server_port
	config.client_port = *client_port
	config.method = *method
	config.key = *key
	if config.run_mode == "server" {
		if config.server_host == "" {
			config.server_host = "0.0.0.0"
		}
	}	else {
		if config.server_host == "" {
			log.Print("ERROR: empty server address, use -h to help")
			return false
		}
		if config.client_host == "" {
			config.client_host = "localhost"
		}
	}
	log.Print("[config] running mode: ", config.run_mode)
	log.Print("[config] server address: ", config.server_host)
	log.Print("[config] server port: ", config.server_port)
	log.Print("[config] client address: ", config.client_host)
	log.Print("[config] client port: ", config.client_port)
	return true
}

//---------------------------------------------------------------------
// handle error
//---------------------------------------------------------------------
func handle_error(err error) bool {
	if err == nil {
		return true
	}
	log.Printf("ERROR: %s", err.Error());
	return false
}

//---------------------------------------------------------------------
// main()
//---------------------------------------------------------------------
func main() {
	if !arguments_init() {
		return
	}

	endpoint := ""

	if config.run_mode == "server" {
		endpoint = config.server_host + ":" + strconv.Itoa(config.server_port)
	}	else {
		endpoint = config.client_host + ":" + strconv.Itoa(config.client_port)
	}

	addr, err := net.ResolveTCPAddr("tcp4", endpoint)
	if !handle_error(err) {
		return
	}

	listener, err := net.ListenTCP("tcp4", addr)
	if !handle_error(err) {
		return
	}

	log.Printf("%s is listening on %s", config.run_mode, endpoint)

	for {
		conn, err := listener.AcceptTCP()
		if err == nil {
			key := []byte(config.key)
			protocol := Protocol {}
			protocol.crypt_recv, _ = rc4.NewCipher(key)
			protocol.crypt_send, _ = rc4.NewCipher(key)
			if config.run_mode == "server" {
			}	else {
			}
		}	else {
			handle_error(err)
		}
	}
}

<commit_msg>update ragesocks<commit_after>package main

import (
	"flag"
	"log"
	"crypto/rc4"
	"strconv"
	"net"
	"io"
	// "os"
)

//---------------------------------------------------------------------
// config
//---------------------------------------------------------------------
type Config struct {
	run_mode string
	server_host string
	client_host string
	server_port int
	client_port int
	key string
	method string
}

var config = Config {}

//---------------------------------------------------------------------
// Protocol
//---------------------------------------------------------------------
type Protocol struct {
	crypt_send *rc4.Cipher
	crypt_recv *rc4.Cipher
	twice_send *rc4.Cipher
	twice_recv *rc4.Cipher
	conn *net.TCPConn
	esock *net.TCPConn
	remote_address string
}

//---------------------------------------------------------------------
// arguments_init()
//---------------------------------------------------------------------
func arguments_init() bool {
	server := flag.Bool("server", false, "Running mode")
	server_host := flag.String("s", "", "Server address")
	client_host := flag.String("b", "", "Local address")
	server_port := flag.Int("p", 1030, "Server port")
	client_port := flag.Int("l", 1030, "Client port")
	method := flag.String("m", "rc4", "Encryption method")
	key := flag.String("k", "", "Key string")
	flag.Parse()
	config.run_mode = "client"
	if *server {
		config.run_mode = "server"
	}
	config.server_host = *server_host
	config.client_host = *client_host
	config.server_port = *server_port
	config.client_port = *client_port
	config.method = *method
	config.key = *key
	if config.run_mode == "server" {
		if config.server_host == "" {
			config.server_host = "0.0.0.0"
		}
	}	else {
		if config.server_host == "" {
			log.Print("ERROR: empty server address, use -h to help")
			return false
		}
		if config.client_host == "" {
			config.client_host = "localhost"
		}
	}
	log.Print("[config] running mode: ", config.run_mode)
	log.Print("[config] server address: ", config.server_host)
	log.Print("[config] server port: ", config.server_port)
	log.Print("[config] client address: ", config.client_host)
	log.Print("[config] client port: ", config.client_port)
	return true
}

//---------------------------------------------------------------------
// handle error
//---------------------------------------------------------------------
func handle_error(err error) bool {
	if err == nil {
		return true
	}
	log.Printf("ERROR: %s", err.Error());
	return false
}

//---------------------------------------------------------------------
// main()
//---------------------------------------------------------------------
func main() {
	if !arguments_init() {
		return
	}

	endpoint := ""

	if config.run_mode == "server" {
		endpoint = config.server_host + ":" + strconv.Itoa(config.server_port)
	}	else {
		endpoint = config.client_host + ":" + strconv.Itoa(config.client_port)
	}

	addr, err := net.ResolveTCPAddr("tcp4", endpoint)
	if !handle_error(err) {
		return
	}

	listener, err := net.ListenTCP("tcp4", addr)
	if !handle_error(err) {
		return
	}

	log.Printf("%s is listening on %s", config.run_mode, endpoint)

	for {
		conn, err := listener.AcceptTCP()
		if err == nil {
			key := []byte(config.key)
			protocol := Protocol {}
			protocol.crypt_recv, _ = rc4.NewCipher(key)
			protocol.crypt_send, _ = rc4.NewCipher(key)
			protocol.twice_recv = nil
			protocol.twice_send = nil
			protocol.conn = conn
			protocol.esock = nil
			protocol.remote_address = ""
			if config.run_mode == "server" {
				protocol.esock = conn
				go func (protocol *Protocol) {
					defer protocol.conn.Close()
					handle_server(protocol)
				}(&protocol)
			}	else {
				go func (protocol *Protocol) {
					defer protocol.conn.Close()
					handle_client(protocol)
				}(&protocol)
			}
		}	else {
			handle_error(err)
		}
	}
}

//---------------------------------------------------------------------
// encryption send
//---------------------------------------------------------------------
func encrypt_send(protocol *Protocol, buf []byte) (int, error) {
	if (protocol.crypt_send != nil) {
		protocol.crypt_send.XORKeyStream(buf, buf)
	}
	if (protocol.twice_send != nil) {
		protocol.twice_send.XORKeyStream(buf, buf)
	}
	return protocol.esock.Write(buf)
}

//---------------------------------------------------------------------
// encryption recv
//---------------------------------------------------------------------
func encrypt_recv(protocol *Protocol, buf []byte) (int, error) {
	n, err := protocol.esock.Read(buf)
	if err != nil {
		return n, err
	}
	if (protocol.twice_recv != nil) {
		protocol.twice_recv.XORKeyStream(buf[:n], buf[:n])
	}
	if (protocol.crypt_recv != nil) {
		protocol.crypt_recv.XORKeyStream(buf[:n], buf[:n])
	}
	return n, nil
}

//---------------------------------------------------------------------
// encryption recv all
//---------------------------------------------------------------------
func encrypt_recv_all(protocol *Protocol, buf []byte) (int, error) {
	n, err := io.ReadFull(protocol.esock, buf)
	if err != nil {
		return n, err
	}
	if (protocol.twice_recv != nil) {
		protocol.twice_recv.XORKeyStream(buf, buf)
	}
	if (protocol.crypt_recv != nil) {
		protocol.crypt_recv.XORKeyStream(buf, buf)
	}
	return n, err
}

//---------------------------------------------------------------------
// server entry
//---------------------------------------------------------------------
func handle_server(protocol *Protocol) {
}

//---------------------------------------------------------------------
// client entry
//---------------------------------------------------------------------
func handle_client(protocol *Protocol) {
}

<|endoftext|>


<commit_before>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012-2014  Casey Marshall

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, version 3.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package openpgp

import (
	"bytes"
	"crypto/md5"
	"io"
	"io/ioutil"
	"sort"
	stdtesting "testing"

	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	gc "gopkg.in/check.v1"

	"github.com/hockeypuck/testing"
)

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

type SamplePacketSuite struct{}

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

func (s *SamplePacketSuite) TestVerifyUserAttributeSig(c *gc.C) {
	key := MustInputAscKey("uat.asc")
	c.Assert(key.UserAttributes, gc.HasLen, 1)
	DropDuplicates(key)
	c.Assert(key.UserAttributes, gc.HasLen, 1)
	uat := key.UserAttributes[0]
	c.Assert(uat.Images, gc.HasLen, 1)
	// TODO: check contents
}

func (s *SamplePacketSuite) TestSksDigest(c *gc.C) {
	key := MustInputAscKey("sksdigest.asc")
	md5, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	c.Assert(key.ShortID(), gc.Equals, "ce353cf4")
	c.Assert(md5, gc.Equals, "da84f40d830a7be2a3c0b7f2e146bfaa")
}

func (s *SamplePacketSuite) TestSksContextualDup(c *gc.C) {
	f := testing.MustInput("sks_fail.asc")

	block, err := armor.Decode(f)
	c.Assert(err, gc.IsNil)
	buf, err := ioutil.ReadAll(block.Body)
	c.Assert(err, gc.IsNil)
	err = f.Close()
	c.Assert(err, gc.IsNil)

	var kr *OpaqueKeyring
	for opkr := range ReadOpaqueKeyrings(bytes.NewBuffer(buf)) {
		c.Assert(kr, gc.IsNil)
		kr = opkr
	}

	var refBuf bytes.Buffer
	for _, op := range kr.Packets {
		err = op.Serialize(&refBuf)
		c.Assert(err, gc.IsNil)
	}
	c.Assert(buf, gc.DeepEquals, refBuf.Bytes())

	pk, err := kr.Parse()
	c.Assert(err, gc.IsNil)
	digest1, err := SksDigest(pk, md5.New())
	c.Assert(err, gc.IsNil)

	err = DropDuplicates(pk)
	c.Assert(err, gc.IsNil)
	digest2, err := SksDigest(pk, md5.New())
	c.Assert(err, gc.IsNil)

	c.Check(digest1, gc.Equals, digest2)

	//sort.Sort(opaquePacketSlice(kr.Packets))
	for _, op := range kr.Packets {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Log("parse primary key")
	key := MustInputAscKey("sks_fail2.asc")
	dupDigest, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	var packetsDup opaquePacketSlice
	for _, node := range key.contents() {
		op, err := node.packet().opaquePacket()
		c.Assert(err, gc.IsNil)
		packetsDup = append(packetsDup, op)
	}
	sort.Sort(packetsDup)
	for _, op := range packetsDup {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Log("deduped primary key")
	key = MustInputAscKey("sks_fail2.asc")
	DropDuplicates(key)
	dedupDigest, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	var packetsDedup opaquePacketSlice
	for _, node := range key.contents() {
		op, err := node.packet().opaquePacket()
		c.Assert(err, gc.IsNil)
		packetsDedup = append(packetsDedup, op)
	}
	sort.Sort(packetsDedup)
	for _, op := range packetsDedup {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Assert(dupDigest, gc.Equals, dedupDigest)
}

func (s *SamplePacketSuite) TestUatRtt(c *gc.C) {
	f := testing.MustInput("uat.asc")
	defer f.Close()
	block, err := armor.Decode(f)
	c.Assert(err, gc.IsNil)
	var p packet.Packet
	for {
		p, err = packet.Read(block.Body)
		if err != nil {
			c.Assert(err, gc.Equals, io.EOF)
			break
		}

		uat, ok := p.(*packet.UserAttribute)
		if ok {
			var buf bytes.Buffer
			uat.Serialize(&buf)
			or := packet.NewOpaqueReader(bytes.NewBuffer(buf.Bytes()))
			op, _ := or.Next()
			c.Assert(buf.Bytes()[3:], gc.DeepEquals, op.Contents)
		}
	}
}

func (s *SamplePacketSuite) TestPacketCounts(c *gc.C) {
	testCases := []struct {
		name                                         string
		nUserID, nUserAttribute, nSubKey, nSignature int
	}{{
		"0ff16c87.asc", 9, 0, 1, 0,
	}, {
		"alice_signed.asc", 1, 0, 1, 0,
	}, {
		"uat.asc", 2, 1, 3, 0,
	}, {
		"252B8B37.dupsig.asc", 3, 0, 2, 1,
	}}
	for i, testCase := range testCases {
		c.Logf("test#%d: %s", i, testCase.name)
		f := testing.MustInput(testCase.name)
		defer f.Close()
		block, err := armor.Decode(f)
		c.Assert(err, gc.IsNil)
		var key *PrimaryKey
		for keyRead := range ReadKeys(block.Body) {
			key = keyRead.PrimaryKey
		}
		c.Assert(key, gc.NotNil)
		c.Assert(key.UserIDs, gc.HasLen, testCase.nUserID)
		c.Assert(key.UserAttributes, gc.HasLen, testCase.nUserAttribute)
		c.Assert(key.SubKeys, gc.HasLen, testCase.nSubKey)
		c.Assert(key.Signatures, gc.HasLen, testCase.nSignature)
	}
}

func (s *SamplePacketSuite) TestDeduplicate(c *gc.C) {
	f := testing.MustInput("d7346e26.asc")
	defer f.Close()
	block, err := armor.Decode(f)
	if err != nil {
		c.Fatal(err)
	}

	// Parse keyring, duplicate all packet types except primary pubkey.
	kr := &OpaqueKeyring{}
	for opkr := range ReadOpaqueKeyrings(block.Body) {
		c.Assert(opkr.Error, gc.IsNil)
		for _, op := range opkr.Packets {
			kr.Packets = append(kr.Packets, op)
			switch op.Tag {
			case 2:
				kr.Packets = append(kr.Packets, op)
				fallthrough
			case 13, 14, 17:
				kr.Packets = append(kr.Packets, op)
			}
		}
	}
	key, err := kr.Parse()
	c.Assert(err, gc.IsNil)

	n := 0
	for _, node := range key.contents() {
		c.Logf("%s", node.uuid())
		n++
	}

	c.Log()
	err = CollectDuplicates(key)
	c.Assert(err, gc.IsNil)

	n2 := 0
	for _, node := range key.contents() {
		c.Logf("%s %d", node.uuid(), node.packet().Count)
		n2++
		switch node.packet().Tag {
		case 2:
			c.Check(node.packet().Count, gc.Equals, 2)
		case 13, 14, 17:
			c.Check(node.packet().Count, gc.Equals, 1)
		case 6:
			c.Check(node.packet().Count, gc.Equals, 0)
		default:
			c.Fatal("should not happen")
		}
	}
	c.Assert(n2 < n, gc.Equals, true)
}

func (s *SamplePacketSuite) TestMerge(c *gc.C) {
	key1 := MustInputAscKey("lp1195901.asc")
	key2 := MustInputAscKey("lp1195901_2.asc")
	err := Merge(key2, key1)
	c.Assert(err, gc.IsNil)
	var matchUID *UserID
	for _, uid := range key2.UserIDs {
		if uid.Keywords == "Phil Pennock <pdp@spodhuis.org>" {
			matchUID = uid
		}
	}
	c.Assert(matchUID, gc.NotNil)
}
<commit_msg>Fix testdata filename.<commit_after>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012-2014  Casey Marshall

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, version 3.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package openpgp

import (
	"bytes"
	"crypto/md5"
	"io"
	"io/ioutil"
	"sort"
	stdtesting "testing"

	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	gc "gopkg.in/check.v1"

	"github.com/hockeypuck/testing"
)

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

type SamplePacketSuite struct{}

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

func (s *SamplePacketSuite) TestVerifyUserAttributeSig(c *gc.C) {
	key := MustInputAscKey("uat.asc")
	c.Assert(key.UserAttributes, gc.HasLen, 1)
	DropDuplicates(key)
	c.Assert(key.UserAttributes, gc.HasLen, 1)
	uat := key.UserAttributes[0]
	c.Assert(uat.Images, gc.HasLen, 1)
	// TODO: check contents
}

func (s *SamplePacketSuite) TestSksDigest(c *gc.C) {
	key := MustInputAscKey("sksdigest.asc")
	md5, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	c.Assert(key.ShortID(), gc.Equals, "ce353cf4")
	c.Assert(md5, gc.Equals, "da84f40d830a7be2a3c0b7f2e146bfaa")
}

func (s *SamplePacketSuite) TestSksContextualDup(c *gc.C) {
	f := testing.MustInput("sks_fail.asc")

	block, err := armor.Decode(f)
	c.Assert(err, gc.IsNil)
	buf, err := ioutil.ReadAll(block.Body)
	c.Assert(err, gc.IsNil)
	err = f.Close()
	c.Assert(err, gc.IsNil)

	var kr *OpaqueKeyring
	for opkr := range ReadOpaqueKeyrings(bytes.NewBuffer(buf)) {
		c.Assert(kr, gc.IsNil)
		kr = opkr
	}

	var refBuf bytes.Buffer
	for _, op := range kr.Packets {
		err = op.Serialize(&refBuf)
		c.Assert(err, gc.IsNil)
	}
	c.Assert(buf, gc.DeepEquals, refBuf.Bytes())

	pk, err := kr.Parse()
	c.Assert(err, gc.IsNil)
	digest1, err := SksDigest(pk, md5.New())
	c.Assert(err, gc.IsNil)

	err = DropDuplicates(pk)
	c.Assert(err, gc.IsNil)
	digest2, err := SksDigest(pk, md5.New())
	c.Assert(err, gc.IsNil)

	c.Check(digest1, gc.Equals, digest2)

	for _, op := range kr.Packets {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Log("parse primary key")
	key := MustInputAscKey("sks_fail.asc")
	dupDigest, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	var packetsDup opaquePacketSlice
	for _, node := range key.contents() {
		op, err := node.packet().opaquePacket()
		c.Assert(err, gc.IsNil)
		packetsDup = append(packetsDup, op)
	}
	sort.Sort(packetsDup)
	for _, op := range packetsDup {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Log("deduped primary key")
	key = MustInputAscKey("sks_fail.asc")
	DropDuplicates(key)
	dedupDigest, err := SksDigest(key, md5.New())
	c.Assert(err, gc.IsNil)
	var packetsDedup opaquePacketSlice
	for _, node := range key.contents() {
		op, err := node.packet().opaquePacket()
		c.Assert(err, gc.IsNil)
		packetsDedup = append(packetsDedup, op)
	}
	sort.Sort(packetsDedup)
	for _, op := range packetsDedup {
		c.Logf("%d %d %s", op.Tag, len(op.Contents), hexmd5(op.Contents))
	}

	c.Assert(dupDigest, gc.Equals, dedupDigest)
}

func (s *SamplePacketSuite) TestUatRtt(c *gc.C) {
	f := testing.MustInput("uat.asc")
	defer f.Close()
	block, err := armor.Decode(f)
	c.Assert(err, gc.IsNil)
	var p packet.Packet
	for {
		p, err = packet.Read(block.Body)
		if err != nil {
			c.Assert(err, gc.Equals, io.EOF)
			break
		}

		uat, ok := p.(*packet.UserAttribute)
		if ok {
			var buf bytes.Buffer
			uat.Serialize(&buf)
			or := packet.NewOpaqueReader(bytes.NewBuffer(buf.Bytes()))
			op, _ := or.Next()
			c.Assert(buf.Bytes()[3:], gc.DeepEquals, op.Contents)
		}
	}
}

func (s *SamplePacketSuite) TestPacketCounts(c *gc.C) {
	testCases := []struct {
		name                                         string
		nUserID, nUserAttribute, nSubKey, nSignature int
	}{{
		"0ff16c87.asc", 9, 0, 1, 0,
	}, {
		"alice_signed.asc", 1, 0, 1, 0,
	}, {
		"uat.asc", 2, 1, 3, 0,
	}, {
		"252B8B37.dupsig.asc", 3, 0, 2, 1,
	}}
	for i, testCase := range testCases {
		c.Logf("test#%d: %s", i, testCase.name)
		f := testing.MustInput(testCase.name)
		defer f.Close()
		block, err := armor.Decode(f)
		c.Assert(err, gc.IsNil)
		var key *PrimaryKey
		for keyRead := range ReadKeys(block.Body) {
			key = keyRead.PrimaryKey
		}
		c.Assert(key, gc.NotNil)
		c.Assert(key.UserIDs, gc.HasLen, testCase.nUserID)
		c.Assert(key.UserAttributes, gc.HasLen, testCase.nUserAttribute)
		c.Assert(key.SubKeys, gc.HasLen, testCase.nSubKey)
		c.Assert(key.Signatures, gc.HasLen, testCase.nSignature)
	}
}

func (s *SamplePacketSuite) TestDeduplicate(c *gc.C) {
	f := testing.MustInput("d7346e26.asc")
	defer f.Close()
	block, err := armor.Decode(f)
	if err != nil {
		c.Fatal(err)
	}

	// Parse keyring, duplicate all packet types except primary pubkey.
	kr := &OpaqueKeyring{}
	for opkr := range ReadOpaqueKeyrings(block.Body) {
		c.Assert(opkr.Error, gc.IsNil)
		for _, op := range opkr.Packets {
			kr.Packets = append(kr.Packets, op)
			switch op.Tag {
			case 2:
				kr.Packets = append(kr.Packets, op)
				fallthrough
			case 13, 14, 17:
				kr.Packets = append(kr.Packets, op)
			}
		}
	}
	key, err := kr.Parse()
	c.Assert(err, gc.IsNil)

	n := 0
	for _, node := range key.contents() {
		c.Logf("%s", node.uuid())
		n++
	}

	c.Log()
	err = CollectDuplicates(key)
	c.Assert(err, gc.IsNil)

	n2 := 0
	for _, node := range key.contents() {
		c.Logf("%s %d", node.uuid(), node.packet().Count)
		n2++
		switch node.packet().Tag {
		case 2:
			c.Check(node.packet().Count, gc.Equals, 2)
		case 13, 14, 17:
			c.Check(node.packet().Count, gc.Equals, 1)
		case 6:
			c.Check(node.packet().Count, gc.Equals, 0)
		default:
			c.Fatal("should not happen")
		}
	}
	c.Assert(n2 < n, gc.Equals, true)
}

func (s *SamplePacketSuite) TestMerge(c *gc.C) {
	key1 := MustInputAscKey("lp1195901.asc")
	key2 := MustInputAscKey("lp1195901_2.asc")
	err := Merge(key2, key1)
	c.Assert(err, gc.IsNil)
	var matchUID *UserID
	for _, uid := range key2.UserIDs {
		if uid.Keywords == "Phil Pennock <pdp@spodhuis.org>" {
			matchUID = uid
		}
	}
	c.Assert(matchUID, gc.NotNil)
}
<|endoftext|>


<commit_before>package irc

import (
	"bufio"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

const delim byte = '\n'
const endline string = "\r\n"

type Connection struct {
	Network   string
	Nick      string
	User      string
	RealName  string
	Input     chan Message
	Output    chan Message
	Reader    *bufio.Reader
	Writer    *bufio.Writer
	conn      net.Conn
	Reconnect chan struct{}
	Quit      chan struct{}
	QuitSend  chan struct{}
	QuitRecv  chan struct{}
	// QuitDispatcher chan struct{}
	L sync.Mutex
}

func (c *Connection) Sender() {
	log.Println(c.Network, "spawned Sender")
	for {
		select {
		case msg := <-c.Input:
			c.Writer.WriteString(msg.String() + endline)
			log.Println(c.Network, "-->", msg.String())
			c.Writer.Flush()
		case <-c.QuitSend:
			log.Println(c.Network, "closing Sender")
			return
		}
	}
}

func (c *Connection) Receiver() {
	log.Println(c.Network, "spawned Receiver")
	for {
		raw, err := c.Reader.ReadString(delim)
		if err != nil {
			log.Println(c.Network, "error reading message", err.Error())
			log.Println(c.Network, "closing Receiver")
			c.Quit <- struct{}{}
			log.Println(c.Network, "sent quit message from Receiver")
			return
		}
		msg, err := ParseMessage(raw)
		if err != nil {
			log.Println(c.Network, "error decoding message", err.Error())
			log.Println(c.Network, "closing Receiver")
			c.Quit <- struct{}{}
			log.Println(c.Network, "sent quit message from Receiver")
			return
		} else {
			log.Println(c.Network, "<--", msg.String())
		}
		select {
		case c.Output <- *msg:
		case <-c.QuitRecv:
			log.Println(c.Network, "closing Receiver")
			return
		}
	}
}

/*
func (c *Connection) Dispatcher() {
	log.Println(c.Network, "spawned Dispatcher")
	for {
		// just sink everything for now
		select {
		case <-c.Output:
		case <-c.QuitDispatcher:
			log.Println(c.Network, "closing Dispatcher")
			return
		}
	}
}
*/

func (c *Connection) Cleaner() {
	log.Println(c.Network, "spawned Cleaner")
	for {
		<-c.Quit
		log.Println(c.Network, "ceceived quit message")
		c.L.Lock()
		log.Println(c.Network, "cleaning up!")
		c.QuitSend <- struct{}{}
		c.QuitRecv <- struct{}{}
		// c.QuitDispatcher <- struct{}{}
		c.Reconnect <- struct{}{}
		c.conn.Close()
		log.Println(c.Network, "closing Cleaner")
		c.L.Unlock()
	}
}

func (c *Connection) Keeper(servers []string) {
	log.Println(c.Network, "spawned Keeper")
	for {
		<-c.Reconnect
		c.L.Lock()
		if c.Input != nil {
			close(c.Input)
			close(c.Output)
			close(c.QuitSend)
			close(c.QuitRecv)
			// close(c.QuitDispatcher)
		}
		c.Input = make(chan Message, 1)
		c.Output = make(chan Message, 1)
		c.QuitSend = make(chan struct{}, 1)
		c.QuitRecv = make(chan struct{}, 1)
		// c.QuitDispatcher = make(chan struct{}, 1)
		server := servers[rand.Intn(len(servers))]
		log.Println(c.Network, "connecting to", server)
		c.Dial(server)
		c.L.Unlock()

		go c.Sender()
		go c.Receiver()
		// go c.Dispatcher()

		log.Println(c.Network, "Initializing IRC connection")
		c.Input <- Message{
			Command:  "NICK",
			Trailing: c.Nick,
		}
		c.Input <- Message{
			Command:  "USER",
			Params:   []string{c.User, "0", "*"},
			Trailing: c.RealName,
		}

	}
}

func (c *Connection) Setup(network string, servers []string, nick string, user string, realname string) {
	rand.Seed(time.Now().UnixNano())

	c.Reconnect = make(chan struct{}, 1)
	c.Quit = make(chan struct{}, 1)
	c.Nick = nick
	c.User = user
	c.RealName = realname
	c.Network = network

	c.Reconnect <- struct{}{}
	go c.Keeper(servers)
	go c.Cleaner()
	return
}

func (c *Connection) Dial(server string) error {

	conn, err := net.Dial("tcp", server)
	if err != nil {
		log.Println(c.Network, "Cannot connect to", server, "error:", err.Error())
		return err
	}
	log.Println(c.Network, "Connected to", server)
	c.Writer = bufio.NewWriter(conn)
	c.Reader = bufio.NewReader(conn)
	c.conn = conn

	return nil
}
<commit_msg>Remove old dispatcher<commit_after>package irc

import (
	"bufio"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

const delim byte = '\n'
const endline string = "\r\n"

type Connection struct {
	Network   string
	Nick      string
	User      string
	RealName  string
	Input     chan Message
	Output    chan Message
	Reader    *bufio.Reader
	Writer    *bufio.Writer
	conn      net.Conn
	Reconnect chan struct{}
	Quit      chan struct{}
	QuitSend  chan struct{}
	QuitRecv  chan struct{}
	L sync.Mutex
}

func (c *Connection) Sender() {
	log.Println(c.Network, "spawned Sender")
	for {
		select {
		case msg := <-c.Input:
			c.Writer.WriteString(msg.String() + endline)
			log.Println(c.Network, "-->", msg.String())
			c.Writer.Flush()
		case <-c.QuitSend:
			log.Println(c.Network, "closing Sender")
			return
		}
	}
}

func (c *Connection) Receiver() {
	log.Println(c.Network, "spawned Receiver")
	for {
		raw, err := c.Reader.ReadString(delim)
		if err != nil {
			log.Println(c.Network, "error reading message", err.Error())
			log.Println(c.Network, "closing Receiver")
			c.Quit <- struct{}{}
			log.Println(c.Network, "sent quit message from Receiver")
			return
		}
		msg, err := ParseMessage(raw)
		if err != nil {
			log.Println(c.Network, "error decoding message", err.Error())
			log.Println(c.Network, "closing Receiver")
			c.Quit <- struct{}{}
			log.Println(c.Network, "sent quit message from Receiver")
			return
		} else {
			log.Println(c.Network, "<--", msg.String())
		}
		select {
		case c.Output <- *msg:
		case <-c.QuitRecv:
			log.Println(c.Network, "closing Receiver")
			return
		}
	}
}

func (c *Connection) Cleaner() {
	log.Println(c.Network, "spawned Cleaner")
	for {
		<-c.Quit
		log.Println(c.Network, "ceceived quit message")
		c.L.Lock()
		log.Println(c.Network, "cleaning up!")
		c.QuitSend <- struct{}{}
		c.QuitRecv <- struct{}{}
		c.Reconnect <- struct{}{}
		c.conn.Close()
		log.Println(c.Network, "closing Cleaner")
		c.L.Unlock()
	}
}

func (c *Connection) Keeper(servers []string) {
	log.Println(c.Network, "spawned Keeper")
	for {
		<-c.Reconnect
		c.L.Lock()
		if c.Input != nil {
			close(c.Input)
			close(c.Output)
			close(c.QuitSend)
			close(c.QuitRecv)
		}
		c.Input = make(chan Message, 1)
		c.Output = make(chan Message, 1)
		c.QuitSend = make(chan struct{}, 1)
		c.QuitRecv = make(chan struct{}, 1)
		server := servers[rand.Intn(len(servers))]
		log.Println(c.Network, "connecting to", server)
		c.Dial(server)
		c.L.Unlock()

		go c.Sender()
		go c.Receiver()

		log.Println(c.Network, "Initializing IRC connection")
		c.Input <- Message{
			Command:  "NICK",
			Trailing: c.Nick,
		}
		c.Input <- Message{
			Command:  "USER",
			Params:   []string{c.User, "0", "*"},
			Trailing: c.RealName,
		}

	}
}

func (c *Connection) Setup(network string, servers []string, nick string, user string, realname string) {
	rand.Seed(time.Now().UnixNano())

	c.Reconnect = make(chan struct{}, 1)
	c.Quit = make(chan struct{}, 1)
	c.Nick = nick
	c.User = user
	c.RealName = realname
	c.Network = network

	c.Reconnect <- struct{}{}
	go c.Keeper(servers)
	go c.Cleaner()
	return
}

func (c *Connection) Dial(server string) error {

	conn, err := net.Dial("tcp", server)
	if err != nil {
		log.Println(c.Network, "Cannot connect to", server, "error:", err.Error())
		return err
	}
	log.Println(c.Network, "Connected to", server)
	c.Writer = bufio.NewWriter(conn)
	c.Reader = bufio.NewReader(conn)
	c.conn = conn

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014-2017 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package device

import (
	"context"
	"flag"
	"strings"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
	"github.com/vmware/govmomi/vim25/types"
)

type boot struct {
	*flags.VirtualMachineFlag

	order string
	types.VirtualMachineBootOptions
}

func init() {
	cli.Register("device.boot", &boot{})
}

func (cmd *boot) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.VirtualMachineFlag, ctx = flags.NewVirtualMachineFlag(ctx)
	cmd.VirtualMachineFlag.Register(ctx, f)

	f.Int64Var(&cmd.BootDelay, "delay", 0, "Delay in ms before starting the boot sequence")
	f.StringVar(&cmd.order, "order", "", "Boot device order [-,floppy,cdrom,ethernet,disk]")
	f.Int64Var(&cmd.BootRetryDelay, "retry-delay", 0, "Delay in ms before a boot retry")

	cmd.BootRetryEnabled = types.NewBool(false)
	f.BoolVar(cmd.BootRetryEnabled, "retry", false, "If true, retry boot after retry-delay")

	cmd.EnterBIOSSetup = types.NewBool(false)
	f.BoolVar(cmd.EnterBIOSSetup, "setup", false, "If true, enter BIOS setup on next boot")
}

func (cmd *boot) Description() string {
	return `Configure VM boot settings.

Examples:
  govc device.boot -vm $vm -delay 1000 -order floppy,cdrom,ethernet,disk
  govc device.boot -vm $vm -order - # reset boot order`
}

func (cmd *boot) Process(ctx context.Context) error {
	if err := cmd.VirtualMachineFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *boot) Run(ctx context.Context, f *flag.FlagSet) error {
	vm, err := cmd.VirtualMachine()
	if err != nil {
		return err
	}

	if vm == nil {
		return flag.ErrHelp
	}

	devices, err := vm.Device(ctx)
	if err != nil {
		return err
	}

	if cmd.order != "" {
		o := strings.Split(cmd.order, ",")
		cmd.BootOrder = devices.BootOrder(o)
	}

	return vm.SetBootOptions(ctx, &cmd.VirtualMachineBootOptions)
}
<commit_msg>govc: add device.boot -secure flag<commit_after>/*
Copyright (c) 2014-2017 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package device

import (
	"context"
	"flag"
	"strings"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
	"github.com/vmware/govmomi/vim25/types"
)

type boot struct {
	*flags.VirtualMachineFlag

	order string
	types.VirtualMachineBootOptions
}

func init() {
	cli.Register("device.boot", &boot{})
}

func (cmd *boot) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.VirtualMachineFlag, ctx = flags.NewVirtualMachineFlag(ctx)
	cmd.VirtualMachineFlag.Register(ctx, f)

	f.Int64Var(&cmd.BootDelay, "delay", 0, "Delay in ms before starting the boot sequence")
	f.StringVar(&cmd.order, "order", "", "Boot device order [-,floppy,cdrom,ethernet,disk]")
	f.Int64Var(&cmd.BootRetryDelay, "retry-delay", 0, "Delay in ms before a boot retry")

	cmd.BootRetryEnabled = types.NewBool(false)
	f.BoolVar(cmd.BootRetryEnabled, "retry", false, "If true, retry boot after retry-delay")

	cmd.EnterBIOSSetup = types.NewBool(false)
	f.BoolVar(cmd.EnterBIOSSetup, "setup", false, "If true, enter BIOS setup on next boot")

	f.Var(flags.NewOptionalBool(&cmd.EfiSecureBootEnabled), "secure", "Enable EFI secure boot")
}

func (cmd *boot) Description() string {
	return `Configure VM boot settings.

Examples:
  govc device.boot -vm $vm -delay 1000 -order floppy,cdrom,ethernet,disk
  govc device.boot -vm $vm -order - # reset boot order
  govc device.boot -vm $vm -secure`
}

func (cmd *boot) Process(ctx context.Context) error {
	if err := cmd.VirtualMachineFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *boot) Run(ctx context.Context, f *flag.FlagSet) error {
	vm, err := cmd.VirtualMachine()
	if err != nil {
		return err
	}

	if vm == nil {
		return flag.ErrHelp
	}

	devices, err := vm.Device(ctx)
	if err != nil {
		return err
	}

	if cmd.order != "" {
		o := strings.Split(cmd.order, ",")
		cmd.BootOrder = devices.BootOrder(o)
	}

	return vm.SetBootOptions(ctx, &cmd.VirtualMachineBootOptions)
}
<|endoftext|>


<commit_before>package main

import (
	//	"errors"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	filepath "path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/rwcarlsen/goexif/exif"
)

// Helper to log an error and then exit
// Helper to log an error and then exit
func handleErr(err error) {
	if err != nil {
		log.Fatal("Error:", err.Error())
	}
}

// Helper to get file modification time, useful as a fallback if file is not a jpg.
func getFileModTime(fileName string) time.Time {
	stat, err := os.Stat(fileName)
	if err != nil {
		log.Error("Unable to get ModTime for file: ", fileName)
		return time.Now()
	}
	return stat.ModTime()
}

// Get date taken of a file. If it is a jpg it will attempt to use EXIF data
func getDateTaken(fileName string) (time.Time, error) {

	if len(fileName) <= 0 {
		log.Warn("Pass filename as parameter.")
		return time.Now(), errors.New("Invalid filename passed.")
	}

	file, err := os.Open(fileName)
	if err != nil {
		return time.Now(), err
	}

	fileExt := strings.ToLower(filepath.Ext(fileName))

	date := time.Now()

	if fileExt == ".jpg" {

		data, err := exif.Decode(file)
		if err != nil {
			// file might not have exif data, use os.Stat
			date = getFileModTime(fileName)
		} else {
			date, _ = data.DateTime()
		}
	} else {
		date = getFileModTime(fileName)
	}

	return date, err
}

// Helper to create a folder
func createDir(dirName string) {
	if _, err := os.Stat(dirName); os.IsNotExist(err) {
		// Ok directory doesn't exist, create it
		err := os.Mkdir(dirName, 0777)
		if err != nil {
			log.Warn("Error happened creating directory:", err.Error())
		}
	}
}

// Helper function to copy a file
func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}

func processPhoto(fileName, outDir, bucketName string, dateTaken time.Time) error {
	outPath := filepath.Join(dateTaken.Format("2006/2006-01-02"), filepath.Base(fileName))
	if len(outDir) > 0 {
		destDir := filepath.Join(outDir, outPath)
		createDir(destDir)
		copyFile(fileName, destDir)
		log.Info("Copied file: " + fileName)
	}
	if len(bucketName) > 0 {
		uploadFile(fileName, bucketName, outPath)
		log.Info("Uploaded file to bucket" + bucketName)
	}
	// TODO! Write index.html file
	return nil
}

// Loops through all files in a dir
func organiseFiles(inDirName, outDirName, bucketName string) {
	files, err := ioutil.ReadDir(inDirName)
	handleErr(err)

	for _, f := range files {
		fileName := inDirName + "/" + f.Name()

		// Get date taken for file
		date, err := getDateTaken(fileName)
		if err != nil {
			log.Warn(err.Error())
		}

		// Organise photo by moving to target folder
		err = processPhoto(fileName, outDirName, bucketName, date)
		if err != nil {
			log.Error(err.Error())
		}
	}
}

func uploadFile(fileName, destName, bucketName string) error {
	// TODO! Upload file to a S3 bucket
	svc := s3.New(session.New(&aws.Config{Region: aws.String("ap-southeast-2")}))

	file, err := os.Open(fileName)

	if err != nil {
		log.Error(err)
	}

	defer file.Close()

	fileInfo, _ := file.Stat()
	size := fileInfo.Size()

	buffer := make([]byte, size)

	// read file content to buffer
	file.Read(buffer)

	fileBytes := bytes.NewReader(buffer) // convert to io.ReadSeeker type

	fileType := http.DetectContentType(buffer)

	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(destName),   // required
		ACL:           aws.String("public-read"),
		Body:          fileBytes,
		ContentLength: aws.Int64(size),
		ContentType:   aws.String(fileType),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err = svc.PutObject(params)

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			// Generic AWS Error with Code, Message, and original error (if any)
			fmt.Println(awsErr.Code(), awsErr.Message(), awsErr.OrigErr())
			if reqErr, ok := err.(awserr.RequestFailure); ok {
				// A service error occurred
				fmt.Println(reqErr.Code(), reqErr.Message(), reqErr.StatusCode(), reqErr.RequestID())
			}
		} else {
			// This case should never be hit, the SDK should always return an
			// error which satisfies the awserr.Error interface.
			fmt.Println(err.Error())
		}
	}
	return nil
}

func main() {

	// Declare a string parameter
	inDirNamePtr := flag.String("i", ".", "input directory")
	outDirNamePtr := flag.String("o", ".", "output directory")
	bucketNamePtr := flag.String("b", "", "bucket name")
	// Parse command line arguments.
	flag.Parse()
	if len(*inDirNamePtr) == 0 {
		log.Fatal("Error, need to define an input directory.")
	}

	organiseFiles(*inDirNamePtr, *outDirNamePtr, *bucketNamePtr)
}
<commit_msg>Add code to handle region.<commit_after>package main

import (
	//	"errors"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	filepath "path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/rwcarlsen/goexif/exif"
)

// Helper to log an error and then exit
// Helper to log an error and then exit
func handleErr(err error) {
	if err != nil {
		log.Fatal("Error:", err.Error())
	}
}

// Helper to get file modification time, useful as a fallback if file is not a jpg.
func getFileModTime(fileName string) time.Time {
	stat, err := os.Stat(fileName)
	if err != nil {
		log.Error("Unable to get ModTime for file: ", fileName)
		return time.Now()
	}
	return stat.ModTime()
}

// Get date taken of a file. If it is a jpg it will attempt to use EXIF data
func getDateTaken(fileName string) (time.Time, error) {

	if len(fileName) <= 0 {
		log.Warn("Pass filename as parameter.")
		return time.Now(), errors.New("Invalid filename passed.")
	}

	file, err := os.Open(fileName)
	if err != nil {
		return time.Now(), err
	}

	fileExt := strings.ToLower(filepath.Ext(fileName))

	date := time.Now()

	if fileExt == ".jpg" {

		data, err := exif.Decode(file)
		if err != nil {
			// file might not have exif data, use os.Stat
			date = getFileModTime(fileName)
		} else {
			date, _ = data.DateTime()
		}
	} else {
		date = getFileModTime(fileName)
	}

	return date, err
}

// Helper to create a folder
func createDir(dirName string) {
	if _, err := os.Stat(dirName); os.IsNotExist(err) {
		// Ok directory doesn't exist, create it
		err := os.Mkdir(dirName, 0777)
		if err != nil {
			log.Warn("Error creating directory:", err.Error())
		}
	}
}

// Helper function to copy a file
func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}

func processPhoto(sourceFile, outDir, bucketName, awsRegion string, dateTaken time.Time) error {
	outPath := dateTaken.Format("2006/2006-01-02")
	if len(outDir) > 0 {
		createDir(filepath.Join(outDir, dateTaken.Format("2006")))
		createDir(filepath.Join(outDir, dateTaken.Format("2006/2006-01-02")))
		fileName := filepath.Base(sourceFile)
		destPath := filepath.Join(outDir, outPath, fileName)

		copyFile(sourceFile, destPath)
		log.Info("Copied file: " + destPath)
	}
	if len(bucketName) > 0 {
		uploadFile(sourceFile, bucketName, outPath, awsRegion)
		log.Info("Uploaded file to bucket: " + bucketName)
	}
	// TODO! Write index.html file
	return nil
}

// Loops through all files in a dir
func organiseFiles(inDirName, outDirName, bucketName, awsRegion string) {
	files, err := ioutil.ReadDir(inDirName)
	handleErr(err)

	for _, f := range files {
		fileName := inDirName + "/" + f.Name()

		// Get date taken for file
		date, err := getDateTaken(fileName)
		if err != nil {
			log.Warn(err.Error())
		}

		// Organise photo by moving to target folder
		err = processPhoto(fileName, outDirName, bucketName, awsRegion, date)
		if err != nil {
			log.Error(err.Error())
		}
		log.Info("Done processing: ", inDirName)
	}
}

func uploadFile(fileName, destName, bucketName, awsRegion string) error {
	svc := s3.New(session.New(&aws.Config{Region: aws.String("ap-southeast-2")}))

	file, err := os.Open(fileName)

	if err != nil {
		log.Error(err)
	}

	defer file.Close()

	fileInfo, _ := file.Stat()
	size := fileInfo.Size()

	buffer := make([]byte, size)

	// read file content to buffer
	file.Read(buffer)

	fileBytes := bytes.NewReader(buffer) // convert to io.ReadSeeker type

	fileType := http.DetectContentType(buffer)

	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(destName),   // required
		ACL:           aws.String("public-read"),
		Body:          fileBytes,
		ContentLength: aws.Int64(size),
		ContentType:   aws.String(fileType),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err = svc.PutObject(params)

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			// Generic AWS Error with Code, Message, and original error (if any)
			fmt.Println(awsErr.Code(), awsErr.Message(), awsErr.OrigErr())
			if reqErr, ok := err.(awserr.RequestFailure); ok {
				// A service error occurred
				fmt.Println(reqErr.Code(), reqErr.Message(), reqErr.StatusCode(), reqErr.RequestID())
			}
		} else {
			// This case should never be hit, the SDK should always return an
			// error which satisfies the awserr.Error interface.
			fmt.Println(err.Error())
		}
	}
	return nil
}

func main() {

	// Declare a string parameter
	inDirNamePtr := flag.String("i", "", "input directory")
	outDirNamePtr := flag.String("o", "", "output directory")
	bucketNamePtr := flag.String("b", "", "bucket name")
	awsRegionNamePtr := flag.String("r", "", "AWS region")
	// Parse command line arguments.
	flag.Parse()
	if len(*inDirNamePtr) == 0 {
		log.Fatal("Error, need to define an input directory.")
	}

	organiseFiles(*inDirNamePtr, *outDirNamePtr, *bucketNamePtr, *awsRegionNamePtr)
	log.Info("Done")
}
<|endoftext|>


<commit_before>// Package tsdb implements a time series storage for float64 sample data.
package tsdb

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"sync"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/cespare/xxhash"
	"github.com/fabxc/tsdb/chunks"
	"github.com/go-kit/kit/log"
)

// DefaultOptions used for the DB. They are sane for setups using
// millisecond precision timestamps.
var DefaultOptions = &Options{
	Retention: 15 * 24 * 3600 * 1000, // 15 days
}

// Options of the DB storage.
type Options struct {
	Retention int64
}

// DB is a time series storage.
type DB struct {
	logger log.Logger
	opts   *Options
	path   string

	shards []*Shard
}

// TODO(fabxc): make configurable
const (
	shardShift   = 2
	numShards    = 1 << shardShift
	maxChunkSize = 1024
)

// Open or create a new DB.
func Open(path string, l log.Logger, opts *Options) (*DB, error) {
	if opts == nil {
		opts = DefaultOptions
	}
	if err := os.MkdirAll(path, 0777); err != nil {
		return nil, err
	}
	if l == nil {
		l = log.NewLogfmtLogger(os.Stdout)
		l = log.NewContext(l).With("ts", log.DefaultTimestampUTC, "caller", log.DefaultCaller)
	}

	c := &DB{
		logger: l,
		opts:   opts,
		path:   path,
	}

	// Initialize vertical shards.
	// TODO(fabxc): validate shard number to be power of 2, which is required
	// for the bitshift-modulo when finding the right shard.
	for i := 0; i < numShards; i++ {
		l := log.NewContext(l).With("shard", i)
		d := shardDir(path, i)

		s, err := OpenShard(d, l)
		if err != nil {
			return nil, fmt.Errorf("initializing shard %q failed: %s", d, err)
		}

		c.shards = append(c.shards, s)
	}

	// TODO(fabxc): run background compaction + GC.

	return c, nil
}

func shardDir(base string, i int) string {
	return filepath.Join(base, strconv.Itoa(i))
}

// Close the database.
func (db *DB) Close() error {
	var g errgroup.Group

	for _, shard := range db.shards {
		// Fix closure argument to goroutine.
		shard := shard
		g.Go(shard.Close)
	}

	return g.Wait()
}

// Appender adds a batch of samples.
type Appender interface {
	// Add adds a sample pair to the appended batch.
	Add(l Labels, t int64, v float64)

	// Commit submits the collected samples.
	Commit() error
}

// Vector is a set of LabelSet associated with one value each.
// Label sets and values must have equal length.
type Vector struct {
	Buckets map[uint16][]Sample
	reused  int
}

type Sample struct {
	Hash   uint64
	Labels Labels
	Value  float64
}

// Reset the vector but keep resources allocated.
func (v *Vector) Reset() {
	// Do a full reset every n-th reusage to avoid memory leaks.
	if v.Buckets == nil || v.reused > 100 {
		v.Buckets = make(map[uint16][]Sample, 0)
		return
	}
	for x, bkt := range v.Buckets {
		v.Buckets[x] = bkt[:0]
	}
	v.reused++
}

// Add a sample to the vector.
func (v *Vector) Add(lset Labels, val float64) {
	h := lset.Hash()
	s := uint16(h >> (64 - shardShift))

	v.Buckets[s] = append(v.Buckets[s], Sample{
		Hash:   h,
		Labels: lset,
		Value:  val,
	})
}

// func (db *DB) Appender() Appender {
// 	return &bucketAppender{
// 		samples: make([]Sample, 1024),
// 	}
// }

// type bucketAppender struct {
// 	db *DB
// 	// buckets []Sam
// }

// func (a *bucketAppender) Add(l Labels, t int64, v float64) {

// }

// func (a *bucketAppender) Commit() error {
// 	// f
// }

// AppendVector adds values for a list of label sets for the given timestamp
// in milliseconds.
func (db *DB) AppendVector(ts int64, v *Vector) error {
	// Sequentially add samples to shards.
	for s, bkt := range v.Buckets {
		shard := db.shards[s]
		if err := shard.appendBatch(ts, bkt); err != nil {
			// TODO(fabxc): handle gracefully and collect multi-error.
			return err
		}
	}

	return nil
}

func (db *DB) AppendSingle(lset Labels, ts int64, v float64) error {
	sort.Sort(lset)
	h := lset.Hash()
	s := uint16(h >> (64 - shardShift))

	return db.shards[s].appendBatch(ts, []Sample{
		{
			Hash:   h,
			Labels: lset,
			Value:  v,
		},
	})
}

const sep = '\xff'

// Shard handles reads and writes of time series falling into
// a hashed shard of a series.
type Shard struct {
	path      string
	persistCh chan struct{}
	logger    log.Logger

	mtx       sync.RWMutex
	persisted persistedBlocks
	head      *HeadBlock
}

// OpenShard returns a new Shard.
func OpenShard(path string, logger log.Logger) (*Shard, error) {
	// Create directory if shard is new.
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := os.MkdirAll(path, 0777); err != nil {
			return nil, err
		}
	}

	// Initialize previously persisted blocks.
	pbs, err := findPersistedBlocks(path)
	if err != nil {
		return nil, err
	}

	s := &Shard{
		path:      path,
		persistCh: make(chan struct{}, 1),
		logger:    logger,
		persisted: pbs,
		// TODO(fabxc): restore from checkpoint.
	}
	// TODO(fabxc): get base time from pre-existing blocks. Otherwise
	// it should come from a user defined start timestamp.
	// Use actual time for now.
	s.head = NewHeadBlock(time.Now().UnixNano() / int64(time.Millisecond))

	return s, nil
}

// Close the shard.
func (s *Shard) Close() error {
	var e MultiError

	for _, pb := range s.persisted {
		e.Add(pb.Close())
	}

	return e.Err()
}

func (s *Shard) appendBatch(ts int64, samples []Sample) error {
	// TODO(fabxc): make configurable.
	const persistenceTimeThreshold = 1000 * 60 * 60 // 1 hour if timestamp in ms

	s.mtx.Lock()
	defer s.mtx.Unlock()

	for _, smpl := range samples {
		if err := s.head.append(smpl.Hash, smpl.Labels, ts, smpl.Value); err != nil {
			// TODO(fabxc): handle gracefully and collect multi-error.
			return err
		}
	}

	if ts > s.head.highTimestamp {
		s.head.highTimestamp = ts
	}

	// TODO(fabxc): randomize over time
	if s.head.stats().samples/uint64(s.head.stats().chunks) > 400 {
		select {
		case s.persistCh <- struct{}{}:
			go s.persist()
		default:
		}
	}

	return nil
}

// blocksForRange returns all blocks within the shard that may contain
// data for the given time range.
func (s *Shard) blocksForRange(mint, maxt int64) (bs []Block) {
	return []Block{s.head}
}

// TODO(fabxc): make configurable.
const shardGracePeriod = 60 * 1000 // 60 seconds for millisecond scale

func (s *Shard) persist() error {
	s.mtx.Lock()

	// Set new head block.
	head := s.head
	s.head = NewHeadBlock(head.highTimestamp)

	s.mtx.Unlock()

	defer func() {
		<-s.persistCh
	}()

	// TODO(fabxc): add grace period where we can still append to old head shard
	// before actually persisting it.
	p := filepath.Join(s.path, fmt.Sprintf("%d", head.baseTimestamp))

	if err := os.MkdirAll(p, 0777); err != nil {
		return err
	}

	sf, err := os.Create(filepath.Join(p, "series"))
	if err != nil {
		return err
	}
	xf, err := os.Create(filepath.Join(p, "index"))
	if err != nil {
		return err
	}

	iw := newIndexWriter(xf)
	sw := newSeriesWriter(sf, iw, s.head.baseTimestamp)

	defer sw.Close()
	defer iw.Close()

	for ref, cd := range head.index.forward {
		if err := sw.WriteSeries(ref, cd.lset, []*chunkDesc{cd}); err != nil {
			return err
		}
	}

	if err := iw.WriteStats(nil); err != nil {
		return err
	}
	for n, v := range head.index.values {
		s := make([]string, 0, len(v))
		for x := range v {
			s = append(s, x)
		}

		if err := iw.WriteLabelIndex([]string{n}, s); err != nil {
			return err
		}
	}

	for t := range head.index.postings.m {
		if err := iw.WritePostings(t.name, t.value, head.index.postings.get(t)); err != nil {
			return err
		}
	}

	sz := fmt.Sprintf("%.2fMiB", float64(sw.Size()+iw.Size())/1024/1024)

	s.logger.Log("size", sz, "samples", head.samples, "chunks", head.stats().chunks, "msg", "persisted head")

	return nil
}

// chunkDesc wraps a plain data chunk and provides cached meta data about it.
type chunkDesc struct {
	lset  Labels
	chunk chunks.Chunk

	// Caching fields.
	lastTimestamp int64
	lastValue     float64

	app chunks.Appender // Current appender for the chunks.
}

func (cd *chunkDesc) append(ts int64, v float64) (err error) {
	if cd.app == nil {
		cd.app, err = cd.chunk.Appender()
		if err != nil {
			return err
		}
	}
	if err := cd.app.Append(ts, v); err != nil {
		return err
	}

	cd.lastTimestamp = ts
	cd.lastValue = v

	return nil
}

// Label is a key/value pair of strings.
type Label struct {
	Name, Value string
}

// Labels is a sorted set of labels. Order has to be guaranteed upon
// instantiation.
type Labels []Label

func (ls Labels) Len() int           { return len(ls) }
func (ls Labels) Swap(i, j int)      { ls[i], ls[j] = ls[j], ls[i] }
func (ls Labels) Less(i, j int) bool { return ls[i].Name < ls[j].Name }

// Hash returns a hash value for the label set.
func (ls Labels) Hash() uint64 {
	b := make([]byte, 0, 1024)

	for _, v := range ls {
		b = append(b, v.Name...)
		b = append(b, sep)
		b = append(b, v.Value...)
		b = append(b, sep)
	}
	return xxhash.Sum64(b)
}

// Get returns the value for the label with the given name.
// Returns an empty string if the label doesn't exist.
func (ls Labels) Get(name string) string {
	for _, l := range ls {
		if l.Name == name {
			return l.Value
		}
	}
	return ""
}

// Equals returns whether the two label sets are equal.
func (ls Labels) Equals(o Labels) bool {
	if len(ls) != len(o) {
		return false
	}
	for i, l := range ls {
		if l.Name != o[i].Name || l.Value != o[i].Value {
			return false
		}
	}
	return true
}

// Map returns a string map of the labels.
func (ls Labels) Map() map[string]string {
	m := make(map[string]string, len(ls))
	for _, l := range ls {
		m[l.Name] = l.Value
	}
	return m
}

// NewLabels returns a sorted Labels from the given labels.
// The caller has to guarantee that all label names are unique.
func NewLabels(ls ...Label) Labels {
	set := make(Labels, 0, len(ls))
	for _, l := range ls {
		set = append(set, l)
	}
	sort.Sort(set)

	return set
}

// LabelsFromMap returns new sorted Labels from the given map.
func LabelsFromMap(m map[string]string) Labels {
	l := make([]Label, 0, len(m))
	for k, v := range m {
		l = append(l, Label{Name: k, Value: v})
	}
	return NewLabels(l...)
}

// The MultiError type implements the error interface, and contains the
// Errors used to construct it.
type MultiError []error

// Returns a concatenated string of the contained errors
func (es MultiError) Error() string {
	var buf bytes.Buffer

	if len(es) > 0 {
		fmt.Fprintf(&buf, "%d errors: ", len(es))
	}

	for i, err := range es {
		if i != 0 {
			buf.WriteString("; ")
		}
		buf.WriteString(err.Error())
	}

	return buf.String()
}

func (es MultiError) Add(err error) {
	if err != nil {
		es = append(es, err)
	}
}

func (es MultiError) Err() error {
	if len(es) == 0 {
		return nil
	}
	return es
}
<commit_msg>Add unsafe string and slice conversions<commit_after>// Package tsdb implements a time series storage for float64 sample data.
package tsdb

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"strconv"
	"sync"
	"time"
	"unsafe"

	"golang.org/x/sync/errgroup"

	"github.com/cespare/xxhash"
	"github.com/fabxc/tsdb/chunks"
	"github.com/go-kit/kit/log"
)

// DefaultOptions used for the DB. They are sane for setups using
// millisecond precision timestamps.
var DefaultOptions = &Options{
	Retention: 15 * 24 * 3600 * 1000, // 15 days
}

// Options of the DB storage.
type Options struct {
	Retention int64
}

// DB is a time series storage.
type DB struct {
	logger log.Logger
	opts   *Options
	path   string

	shards []*Shard
}

// TODO(fabxc): make configurable
const (
	shardShift   = 2
	numShards    = 1 << shardShift
	maxChunkSize = 1024
)

// Open or create a new DB.
func Open(path string, l log.Logger, opts *Options) (*DB, error) {
	if opts == nil {
		opts = DefaultOptions
	}
	if err := os.MkdirAll(path, 0777); err != nil {
		return nil, err
	}
	if l == nil {
		l = log.NewLogfmtLogger(os.Stdout)
		l = log.NewContext(l).With("ts", log.DefaultTimestampUTC, "caller", log.DefaultCaller)
	}

	c := &DB{
		logger: l,
		opts:   opts,
		path:   path,
	}

	// Initialize vertical shards.
	// TODO(fabxc): validate shard number to be power of 2, which is required
	// for the bitshift-modulo when finding the right shard.
	for i := 0; i < numShards; i++ {
		l := log.NewContext(l).With("shard", i)
		d := shardDir(path, i)

		s, err := OpenShard(d, l)
		if err != nil {
			return nil, fmt.Errorf("initializing shard %q failed: %s", d, err)
		}

		c.shards = append(c.shards, s)
	}

	// TODO(fabxc): run background compaction + GC.

	return c, nil
}

func shardDir(base string, i int) string {
	return filepath.Join(base, strconv.Itoa(i))
}

// Close the database.
func (db *DB) Close() error {
	var g errgroup.Group

	for _, shard := range db.shards {
		// Fix closure argument to goroutine.
		shard := shard
		g.Go(shard.Close)
	}

	return g.Wait()
}

// Appender adds a batch of samples.
type Appender interface {
	// Add adds a sample pair to the appended batch.
	Add(l Labels, t int64, v float64)

	// Commit submits the collected samples.
	Commit() error
}

// Vector is a set of LabelSet associated with one value each.
// Label sets and values must have equal length.
type Vector struct {
	Buckets map[uint16][]Sample
	reused  int
}

type Sample struct {
	Hash   uint64
	Labels Labels
	Value  float64
}

// Reset the vector but keep resources allocated.
func (v *Vector) Reset() {
	// Do a full reset every n-th reusage to avoid memory leaks.
	if v.Buckets == nil || v.reused > 100 {
		v.Buckets = make(map[uint16][]Sample, 0)
		return
	}
	for x, bkt := range v.Buckets {
		v.Buckets[x] = bkt[:0]
	}
	v.reused++
}

// Add a sample to the vector.
func (v *Vector) Add(lset Labels, val float64) {
	h := lset.Hash()
	s := uint16(h >> (64 - shardShift))

	v.Buckets[s] = append(v.Buckets[s], Sample{
		Hash:   h,
		Labels: lset,
		Value:  val,
	})
}

// func (db *DB) Appender() Appender {
// 	return &bucketAppender{
// 		samples: make([]Sample, 1024),
// 	}
// }

// type bucketAppender struct {
// 	db *DB
// 	// buckets []Sam
// }

// func (a *bucketAppender) Add(l Labels, t int64, v float64) {

// }

// func (a *bucketAppender) Commit() error {
// 	// f
// }

// AppendVector adds values for a list of label sets for the given timestamp
// in milliseconds.
func (db *DB) AppendVector(ts int64, v *Vector) error {
	// Sequentially add samples to shards.
	for s, bkt := range v.Buckets {
		shard := db.shards[s]
		if err := shard.appendBatch(ts, bkt); err != nil {
			// TODO(fabxc): handle gracefully and collect multi-error.
			return err
		}
	}

	return nil
}

func (db *DB) AppendSingle(lset Labels, ts int64, v float64) error {
	sort.Sort(lset)
	h := lset.Hash()
	s := uint16(h >> (64 - shardShift))

	return db.shards[s].appendBatch(ts, []Sample{
		{
			Hash:   h,
			Labels: lset,
			Value:  v,
		},
	})
}

const sep = '\xff'

// Shard handles reads and writes of time series falling into
// a hashed shard of a series.
type Shard struct {
	path      string
	persistCh chan struct{}
	logger    log.Logger

	mtx       sync.RWMutex
	persisted persistedBlocks
	head      *HeadBlock
}

// OpenShard returns a new Shard.
func OpenShard(path string, logger log.Logger) (*Shard, error) {
	// Create directory if shard is new.
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := os.MkdirAll(path, 0777); err != nil {
			return nil, err
		}
	}

	// Initialize previously persisted blocks.
	pbs, err := findPersistedBlocks(path)
	if err != nil {
		return nil, err
	}

	s := &Shard{
		path:      path,
		persistCh: make(chan struct{}, 1),
		logger:    logger,
		persisted: pbs,
		// TODO(fabxc): restore from checkpoint.
	}
	// TODO(fabxc): get base time from pre-existing blocks. Otherwise
	// it should come from a user defined start timestamp.
	// Use actual time for now.
	s.head = NewHeadBlock(time.Now().UnixNano() / int64(time.Millisecond))

	return s, nil
}

// Close the shard.
func (s *Shard) Close() error {
	var e MultiError

	for _, pb := range s.persisted {
		e.Add(pb.Close())
	}

	return e.Err()
}

func (s *Shard) appendBatch(ts int64, samples []Sample) error {
	// TODO(fabxc): make configurable.
	const persistenceTimeThreshold = 1000 * 60 * 60 // 1 hour if timestamp in ms

	s.mtx.Lock()
	defer s.mtx.Unlock()

	for _, smpl := range samples {
		if err := s.head.append(smpl.Hash, smpl.Labels, ts, smpl.Value); err != nil {
			// TODO(fabxc): handle gracefully and collect multi-error.
			return err
		}
	}

	if ts > s.head.highTimestamp {
		s.head.highTimestamp = ts
	}

	// TODO(fabxc): randomize over time
	if s.head.stats().samples/uint64(s.head.stats().chunks) > 400 {
		select {
		case s.persistCh <- struct{}{}:
			go s.persist()
		default:
		}
	}

	return nil
}

// blocksForRange returns all blocks within the shard that may contain
// data for the given time range.
func (s *Shard) blocksForRange(mint, maxt int64) (bs []Block) {
	return []Block{s.head}
}

// TODO(fabxc): make configurable.
const shardGracePeriod = 60 * 1000 // 60 seconds for millisecond scale

func (s *Shard) persist() error {
	s.mtx.Lock()

	// Set new head block.
	head := s.head
	s.head = NewHeadBlock(head.highTimestamp)

	s.mtx.Unlock()

	defer func() {
		<-s.persistCh
	}()

	// TODO(fabxc): add grace period where we can still append to old head shard
	// before actually persisting it.
	p := filepath.Join(s.path, fmt.Sprintf("%d", head.baseTimestamp))

	if err := os.MkdirAll(p, 0777); err != nil {
		return err
	}

	sf, err := os.Create(filepath.Join(p, "series"))
	if err != nil {
		return err
	}
	xf, err := os.Create(filepath.Join(p, "index"))
	if err != nil {
		return err
	}

	iw := newIndexWriter(xf)
	sw := newSeriesWriter(sf, iw, s.head.baseTimestamp)

	defer sw.Close()
	defer iw.Close()

	for ref, cd := range head.index.forward {
		if err := sw.WriteSeries(ref, cd.lset, []*chunkDesc{cd}); err != nil {
			return err
		}
	}

	if err := iw.WriteStats(nil); err != nil {
		return err
	}
	for n, v := range head.index.values {
		s := make([]string, 0, len(v))
		for x := range v {
			s = append(s, x)
		}

		if err := iw.WriteLabelIndex([]string{n}, s); err != nil {
			return err
		}
	}

	for t := range head.index.postings.m {
		if err := iw.WritePostings(t.name, t.value, head.index.postings.get(t)); err != nil {
			return err
		}
	}

	sz := fmt.Sprintf("%.2fMiB", float64(sw.Size()+iw.Size())/1024/1024)

	s.logger.Log("size", sz, "samples", head.samples, "chunks", head.stats().chunks, "msg", "persisted head")

	return nil
}

// chunkDesc wraps a plain data chunk and provides cached meta data about it.
type chunkDesc struct {
	lset  Labels
	chunk chunks.Chunk

	// Caching fields.
	lastTimestamp int64
	lastValue     float64

	app chunks.Appender // Current appender for the chunks.
}

func (cd *chunkDesc) append(ts int64, v float64) (err error) {
	if cd.app == nil {
		cd.app, err = cd.chunk.Appender()
		if err != nil {
			return err
		}
	}
	if err := cd.app.Append(ts, v); err != nil {
		return err
	}

	cd.lastTimestamp = ts
	cd.lastValue = v

	return nil
}

// Label is a key/value pair of strings.
type Label struct {
	Name, Value string
}

// Labels is a sorted set of labels. Order has to be guaranteed upon
// instantiation.
type Labels []Label

func (ls Labels) Len() int           { return len(ls) }
func (ls Labels) Swap(i, j int)      { ls[i], ls[j] = ls[j], ls[i] }
func (ls Labels) Less(i, j int) bool { return ls[i].Name < ls[j].Name }

// Hash returns a hash value for the label set.
func (ls Labels) Hash() uint64 {
	b := make([]byte, 0, 1024)

	for _, v := range ls {
		b = append(b, v.Name...)
		b = append(b, sep)
		b = append(b, v.Value...)
		b = append(b, sep)
	}
	return xxhash.Sum64(b)
}

// Get returns the value for the label with the given name.
// Returns an empty string if the label doesn't exist.
func (ls Labels) Get(name string) string {
	for _, l := range ls {
		if l.Name == name {
			return l.Value
		}
	}
	return ""
}

// Equals returns whether the two label sets are equal.
func (ls Labels) Equals(o Labels) bool {
	if len(ls) != len(o) {
		return false
	}
	for i, l := range ls {
		if l.Name != o[i].Name || l.Value != o[i].Value {
			return false
		}
	}
	return true
}

// Map returns a string map of the labels.
func (ls Labels) Map() map[string]string {
	m := make(map[string]string, len(ls))
	for _, l := range ls {
		m[l.Name] = l.Value
	}
	return m
}

// NewLabels returns a sorted Labels from the given labels.
// The caller has to guarantee that all label names are unique.
func NewLabels(ls ...Label) Labels {
	set := make(Labels, 0, len(ls))
	for _, l := range ls {
		set = append(set, l)
	}
	sort.Sort(set)

	return set
}

// LabelsFromMap returns new sorted Labels from the given map.
func LabelsFromMap(m map[string]string) Labels {
	l := make([]Label, 0, len(m))
	for k, v := range m {
		l = append(l, Label{Name: k, Value: v})
	}
	return NewLabels(l...)
}

// The MultiError type implements the error interface, and contains the
// Errors used to construct it.
type MultiError []error

// Returns a concatenated string of the contained errors
func (es MultiError) Error() string {
	var buf bytes.Buffer

	if len(es) > 0 {
		fmt.Fprintf(&buf, "%d errors: ", len(es))
	}

	for i, err := range es {
		if i != 0 {
			buf.WriteString("; ")
		}
		buf.WriteString(err.Error())
	}

	return buf.String()
}

// Add adds the error to the error list if it is not nil.
func (es MultiError) Add(err error) {
	if err != nil {
		es = append(es, err)
	}
}

// Err returns the error list as an error or nil if it is empty.
func (es MultiError) Err() error {
	if len(es) == 0 {
		return nil
	}
	return es
}

func yoloString(b []byte) string {
	h := reflect.StringHeader{
		Data: uintptr(unsafe.Pointer(&b[0])),
		Len:  len(b),
	}
	return *((*string)(unsafe.Pointer(&h)))
}

func yoloBytes(s string) []byte {
	sh := (*reflect.StringHeader)(unsafe.Pointer(&s))

	h := reflect.SliceHeader{
		Cap:  sh.Len,
		Len:  sh.Len,
		Data: sh.Data,
	}
	return *((*[]byte)(unsafe.Pointer(&h)))
}
<|endoftext|>


<commit_before>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/hashicorp/go-version"
	sq "github.com/lann/squirrel"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
	"github.com/koding/multiconfig"
)

var (
	ErrQueryFieldsEmpty = errors.New("all query fields are empty")
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string `default:"localhost"`
	Port     int    `default:"5432"`
	Username string `required:"true"`
	Password string
	DBName   string `required:"true" `
}

type Postgres struct {
	DB  *sql.DB
	Log kite.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf == nil {
		conf = new(PostgresConfig)

		envLoader := &multiconfig.EnvironmentLoader{Prefix: "kontrol_postgres"}
		configLoader := multiconfig.MultiLoader(
			&multiconfig.TagLoader{}, envLoader,
		)

		if err := configLoader.Load(conf); err != nil {
			fmt.Println("Valid environment variables are: ")
			envLoader.PrintEnvs(conf)
			panic(err)
		}

		err := multiconfig.MultiValidator(&multiconfig.RequiredValidator{}).Validate(conf)
		if err != nil {
			fmt.Println("Valid environment variables are: ")
			envLoader.PrintEnvs(conf)
			panic(err)
		}
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s user=%s password=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName, conf.Username, conf.Password,
	)

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	p := &Postgres{
		DB:  db,
		Log: log,
	}

	cleanInterval := 120 * time.Second // clean every 120 second
	go p.RunCleaner(cleanInterval, KeyTTL)

	return p
}

// RunCleaner deletes every "interval" duration rows which are older than
// "expire" duration based on the "updated_at" field. For more info check
// CleanExpireRows which is used to delete old rows.
func (p *Postgres) RunCleaner(interval, expire time.Duration) {
	cleanFunc := func() {
		affectedRows, err := p.CleanExpiredRows(expire)
		if err != nil {
			p.Log.Warning("postgres: cleaning old rows failed: %s", err)
		} else if affectedRows != 0 {
			p.Log.Debug("postgres: cleaned up %d rows", affectedRows)
		}
	}

	for _ = range time.Tick(interval) {
		cleanFunc()
	}
}

// CleanExpiredRows deletes rows that are at least "expire" duration old. So if
// say an expire duration of 10 second is given, it will delete all rows that
// were updated 10 seconds ago
func (p *Postgres) CleanExpiredRows(expire time.Duration) (int64, error) {
	// See: http://stackoverflow.com/questions/14465727/how-to-insert-things-like-now-interval-2-minutes-into-php-pdo-query
	// basically by passing an integer to INTERVAL is not possible, we need to
	// cast it. However there is a more simpler way, we can multiply INTERVAL
	// with an integer so we just declare a one second INTERVAL and multiply it
	// with the amount we want.
	cleanOldRows := `DELETE FROM kite.kite WHERE updated_at < (now() at time zone 'utc') - ((INTERVAL '1 second') * $1)`

	rows, err := p.DB.Exec(cleanOldRows, int64(expire/time.Second))
	if err != nil {
		return 0, err
	}

	return rows.RowsAffected()
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	sqlQuery, args, err := selectQuery(query)
	if err != nil {
		return nil, err
	}

	var hasVersionConstraint bool // does query contains a constraint on version?
	var keyRest string            // query key after the version field
	var versionConstraint version.Constraints
	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err = version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		// now parse our constraint
		versionConstraint, err = version.NewConstraint(query.Version)
		if err != nil {
			// version is a malformed, just return the error
			return nil, err
		}

		hasVersionConstraint = true
		nameQuery := &protocol.KontrolQuery{
			Username:    query.Username,
			Environment: query.Environment,
			Name:        query.Name,
		}

		// We will make a get request to all nodes under this name
		// and filter the result later.
		sqlQuery, args, err = selectQuery(nameQuery)
		if err != nil {
			return nil, err
		}

		// Rest of the key after version field
		keyRest = "/" + strings.TrimRight(
			query.Region+"/"+query.Hostname+"/"+query.ID, "/")
	}

	rows, err := p.DB.Query(sqlQuery, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var (
		username    string
		environment string
		kitename    string
		version     string
		region      string
		hostname    string
		id          string
		url         string
		updated_at  time.Time
		created_at  time.Time
		keyId       string
	)

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(
			&username,
			&environment,
			&kitename,
			&version,
			&region,
			&hostname,
			&id,
			&url,
			&updated_at,
			&created_at,
			&keyId,
		)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: protocol.Kite{
				Username:    username,
				Environment: environment,
				Name:        kitename,
				Version:     version,
				Region:      region,
				Hostname:    hostname,
				ID:          id,
			},
			URL:   url,
			KeyID: keyId,
		})
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	// if it's just single result there is no need to shuffle or filter
	// according to the version constraint
	if len(kites) == 1 {
		return kites, nil
	}

	// Filter kites by version constraint
	if hasVersionConstraint {
		kites.Filter(versionConstraint, keyRest)
	}

	// randomize the result
	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Upsert(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) (err error) {
	// check that the incoming URL is valid to prevent malformed input
	_, err = url.Parse(value.URL)
	if err != nil {
		return err
	}

	if value.KeyID == "" {
		return errors.New("postgres: keyId is empty. Aborting upsert")
	}

	// we are going to try an UPDATE, if it's not successfull we are going to
	// INSERT the document, all ine one single transaction
	tx, err := p.DB.Begin()
	if err != nil {
		return err
	}

	defer func() {
		if err != nil {
			err = tx.Rollback()
		} else {
			// it calls Rollback inside if it fails again :)
			err = tx.Commit()
		}
	}()

	res, err := tx.Exec(`UPDATE kite.kite SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE id = $2`, value.URL, kiteProt.ID)
	if err != nil {
		return err
	}

	rowAffected, err := res.RowsAffected()
	if err != nil {
		return err
	}

	// we got an update! so this was successfull, just return without an error
	if rowAffected != 0 {
		return nil
	}

	insertSQL, args, err := insertKiteQuery(kiteProt, value.URL, value.KeyID)
	if err != nil {
		return err
	}

	_, err = tx.Exec(insertSQL, args...)
	return err
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	sqlQuery, args, err := insertKiteQuery(kiteProt, value.URL, value.KeyID)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(sqlQuery, args...)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kite.kite SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE id = $2`,
		value.URL, kiteProt.ID)

	return err
}

func (p *Postgres) Delete(kiteProt *protocol.Kite) error {
	deleteKite := `DELETE FROM kite.kite WHERE id = $1`
	_, err := p.DB.Exec(deleteKite, kiteProt.ID)
	return err
}

// selectQuery returns a SQL query for the given query
func selectQuery(query *protocol.KontrolQuery) (string, []interface{}, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

	kites := psql.Select("*").From("kite.kite")
	fields := query.Fields()
	andQuery := sq.And{}

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			continue
		}

		// we are using "kitename" as the columname
		if key == "name" {
			key = "kitename"
		}

		andQuery = append(andQuery, sq.Eq{key: v})
	}

	if len(andQuery) == 0 {
		return "", nil, ErrQueryFieldsEmpty
	}

	return kites.Where(andQuery).ToSql()
}

// inseryKiteQuery inserts the given kite, url and key to the kite.kite table
func insertKiteQuery(kiteProt *protocol.Kite, url, keyId string) (string, []interface{}, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

	kiteValues := kiteProt.Values()
	values := make([]interface{}, len(kiteValues))

	for i, kiteVal := range kiteValues {
		values[i] = kiteVal
	}

	values = append(values, url)
	values = append(values, keyId)

	return psql.Insert("kite.kite").Columns(
		"username",
		"environment",
		"kitename",
		"version",
		"region",
		"hostname",
		"id",
		"url",
		"key_id",
	).Values(values...).ToSql()
}

/*

--- Key Pair -----------------

*/

func (p *Postgres) AddKey(keyPair *KeyPair) error {
	if err := keyPair.Validate(); err != nil {
		return err
	}

	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.Insert("kite.key").Columns(
		"id",
		"public",
		"private",
	).Values(keyPair.ID, keyPair.Public, keyPair.Private).ToSql()
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(sqlQuery, args...)
	return err
}

func (p *Postgres) DeleteKey(keyPair *KeyPair) error {
	res, err := p.DB.Exec(`UPDATE kite.key SET deleted_at = (now() at time zone 'utc')`)
	if err != nil {
		return err
	}

	_, err = res.RowsAffected()
	return err
}

func (p *Postgres) GetKeyFromID(id string) (*KeyPair, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.
		Select("id", "public", "private").
		From("kite.key").
		Where(map[string]interface{}{"id": id}).
		ToSql()
	if err != nil {
		return nil, err
	}

	keyPair := &KeyPair{}
	err = p.DB.QueryRow(sqlQuery, args...).Scan(&keyPair.ID, &keyPair.Public, &keyPair.Private)
	if err != nil {
		return nil, err
	}

	return keyPair, nil
}

func (p *Postgres) GetKeyFromPublic(public string) (*KeyPair, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.
		Select("id", "public", "private").
		From("kite.key").
		Where(map[string]interface{}{"public": public}).
		ToSql()
	if err != nil {
		return nil, err
	}

	keyPair := &KeyPair{}
	err = p.DB.QueryRow(sqlQuery, args...).Scan(&keyPair.ID, &keyPair.Public, &keyPair.Private)
	if err != nil {
		return nil, err
	}

	return keyPair, nil
}
<commit_msg>postgres: only select valid keys<commit_after>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/hashicorp/go-version"
	sq "github.com/lann/squirrel"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
	"github.com/koding/multiconfig"
)

var (
	ErrQueryFieldsEmpty = errors.New("all query fields are empty")
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string `default:"localhost"`
	Port     int    `default:"5432"`
	Username string `required:"true"`
	Password string
	DBName   string `required:"true" `
}

type Postgres struct {
	DB  *sql.DB
	Log kite.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf == nil {
		conf = new(PostgresConfig)

		envLoader := &multiconfig.EnvironmentLoader{Prefix: "kontrol_postgres"}
		configLoader := multiconfig.MultiLoader(
			&multiconfig.TagLoader{}, envLoader,
		)

		if err := configLoader.Load(conf); err != nil {
			fmt.Println("Valid environment variables are: ")
			envLoader.PrintEnvs(conf)
			panic(err)
		}

		err := multiconfig.MultiValidator(&multiconfig.RequiredValidator{}).Validate(conf)
		if err != nil {
			fmt.Println("Valid environment variables are: ")
			envLoader.PrintEnvs(conf)
			panic(err)
		}
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s user=%s password=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName, conf.Username, conf.Password,
	)

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	p := &Postgres{
		DB:  db,
		Log: log,
	}

	cleanInterval := 120 * time.Second // clean every 120 second
	go p.RunCleaner(cleanInterval, KeyTTL)

	return p
}

// RunCleaner deletes every "interval" duration rows which are older than
// "expire" duration based on the "updated_at" field. For more info check
// CleanExpireRows which is used to delete old rows.
func (p *Postgres) RunCleaner(interval, expire time.Duration) {
	cleanFunc := func() {
		affectedRows, err := p.CleanExpiredRows(expire)
		if err != nil {
			p.Log.Warning("postgres: cleaning old rows failed: %s", err)
		} else if affectedRows != 0 {
			p.Log.Debug("postgres: cleaned up %d rows", affectedRows)
		}
	}

	for _ = range time.Tick(interval) {
		cleanFunc()
	}
}

// CleanExpiredRows deletes rows that are at least "expire" duration old. So if
// say an expire duration of 10 second is given, it will delete all rows that
// were updated 10 seconds ago
func (p *Postgres) CleanExpiredRows(expire time.Duration) (int64, error) {
	// See: http://stackoverflow.com/questions/14465727/how-to-insert-things-like-now-interval-2-minutes-into-php-pdo-query
	// basically by passing an integer to INTERVAL is not possible, we need to
	// cast it. However there is a more simpler way, we can multiply INTERVAL
	// with an integer so we just declare a one second INTERVAL and multiply it
	// with the amount we want.
	cleanOldRows := `DELETE FROM kite.kite WHERE updated_at < (now() at time zone 'utc') - ((INTERVAL '1 second') * $1)`

	rows, err := p.DB.Exec(cleanOldRows, int64(expire/time.Second))
	if err != nil {
		return 0, err
	}

	return rows.RowsAffected()
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	sqlQuery, args, err := selectQuery(query)
	if err != nil {
		return nil, err
	}

	var hasVersionConstraint bool // does query contains a constraint on version?
	var keyRest string            // query key after the version field
	var versionConstraint version.Constraints
	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err = version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		// now parse our constraint
		versionConstraint, err = version.NewConstraint(query.Version)
		if err != nil {
			// version is a malformed, just return the error
			return nil, err
		}

		hasVersionConstraint = true
		nameQuery := &protocol.KontrolQuery{
			Username:    query.Username,
			Environment: query.Environment,
			Name:        query.Name,
		}

		// We will make a get request to all nodes under this name
		// and filter the result later.
		sqlQuery, args, err = selectQuery(nameQuery)
		if err != nil {
			return nil, err
		}

		// Rest of the key after version field
		keyRest = "/" + strings.TrimRight(
			query.Region+"/"+query.Hostname+"/"+query.ID, "/")
	}

	rows, err := p.DB.Query(sqlQuery, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var (
		username    string
		environment string
		kitename    string
		version     string
		region      string
		hostname    string
		id          string
		url         string
		updated_at  time.Time
		created_at  time.Time
		keyId       string
	)

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(
			&username,
			&environment,
			&kitename,
			&version,
			&region,
			&hostname,
			&id,
			&url,
			&updated_at,
			&created_at,
			&keyId,
		)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: protocol.Kite{
				Username:    username,
				Environment: environment,
				Name:        kitename,
				Version:     version,
				Region:      region,
				Hostname:    hostname,
				ID:          id,
			},
			URL:   url,
			KeyID: keyId,
		})
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	// if it's just single result there is no need to shuffle or filter
	// according to the version constraint
	if len(kites) == 1 {
		return kites, nil
	}

	// Filter kites by version constraint
	if hasVersionConstraint {
		kites.Filter(versionConstraint, keyRest)
	}

	// randomize the result
	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Upsert(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) (err error) {
	// check that the incoming URL is valid to prevent malformed input
	_, err = url.Parse(value.URL)
	if err != nil {
		return err
	}

	if value.KeyID == "" {
		return errors.New("postgres: keyId is empty. Aborting upsert")
	}

	// we are going to try an UPDATE, if it's not successfull we are going to
	// INSERT the document, all ine one single transaction
	tx, err := p.DB.Begin()
	if err != nil {
		return err
	}

	defer func() {
		if err != nil {
			err = tx.Rollback()
		} else {
			// it calls Rollback inside if it fails again :)
			err = tx.Commit()
		}
	}()

	res, err := tx.Exec(`UPDATE kite.kite SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE id = $2`, value.URL, kiteProt.ID)
	if err != nil {
		return err
	}

	rowAffected, err := res.RowsAffected()
	if err != nil {
		return err
	}

	// we got an update! so this was successfull, just return without an error
	if rowAffected != 0 {
		return nil
	}

	insertSQL, args, err := insertKiteQuery(kiteProt, value.URL, value.KeyID)
	if err != nil {
		return err
	}

	_, err = tx.Exec(insertSQL, args...)
	return err
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	sqlQuery, args, err := insertKiteQuery(kiteProt, value.URL, value.KeyID)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(sqlQuery, args...)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kite.kite SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE id = $2`,
		value.URL, kiteProt.ID)

	return err
}

func (p *Postgres) Delete(kiteProt *protocol.Kite) error {
	deleteKite := `DELETE FROM kite.kite WHERE id = $1`
	_, err := p.DB.Exec(deleteKite, kiteProt.ID)
	return err
}

// selectQuery returns a SQL query for the given query
func selectQuery(query *protocol.KontrolQuery) (string, []interface{}, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

	kites := psql.Select("*").From("kite.kite")
	fields := query.Fields()
	andQuery := sq.And{}

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			continue
		}

		// we are using "kitename" as the columname
		if key == "name" {
			key = "kitename"
		}

		andQuery = append(andQuery, sq.Eq{key: v})
	}

	if len(andQuery) == 0 {
		return "", nil, ErrQueryFieldsEmpty
	}

	return kites.Where(andQuery).ToSql()
}

// inseryKiteQuery inserts the given kite, url and key to the kite.kite table
func insertKiteQuery(kiteProt *protocol.Kite, url, keyId string) (string, []interface{}, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

	kiteValues := kiteProt.Values()
	values := make([]interface{}, len(kiteValues))

	for i, kiteVal := range kiteValues {
		values[i] = kiteVal
	}

	values = append(values, url)
	values = append(values, keyId)

	return psql.Insert("kite.kite").Columns(
		"username",
		"environment",
		"kitename",
		"version",
		"region",
		"hostname",
		"id",
		"url",
		"key_id",
	).Values(values...).ToSql()
}

/*

--- Key Pair -----------------

*/

func (p *Postgres) AddKey(keyPair *KeyPair) error {
	if err := keyPair.Validate(); err != nil {
		return err
	}

	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.Insert("kite.key").Columns(
		"id",
		"public",
		"private",
	).Values(keyPair.ID, keyPair.Public, keyPair.Private).ToSql()
	if err != nil {
		return err
	}

	_, err = p.DB.Exec(sqlQuery, args...)
	return err
}

func (p *Postgres) DeleteKey(keyPair *KeyPair) error {
	res, err := p.DB.Exec(`UPDATE kite.key SET deleted_at = (now() at time zone 'utc')`)
	if err != nil {
		return err
	}

	_, err = res.RowsAffected()
	return err
}

func (p *Postgres) GetKeyFromID(id string) (*KeyPair, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.
		Select("id", "public", "private").
		From("kite.key").
		Where(map[string]interface{}{
		"id":         id,
		"deleted_at": nil,
	}).ToSql()
	if err != nil {
		return nil, err
	}

	keyPair := &KeyPair{}
	err = p.DB.QueryRow(sqlQuery, args...).Scan(&keyPair.ID, &keyPair.Public, &keyPair.Private)
	if err != nil {
		return nil, err
	}

	return keyPair, nil
}

func (p *Postgres) GetKeyFromPublic(public string) (*KeyPair, error) {
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	sqlQuery, args, err := psql.
		Select("id", "public", "private").
		From("kite.key").
		Where(map[string]interface{}{
		"public":     public,
		"deleted_at": nil,
	}).ToSql()
	if err != nil {
		return nil, err
	}

	keyPair := &KeyPair{}
	err = p.DB.QueryRow(sqlQuery, args...).Scan(&keyPair.ID, &keyPair.Public, &keyPair.Private)
	if err != nil {
		return nil, err
	}

	return keyPair, nil
}
<|endoftext|>