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

package github

import (
	"encoding/json"
	"fmt"
)

// Deployment represents a deployment in a repo
type Deployment struct {
	Url         *string         `json:"url,omitempty"`
	ID          *int            `json:"id,omitempty"`
	SHA         *string         `json:"sha,omitempty"`
	Ref         *string         `json:"ref,omitempty"`
	Task        *string         `json:"task,omitempty"`
	Payload     json.RawMessage `json:"payload,omitempty"`
	Environment *string         `json:"environment,omitempty"`
	Description *string         `json:"description,omitempty"`
	Creator     *User           `json:"creator,omitempty"`
	CreatedAt   *Timestamp      `json:"created_at,omitempty"`
	UpdatedAt   *Timestamp      `json:"pushed_at,omitempty"`
}

// DeploymentRequest represents a deployment request
type DeploymentRequest struct {
	Ref              *string  `json:"ref,omitempty"`
	Task             *string  `json:"task,omitempty"`
	AutoMerge        *bool    `json:"auto_merge,omitempty"`
	RequiredContexts []string `json:"required_contexts,omitempty"`
	Payload          *string  `json:"payload,omitempty"`
	Environment      *string  `json:"environment,omitempty"`
	Description      *string  `json:"description,omitempty"`
}

// DeploymentsListOptions specifies the optional parameters to the
// RepositoriesService.ListDeployments method.
type DeploymentsListOptions struct {
	// SHA of the Deployment.
	SHA string `url:"sha,omitempty"`

	// List deployments for a given ref.
	Ref string `url:"ref,omitempty"`

	// List deployments for a given task.
	Task string `url:"task,omitempty"`

	// List deployments for a given environment.
	Environment string `url:"environment,omitempty"`

	ListOptions
}

// ListDeployments lists the deployments of a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#list-deployments
func (s *RepositoriesService) ListDeployments(owner, repo string, opt *DeploymentsListOptions) ([]Deployment, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments", owner, repo)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	deployments := new([]Deployment)
	resp, err := s.client.Do(req, deployments)
	if err != nil {
		return nil, resp, err
	}

	return *deployments, resp, err
}

// CreateDeployment creates a new deployment for a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#create-a-deployment
func (s *RepositoriesService) CreateDeployment(owner, repo string, request *DeploymentRequest) (*Deployment, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments", owner, repo)

	req, err := s.client.NewRequest("POST", u, request)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	d := new(Deployment)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// DeploymentStatus represents the status of a
// particular deployment.
type DeploymentStatus struct {
	ID          *int       `json:"id,omitempty"`
	State       *string    `json:"state,omitempty"`
	Creator     *User      `json:"creator,omitempty"`
	Description *string    `json:"description,omitempty"`
	TargetUrl   *string    `json:"target_url,omitempty"`
	CreatedAt   *Timestamp `json:"created_at,omitempty"`
	UpdatedAt   *Timestamp `json:"pushed_at,omitempty"`
}

// DeploymentRequest represents a deployment request
type DeploymentStatusRequest struct {
	State       *string `json:"state,omitempty"`
	TargetUrl   *string `json:"target_url,omitempty"`
	Description *string `json:"description,omitempty"`
}

// ListDeploymentStatuses lists the statuses of a given deployment of a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#list-deployment-statuses
func (s *RepositoriesService) ListDeploymentStatuses(owner, repo string, deployment int, opt *ListOptions) ([]DeploymentStatus, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments/%v/statuses", owner, repo, deployment)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	statuses := new([]DeploymentStatus)
	resp, err := s.client.Do(req, statuses)
	if err != nil {
		return nil, resp, err
	}

	return *statuses, resp, err
}

// CreateDeploymentStatus creates a new status for a deployment.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#create-a-deployment-status
func (s *RepositoriesService) CreateDeploymentStatus(owner, repo string, deployment int, request *DeploymentStatusRequest) (*DeploymentStatus, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments/%v/statuses", owner, repo, deployment)

	req, err := s.client.NewRequest("POST", u, request)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	d := new(DeploymentStatus)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}
<commit_msg>fix minor golint issues<commit_after>// Copyright 2014 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"encoding/json"
	"fmt"
)

// Deployment represents a deployment in a repo
type Deployment struct {
	URL         *string         `json:"url,omitempty"`
	ID          *int            `json:"id,omitempty"`
	SHA         *string         `json:"sha,omitempty"`
	Ref         *string         `json:"ref,omitempty"`
	Task        *string         `json:"task,omitempty"`
	Payload     json.RawMessage `json:"payload,omitempty"`
	Environment *string         `json:"environment,omitempty"`
	Description *string         `json:"description,omitempty"`
	Creator     *User           `json:"creator,omitempty"`
	CreatedAt   *Timestamp      `json:"created_at,omitempty"`
	UpdatedAt   *Timestamp      `json:"pushed_at,omitempty"`
}

// DeploymentRequest represents a deployment request
type DeploymentRequest struct {
	Ref              *string  `json:"ref,omitempty"`
	Task             *string  `json:"task,omitempty"`
	AutoMerge        *bool    `json:"auto_merge,omitempty"`
	RequiredContexts []string `json:"required_contexts,omitempty"`
	Payload          *string  `json:"payload,omitempty"`
	Environment      *string  `json:"environment,omitempty"`
	Description      *string  `json:"description,omitempty"`
}

// DeploymentsListOptions specifies the optional parameters to the
// RepositoriesService.ListDeployments method.
type DeploymentsListOptions struct {
	// SHA of the Deployment.
	SHA string `url:"sha,omitempty"`

	// List deployments for a given ref.
	Ref string `url:"ref,omitempty"`

	// List deployments for a given task.
	Task string `url:"task,omitempty"`

	// List deployments for a given environment.
	Environment string `url:"environment,omitempty"`

	ListOptions
}

// ListDeployments lists the deployments of a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#list-deployments
func (s *RepositoriesService) ListDeployments(owner, repo string, opt *DeploymentsListOptions) ([]Deployment, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments", owner, repo)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	deployments := new([]Deployment)
	resp, err := s.client.Do(req, deployments)
	if err != nil {
		return nil, resp, err
	}

	return *deployments, resp, err
}

// CreateDeployment creates a new deployment for a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#create-a-deployment
func (s *RepositoriesService) CreateDeployment(owner, repo string, request *DeploymentRequest) (*Deployment, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments", owner, repo)

	req, err := s.client.NewRequest("POST", u, request)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	d := new(Deployment)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// DeploymentStatus represents the status of a
// particular deployment.
type DeploymentStatus struct {
	ID          *int       `json:"id,omitempty"`
	State       *string    `json:"state,omitempty"`
	Creator     *User      `json:"creator,omitempty"`
	Description *string    `json:"description,omitempty"`
	TargetURL   *string    `json:"target_url,omitempty"`
	CreatedAt   *Timestamp `json:"created_at,omitempty"`
	UpdatedAt   *Timestamp `json:"pushed_at,omitempty"`
}

// DeploymentStatusRequest represents a deployment request
type DeploymentStatusRequest struct {
	State       *string `json:"state,omitempty"`
	TargetURL   *string `json:"target_url,omitempty"`
	Description *string `json:"description,omitempty"`
}

// ListDeploymentStatuses lists the statuses of a given deployment of a repository.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#list-deployment-statuses
func (s *RepositoriesService) ListDeploymentStatuses(owner, repo string, deployment int, opt *ListOptions) ([]DeploymentStatus, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments/%v/statuses", owner, repo, deployment)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	statuses := new([]DeploymentStatus)
	resp, err := s.client.Do(req, statuses)
	if err != nil {
		return nil, resp, err
	}

	return *statuses, resp, err
}

// CreateDeploymentStatus creates a new status for a deployment.
//
// GitHub API docs: https://developer.github.com/v3/repos/deployments/#create-a-deployment-status
func (s *RepositoriesService) CreateDeploymentStatus(owner, repo string, deployment int, request *DeploymentStatusRequest) (*DeploymentStatus, *Response, error) {
	u := fmt.Sprintf("repos/%v/%v/deployments/%v/statuses", owner, repo, deployment)

	req, err := s.client.NewRequest("POST", u, request)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches
	req.Header.Set("Accept", mediaTypeDeploymentPreview)

	d := new(DeploymentStatus)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}
<|endoftext|>


<commit_before>package routes

import (
	"net/http"

	"github.com/generationtux/brizo/app/handlers/web"
	"github.com/go-zoo/bone"
)

// BuildRouter configures the application router
func BuildRouter() *bone.Mux {
	r := mainRoutes()
	r.SubRoute("/api", apiRoutes())

	return r
}

// mainRoutes registers the routes for the user interface
func mainRoutes() *bone.Mux {
	router := bone.New()

	// Home page
	router.GetFunc("/", web.RootHandler)

	// Javascript UI
	router.GetFunc("/app", web.UIHandler)
	router.GetFunc("/app/*", web.UIHandler)
	router.Get("/assets/", http.StripPrefix("/assets/", http.FileServer(http.Dir("ui/dist/"))))

	// Healthz endpoint
	router.GetFunc("/healthz", web.HealthzHandler)

	// OAuth endpoints
	router.GetFunc("/o/auth/login/github", web.AuthGithubHandler)
	router.GetFunc("/o/auth/callback/github", web.AuthGithubCallbackHandler)

	// @todo move to JS UI
	router.GetFunc("/users", web.AuthAddNewUser)
	router.GetFunc("/login", web.AuthMainHandler)

	return router
}
<commit_msg>v1 prefix<commit_after>package routes

import (
	"net/http"

	"github.com/generationtux/brizo/app/handlers/web"
	"github.com/go-zoo/bone"
)

// BuildRouter configures the application router
func BuildRouter() *bone.Mux {
	r := mainRoutes()
	r.SubRoute("/api/v1", apiRoutes())

	return r
}

// mainRoutes registers the routes for the user interface
func mainRoutes() *bone.Mux {
	router := bone.New()

	// Home page
	router.GetFunc("/", web.RootHandler)

	// Javascript UI
	router.GetFunc("/app", web.UIHandler)
	router.GetFunc("/app/*", web.UIHandler)
	router.Get("/assets/", http.StripPrefix("/assets/", http.FileServer(http.Dir("ui/dist/"))))

	// Healthz endpoint
	router.GetFunc("/healthz", web.HealthzHandler)

	// OAuth endpoints
	router.GetFunc("/o/auth/login/github", web.AuthGithubHandler)
	router.GetFunc("/o/auth/callback/github", web.AuthGithubCallbackHandler)

	// @todo move to JS UI
	router.GetFunc("/users", web.AuthAddNewUser)
	router.GetFunc("/login", web.AuthMainHandler)

	return router
}
<|endoftext|>


<commit_before>package chat

import (
	"crypto/rand"
	"errors"
	"io"

	"github.com/agl/ed25519"
	"github.com/keybase/client/go/chat/signencrypt"
	"github.com/keybase/client/go/libkb"
)

type Encrypter interface {
	// EncryptedLen returns the number of bytes that the ciphertext of
	// size plaintext bytes will be.
	EncryptedLen(size int) int

	// Encrypt takes a plaintext reader and returns a ciphertext reader.
	Encrypt(plaintext io.Reader) (ciphertext io.Reader)

	// EncryptKey returns the ephemeral key that was used during Encrypt.
	EncryptKey() []byte

	// VerifyKey returns the public portion of the signing key that
	// can be used for signature verification.
	VerifyKey() []byte
}

type Decrypter interface {
	// Decrypt takes a ciphertext reader, encryption and verify keys.
	// It returns a plaintext reader.
	Decrypt(ciphertext io.Reader, encKey, verifyKey []byte) (plaintext io.Reader)
}

var nonce = &[signencrypt.NonceSize]byte{
	0x10, 0x10, 0x10, 0x10, 0x20, 0x20, 0x20, 0x20,
	0x30, 0x30, 0x30, 0x30, 0x40, 0x40, 0x40, 0x40,
}

type SignEncrypter struct {
	encKey    signencrypt.SecretboxKey
	signKey   signencrypt.SignKey
	verifyKey signencrypt.VerifyKey
}

func NewSignEncrypter() (*SignEncrypter, error) {
	var encKey [signencrypt.SecretboxKeySize]byte
	n, err := rand.Read(encKey[:])
	if err != nil {
		return nil, err
	}
	if n != signencrypt.SecretboxKeySize {
		return nil, errors.New("failed to rand.Read the correct number of bytes")
	}

	sign, err := libkb.GenerateNaclSigningKeyPair()
	if err != nil {
		return nil, err
	}

	var signKey [ed25519.PrivateKeySize]byte
	copy(signKey[:], (*sign.Private)[:])
	var verifyKey [ed25519.PublicKeySize]byte
	copy(verifyKey[:], sign.Public[:])

	res := SignEncrypter{
		encKey:    &encKey,
		signKey:   &signKey,
		verifyKey: &verifyKey,
	}

	return &res, nil
}

func (s *SignEncrypter) EncryptedLen(size int) int {
	return signencrypt.GetSealedSize(size)
}

func (s *SignEncrypter) Encrypt(r io.Reader) io.Reader {
	return signencrypt.NewEncodingReader(s.encKey, s.signKey, nonce, r)
}

func (s *SignEncrypter) EncryptKey() []byte {
	return []byte((*s.encKey)[:])
}

func (s *SignEncrypter) VerifyKey() []byte {
	return []byte((*s.verifyKey)[:])
}

type SignDecrypter struct{}

func NewSignDecrypter() *SignDecrypter {
	return &SignDecrypter{}
}

func (s *SignDecrypter) Decrypt(r io.Reader, encKey, verifyKey []byte) io.Reader {
	var xencKey [signencrypt.SecretboxKeySize]byte
	copy(xencKey[:], encKey)
	var xverifyKey [ed25519.PublicKeySize]byte
	copy(xverifyKey[:], verifyKey)
	return signencrypt.NewDecodingReader(&xencKey, &xverifyKey, nonce, r)
}
<commit_msg>Change nonce<commit_after>package chat

import (
	"crypto/rand"
	"errors"
	"io"

	"github.com/agl/ed25519"
	"github.com/keybase/client/go/chat/signencrypt"
	"github.com/keybase/client/go/libkb"
)

type Encrypter interface {
	// EncryptedLen returns the number of bytes that the ciphertext of
	// size plaintext bytes will be.
	EncryptedLen(size int) int

	// Encrypt takes a plaintext reader and returns a ciphertext reader.
	Encrypt(plaintext io.Reader) (ciphertext io.Reader)

	// EncryptKey returns the ephemeral key that was used during Encrypt.
	EncryptKey() []byte

	// VerifyKey returns the public portion of the signing key that
	// can be used for signature verification.
	VerifyKey() []byte
}

type Decrypter interface {
	// Decrypt takes a ciphertext reader, encryption and verify keys.
	// It returns a plaintext reader.
	Decrypt(ciphertext io.Reader, encKey, verifyKey []byte) (plaintext io.Reader)
}

var nonce signencrypt.Nonce

func init() {
	var n [signencrypt.NonceSize]byte
	copy(n[:], "kbchatattachment")
	nonce = &n
}

type SignEncrypter struct {
	encKey    signencrypt.SecretboxKey
	signKey   signencrypt.SignKey
	verifyKey signencrypt.VerifyKey
}

func NewSignEncrypter() (*SignEncrypter, error) {
	var encKey [signencrypt.SecretboxKeySize]byte
	n, err := rand.Read(encKey[:])
	if err != nil {
		return nil, err
	}
	if n != signencrypt.SecretboxKeySize {
		return nil, errors.New("failed to rand.Read the correct number of bytes")
	}

	sign, err := libkb.GenerateNaclSigningKeyPair()
	if err != nil {
		return nil, err
	}

	var signKey [ed25519.PrivateKeySize]byte
	copy(signKey[:], (*sign.Private)[:])
	var verifyKey [ed25519.PublicKeySize]byte
	copy(verifyKey[:], sign.Public[:])

	res := SignEncrypter{
		encKey:    &encKey,
		signKey:   &signKey,
		verifyKey: &verifyKey,
	}

	return &res, nil
}

func (s *SignEncrypter) EncryptedLen(size int) int {
	return signencrypt.GetSealedSize(size)
}

func (s *SignEncrypter) Encrypt(r io.Reader) io.Reader {
	return signencrypt.NewEncodingReader(s.encKey, s.signKey, nonce, r)
}

func (s *SignEncrypter) EncryptKey() []byte {
	return []byte((*s.encKey)[:])
}

func (s *SignEncrypter) VerifyKey() []byte {
	return []byte((*s.verifyKey)[:])
}

type SignDecrypter struct{}

func NewSignDecrypter() *SignDecrypter {
	return &SignDecrypter{}
}

func (s *SignDecrypter) Decrypt(r io.Reader, encKey, verifyKey []byte) io.Reader {
	var xencKey [signencrypt.SecretboxKeySize]byte
	copy(xencKey[:], encKey)
	var xverifyKey [ed25519.PublicKeySize]byte
	copy(xverifyKey[:], verifyKey)
	return signencrypt.NewDecodingReader(&xencKey, &xverifyKey, nonce, r)
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

// CmdDeviceAdd is the 'device add' command.  It is used for
// device provisioning on the provisioner/device X/C1.
type CmdDeviceAdd struct {
	libkb.Contextified
}

const cmdDevAddDesc = `When you are adding a new device to your account and you have an
existing device, you will be prompted to use this command on your
existing device to authorize the new device.`

// NewCmdDeviceAdd creates a new cli.Command.
func NewCmdDeviceAdd(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:        "add",
		Usage:       "Authorize a new device",
		Description: cmdDevAddDesc,
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdDeviceAdd{Contextified: libkb.NewContextified(g)}, "add", c)
		},
	}
}

// RunClient runs the command in client/server mode.
func (c *CmdDeviceAdd) Run() error {
	var err error
	cli, err := GetDeviceClient()
	if err != nil {
		return err
	}
	protocols := []rpc.Protocol{
		NewProvisionUIProtocol(c.G(), libkb.KexRoleProvisioner),
		NewSecretUIProtocol(c.G()),
	}
	if err := RegisterProtocols(protocols); err != nil {
		return err
	}

	return cli.DeviceXAdd(context.TODO(), 0)
}

// ParseArgv gets the secret phrase from the command args.
func (c *CmdDeviceAdd) ParseArgv(ctx *cli.Context) error {
	if len(ctx.Args()) != 0 {
		return fmt.Errorf("device xadd takes zero arguments")
	}
	return nil
}

// GetUsage says what this command needs to operate.
func (c *CmdDeviceAdd) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:    true,
		KbKeyring: true,
		API:       true,
	}
}
<commit_msg>Fix xadd -> add in error<commit_after>package client

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

// CmdDeviceAdd is the 'device add' command.  It is used for
// device provisioning on the provisioner/device X/C1.
type CmdDeviceAdd struct {
	libkb.Contextified
}

const cmdDevAddDesc = `When you are adding a new device to your account and you have an
existing device, you will be prompted to use this command on your
existing device to authorize the new device.`

// NewCmdDeviceAdd creates a new cli.Command.
func NewCmdDeviceAdd(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:        "add",
		Usage:       "Authorize a new device",
		Description: cmdDevAddDesc,
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdDeviceAdd{Contextified: libkb.NewContextified(g)}, "add", c)
		},
	}
}

// RunClient runs the command in client/server mode.
func (c *CmdDeviceAdd) Run() error {
	var err error
	cli, err := GetDeviceClient()
	if err != nil {
		return err
	}
	protocols := []rpc.Protocol{
		NewProvisionUIProtocol(c.G(), libkb.KexRoleProvisioner),
		NewSecretUIProtocol(c.G()),
	}
	if err := RegisterProtocols(protocols); err != nil {
		return err
	}

	return cli.DeviceXAdd(context.TODO(), 0)
}

// ParseArgv gets the secret phrase from the command args.
func (c *CmdDeviceAdd) ParseArgv(ctx *cli.Context) error {
	if len(ctx.Args()) != 0 {
		return fmt.Errorf("device add takes zero arguments")
	}
	return nil
}

// GetUsage says what this command needs to operate.
func (c *CmdDeviceAdd) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:    true,
		KbKeyring: true,
		API:       true,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"sort"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/olekukonko/tablewriter"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/vitessdriver"
	"github.com/youtube/vitess/go/vt/vterrors"
	"github.com/youtube/vitess/go/vt/vtgate/vtgateconn"

	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
)

var (
	usage = `
vtclient connects to a vtgate server using the standard go driver API.
Version 3 of the API is used, we do not send any hint to the server.

For query bound variables, we assume place-holders in the query string
in the form of :v1, :v2, etc.

Examples:

  $ vtclient -server vtgate:15991 "SELECT * FROM messages"

  $ vtclient -server vtgate:15991 -target '@master' -bind_variables '[ 12345, 1, "msg 12345" ]' "INSERT INTO messages (page,time_created_ns,message) VALUES (:v1, :v2, :v3)"

`
	server        = flag.String("server", "", "vtgate server to connect to")
	timeout       = flag.Duration("timeout", 30*time.Second, "timeout for queries")
	streaming     = flag.Bool("streaming", false, "use a streaming query")
	bindVariables = newBindvars("bind_variables", "bind variables as a json list")
	targetString  = flag.String("target", "", "keyspace:shard@tablet_type")
	jsonOutput    = flag.Bool("json", false, "Output JSON instead of human-readable table")
	parallel      = flag.Int("parallel", 1, "DMLs only: Number of threads executing the same query in parallel. Useful for simple load testing.")
	count         = flag.Int("count", 1, "DMLs only: Number of times each thread executes the query. Useful for simple, sustained load testing.")
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, usage)
	}
}

type bindvars []interface{}

func (bv *bindvars) String() string {
	b, err := json.Marshal(bv)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

func (bv *bindvars) Set(s string) (err error) {
	err = json.Unmarshal([]byte(s), &bv)
	if err != nil {
		return err
	}
	// json reads all numbers as float64
	// So, we just ditch floats for bindvars
	for i, v := range *bv {
		if f, ok := v.(float64); ok {
			if f > 0 {
				(*bv)[i] = uint64(f)
			} else {
				(*bv)[i] = int64(f)
			}
		}
	}

	return nil
}

// For internal flag compatibility
func (bv *bindvars) Get() interface{} {
	return bv
}

func newBindvars(name, usage string) *bindvars {
	var bv bindvars
	flag.Var(&bv, name, usage)
	return &bv
}

func main() {
	defer logutil.Flush()

	qr, err := run()
	if *jsonOutput && qr != nil {
		data, err := json.MarshalIndent(qr, "", "  ")
		if err != nil {
			log.Exitf("cannot marshal data: %v", err)
		}
		fmt.Print(string(data))
		return
	}

	qr.print()

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

func run() (*results, error) {
	flag.Parse()
	args := flag.Args()

	if len(args) == 0 {
		flag.Usage()
		return nil, errors.New("no arguments provided. See usage above")
	}
	if len(args) > 1 {
		return nil, errors.New("no additional arguments after the query allowed")
	}

	c := vitessdriver.Configuration{
		Protocol:  *vtgateconn.VtgateProtocol,
		Address:   *server,
		Target:    *targetString,
		Timeout:   *timeout,
		Streaming: *streaming,
	}
	db, err := vitessdriver.OpenWithConfiguration(c)
	if err != nil {
		return nil, fmt.Errorf("client error: %v", err)
	}

	log.Infof("Sending the query...")

	return execMulti(db, args[0])
}

func execMulti(db *sql.DB, sql string) (*results, error) {
	all := newResults()
	ec := concurrency.FirstErrorRecorder{}
	wg := sync.WaitGroup{}
	isDML := sqlparser.IsDML(sql)

	start := time.Now()
	for i := 0; i < *parallel; i++ {
		wg.Add(1)

		go func() {
			defer wg.Done()

			for j := 0; j < *count; j++ {
				var qr *results
				var err error
				if isDML {
					qr, err = execDml(db, sql)
				} else {
					qr, err = execNonDml(db, sql)
				}
				if *count == 1 {
					all = qr
				} else {
					all.merge(qr)
				}
				if err != nil {
					ec.RecordError(err)
					all.recordError(err)
					// We keep going and do not return early purpose.
				}
			}
		}()
	}
	wg.Wait()
	all.duration = time.Since(start)

	return all, ec.Error()
}

func execDml(db *sql.DB, sql string) (*results, error) {
	start := time.Now()
	tx, err := db.Begin()
	if err != nil {
		return nil, vterrors.Wrap(err, "BEGIN failed")
	}

	result, err := tx.Exec(sql, []interface{}(*bindVariables)...)
	if err != nil {
		return nil, vterrors.Wrap(err, "failed to execute DML")
	}

	err = tx.Commit()
	if err != nil {
		return nil, vterrors.Wrap(err, "COMMIT failed")
	}

	rowsAffected, err := result.RowsAffected()
	lastInsertID, err := result.LastInsertId()
	return &results{
		rowsAffected: rowsAffected,
		lastInsertID: lastInsertID,
		duration:     time.Since(start),
	}, nil
}

func execNonDml(db *sql.DB, sql string) (*results, error) {
	start := time.Now()
	rows, err := db.Query(sql, []interface{}(*bindVariables)...)
	if err != nil {
		return nil, vterrors.Wrap(err, "client error")
	}
	defer rows.Close()

	// get the headers
	var qr results
	cols, err := rows.Columns()
	if err != nil {
		return nil, fmt.Errorf("client error: %v", err)
	}
	qr.Fields = cols

	// get the rows
	for rows.Next() {
		row := make([]interface{}, len(cols))
		for i := range row {
			var col string
			row[i] = &col
		}
		if err := rows.Scan(row...); err != nil {
			return nil, fmt.Errorf("client error: %v", err)
		}

		// unpack []*string into []string
		vals := make([]string, 0, len(row))
		for _, value := range row {
			vals = append(vals, *(value.(*string)))
		}
		qr.Rows = append(qr.Rows, vals)
	}
	qr.rowsAffected = int64(len(qr.Rows))

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("Vitess returned an error: %v", err)
	}

	qr.duration = time.Since(start)
	return &qr, nil
}

type results struct {
	mu                 sync.Mutex
	Fields             []string   `json:"fields"`
	Rows               [][]string `json:"rows"`
	rowsAffected       int64
	lastInsertID       int64
	duration           time.Duration
	cumulativeDuration time.Duration

	// Multi DML mode: Track total error count, error count per code and the first error.
	totalErrorCount int
	errorCount      map[vtrpcpb.Code]int
	firstError      map[vtrpcpb.Code]error
}

func newResults() *results {
	return &results{
		errorCount: make(map[vtrpcpb.Code]int),
		firstError: make(map[vtrpcpb.Code]error),
	}
}

// merge aggregates "other" into "r".
// This is only used for executing DMLs concurrently and repeatedly.
// Therefore, "Fields" and "Rows" are not merged.
func (r *results) merge(other *results) {
	if other == nil {
		return
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	r.rowsAffected += other.rowsAffected
	if other.lastInsertID > r.lastInsertID {
		r.lastInsertID = other.lastInsertID
	}
	r.cumulativeDuration += other.duration
}

func (r *results) recordError(err error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.totalErrorCount++
	code := vterrors.Code(err)
	r.errorCount[code]++

	if r.errorCount[code] == 1 {
		r.firstError[code] = err
	}
}

func (r *results) print() {
	if r == nil {
		return
	}

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader(r.Fields)
	table.SetAutoFormatHeaders(false)
	table.AppendBulk(r.Rows)
	table.Render()
	fmt.Printf("%v row(s) affected (%v, cum: %v)\n", r.rowsAffected, r.duration, r.cumulativeDuration)
	if r.lastInsertID != 0 {
		fmt.Printf("Last insert ID: %v\n", r.lastInsertID)
	}

	if r.totalErrorCount == 0 {
		return
	}

	fmt.Printf("%d error(s) were returned. Number of errors by error code:\n\n", r.totalErrorCount)
	// Sort different error codes by count (descending).
	type errorCounts struct {
		code  vtrpcpb.Code
		count int
	}
	var counts []errorCounts
	for code, count := range r.errorCount {
		counts = append(counts, errorCounts{code, count})
	}
	sort.Slice(counts, func(i, j int) bool { return counts[i].count >= counts[j].count })
	for _, c := range counts {
		fmt.Printf("%- 30v= % 5d\n", c.code, c.count)
	}

	fmt.Printf("\nFirst error per code:\n\n")
	for code, err := range r.firstError {
		fmt.Printf("Code:  %v\nError: %v\n\n", code, err)
	}
}
<commit_msg>vtclient: Feature to generate queries with random IDs<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"sort"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/olekukonko/tablewriter"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/vitessdriver"
	"github.com/youtube/vitess/go/vt/vterrors"
	"github.com/youtube/vitess/go/vt/vtgate/vtgateconn"

	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
)

var (
	usage = `
vtclient connects to a vtgate server using the standard go driver API.
Version 3 of the API is used, we do not send any hint to the server.

For query bound variables, we assume place-holders in the query string
in the form of :v1, :v2, etc.

Examples:

  $ vtclient -server vtgate:15991 "SELECT * FROM messages"

  $ vtclient -server vtgate:15991 -target '@master' -bind_variables '[ 12345, 1, "msg 12345" ]' "INSERT INTO messages (page,time_created_ns,message) VALUES (:v1, :v2, :v3)"

`
	server        = flag.String("server", "", "vtgate server to connect to")
	timeout       = flag.Duration("timeout", 30*time.Second, "timeout for queries")
	streaming     = flag.Bool("streaming", false, "use a streaming query")
	bindVariables = newBindvars("bind_variables", "bind variables as a json list")
	targetString  = flag.String("target", "", "keyspace:shard@tablet_type")
	jsonOutput    = flag.Bool("json", false, "Output JSON instead of human-readable table")
	parallel      = flag.Int("parallel", 1, "DMLs only: Number of threads executing the same query in parallel. Useful for simple load testing.")
	count         = flag.Int("count", 1, "DMLs only: Number of times each thread executes the query. Useful for simple, sustained load testing.")
	minRandomID   = flag.Int("min_random_id", 0, "min random ID to generate. When max_random_id > min_random_id, for each query, a random number is generated in [min_random_id, max_random_id) and attached to the end of the bind variables.")
	maxRandomID   = flag.Int("max_random_id", 0, "max random ID.")
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, usage)
	}
}

type bindvars []interface{}

func (bv *bindvars) String() string {
	b, err := json.Marshal(bv)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

func (bv *bindvars) Set(s string) (err error) {
	err = json.Unmarshal([]byte(s), &bv)
	if err != nil {
		return err
	}
	// json reads all numbers as float64
	// So, we just ditch floats for bindvars
	for i, v := range *bv {
		if f, ok := v.(float64); ok {
			if f > 0 {
				(*bv)[i] = uint64(f)
			} else {
				(*bv)[i] = int64(f)
			}
		}
	}

	return nil
}

// For internal flag compatibility
func (bv *bindvars) Get() interface{} {
	return bv
}

func newBindvars(name, usage string) *bindvars {
	var bv bindvars
	flag.Var(&bv, name, usage)
	return &bv
}

func main() {
	defer logutil.Flush()

	qr, err := run()
	if *jsonOutput && qr != nil {
		data, err := json.MarshalIndent(qr, "", "  ")
		if err != nil {
			log.Exitf("cannot marshal data: %v", err)
		}
		fmt.Print(string(data))
		return
	}

	qr.print()

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

func run() (*results, error) {
	flag.Parse()
	args := flag.Args()

	if len(args) == 0 {
		flag.Usage()
		return nil, errors.New("no arguments provided. See usage above")
	}
	if len(args) > 1 {
		return nil, errors.New("no additional arguments after the query allowed")
	}

	c := vitessdriver.Configuration{
		Protocol:  *vtgateconn.VtgateProtocol,
		Address:   *server,
		Target:    *targetString,
		Timeout:   *timeout,
		Streaming: *streaming,
	}
	db, err := vitessdriver.OpenWithConfiguration(c)
	if err != nil {
		return nil, fmt.Errorf("client error: %v", err)
	}

	log.Infof("Sending the query...")

	return execMulti(db, args[0])
}

func prepareBindVariables() []interface{} {
	bv := *bindVariables
	if *maxRandomID > *minRandomID {
		bv = append(bv, rand.Intn(*maxRandomID-*minRandomID)+*minRandomID)
	}
	return bv
}

func execMulti(db *sql.DB, sql string) (*results, error) {
	all := newResults()
	ec := concurrency.FirstErrorRecorder{}
	wg := sync.WaitGroup{}
	isDML := sqlparser.IsDML(sql)

	start := time.Now()
	for i := 0; i < *parallel; i++ {
		wg.Add(1)

		go func() {
			defer wg.Done()

			for j := 0; j < *count; j++ {
				var qr *results
				var err error
				if isDML {
					qr, err = execDml(db, sql)
				} else {
					qr, err = execNonDml(db, sql)
				}
				if *count == 1 {
					all = qr
				} else {
					all.merge(qr)
				}
				if err != nil {
					ec.RecordError(err)
					all.recordError(err)
					// We keep going and do not return early purpose.
				}
			}
		}()
	}
	wg.Wait()
	all.duration = time.Since(start)

	return all, ec.Error()
}

func execDml(db *sql.DB, sql string) (*results, error) {
	start := time.Now()
	tx, err := db.Begin()
	if err != nil {
		return nil, vterrors.Wrap(err, "BEGIN failed")
	}

	result, err := tx.Exec(sql, []interface{}(prepareBindVariables())...)
	if err != nil {
		return nil, vterrors.Wrap(err, "failed to execute DML")
	}

	err = tx.Commit()
	if err != nil {
		return nil, vterrors.Wrap(err, "COMMIT failed")
	}

	rowsAffected, err := result.RowsAffected()
	lastInsertID, err := result.LastInsertId()
	return &results{
		rowsAffected: rowsAffected,
		lastInsertID: lastInsertID,
		duration:     time.Since(start),
	}, nil
}

func execNonDml(db *sql.DB, sql string) (*results, error) {
	start := time.Now()
	rows, err := db.Query(sql, []interface{}(prepareBindVariables())...)
	if err != nil {
		return nil, vterrors.Wrap(err, "client error")
	}
	defer rows.Close()

	// get the headers
	var qr results
	cols, err := rows.Columns()
	if err != nil {
		return nil, fmt.Errorf("client error: %v", err)
	}
	qr.Fields = cols

	// get the rows
	for rows.Next() {
		row := make([]interface{}, len(cols))
		for i := range row {
			var col string
			row[i] = &col
		}
		if err := rows.Scan(row...); err != nil {
			return nil, fmt.Errorf("client error: %v", err)
		}

		// unpack []*string into []string
		vals := make([]string, 0, len(row))
		for _, value := range row {
			vals = append(vals, *(value.(*string)))
		}
		qr.Rows = append(qr.Rows, vals)
	}
	qr.rowsAffected = int64(len(qr.Rows))

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("Vitess returned an error: %v", err)
	}

	qr.duration = time.Since(start)
	return &qr, nil
}

type results struct {
	mu                 sync.Mutex
	Fields             []string   `json:"fields"`
	Rows               [][]string `json:"rows"`
	rowsAffected       int64
	lastInsertID       int64
	duration           time.Duration
	cumulativeDuration time.Duration

	// Multi DML mode: Track total error count, error count per code and the first error.
	totalErrorCount int
	errorCount      map[vtrpcpb.Code]int
	firstError      map[vtrpcpb.Code]error
}

func newResults() *results {
	return &results{
		errorCount: make(map[vtrpcpb.Code]int),
		firstError: make(map[vtrpcpb.Code]error),
	}
}

// merge aggregates "other" into "r".
// This is only used for executing DMLs concurrently and repeatedly.
// Therefore, "Fields" and "Rows" are not merged.
func (r *results) merge(other *results) {
	if other == nil {
		return
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	r.rowsAffected += other.rowsAffected
	if other.lastInsertID > r.lastInsertID {
		r.lastInsertID = other.lastInsertID
	}
	r.cumulativeDuration += other.duration
}

func (r *results) recordError(err error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.totalErrorCount++
	code := vterrors.Code(err)
	r.errorCount[code]++

	if r.errorCount[code] == 1 {
		r.firstError[code] = err
	}
}

func (r *results) print() {
	if r == nil {
		return
	}

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader(r.Fields)
	table.SetAutoFormatHeaders(false)
	table.AppendBulk(r.Rows)
	table.Render()
	fmt.Printf("%v row(s) affected (%v, cum: %v)\n", r.rowsAffected, r.duration, r.cumulativeDuration)
	if r.lastInsertID != 0 {
		fmt.Printf("Last insert ID: %v\n", r.lastInsertID)
	}

	if r.totalErrorCount == 0 {
		return
	}

	fmt.Printf("%d error(s) were returned. Number of errors by error code:\n\n", r.totalErrorCount)
	// Sort different error codes by count (descending).
	type errorCounts struct {
		code  vtrpcpb.Code
		count int
	}
	var counts []errorCounts
	for code, count := range r.errorCount {
		counts = append(counts, errorCounts{code, count})
	}
	sort.Slice(counts, func(i, j int) bool { return counts[i].count >= counts[j].count })
	for _, c := range counts {
		fmt.Printf("%- 30v= % 5d\n", c.code, c.count)
	}

	fmt.Printf("\nFirst error per code:\n\n")
	for code, err := range r.firstError {
		fmt.Printf("Code:  %v\nError: %v\n\n", code, err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

// vt tablet server: Serves queries and performs housekeeping jobs.
package main

import (
	"flag"

	"golang.org/x/net/context"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"
	"vitess.io/vitess/go/vt/servenv"
	"vitess.io/vitess/go/vt/tableacl"
	"vitess.io/vitess/go/vt/tableacl/simpleacl"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager"
	"vitess.io/vitess/go/vt/vttablet/tabletserver"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
)

var (
	enforceTableACLConfig        = flag.Bool("enforce-tableacl-config", false, "if this flag is true, vttablet will fail to start if a valid tableacl config does not exist")
	tableACLConfig               = flag.String("table-acl-config", "", "path to table access checker config file; send SIGHUP to reload this file")
	tableACLConfigReloadInterval = flag.Duration("table-acl-config-reload-interval", 0, "Ticker to reload ACLs")
	tabletPath                   = flag.String("tablet-path", "", "tablet alias")

	agent *tabletmanager.ActionAgent
)

func init() {
	servenv.RegisterDefaultFlags()
}

func main() {
	dbconfigs.RegisterFlags(dbconfigs.All...)
	mysqlctl.RegisterFlags()

	servenv.ParseFlags("vttablet")

	if err := tabletenv.NewCurrentConfig().Verify(); err != nil {
		log.Exitf("invalid config: %v", err)
	}

	tabletenv.Init()

	servenv.Init()

	if *tabletPath == "" {
		log.Exit("-tablet-path required")
	}
	tabletAlias, err := topoproto.ParseTabletAlias(*tabletPath)
	if err != nil {
		log.Exitf("failed to parse -tablet-path: %v", err)
	}

	var mycnf *mysqlctl.Mycnf
	var socketFile string
	// If no connection parameters were specified, load the mycnf file
	// and use the socket from it. If connection parameters were specified,
	// we assume that the mysql is not local, and we skip loading mycnf.
	// This also means that backup and restore will not be allowed.
	if !dbconfigs.HasConnectionParams() {
		var err error
		if mycnf, err = mysqlctl.NewMycnfFromFlags(tabletAlias.Uid); err != nil {
			log.Exitf("mycnf read failed: %v", err)
		}
		socketFile = mycnf.SocketFile
	} else {
		log.Info("connection parameters were specified. Not loading my.cnf.")
	}

	// If connection parameters were specified, socketFile will be empty.
	// Otherwise, the socketFile (read from mycnf) will be used to initialize
	// dbconfigs.
	dbcfgs, err := dbconfigs.Init(socketFile)
	if err != nil {
		log.Warning(err)
	}

	if *tableACLConfig != "" {
		// To override default simpleacl, other ACL plugins must set themselves to be default ACL factory
		tableacl.Register("simpleacl", &simpleacl.Factory{})
	} else if *enforceTableACLConfig {
		log.Exit("table acl config has to be specified with table-acl-config flag because enforce-tableacl-config is set.")
	}

	// creates and registers the query service
	ts := topo.Open()
	qsc := tabletserver.NewServer("", ts, *tabletAlias)
	servenv.OnRun(func() {
		qsc.Register()
		addStatusParts(qsc)
	})
	servenv.OnClose(func() {
		// We now leave the queryservice running during lameduck,
		// so stop it in OnClose(), after lameduck is over.
		qsc.StopService()
	})

	qsc.InitACL(*tableACLConfig, *enforceTableACLConfig, *tableACLConfigReloadInterval)

	// Create mysqld and register the health reporter (needs to be done
	// before initializing the agent, so the initial health check
	// done by the agent has the right reporter)
	mysqld := mysqlctl.NewMysqld(dbcfgs)
	servenv.OnClose(mysqld.Close)

	// Depends on both query and updateStream.
	gRPCPort := int32(0)
	if servenv.GRPCPort != nil {
		gRPCPort = int32(*servenv.GRPCPort)
	}
	agent, err = tabletmanager.NewActionAgent(context.Background(), ts, mysqld, qsc, tabletAlias, dbcfgs, mycnf, int32(*servenv.Port), gRPCPort)
	if err != nil {
		log.Exitf("NewActionAgent() failed: %v", err)
	}

	servenv.OnClose(func() {
		// Close the agent so that our topo entry gets pruned properly and any
		// background goroutines that use the topo connection are stopped.
		agent.Close()

		// We will still use the topo server during lameduck period
		// to update our state, so closing it in OnClose()
		ts.Close()
	})

	servenv.RunDefault()
}
<commit_msg>yaml: add -tablet_config to read yaml<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

// vt tablet server: Serves queries and performs housekeeping jobs.
package main

import (
	"flag"
	"io/ioutil"

	"golang.org/x/net/context"
	"sigs.k8s.io/yaml"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"
	"vitess.io/vitess/go/vt/servenv"
	"vitess.io/vitess/go/vt/tableacl"
	"vitess.io/vitess/go/vt/tableacl/simpleacl"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager"
	"vitess.io/vitess/go/vt/vttablet/tabletserver"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
)

var (
	enforceTableACLConfig        = flag.Bool("enforce-tableacl-config", false, "if this flag is true, vttablet will fail to start if a valid tableacl config does not exist")
	tableACLConfig               = flag.String("table-acl-config", "", "path to table access checker config file; send SIGHUP to reload this file")
	tableACLConfigReloadInterval = flag.Duration("table-acl-config-reload-interval", 0, "Ticker to reload ACLs")
	tabletPath                   = flag.String("tablet-path", "", "tablet alias")
	tabletConfig                 = flag.String("tablet_config", "", "YAML file config for tablet")

	agent *tabletmanager.ActionAgent
)

func init() {
	servenv.RegisterDefaultFlags()
}

func main() {
	dbconfigs.RegisterFlags(dbconfigs.All...)
	mysqlctl.RegisterFlags()

	servenv.ParseFlags("vttablet")

	config := tabletenv.NewCurrentConfig()
	if err := config.Verify(); err != nil {
		log.Exitf("invalid config: %v", err)
	}

	tabletenv.Init()
	if *tabletConfig != "" {
		bytes, err := ioutil.ReadFile(*tabletConfig)
		if err != nil {
			log.Exitf("error reading config file %s: %v", *tabletConfig, err)
		}
		if err := yaml.Unmarshal(bytes, config); err != nil {
			log.Exitf("error parsing config file %s: %v", bytes, err)
		}
		gotBytes, _ := yaml.Marshal(config)
		log.Infof("Loaded config file %s successfully:\n%s", *tabletConfig, gotBytes)
	}

	servenv.Init()

	if *tabletPath == "" {
		log.Exit("-tablet-path required")
	}
	tabletAlias, err := topoproto.ParseTabletAlias(*tabletPath)
	if err != nil {
		log.Exitf("failed to parse -tablet-path: %v", err)
	}

	var mycnf *mysqlctl.Mycnf
	var socketFile string
	// If no connection parameters were specified, load the mycnf file
	// and use the socket from it. If connection parameters were specified,
	// we assume that the mysql is not local, and we skip loading mycnf.
	// This also means that backup and restore will not be allowed.
	if !dbconfigs.HasConnectionParams() {
		var err error
		if mycnf, err = mysqlctl.NewMycnfFromFlags(tabletAlias.Uid); err != nil {
			log.Exitf("mycnf read failed: %v", err)
		}
		socketFile = mycnf.SocketFile
	} else {
		log.Info("connection parameters were specified. Not loading my.cnf.")
	}

	// If connection parameters were specified, socketFile will be empty.
	// Otherwise, the socketFile (read from mycnf) will be used to initialize
	// dbconfigs.
	dbcfgs, err := dbconfigs.Init(socketFile)
	if err != nil {
		log.Warning(err)
	}

	if *tableACLConfig != "" {
		// To override default simpleacl, other ACL plugins must set themselves to be default ACL factory
		tableacl.Register("simpleacl", &simpleacl.Factory{})
	} else if *enforceTableACLConfig {
		log.Exit("table acl config has to be specified with table-acl-config flag because enforce-tableacl-config is set.")
	}

	// creates and registers the query service
	ts := topo.Open()
	qsc := tabletserver.NewTabletServer("", config, ts, *tabletAlias)
	servenv.OnRun(func() {
		qsc.Register()
		addStatusParts(qsc)
	})
	servenv.OnClose(func() {
		// We now leave the queryservice running during lameduck,
		// so stop it in OnClose(), after lameduck is over.
		qsc.StopService()
	})

	qsc.InitACL(*tableACLConfig, *enforceTableACLConfig, *tableACLConfigReloadInterval)

	// Create mysqld and register the health reporter (needs to be done
	// before initializing the agent, so the initial health check
	// done by the agent has the right reporter)
	mysqld := mysqlctl.NewMysqld(dbcfgs)
	servenv.OnClose(mysqld.Close)

	// Depends on both query and updateStream.
	gRPCPort := int32(0)
	if servenv.GRPCPort != nil {
		gRPCPort = int32(*servenv.GRPCPort)
	}
	agent, err = tabletmanager.NewActionAgent(context.Background(), ts, mysqld, qsc, tabletAlias, dbcfgs, mycnf, int32(*servenv.Port), gRPCPort)
	if err != nil {
		log.Exitf("NewActionAgent() failed: %v", err)
	}

	servenv.OnClose(func() {
		// Close the agent so that our topo entry gets pruned properly and any
		// background goroutines that use the topo connection are stopped.
		agent.Close()

		// We will still use the topo server during lameduck period
		// to update our state, so closing it in OnClose()
		ts.Close()
	})

	servenv.RunDefault()
}
<|endoftext|>


<commit_before>package systests

import (
	"bytes"
	"net/http"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/teams"
	"github.com/keybase/stellarnet"
	"github.com/stellar/go/build"
	"github.com/stellar/go/clients/horizon"
	"github.com/stretchr/testify/require"
)

func TestStellarNoteRoundtripAndResets(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	// Sign up two users, bob and alice.
	alice := ctx.installKeybaseForUser("alice", 10)
	alice.signup()
	divDebug(ctx, "Signed up alice (%s)", alice.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)

	t.Logf("note to self")
	encB64, err := stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), nil)
	require.NoError(t, err)
	note, err := stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("note to both users")
	other := bob.userVersion()
	encB64, err = stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), &other)
	require.NoError(t, err)

	t.Logf("decrypt as self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("reset sender")
	alice.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)
	alice.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	t.Logf("fail to decrypt as post-reset self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.Error(t, err)
	require.Equal(t, "note not encrypted for logged-in user", err.Error())

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)
}

// Test took 38s on a dev server 2018-06-07
func TestStellarRelayAutoClaims(t *testing.T) {
	testStellarRelayAutoClaims(t, false, false)
}

// Test took 29s on a dev server 2018-06-07
func TestStellarRelayAutoClaimsWithPUK(t *testing.T) {
	testStellarRelayAutoClaims(t, true, true)
}

// Part 1:
// XLM is sent to a user before they have a [PUK / wallet].
// In the form of multiple relay payments.
// They then [get a PUK,] add a wallet, and enter the impteam,
// which all kick the autoclaim into gear.
//
// Part 2:
// A relay payment is sent to the user who already has a wallet.
// The funds should be claimed asap.
//
// To debug this test use log filter "stellar_test|poll-|AutoClaim|stellar.claim|pollfor"
func testStellarRelayAutoClaims(t *testing.T, startWithPUK, skipPart2 bool) {
	tt := newTeamTester(t)
	defer tt.cleanup()
	useStellarTestNet(t)

	alice := tt.addUser("alice")
	var bob *userPlusDevice
	if startWithPUK {
		bob = tt.addWalletlessUser("bob")
	} else {
		bob = tt.addPuklessUser("bob")
	}
	alice.kickTeamRekeyd()

	t.Logf("alice gets funded")
	res, err := alice.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
	require.NoError(t, err)
	gift(t, res[0].AccountID)

	t.Logf("alice sends a first relay payment to bob P1")
	attachIdentifyUI(t, alice.tc.G, newSimpleIdentifyUI())
	cmd := client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "50",
	}
	require.NoError(t, cmd.Run())

	t.Logf("alice sends a second relay payment to bob P2")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "30",
	}
	require.NoError(t, cmd.Run())

	t.Logf("get the impteam seqno to wait on later")
	team, _, _, err := teams.LookupImplicitTeam(context.Background(), alice.tc.G, alice.username+","+bob.username, false)
	require.NoError(t, err)
	nextSeqno := team.NextSeqno()

	if startWithPUK {
		t.Logf("bob gets a wallet")
		bob.tc.Tp.DisableAutoWallet = false
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())
	} else {
		t.Logf("bob gets a PUK and wallet")
		bob.perUserKeyUpgrade()
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())

		t.Logf("wait for alice to add bob to their impteam")
		alice.pollForTeamSeqnoLinkWithLoadArgs(keybase1.LoadTeamArg{ID: team.ID}, nextSeqno)
	}

	pollTime := 10 * time.Second
	if libkb.UseCITime(bob.tc.G) {
		// This test is especially slow.
		pollTime = 15 * time.Second
	}

	pollFor(t, "claims to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-1-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "0 XLM" {
			return false
		}
		if res[0].BalanceDescription == "49.9999800 XLM" {
			t.Logf("poll-1-%v: received T1 but not T2", i)
			return false
		}
		if res[0].BalanceDescription == "29.9999800 XLM" {
			t.Logf("poll-1-%v: received T2 but not T1", i)
			return false
		}
		t.Logf("poll-1-%v: received both payments", i)
		require.Equal(t, "79.9999700 XLM", res[0].BalanceDescription)
		return true
	})

	if skipPart2 {
		t.Logf("Skipping part 2")
		return
	}

	t.Logf("--------------------")
	t.Logf("Part 2: Alice sends a relay payment to bob who now already has a wallet")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "10",
		ForceRelay:   true,
	}
	require.NoError(t, cmd.Run())

	pollFor(t, "final claim to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-2-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "79.9999700 XLM" {
			return false
		}
		t.Logf("poll-1-%v: received final payment", i)
		require.Equal(t, "89.9999600 XLM", res[0].BalanceDescription)
		return true
	})

}

func sampleNote() stellar1.NoteContents {
	return stellar1.NoteContents{
		Note:      "wizbang",
		StellarID: stellar1.TransactionID("6653fc2fdbc42ad51ccbe77ee0a3c29e258a5513c62fdc532cbfff91ab101abf"),
	}
}

// Friendbot sends someone XLM
func gift(t testing.TB, accountID stellar1.AccountID) {
	t.Logf("gift -> %v", accountID)
	url := "https://friendbot.stellar.org/?addr=" + accountID.String()
	t.Logf("gift url: %v", url)
	res, err := http.Get(url)
	require.NoError(t, err, "friendbot request error")
	bodyBuf := new(bytes.Buffer)
	bodyBuf.ReadFrom(res.Body)
	t.Logf("gift res: %v", bodyBuf.String())
	require.Equal(t, 200, res.StatusCode, "friendbot response status code")
}

func useStellarTestNet(t testing.TB) {
	stellarnet.SetClientAndNetwork(horizon.DefaultTestNetClient, build.TestNetwork)
}
<commit_msg>skip test<commit_after>package systests

import (
	"bytes"
	"net/http"
	"os"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/teams"
	"github.com/keybase/stellarnet"
	"github.com/stellar/go/build"
	"github.com/stellar/go/clients/horizon"
	"github.com/stretchr/testify/require"
)

func TestStellarNoteRoundtripAndResets(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	// Sign up two users, bob and alice.
	alice := ctx.installKeybaseForUser("alice", 10)
	alice.signup()
	divDebug(ctx, "Signed up alice (%s)", alice.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)

	t.Logf("note to self")
	encB64, err := stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), nil)
	require.NoError(t, err)
	note, err := stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("note to both users")
	other := bob.userVersion()
	encB64, err = stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), &other)
	require.NoError(t, err)

	t.Logf("decrypt as self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("reset sender")
	alice.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)
	alice.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	t.Logf("fail to decrypt as post-reset self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.Error(t, err)
	require.Equal(t, "note not encrypted for logged-in user", err.Error())

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)
}

// Test took 38s on a dev server 2018-06-07
func TestStellarRelayAutoClaims(t *testing.T) {
	testStellarRelayAutoClaims(t, false, false)
}

// Test took 29s on a dev server 2018-06-07
func TestStellarRelayAutoClaimsWithPUK(t *testing.T) {
	testStellarRelayAutoClaims(t, true, true)
}

// Part 1:
// XLM is sent to a user before they have a [PUK / wallet].
// In the form of multiple relay payments.
// They then [get a PUK,] add a wallet, and enter the impteam,
// which all kick the autoclaim into gear.
//
// Part 2:
// A relay payment is sent to the user who already has a wallet.
// The funds should be claimed asap.
//
// To debug this test use log filter "stellar_test|poll-|AutoClaim|stellar.claim|pollfor"
func testStellarRelayAutoClaims(t *testing.T, startWithPUK, skipPart2 bool) {
	if os.Getenv("UNSKIP_CORE_8044") != "1" {
		t.Skip("CORE-8044")
	}
	tt := newTeamTester(t)
	defer tt.cleanup()
	useStellarTestNet(t)

	alice := tt.addUser("alice")
	var bob *userPlusDevice
	if startWithPUK {
		bob = tt.addWalletlessUser("bob")
	} else {
		bob = tt.addPuklessUser("bob")
	}
	alice.kickTeamRekeyd()

	t.Logf("alice gets funded")
	res, err := alice.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
	require.NoError(t, err)
	gift(t, res[0].AccountID)

	t.Logf("alice sends a first relay payment to bob P1")
	attachIdentifyUI(t, alice.tc.G, newSimpleIdentifyUI())
	cmd := client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "50",
	}
	require.NoError(t, cmd.Run())

	t.Logf("alice sends a second relay payment to bob P2")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "30",
	}
	require.NoError(t, cmd.Run())

	t.Logf("get the impteam seqno to wait on later")
	team, _, _, err := teams.LookupImplicitTeam(context.Background(), alice.tc.G, alice.username+","+bob.username, false)
	require.NoError(t, err)
	nextSeqno := team.NextSeqno()

	if startWithPUK {
		t.Logf("bob gets a wallet")
		bob.tc.Tp.DisableAutoWallet = false
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())
	} else {
		t.Logf("bob gets a PUK and wallet")
		bob.perUserKeyUpgrade()
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())

		t.Logf("wait for alice to add bob to their impteam")
		alice.pollForTeamSeqnoLinkWithLoadArgs(keybase1.LoadTeamArg{ID: team.ID}, nextSeqno)
	}

	pollTime := 10 * time.Second
	if libkb.UseCITime(bob.tc.G) {
		// This test is especially slow.
		pollTime = 15 * time.Second
	}

	pollFor(t, "claims to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-1-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "0 XLM" {
			return false
		}
		if res[0].BalanceDescription == "49.9999800 XLM" {
			t.Logf("poll-1-%v: received T1 but not T2", i)
			return false
		}
		if res[0].BalanceDescription == "29.9999800 XLM" {
			t.Logf("poll-1-%v: received T2 but not T1", i)
			return false
		}
		t.Logf("poll-1-%v: received both payments", i)
		require.Equal(t, "79.9999700 XLM", res[0].BalanceDescription)
		return true
	})

	if skipPart2 {
		t.Logf("Skipping part 2")
		return
	}

	t.Logf("--------------------")
	t.Logf("Part 2: Alice sends a relay payment to bob who now already has a wallet")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "10",
		ForceRelay:   true,
	}
	require.NoError(t, cmd.Run())

	pollFor(t, "final claim to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-2-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "79.9999700 XLM" {
			return false
		}
		t.Logf("poll-1-%v: received final payment", i)
		require.Equal(t, "89.9999600 XLM", res[0].BalanceDescription)
		return true
	})

}

func sampleNote() stellar1.NoteContents {
	return stellar1.NoteContents{
		Note:      "wizbang",
		StellarID: stellar1.TransactionID("6653fc2fdbc42ad51ccbe77ee0a3c29e258a5513c62fdc532cbfff91ab101abf"),
	}
}

// Friendbot sends someone XLM
func gift(t testing.TB, accountID stellar1.AccountID) {
	t.Logf("gift -> %v", accountID)
	url := "https://friendbot.stellar.org/?addr=" + accountID.String()
	t.Logf("gift url: %v", url)
	res, err := http.Get(url)
	require.NoError(t, err, "friendbot request error")
	bodyBuf := new(bytes.Buffer)
	bodyBuf.ReadFrom(res.Body)
	t.Logf("gift res: %v", bodyBuf.String())
	require.Equal(t, 200, res.StatusCode, "friendbot response status code")
}

func useStellarTestNet(t testing.TB) {
	stellarnet.SetClientAndNetwork(horizon.DefaultTestNetClient, build.TestNetwork)
}
<|endoftext|>


<commit_before>// Copyright 2014, Truveris Inc. All Rights Reserved.
// Use of this source code is governed by the ISC license in the LICENSE file.
//
// This is ygord's tokenizer. It takes strings of characters such as:
//
//    unalias foo;alias foo "play http://foo/music.mp3; image bar.jpg"
//
// And creates a list of tokens:
//
//    - ["unalias", "foo", ";", "alias", "foo", "play http://foo/music.mp3; image bar.jpg"]
//
// Note that unescaped semi-colons are singled out as tokens in a shell
// fashion.
//

package main

import (
	"bufio"
	"errors"
	"io"
	"strings"
)

var (
)

const (
	STATE_START TokenizerState = iota
	STATE_INWORD TokenizerState = iota
	STATE_QUOTED TokenizerState = iota
	STATE_QUOTED_ESCAPED TokenizerState = iota
	STATE_ESCAPED TokenizerState = iota
	STATE_INWORD_ESCAPED TokenizerState = iota

	RUNETYPE_EOF RuneType = iota
	RUNETYPE_EOS RuneType = iota
	RUNETYPE_CHAR RuneType = iota
	RUNETYPE_SPACE RuneType = iota
	RUNETYPE_QUOTE RuneType = iota
	RUNETYPE_BACKSLASH RuneType = iota
)

type RuneType int
type TokenizerState int

type Tokenizer struct {
	input *bufio.Reader
}

func NewTokenizer(input io.Reader) *Tokenizer {
	t := &Tokenizer{bufio.NewReader(input)}
	return t
}

func (t *Tokenizer) GetRuneType(r rune) RuneType {
	switch {
	case r == ';':
		return RUNETYPE_EOS
	case r == '"':
		return RUNETYPE_QUOTE
	case r == '\\':
		return RUNETYPE_BACKSLASH
	case strings.ContainsRune(" \t\n\r", r):
		return RUNETYPE_SPACE
	}

	return RUNETYPE_CHAR
}

func (t *Tokenizer) NextToken() (string, error) {
	state := STATE_START
	token := make([]rune, 0)
	for {
		nextRune, _, err := t.input.ReadRune()
		nextRuneType := t.GetRuneType(nextRune)
		if err != nil {
			if err == io.EOF {
				nextRuneType = RUNETYPE_EOF
			} else {
				return "", err
			}
		}

		switch state {
		case STATE_START:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return "", io.EOF
			case RUNETYPE_CHAR:
				state = STATE_INWORD
				token = append(token, nextRune)
			case RUNETYPE_EOS:
				return ";", nil
			case RUNETYPE_SPACE:
				continue
			case RUNETYPE_BACKSLASH:
				state = STATE_INWORD_ESCAPED
			case RUNETYPE_QUOTE:
				state = STATE_QUOTED
			}
		case STATE_INWORD:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), io.EOF
			case RUNETYPE_CHAR:
				token = append(token, nextRune)
			case RUNETYPE_SPACE:
				return string(token), nil
			case RUNETYPE_BACKSLASH:
				state = STATE_INWORD_ESCAPED
			case RUNETYPE_EOS:
				t.input.UnreadRune()
				return string(token), nil
			case RUNETYPE_QUOTE:
				state = STATE_QUOTED
			}
		case STATE_QUOTED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("missing quote termination")
			case RUNETYPE_CHAR, RUNETYPE_SPACE, RUNETYPE_EOS:
				token = append(token, nextRune)
			case RUNETYPE_BACKSLASH:
				state = STATE_QUOTED_ESCAPED
			case RUNETYPE_QUOTE:
				state = STATE_INWORD
			}
		case STATE_INWORD_ESCAPED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("unterminated escape character")
			case RUNETYPE_QUOTE:
				token = append(token, '"')
				state = STATE_INWORD
			case RUNETYPE_BACKSLASH:
				token = append(token, '\\')
				state = STATE_INWORD
			default:
				return string(token), errors.New("unknown escape character")
			}
		case STATE_QUOTED_ESCAPED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("unterminated escape character")
			case RUNETYPE_QUOTE:
				token = append(token, '"')
				state = STATE_QUOTED
			case RUNETYPE_BACKSLASH:
				token = append(token, '\\')
				state = STATE_QUOTED
			default:
				return string(token), errors.New("unknown escape character")
			}
		}
	}
	return "", nil
}
<commit_msg>Don't look for magic in escaped letters.<commit_after>// Copyright 2014, Truveris Inc. All Rights Reserved.
// Use of this source code is governed by the ISC license in the LICENSE file.
//
// This is ygord's tokenizer. It takes strings of characters such as:
//
//    unalias foo;alias foo "play http://foo/music.mp3; image bar.jpg"
//
// And creates a list of tokens:
//
//    - ["unalias", "foo", ";", "alias", "foo", "play http://foo/music.mp3; image bar.jpg"]
//
// Note that unescaped semi-colons are singled out as tokens in a shell
// fashion.
//

package main

import (
	"bufio"
	"errors"
	"io"
	"strings"
)

var (
)

const (
	STATE_START TokenizerState = iota
	STATE_INWORD TokenizerState = iota
	STATE_QUOTED TokenizerState = iota
	STATE_QUOTED_ESCAPED TokenizerState = iota
	STATE_ESCAPED TokenizerState = iota
	STATE_INWORD_ESCAPED TokenizerState = iota

	RUNETYPE_EOF RuneType = iota
	RUNETYPE_EOS RuneType = iota
	RUNETYPE_CHAR RuneType = iota
	RUNETYPE_SPACE RuneType = iota
	RUNETYPE_QUOTE RuneType = iota
	RUNETYPE_BACKSLASH RuneType = iota
)

type RuneType int
type TokenizerState int

type Tokenizer struct {
	input *bufio.Reader
}

func NewTokenizer(input io.Reader) *Tokenizer {
	t := &Tokenizer{bufio.NewReader(input)}
	return t
}

func (t *Tokenizer) GetRuneType(r rune) RuneType {
	switch {
	case r == ';':
		return RUNETYPE_EOS
	case r == '"':
		return RUNETYPE_QUOTE
	case r == '\\':
		return RUNETYPE_BACKSLASH
	case strings.ContainsRune(" \t\n\r", r):
		return RUNETYPE_SPACE
	}

	return RUNETYPE_CHAR
}

func (t *Tokenizer) NextToken() (string, error) {
	state := STATE_START
	token := make([]rune, 0)
	for {
		nextRune, _, err := t.input.ReadRune()
		nextRuneType := t.GetRuneType(nextRune)
		if err != nil {
			if err == io.EOF {
				nextRuneType = RUNETYPE_EOF
			} else {
				return "", err
			}
		}

		switch state {
		case STATE_START:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return "", io.EOF
			case RUNETYPE_CHAR:
				state = STATE_INWORD
				token = append(token, nextRune)
			case RUNETYPE_EOS:
				return ";", nil
			case RUNETYPE_SPACE:
				continue
			case RUNETYPE_BACKSLASH:
				state = STATE_INWORD_ESCAPED
			case RUNETYPE_QUOTE:
				state = STATE_QUOTED
			}
		case STATE_INWORD:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), io.EOF
			case RUNETYPE_CHAR:
				token = append(token, nextRune)
			case RUNETYPE_SPACE:
				return string(token), nil
			case RUNETYPE_BACKSLASH:
				state = STATE_INWORD_ESCAPED
			case RUNETYPE_EOS:
				t.input.UnreadRune()
				return string(token), nil
			case RUNETYPE_QUOTE:
				state = STATE_QUOTED
			}
		case STATE_QUOTED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("missing quote termination")
			case RUNETYPE_CHAR, RUNETYPE_SPACE, RUNETYPE_EOS:
				token = append(token, nextRune)
			case RUNETYPE_BACKSLASH:
				state = STATE_QUOTED_ESCAPED
			case RUNETYPE_QUOTE:
				state = STATE_INWORD
			}
		case STATE_INWORD_ESCAPED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("unterminated escape character")
			case RUNETYPE_CHAR:
				token = append(token, nextRune)
			case RUNETYPE_QUOTE:
				token = append(token, '"')
				state = STATE_INWORD
			case RUNETYPE_BACKSLASH:
				token = append(token, '\\')
				state = STATE_INWORD
			default:
				return string(token), errors.New("unknown escape character")
			}
		case STATE_QUOTED_ESCAPED:
			switch nextRuneType {
			case RUNETYPE_EOF:
				return string(token), errors.New("unterminated escape character")
			case RUNETYPE_CHAR:
				token = append(token, nextRune)
			case RUNETYPE_QUOTE:
				token = append(token, '"')
				state = STATE_QUOTED
			case RUNETYPE_BACKSLASH:
				token = append(token, '\\')
				state = STATE_QUOTED
			default:
				return string(token), errors.New("unknown escape character")
			}
		}
	}
	return "", nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/golang/glog"
	"github.com/miekg/dns"
	"github.com/shane-kerr/ymmv/dnsstub"
	"math/rand"
	"net"
	"sort"
	"strings"
	"sync"
	"time"
)

// starting point to find Yeti root servers
var yeti_root_hints = []string{
	"bii.dns-lab.net.",
	"yeti-ns.wide.ad.jp.",
	"yeti-ns.tisf.net.",
}

// allowed server-selection algorithms
var server_algorithms = map[string]bool{
	"rtt":         true,
	"round-robin": true,
	"random":      true,
	"all":         true,
}

// maximum TTL to use
//const MAX_TTL = 30
//const MAX_TTL = 300
const MAX_TTL = 86400

// convert a TTL into a bounded duration
func use_ttl(ttl uint32) time.Duration {
	if ttl > MAX_TTL {
		ttl = MAX_TTL
	}
	result := time.Second * time.Duration(ttl+1)
	return result
}

// information about each IP address for each Yeti name server
type ip_info struct {
	// IP address
	ip net.IP
	// smoothed round-trip time (SRTT) for this IP address
	srtt time.Duration
}

// information about each Yeti name server
type ns_info struct {
	// name of the name server, like "bii.dns-lab.net" (may be "")
	name string

	// IP addresses of the name server
	ip_info []*ip_info

	// timer for seeing if the IP address changed (may be nil)
	timer *time.Timer

	// server set we belong to
	srvs *yeti_server_set
}

type yeti_server_set struct {
	// prevent use of data during modification
	lock sync.Mutex

	// timer for refreshing the root zone NS RRset
	root_ns_timer *time.Timer

	// information about each IP address
	ns []*ns_info

	// algorithm used to pick server
	algorithm string

	// next server to use (for round-robin)
	next_server int
	next_ip     int

	// resolver for lookups
	resolver *dnsstub.StubResolver
}

// does a lookup and gets the root NS RRset
func get_root_ns(hints []string) (ns_names []string, ns_ttl uint32) {
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	ns_query.RecursionDesired = true

	// Try each of our servers listed in our hints until one works.
	// Use the rand.Perm() function so we try them in a random order.
	var ns_response *dns.Msg
	for _, n := range rand.Perm(len(hints)) {
		root := hints[n]
		if !strings.ContainsRune(root, ':') {
			root = root + ":53"
		}
		var err error
		ns_response, _, err = dnsstub.DnsQuery(root, ns_query)
		if err != nil {
			glog.Warningf("Error looking up Yeti root server NS from %s; %s", root, err)
		}
		if ns_response != nil {
			break
		}
	}
	if ns_response == nil {
		glog.Fatalf("Unable to get NS from any Yeti root server")
	}

	// extract out the names from the NS RRset, and also save the TTL
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns_names = append(ns_names, root_server.(*dns.NS).Ns)
			ns_ttl = root_server.Header().Ttl
		}
	}

	if len(ns_names) == 0 {
		glog.Fatalf("No NS found for Yeti root")
	}

	// insure our order is repeatable
	sort.Strings(ns_names)

	return ns_names, ns_ttl
}

func refresh_ns(srvs *yeti_server_set) {
	glog.V(1).Infof("refreshing root NS list")

	ns_names, ns_ttl := get_root_ns(yeti_root_hints)

	srvs.lock.Lock()
	n := 0
	for _, name := range ns_names {
		// remove any names that are no longer present
		for (len(srvs.ns) > n) && (srvs.ns[n].name < name) {
			if !srvs.ns[n].timer.Stop() {
				<-srvs.ns[n].timer.C
			}
			srvs.ns = append(srvs.ns[:n], srvs.ns[n+1:]...)
		}
		// if name is the same, we are good, advance to next name
		if srvs.ns[n].name == name {
			n++
			// otherwise we have to add it
		} else {
			new_ns := &ns_info{name: name, srvs: srvs}
			srvs.ns = append(srvs.ns, nil)
			copy(srvs.ns[n+1:], srvs.ns[n:])
			srvs.ns[n] = new_ns
			go refresh_aaaa(new_ns, nil)
		}
	}
	srvs.lock.Unlock()

	// set timer to refresh our NS RRset
	srvs.root_ns_timer = time.AfterFunc(use_ttl(ns_ttl), func() { refresh_ns(srvs) })
}

// Note that on error we don't use this name server until we can find out
// the IP address for it. In theory we could use the name server until the
// TTL for the address expires, but rather than track that we just
// temporarily stop using it.
func refresh_aaaa(ns *ns_info, done chan int) {
	glog.V(1).Infof("refreshing %s", ns.name)

	var new_ip []net.IP
	var this_ttl uint32

	answer, _, err := ns.srvs.resolver.SyncQuery(ns.name, dns.TypeAAAA)
	if err != nil {
		glog.Warningf("Error looking up %s: %s\n", ns.name, err)
	}

	if answer != nil {
		for _, root_address := range answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa := root_address.(*dns.AAAA).AAAA
				glog.V(2).Infof("AAAA for %s = %s", ns.name, aaaa)
				new_ip = append(new_ip, aaaa)
				this_ttl = root_address.Header().Ttl
			}
		}
	}

	if len(new_ip) == 0 {
		glog.Warningf("No AAAA for %s, checking again in 300 seconds", ns.name)
		this_ttl = 300
	}

	// make a set of IP information, copying old SRTT if present in the old set
	var new_ip_info []*ip_info
	for _, ip := range new_ip {
		found := false
		for _, info := range ns.ip_info {
			if ip.Equal(info.ip) {
				glog.V(1).Infof("%s: copying old IP information ip=%s, srtt=%s", ns.name, info.ip, info.srtt)
				new_ip_info = append(new_ip_info, info)
				found = true
				break
			}
		}
		if !found {
			// if we are updating a name server that already had an IP, warn about this
			if len(ns.ip_info) > 0 {
				glog.Warningf("%s: adding new IP information ip=%s, srtt=0", ns.name, ip)
			}
			new_ip_info = append(new_ip_info, &ip_info{ip: ip, srtt: 0})
		}
	}

	ns.srvs.lock.Lock()
	ns.ip_info = new_ip_info
	when := use_ttl(this_ttl)
	glog.V(1).Infof("scheduling refresh of AAAA for %s in %s\n", ns.name, when)
	ns.timer = time.AfterFunc(when, func() { refresh_aaaa(ns, nil) })
	ns.srvs.lock.Unlock()

	done <- len(new_ip)
}

// get the list of root servers from known Yeti root servers
func yeti_priming(srvs *yeti_server_set) {
	// get the names from the NS RRset
	glog.V(1).Infof("getting root NS RRset")
	ns_names, ns_ttl := get_root_ns(yeti_root_hints)

	aaaa_done := make(chan int)
	for _, ns_name := range ns_names {
		this_ns := &ns_info{name: ns_name, srvs: srvs}
		// we want to complete at least one lookup before we return from priming
		go refresh_aaaa(this_ns, aaaa_done)
		srvs.ns = append(srvs.ns, this_ns)
	}

	found_aaaa := false
	for _ = range ns_names {
		num_ip := <-aaaa_done
		if num_ip > 0 {
			found_aaaa = true
			break
		}
	}
	if !found_aaaa {
		glog.Fatalf("No AAAA found for Yeti root NS")
	}

	// set timer to refresh our NS RRset
	when := use_ttl(ns_ttl)
	glog.V(1).Infof("scheduling refresh of NS in %s\n", when)
	srvs.root_ns_timer = time.AfterFunc(when, func() { refresh_ns(srvs) })
}

func init_yeti_server_set(ips []net.IP, algo string) (srvs *yeti_server_set) {
	srvs = new(yeti_server_set)

	if len(ips) == 0 {
		glog.Infof("no IP's passed, performing Yeti priming\n")
		var err error
		srvs.resolver, err = dnsstub.Init(4, nil)
		if err != nil {
			glog.Fatalf("Error setting up DNS stub resolver: %s\n", err)
		}
		yeti_priming(srvs)
		glog.V(1).Infof("priming done\n")
	} else {
		var ns_ip_info []*ip_info
		for _, ip := range ips {
			ns_ip_info = append(ns_ip_info, &ip_info{ip: ip, srtt: 0})
		}
		srvs.ns = append(srvs.ns, &ns_info{ip_info: ns_ip_info})
	}

	srvs.algorithm = algo
	srvs.next_server = 0
	srvs.next_ip = 0

	return srvs
}

type query_target struct {
	ip      net.IP
	ns_name string
}

// Get the next set of IP addresses to query.
// For most algorithms this is a single address, but it may be more (for "all").
func (srvs *yeti_server_set) next() (targets []*query_target) {
	srvs.lock.Lock()
	defer srvs.lock.Unlock()

	if srvs.algorithm == "round-robin" {
		for srvs.next_ip >= len(srvs.ns[srvs.next_server].ip_info) {
			srvs.next_server = (srvs.next_server + 1) % len(srvs.ns)
			srvs.next_ip = 0
		}
		ns := srvs.ns[srvs.next_server]
		ip := ns.ip_info[srvs.next_ip].ip
		targets = append(targets, &query_target{ip: ip, ns_name: ns.name})
		srvs.next_ip = srvs.next_ip + 1
	} else if srvs.algorithm == "rtt" {
		var lowest_ip_info *ip_info = nil
		var ns_name string
		for _, ns := range srvs.ns {
			for _, info := range ns.ip_info {
				if (lowest_ip_info == nil) || (lowest_ip_info.srtt > info.srtt) {
					lowest_ip_info = info
					ns_name = ns.name
				}
			}
		}
		targets = append(targets, &query_target{ip: lowest_ip_info.ip, ns_name: ns_name})
	} else {
		var all_targets []*query_target
		for _, ns := range srvs.ns {
			for _, info := range ns.ip_info {
				all_targets = append(all_targets, &query_target{ip: info.ip, ns_name: ns.name})
			}
		}
		if srvs.algorithm == "all" {
			targets = all_targets
		} else if srvs.algorithm == "random" {
			targets = append(targets, all_targets[rand.Intn(len(all_targets))])
		}
	}
	return targets
}

func (srvs *yeti_server_set) update_srtt(ip net.IP, rtt time.Duration) {
	glog.V(3).Infof("update_srtt ip=%s, rtt=%s", ip, rtt)
	srvs.lock.Lock()
	defer srvs.lock.Unlock()

	for _, ns_info := range srvs.ns {
		for _, ip_info := range ns_info.ip_info {
			// update the time for the IP that we just queried
			if ip_info.ip.Equal(ip) {
				if ip_info.srtt == 0 {
					ip_info.srtt = rtt
				} else {
					ip_info.srtt = ((ip_info.srtt * 7) + (rtt * 3)) / 10
				}
				glog.V(2).Infof("%s: update SRTT ip=%s, srtt=%s", ns_info.name, ip_info.ip, ip_info.srtt)
				// all other IP have their time decayed a bit
			} else {
				// There may be overflow issues to worry about here. Durations
				// are 64-bit nanoseconds, so we should be able to handle any
				//
				ip_info.srtt = (ip_info.srtt * 49) / 50
				glog.V(3).Infof("%s: decay SRTT ip=%s, srtt=%s", ns_info.name, ip_info.ip, ip_info.srtt)
			}
		}
	}
}
<commit_msg>Change priming to work with IPv6-only servers<commit_after>package main

import (
	"github.com/golang/glog"
	"github.com/miekg/dns"
	"github.com/shane-kerr/ymmv/dnsstub"
	"math/rand"
	"net"
	"sort"
	"strings"
	"sync"
	"time"
)

// starting point to find Yeti root servers
var yeti_root_hints = []string{
	// we use IPv6 literals here, since for some reason the Go
	// net.DialTimeout() function seems to have stopped working
	// for IPv6-only hostnames
	"[240c:f:1:22::6]:53",   // bii.dns-lab.net
	"[2001:200:1d9::35]:53", // yeti-ns.wide.ad.jp
	"[2001:559:8000::6]:53", // yeti-ns.tisf.net
}

// allowed server-selection algorithms
var server_algorithms = map[string]bool{
	"rtt":         true,
	"round-robin": true,
	"random":      true,
	"all":         true,
}

// maximum TTL to use
//const MAX_TTL = 30
//const MAX_TTL = 300
const MAX_TTL = 86400

// convert a TTL into a bounded duration
func use_ttl(ttl uint32) time.Duration {
	if ttl > MAX_TTL {
		ttl = MAX_TTL
	}
	result := time.Second * time.Duration(ttl+1)
	return result
}

// information about each IP address for each Yeti name server
type ip_info struct {
	// IP address
	ip net.IP
	// smoothed round-trip time (SRTT) for this IP address
	srtt time.Duration
}

// information about each Yeti name server
type ns_info struct {
	// name of the name server, like "bii.dns-lab.net" (may be "")
	name string

	// IP addresses of the name server
	ip_info []*ip_info

	// timer for seeing if the IP address changed (may be nil)
	timer *time.Timer

	// server set we belong to
	srvs *yeti_server_set
}

type yeti_server_set struct {
	// prevent use of data during modification
	lock sync.Mutex

	// timer for refreshing the root zone NS RRset
	root_ns_timer *time.Timer

	// information about each IP address
	ns []*ns_info

	// algorithm used to pick server
	algorithm string

	// next server to use (for round-robin)
	next_server int
	next_ip     int

	// resolver for lookups
	resolver *dnsstub.StubResolver
}

// does a lookup and gets the root NS RRset
func get_root_ns(hints []string) (ns_names []string, ns_ttl uint32) {
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	ns_query.RecursionDesired = true

	// Try each of our servers listed in our hints until one works.
	// Use the rand.Perm() function so we try them in a random order.
	var ns_response *dns.Msg
	for _, n := range rand.Perm(len(hints)) {
		root := hints[n]
		if !strings.ContainsRune(root, ':') {
			root = root + ":53"
		}
		var err error
		ns_response, _, err = dnsstub.DnsQuery(root, ns_query)
		if err != nil {
			glog.Warningf("Error looking up Yeti root server NS from %s; %s", root, err)
		}
		if ns_response != nil {
			break
		}
	}
	if ns_response == nil {
		glog.Fatalf("Unable to get NS from any Yeti root server")
	}

	// extract out the names from the NS RRset, and also save the TTL
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns_names = append(ns_names, root_server.(*dns.NS).Ns)
			ns_ttl = root_server.Header().Ttl
		}
	}

	if len(ns_names) == 0 {
		glog.Fatalf("No NS found for Yeti root")
	}

	// insure our order is repeatable
	sort.Strings(ns_names)

	return ns_names, ns_ttl
}

func refresh_ns(srvs *yeti_server_set) {
	glog.V(1).Infof("refreshing root NS list")

	ns_names, ns_ttl := get_root_ns(yeti_root_hints)

	srvs.lock.Lock()
	n := 0
	for _, name := range ns_names {
		// remove any names that are no longer present
		for (len(srvs.ns) > n) && (srvs.ns[n].name < name) {
			if !srvs.ns[n].timer.Stop() {
				<-srvs.ns[n].timer.C
			}
			srvs.ns = append(srvs.ns[:n], srvs.ns[n+1:]...)
		}
		// if name is the same, we are good, advance to next name
		if srvs.ns[n].name == name {
			n++
			// otherwise we have to add it
		} else {
			new_ns := &ns_info{name: name, srvs: srvs}
			srvs.ns = append(srvs.ns, nil)
			copy(srvs.ns[n+1:], srvs.ns[n:])
			srvs.ns[n] = new_ns
			go refresh_aaaa(new_ns, nil)
		}
	}
	srvs.lock.Unlock()

	// set timer to refresh our NS RRset
	srvs.root_ns_timer = time.AfterFunc(use_ttl(ns_ttl), func() { refresh_ns(srvs) })
}

// Note that on error we don't use this name server until we can find out
// the IP address for it. In theory we could use the name server until the
// TTL for the address expires, but rather than track that we just
// temporarily stop using it.
func refresh_aaaa(ns *ns_info, done chan int) {
	glog.V(1).Infof("refreshing %s", ns.name)

	var new_ip []net.IP
	var this_ttl uint32

	answer, _, err := ns.srvs.resolver.SyncQuery(ns.name, dns.TypeAAAA)
	if err != nil {
		glog.Warningf("Error looking up %s: %s\n", ns.name, err)
	}

	if answer != nil {
		for _, root_address := range answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa := root_address.(*dns.AAAA).AAAA
				glog.V(2).Infof("AAAA for %s = %s", ns.name, aaaa)
				new_ip = append(new_ip, aaaa)
				this_ttl = root_address.Header().Ttl
			}
		}
	}

	if len(new_ip) == 0 {
		glog.Warningf("No AAAA for %s, checking again in 300 seconds", ns.name)
		this_ttl = 300
	}

	// make a set of IP information, copying old SRTT if present in the old set
	var new_ip_info []*ip_info
	for _, ip := range new_ip {
		found := false
		for _, info := range ns.ip_info {
			if ip.Equal(info.ip) {
				glog.V(1).Infof("%s: copying old IP information ip=%s, srtt=%s", ns.name, info.ip, info.srtt)
				new_ip_info = append(new_ip_info, info)
				found = true
				break
			}
		}
		if !found {
			// if we are updating a name server that already had an IP, warn about this
			if len(ns.ip_info) > 0 {
				glog.Warningf("%s: adding new IP information ip=%s, srtt=0", ns.name, ip)
			}
			new_ip_info = append(new_ip_info, &ip_info{ip: ip, srtt: 0})
		}
	}

	ns.srvs.lock.Lock()
	ns.ip_info = new_ip_info
	when := use_ttl(this_ttl)
	glog.V(1).Infof("scheduling refresh of AAAA for %s in %s\n", ns.name, when)
	ns.timer = time.AfterFunc(when, func() { refresh_aaaa(ns, nil) })
	ns.srvs.lock.Unlock()

	done <- len(new_ip)
}

// get the list of root servers from known Yeti root servers
func yeti_priming(srvs *yeti_server_set) {
	// get the names from the NS RRset
	glog.V(1).Infof("getting root NS RRset")
	ns_names, ns_ttl := get_root_ns(yeti_root_hints)

	aaaa_done := make(chan int)
	for _, ns_name := range ns_names {
		this_ns := &ns_info{name: ns_name, srvs: srvs}
		// we want to complete at least one lookup before we return from priming
		go refresh_aaaa(this_ns, aaaa_done)
		srvs.ns = append(srvs.ns, this_ns)
	}

	found_aaaa := false
	for _ = range ns_names {
		num_ip := <-aaaa_done
		if num_ip > 0 {
			found_aaaa = true
			break
		}
	}
	if !found_aaaa {
		glog.Fatalf("No AAAA found for Yeti root NS")
	}

	// set timer to refresh our NS RRset
	when := use_ttl(ns_ttl)
	glog.V(1).Infof("scheduling refresh of NS in %s\n", when)
	srvs.root_ns_timer = time.AfterFunc(when, func() { refresh_ns(srvs) })
}

func init_yeti_server_set(ips []net.IP, algo string) (srvs *yeti_server_set) {
	srvs = new(yeti_server_set)

	if len(ips) == 0 {
		glog.Infof("no IP's passed, performing Yeti priming\n")
		var err error
		srvs.resolver, err = dnsstub.Init(4, nil)
		if err != nil {
			glog.Fatalf("Error setting up DNS stub resolver: %s\n", err)
		}
		yeti_priming(srvs)
		glog.V(1).Infof("priming done\n")
	} else {
		var ns_ip_info []*ip_info
		for _, ip := range ips {
			ns_ip_info = append(ns_ip_info, &ip_info{ip: ip, srtt: 0})
		}
		srvs.ns = append(srvs.ns, &ns_info{ip_info: ns_ip_info})
	}

	srvs.algorithm = algo
	srvs.next_server = 0
	srvs.next_ip = 0

	return srvs
}

type query_target struct {
	ip      net.IP
	ns_name string
}

// Get the next set of IP addresses to query.
// For most algorithms this is a single address, but it may be more (for "all").
func (srvs *yeti_server_set) next() (targets []*query_target) {
	srvs.lock.Lock()
	defer srvs.lock.Unlock()

	if srvs.algorithm == "round-robin" {
		for srvs.next_ip >= len(srvs.ns[srvs.next_server].ip_info) {
			srvs.next_server = (srvs.next_server + 1) % len(srvs.ns)
			srvs.next_ip = 0
		}
		ns := srvs.ns[srvs.next_server]
		ip := ns.ip_info[srvs.next_ip].ip
		targets = append(targets, &query_target{ip: ip, ns_name: ns.name})
		srvs.next_ip = srvs.next_ip + 1
	} else if srvs.algorithm == "rtt" {
		var lowest_ip_info *ip_info = nil
		var ns_name string
		for _, ns := range srvs.ns {
			for _, info := range ns.ip_info {
				if (lowest_ip_info == nil) || (lowest_ip_info.srtt > info.srtt) {
					lowest_ip_info = info
					ns_name = ns.name
				}
			}
		}
		targets = append(targets, &query_target{ip: lowest_ip_info.ip, ns_name: ns_name})
	} else {
		var all_targets []*query_target
		for _, ns := range srvs.ns {
			for _, info := range ns.ip_info {
				all_targets = append(all_targets, &query_target{ip: info.ip, ns_name: ns.name})
			}
		}
		if srvs.algorithm == "all" {
			targets = all_targets
		} else if srvs.algorithm == "random" {
			targets = append(targets, all_targets[rand.Intn(len(all_targets))])
		}
	}
	return targets
}

func (srvs *yeti_server_set) update_srtt(ip net.IP, rtt time.Duration) {
	glog.V(3).Infof("update_srtt ip=%s, rtt=%s", ip, rtt)
	srvs.lock.Lock()
	defer srvs.lock.Unlock()

	for _, ns_info := range srvs.ns {
		for _, ip_info := range ns_info.ip_info {
			// update the time for the IP that we just queried
			if ip_info.ip.Equal(ip) {
				if ip_info.srtt == 0 {
					ip_info.srtt = rtt
				} else {
					ip_info.srtt = ((ip_info.srtt * 7) + (rtt * 3)) / 10
				}
				glog.V(2).Infof("%s: update SRTT ip=%s, srtt=%s", ns_info.name, ip_info.ip, ip_info.srtt)
				// all other IP have their time decayed a bit
			} else {
				// There may be overflow issues to worry about here. Durations
				// are 64-bit nanoseconds, so we should be able to handle any
				//
				ip_info.srtt = (ip_info.srtt * 49) / 50
				glog.V(3).Infof("%s: decay SRTT ip=%s, srtt=%s", ns_info.name, ip_info.ip, ip_info.srtt)
			}
		}
	}
}
<|endoftext|>


<commit_before>package screenshot

import (
	"image"
)

func ScreenRect() (image.Rectangle, error) {
	panic("ScreenRect is not supported on this platform.")
}

func CaptureScreen() (*image.RGBA, error) {
	panic("CaptureScreen is not supported on this platform.")
}

func CaptureRect(rect image.Rectangle) (*image.RGBA, error) {
	panic("CaptureRect is not supported on this platform.")
}
<commit_msg>Implemented all functions using native frameworks<commit_after>package screenshot

import (
	// #cgo LDFLAGS: -framework CoreGraphics
	// #cgo LDFLAGS: -framework CoreFoundation
	// #include <CoreGraphics/CoreGraphics.h>
	// #include <CoreFoundation/CoreFoundation.h>
	"C"
	"image"
	"reflect"
	"unsafe"
)

func ScreenRect() (image.Rectangle, error) {
	displayID := C.CGMainDisplayID()
	width := int(C.CGDisplayPixelsWide(displayID))
	height := int(C.CGDisplayPixelsHigh(displayID))
	return image.Rect(0, 0, width, height), nil
}

func CaptureScreen() (*image.RGBA, error) {
	rect, err := ScreenRect()
	if err != nil {
		return nil, err
	}
	return CaptureRect(rect)
}

func CaptureRect(rect image.Rectangle) (*image.RGBA, error) {
	displayID := C.CGMainDisplayID()
	width := int(C.CGDisplayPixelsWide(displayID))
	rawData := C.CGDataProviderCopyData(C.CGImageGetDataProvider(C.CGDisplayCreateImage(displayID)))

	length := int(C.CFDataGetLength(rawData))
	ptr := unsafe.Pointer(C.CFDataGetBytePtr(rawData))

	var slice []byte
	hdrp := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
	hdrp.Data = uintptr(ptr)
	hdrp.Len = length
	hdrp.Cap = length

	imageBytes := make([]byte, length)

	for i := 0; i < length; i += 4 {
		imageBytes[i], imageBytes[i+2], imageBytes[i+1], imageBytes[i+3] = slice[i+2], slice[i], slice[i+1], slice[i+3]
	}

	C.CFRelease(rawData)

	img := &image.RGBA{Pix: imageBytes, Stride: 4 * width, Rect: rect}
	return img, nil
}
<|endoftext|>