<commit_before>package bdconfig

import (
	"jrubin.io/blamedns/bdconfig/bdtype"
	"jrubin.io/blamedns/dnscache"
	"jrubin.io/blamedns/dnsserver"
)

type DNSConfig struct {
	Forward            []string             `toml:"forward" cli:",dns server(s) to forward requests to"`
	Listen             []string             `toml:"listen" cli:",url(s) to listen for dns requests on"`
	Block              BlockConfig          `toml:"block"`
	ClientTimeout      bdtype.Duration      `toml:"client_timeout" cli:",time to wait for remote servers to respond to queries"`
	ServerTimeout      bdtype.Duration      `toml:"server_timeout" cli:",time for server to wait for remote clients to respond to queries"`
	DialTimeout        bdtype.Duration      `toml:"dial_timeout" cli:",time to wait to establish connection to remote server"`
	Interval           bdtype.Duration      `toml:"interval" cli:",concurrency interval for lookups in miliseconds"`
	Server             *dnsserver.DNSServer `toml:"-" cli:"-"`
	CachePruneInterval bdtype.Duration      `toml:"cache_prune_interval" cli:",how often to run cache expiration"`
	DisableDNSSEC      bool                 `toml:"disable_dnssec" cli:",disable dnssec validation"`
	DisableRecursion   bool                 `toml:"disable_recursion" cli:",refuse recursive queries (you probably don't want to set this)"`
}

func defaultDNSConfig() DNSConfig {
	return DNSConfig{
		Forward: []string{
			"8.8.8.8",
			"8.8.4.4",
		},
		Listen: []string{
			"udp://[::]:53",
			"tcp://[::]:53",
		},
		Block: defaultBlockConfig(),
	}
}

func (cfg *DNSConfig) Init(root *Config) error {
	if err := cfg.Block.Init(root); err != nil {
		return err
	}

	cfg.Server = &dnsserver.DNSServer{
		Forward:            cfg.Forward,
		Listen:             cfg.Listen,
		Block:              cfg.Block.Block(),
		ClientTimeout:      cfg.ClientTimeout.Duration(),
		ServerTimeout:      cfg.ServerTimeout.Duration(),
		DialTimeout:        cfg.DialTimeout.Duration(),
		Interval:           cfg.Interval.Duration(),
		Logger:             root.Logger,
		Cache:              dnscache.NewMemory(root.Logger),
		CachePruneInterval: cfg.CachePruneInterval.Duration(),
		NotifyStartedFunc:  cfg.Block.Start,
		DisableDNSSEC:      cfg.DisableDNSSEC,
		DisableRecursion:   cfg.DisableRecursion,
	}

	return nil
}

func (cfg DNSConfig) Start() error {
	return cfg.Server.ListenAndServe()
}

func (cfg DNSConfig) Shutdown() {
	cfg.Block.Shutdown()
	cfg.Server.Shutdown()
}
<commit_msg>add option to disable cache<commit_after>package bdconfig

import (
	"jrubin.io/blamedns/bdconfig/bdtype"
	"jrubin.io/blamedns/dnscache"
	"jrubin.io/blamedns/dnsserver"
)

type DNSConfig struct {
	Forward            []string             `toml:"forward" cli:",dns server(s) to forward requests to"`
	Listen             []string             `toml:"listen" cli:",url(s) to listen for dns requests on"`
	Block              BlockConfig          `toml:"block"`
	ClientTimeout      bdtype.Duration      `toml:"client_timeout" cli:",time to wait for remote servers to respond to queries"`
	ServerTimeout      bdtype.Duration      `toml:"server_timeout" cli:",time for server to wait for remote clients to respond to queries"`
	DialTimeout        bdtype.Duration      `toml:"dial_timeout" cli:",time to wait to establish connection to remote server"`
	Interval           bdtype.Duration      `toml:"interval" cli:",concurrency interval for lookups in miliseconds"`
	Server             *dnsserver.DNSServer `toml:"-" cli:"-"`
	CachePruneInterval bdtype.Duration      `toml:"cache_prune_interval" cli:",how often to run cache expiration"`
	DisableDNSSEC      bool                 `toml:"disable_dnssec" cli:",disable dnssec validation"`
	DisableRecursion   bool                 `toml:"disable_recursion" cli:",refuse recursive queries (you probably don't want to set this)"`
	DisableCache       bool                 `toml:"disable_cache" cli:",disable dns lookup cache"`
}

func defaultDNSConfig() DNSConfig {
	return DNSConfig{
		Forward: []string{
			"8.8.8.8",
			"8.8.4.4",
		},
		Listen: []string{
			"udp://[::]:53",
			"tcp://[::]:53",
		},
		Block: defaultBlockConfig(),
	}
}

func (cfg *DNSConfig) Init(root *Config) error {
	if err := cfg.Block.Init(root); err != nil {
		return err
	}

	cfg.Server = &dnsserver.DNSServer{
		Forward:            cfg.Forward,
		Listen:             cfg.Listen,
		Block:              cfg.Block.Block(),
		ClientTimeout:      cfg.ClientTimeout.Duration(),
		ServerTimeout:      cfg.ServerTimeout.Duration(),
		DialTimeout:        cfg.DialTimeout.Duration(),
		Interval:           cfg.Interval.Duration(),
		Logger:             root.Logger,
		CachePruneInterval: cfg.CachePruneInterval.Duration(),
		NotifyStartedFunc:  cfg.Block.Start,
		DisableDNSSEC:      cfg.DisableDNSSEC,
		DisableRecursion:   cfg.DisableRecursion,
	}

	if !cfg.DisableCache {
		cfg.Server.Cache = dnscache.NewMemory(root.Logger)
	}

	return nil
}

func (cfg DNSConfig) Start() error {
	return cfg.Server.ListenAndServe()
}

func (cfg DNSConfig) Shutdown() {
	cfg.Block.Shutdown()
	cfg.Server.Shutdown()
}
<|endoftext|>


<commit_before>package deploy_docus

import (
	"code.google.com/p/go.crypto/ssh"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Rsa struct {
	Repository *Repository
	Private    *rsa.PrivateKey
}

func (r *Rsa) Encrypt(value []byte) (string, error) {
	hasher := sha1.New()
	key, err := r.PrivateKey()
	if err != nil {
		return "", err
	}

	hasher.Write([]byte(fmt.Sprintf("%s-%s", key, value)))
	return base64.URLEncoding.EncodeToString(hasher.Sum(nil)), nil
}

func (r *Rsa) WriteKey() error {
	path := r.KeyPath()

	err := os.MkdirAll(filepath.Dir(path), 0700)
	if err != nil {
		return err
	}

	key, _ := r.PrivateKey()
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(path, key, 0600)
	return err
}

func (r *Rsa) PrivateKey() ([]byte, error) {
	cert := x509.MarshalPKCS1PrivateKey(r.Private)
	blk := new(pem.Block)
	blk.Type = "RSA PRIVATE KEY"
	blk.Bytes = cert
	content := pem.EncodeToMemory(blk)

	return content, nil
}

func (r *Rsa) PublicKey() string {
	key, _ := ssh.NewPublicKey(&r.Private.PublicKey)
	marshalled := ssh.MarshalPublicKey(key)
	encoded := base64.StdEncoding.EncodeToString(marshalled) + "\n"

	return fmt.Sprintf("ssh-rsa %s", encoded)
}

func (r *Rsa) KeyPath() string {
	return fmt.Sprintf("/tmp/deploy_docus/keys/%s", r.Repository.Name())
}

func BuildPrivateKey(content []byte) (*rsa.PrivateKey, error) {
	if content == nil {
		return rsa.GenerateKey(rand.Reader, 1024)
	} else {
		block, _ := pem.Decode(content)
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	}
}

func BuildRsa(repository *Repository, ssh_key []byte) *Rsa {
	key, err := BuildPrivateKey(ssh_key)
	if err != nil {
		panic(err)
	}
	rsa := &Rsa{repository, key}

	return rsa
}
<commit_msg>check that the file exists after creating it<commit_after>package deploy_docus

import (
	"code.google.com/p/go.crypto/ssh"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Rsa struct {
	Repository *Repository
	Private    *rsa.PrivateKey
}

func (r *Rsa) Encrypt(value []byte) (string, error) {
	hasher := sha1.New()
	key, err := r.PrivateKey()
	if err != nil {
		return "", err
	}

	hasher.Write([]byte(fmt.Sprintf("%s-%s", key, value)))
	return base64.URLEncoding.EncodeToString(hasher.Sum(nil)), nil
}

func (r *Rsa) WriteKey() error {
	path := r.KeyPath()

	err := os.MkdirAll(filepath.Dir(path), 0700)
	if err != nil {
		return err
	}

	key, _ := r.PrivateKey()
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(path, key, 0600)
	if err != nil {
		return err
	}

	_, err = os.Stat(path)
	return err
}

func (r *Rsa) PrivateKey() ([]byte, error) {
	cert := x509.MarshalPKCS1PrivateKey(r.Private)
	blk := new(pem.Block)
	blk.Type = "RSA PRIVATE KEY"
	blk.Bytes = cert
	content := pem.EncodeToMemory(blk)

	return content, nil
}

func (r *Rsa) PublicKey() string {
	key, _ := ssh.NewPublicKey(&r.Private.PublicKey)
	marshalled := ssh.MarshalPublicKey(key)
	encoded := base64.StdEncoding.EncodeToString(marshalled) + "\n"

	return fmt.Sprintf("ssh-rsa %s", encoded)
}

func (r *Rsa) KeyPath() string {
	return fmt.Sprintf("/tmp/deploy_docus/keys/%s", r.Repository.Name())
}

func BuildPrivateKey(content []byte) (*rsa.PrivateKey, error) {
	if content == nil {
		return rsa.GenerateKey(rand.Reader, 1024)
	} else {
		block, _ := pem.Decode(content)
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	}
}

func BuildRsa(repository *Repository, ssh_key []byte) *Rsa {
	key, err := BuildPrivateKey(ssh_key)
	if err != nil {
		panic(err)
	}
	rsa := &Rsa{repository, key}

	return rsa
}
<|endoftext|>


<commit_before>// Package passenger revolves around the User associated with the
// request in flight.
//
// It establishes this link via the HTTP authorization header,
// which can be used to authenticate with basic authorization or
// token based authorization.
package passenger

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/subtle"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"fmt"
	mathrand "math/rand"
	"net/http"
	"strings"
	"time"

	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/memcache"

	"github.com/coduno/api/model"
	"github.com/coduno/api/util/password"

	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"

	// Indirect import of SHA2, will be called by crypto.Hash.New()
	_ "crypto/sha256"
	_ "crypto/sha512"

	// Indirect import of SHA3, will be called by crypto.Hash.New()
	_ "golang.org/x/crypto/sha3"
)

const maxValidity = time.Hour * 24 * 365 * 2
const minValidity = time.Minute * 30
const defaultValidity = time.Hour * 24 * 14

const defaultHash = crypto.SHA3_256

type key int64

var passengerKey key

// ErrNoAuthHeader signals that the request did not carry an Authorization
// header.
var ErrNoAuthHeader = errors.New("passenger: no authorization header present")

// ErrUnkAuthHeader is returned if the request contained an Authorization
// header, but it could not be parsed. The two accepted authorization
// methods are "Basic" and "Token".
var ErrUnkAuthHeader = errors.New("passenger: cannot interpret authorization header")

// ErrDigestMismatch is returned if the token provided by the user did not
// result in the same digest value after hashing as the one stored on the
// server.
var ErrDigestMismatch = errors.New("passenger: digest mismatch")

// ErrTokenNotAssociated can be returned if this package encounters invalid
// data during authentication. An AccessToken must be child of a User.
var ErrTokenNotAssociated = errors.New("passenger: token not associated to any user")

// ErrTokenExpired is returned if the token is not valid anymore.
var ErrTokenExpired = errors.New("passenger: token expired")

// ErrTokenNotMatchingUser is returned if basic auth with token password
// was attempted, and the token could be found but does not match the user.
type ErrTokenNotMatchingUser struct {
	Parent, Actual *datastore.Key
}

func (e ErrTokenNotMatchingUser) Error() string {
	return fmt.Sprintf("passenger: parent of token %+v does not match with actual user %+v", e.Parent, e.Actual)
}

func init() {
	passengerKey = key(mathrand.Int63())
}

// Passenger holds the currently authenticated user
// together with the access token (if relevant).
type Passenger struct {
	UserKey     *datastore.Key
	AccessToken model.AccessToken
}

// HasScope can be used to check whether this Passenger was
// granted access to a given scope.
// Note that a user authenticated via username and password
// (not via access token) will have access to all scopes by
// default.
func (p *Passenger) HasScope(scope string) (has bool) {
	for _, grantedScope := range p.AccessToken.Scopes {
		if scope == grantedScope {
			return true
		}
	}
	return
}

// Save will persist the passenger to Datastore and send it to Memcache.
func (p *Passenger) Save(ctx context.Context) (*datastore.Key, error) {
	now := time.Now()

	key, err := p.AccessToken.SaveWithParent(ctx, p.UserKey)
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	if err = gob.NewEncoder(buf).Encode(p); err != nil {
		return nil, err
	}

	item := &memcache.Item{
		Key:        key.Encode(),
		Value:      buf.Bytes(),
		Expiration: p.AccessToken.Expiry.Sub(now) + 10*time.Second,
	}

	if err = memcache.Set(ctx, item); err != nil {
		return nil, err
	}

	return key, nil
}

func (p *Passenger) check(raw []byte) error {
	digest := crypto.Hash(p.AccessToken.Hash).New().Sum(raw)

	if subtle.ConstantTimeCompare(digest, p.AccessToken.Digest) != 1 {
		return ErrDigestMismatch
	}

	if p.AccessToken.Expiry.Before(time.Now()) {
		return ErrTokenExpired
	}

	return nil
}

// IssueToken creates a new Token for the authenticated user. Callers should
// prefill the accessToken with whatever values they like and leave zero values
// to be set.
// The generated token will also be persisted and can be handed to the client
// with no more handling.
func (p *Passenger) IssueToken(ctx context.Context, token *model.AccessToken) (string, error) {
	now := time.Now()

	if token.Expiry == (time.Time{}) {
		token.Expiry = now.Add(defaultValidity)
	} else {
		if token.Expiry.Before(now.Add(minValidity)) {
			return "", fmt.Errorf("passenger: token must be valid at least %s", minValidity)
		}

		if token.Expiry.After(now.Add(maxValidity)) {
			return "", fmt.Errorf("passenger: token must be valid at most %s", maxValidity)
		}
	}

	// TODO(flowlo): This will reject all scopes for now, as we are not using them.
	// As soon as we introduce scopes, this check must be rewritten accordingly.
	if len(token.Scopes) > 0 {
		return "", fmt.Errorf("passenger: unknown scopes: %s", strings.Join(token.Scopes, ", "))
	}

	if len(token.Description) > 512 || len(token.Description) < 4 {
		return "", fmt.Errorf("passenger: description has bad len %d", len(token.Description))
	}

	raw := make([]byte, 16)

	if _, err := rand.Read(raw); err != nil {
		return "", err
	}

	if token.Hash == 0 {
		token.Hash = int(defaultHash)
	}
	token.Digest = crypto.Hash(token.Hash).New().Sum(raw)

	clone := Passenger{
		UserKey:     p.UserKey,
		AccessToken: *token,
	}

	key, err := clone.Save(ctx)
	if err != nil {
		return "", err
	}

	return encodeToken(key, raw)
}

// FromAccessToken tries do identify a Passenger by the access token he gave us.
// It will look up the AccessToken and consequently the corresponding User.
func FromAccessToken(ctx context.Context, accessToken string) (*Passenger, error) {
	key, raw, err := decodeToken(accessToken)
	if err != nil {
		return nil, err
	}

	var p *Passenger
	if p, err = fromCache(ctx, key); err != nil {
		if p, err = fromDatastore(ctx, key); err != nil {
			return nil, err
		}
	}

	return p, p.check(raw)
}

func fromCache(ctx context.Context, key *datastore.Key) (p *Passenger, err error) {
	item, err := memcache.Get(ctx, key.Encode())
	if err != nil {
		return nil, err
	}
	p = new(Passenger)
	err = gob.NewDecoder(bytes.NewReader(item.Value)).Decode(&p)
	return
}

func fromDatastore(ctx context.Context, key *datastore.Key) (p *Passenger, err error) {
	p = new(Passenger)
	if err = datastore.Get(ctx, key, p.AccessToken); err != nil {
		return
	}
	if p.UserKey = key.Parent(); p.UserKey == nil {
		return nil, ErrTokenNotAssociated
	}
	return
}

// FromBasicAuth tries do identify a Passenger by the access token he gave us.
// It will look up the the user by username and try to match password.
func FromBasicAuth(ctx context.Context, username, pw string) (p *Passenger, err error) {
	p = new(Passenger)
	var user model.User
	p.UserKey, err = model.NewQueryForUser().
		Filter("Nick=", username).
		Limit(1).
		Run(ctx).
		Next(&user)

	if err != nil {
		return
	}
	err = password.Check([]byte(pw), user.HashedPassword)

	// TODO(flowlo): Depending on bcrypt is very fragile. We
	// should encapsulate that.
	if err == bcrypt.ErrMismatchedHashAndPassword {
		userKey := p.UserKey
		p, err = FromAccessToken(ctx, pw)
		if err != nil {
			return
		}
		if !p.UserKey.Equal(userKey) {
			return nil, ErrTokenNotMatchingUser{Parent: p.UserKey, Actual: userKey}
		}
	}
	return
}

// FromRequest inspects the HTTP Authorization header of the given request
// and tries to identify a passenger.
func FromRequest(ctx context.Context, r *http.Request) (p *Passenger, err error) {
	auth := ""
	if auth = r.Header.Get("Authorization"); auth == "" {
		return nil, ErrNoAuthHeader
	}

	if strings.HasPrefix(auth, "Token ") {
		return FromAccessToken(ctx, auth[6:])
	}

	username, password, ok := "", "", false
	if username, password, ok = r.BasicAuth(); !ok {
		return nil, ErrUnkAuthHeader
	}

	p, err = FromBasicAuth(ctx, username, password)
	return
}

// FromContext returns the Passenger value stored in ctx if any.
func FromContext(ctx context.Context) (p *Passenger, ok bool) {
	p, ok = ctx.Value(passengerKey).(*Passenger)
	return
}

// NewContext returns a new Context that carries value p.
func NewContext(ctx context.Context, p *Passenger) context.Context {
	return context.WithValue(ctx, passengerKey, p)
}

// NewContextFromRequest wraps FromRequest and NewContext.
func NewContextFromRequest(ctx context.Context, r *http.Request) (context.Context, error) {
	p, err := FromRequest(ctx, r)
	if err != nil {
		return ctx, err
	}
	return NewContext(ctx, p), nil
}

// DecodeToken will take a token as sent by a client and translate it into a
// key to look up the full token on the server side and the raw secret.
func decodeToken(token string) (*datastore.Key, []byte, error) {
	b, err := hex.DecodeString(token)
	if err != nil {
		return nil, nil, err
	}

	dec := gob.NewDecoder(bytes.NewReader(b))

	var key *datastore.Key
	if err := dec.Decode(&key); err != nil {
		return nil, nil, err
	}
	var raw []byte
	if err := dec.Decode(&raw); err != nil {
		return nil, nil, err
	}
	return key, raw, nil
}

// EncodeToken translates the key and raw secret of a newly generated token to
// a form suitable for the client.
func encodeToken(key *datastore.Key, raw []byte) (string, error) {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	if err := enc.Encode(key); err != nil {
		return "", err
	}
	if err := enc.Encode(raw); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf.Bytes()), nil
}
<commit_msg>perf: Implement shorter and fixed size tokens<commit_after>// Package passenger revolves around the User associated with the
// request in flight.
//
// It establishes this link via the HTTP authorization header,
// which can be used to authenticate with basic authorization or
// token based authorization.
package passenger

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"encoding/gob"
	"errors"
	"fmt"
	mathrand "math/rand"
	"net/http"
	"strings"
	"time"

	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/memcache"

	"github.com/coduno/api/model"
	"github.com/coduno/api/util/password"

	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"

	// Indirect import of SHA2, will be called by crypto.Hash.New()
	_ "crypto/sha256"
	_ "crypto/sha512"

	// Indirect import of SHA3, will be called by crypto.Hash.New()
	_ "golang.org/x/crypto/sha3"
)

const maxValidity = time.Hour * 24 * 365 * 2
const minValidity = time.Minute * 30
const defaultValidity = time.Hour * 24 * 14

const defaultHash = crypto.SHA3_256

const tokenLength = 16

var endianness = binary.LittleEndian
var encoding = base64.StdEncoding

type key int64

var passengerKey key

// ErrNoAuthHeader signals that the request did not carry an Authorization
// header.
var ErrNoAuthHeader = errors.New("passenger: no authorization header present")

// ErrUnkAuthHeader is returned if the request contained an Authorization
// header, but it could not be parsed. The two accepted authorization
// methods are "Basic" and "Token".
var ErrUnkAuthHeader = errors.New("passenger: cannot interpret authorization header")

// ErrDigestMismatch is returned if the token provided by the user did not
// result in the same digest value after hashing as the one stored on the
// server.
var ErrDigestMismatch = errors.New("passenger: digest mismatch")

// ErrTokenNotAssociated can be returned if this package encounters invalid
// data during authentication. An AccessToken must be child of a User.
var ErrTokenNotAssociated = errors.New("passenger: token not associated to any user")

// ErrTokenExpired is returned if the token is not valid anymore.
var ErrTokenExpired = errors.New("passenger: token expired")

// ErrTokenNotMatchingUser is returned if basic auth with token password
// was attempted, and the token could be found but does not match the user.
type ErrTokenNotMatchingUser struct {
	Parent, Actual *datastore.Key
}

func (e ErrTokenNotMatchingUser) Error() string {
	return fmt.Sprintf("passenger: parent of token %+v does not match with actual user %+v", e.Parent, e.Actual)
}

func init() {
	passengerKey = key(mathrand.Int63())
}

// Passenger holds the currently authenticated user
// together with the access token (if relevant).
type Passenger struct {
	UserKey     *datastore.Key
	AccessToken model.AccessToken
}

// HasScope can be used to check whether this Passenger was
// granted access to a given scope.
// Note that a user authenticated via username and password
// (not via access token) will have access to all scopes by
// default.
func (p *Passenger) HasScope(scope string) (has bool) {
	for _, grantedScope := range p.AccessToken.Scopes {
		if scope == grantedScope {
			return true
		}
	}
	return
}

// Save will persist the passenger to Datastore and send it to Memcache.
func (p *Passenger) Save(ctx context.Context) (*datastore.Key, error) {
	now := time.Now()

	key, err := p.AccessToken.SaveWithParent(ctx, p.UserKey)
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	if err = gob.NewEncoder(buf).Encode(p); err != nil {
		return nil, err
	}

	item := &memcache.Item{
		Key:        key.Encode(),
		Value:      buf.Bytes(),
		Expiration: p.AccessToken.Expiry.Sub(now) + 10*time.Second,
	}

	if err = memcache.Set(ctx, item); err != nil {
		return nil, err
	}

	return key, nil
}

func (p *Passenger) check(raw []byte) error {
	digest := crypto.Hash(p.AccessToken.Hash).New().Sum(raw)

	if subtle.ConstantTimeCompare(digest, p.AccessToken.Digest) != 1 {
		return ErrDigestMismatch
	}

	if p.AccessToken.Expiry.Before(time.Now()) {
		return ErrTokenExpired
	}

	return nil
}

// IssueToken creates a new Token for the authenticated user. Callers should
// prefill the accessToken with whatever values they like and leave zero values
// to be set.
// The generated token will also be persisted and can be handed to the client
// with no more handling.
func (p *Passenger) IssueToken(ctx context.Context, token *model.AccessToken) (string, error) {
	now := time.Now()

	if token.Expiry == (time.Time{}) {
		token.Expiry = now.Add(defaultValidity)
	} else {
		if token.Expiry.Before(now.Add(minValidity)) {
			return "", fmt.Errorf("passenger: token must be valid at least %s", minValidity)
		}

		if token.Expiry.After(now.Add(maxValidity)) {
			return "", fmt.Errorf("passenger: token must be valid at most %s", maxValidity)
		}
	}

	// TODO(flowlo): This will reject all scopes for now, as we are not using them.
	// As soon as we introduce scopes, this check must be rewritten accordingly.
	if len(token.Scopes) > 0 {
		return "", fmt.Errorf("passenger: unknown scopes: %s", strings.Join(token.Scopes, ", "))
	}

	if len(token.Description) > 512 || len(token.Description) < 4 {
		return "", fmt.Errorf("passenger: description has bad len %d", len(token.Description))
	}

	var raw [tokenLength]byte

	if _, err := rand.Read(raw[:]); err != nil {
		return "", err
	}

	if token.Hash == 0 {
		token.Hash = int(defaultHash)
	}
	token.Digest = crypto.Hash(token.Hash).New().Sum(raw[:])

	clone := Passenger{
		UserKey:     p.UserKey,
		AccessToken: *token,
	}

	key, err := clone.Save(ctx)
	if err != nil {
		return "", err
	}

	return encodeToken(key, &raw)
}

// FromAccessToken tries do identify a Passenger by the access token he gave us.
// It will look up the AccessToken and consequently the corresponding User.
func FromAccessToken(ctx context.Context, accessToken string) (*Passenger, error) {
	key, raw, err := decodeToken(ctx, accessToken)
	if err != nil {
		return nil, err
	}

	var p *Passenger
	if p, err = fromCache(ctx, key); err != nil {
		if p, err = fromDatastore(ctx, key); err != nil {
			return nil, err
		}
	}

	return p, p.check(raw)
}

func fromCache(ctx context.Context, key *datastore.Key) (p *Passenger, err error) {
	item, err := memcache.Get(ctx, key.Encode())
	if err != nil {
		return nil, err
	}
	p = new(Passenger)
	err = gob.NewDecoder(bytes.NewReader(item.Value)).Decode(&p)
	return
}

func fromDatastore(ctx context.Context, key *datastore.Key) (p *Passenger, err error) {
	p = new(Passenger)
	if err = datastore.Get(ctx, key, p.AccessToken); err != nil {
		return
	}
	if p.UserKey = key.Parent(); p.UserKey == nil {
		return nil, ErrTokenNotAssociated
	}
	return
}

// FromBasicAuth tries do identify a Passenger by the access token he gave us.
// It will look up the the user by username and try to match password.
func FromBasicAuth(ctx context.Context, username, pw string) (p *Passenger, err error) {
	p = new(Passenger)
	var user model.User
	p.UserKey, err = model.NewQueryForUser().
		Filter("Nick=", username).
		Limit(1).
		Run(ctx).
		Next(&user)

	if err != nil {
		return
	}
	err = password.Check([]byte(pw), user.HashedPassword)

	// TODO(flowlo): Depending on bcrypt is very fragile. We
	// should encapsulate that.
	if err == bcrypt.ErrMismatchedHashAndPassword {
		userKey := p.UserKey
		p, err = FromAccessToken(ctx, pw)
		if err != nil {
			return
		}
		if !p.UserKey.Equal(userKey) {
			return nil, ErrTokenNotMatchingUser{Parent: p.UserKey, Actual: userKey}
		}
	}
	return
}

// FromRequest inspects the HTTP Authorization header of the given request
// and tries to identify a passenger.
func FromRequest(ctx context.Context, r *http.Request) (p *Passenger, err error) {
	auth := ""
	if auth = r.Header.Get("Authorization"); auth == "" {
		return nil, ErrNoAuthHeader
	}

	if strings.HasPrefix(auth, "Token ") {
		return FromAccessToken(ctx, auth[6:])
	}

	username, password, ok := "", "", false
	if username, password, ok = r.BasicAuth(); !ok {
		return nil, ErrUnkAuthHeader
	}

	p, err = FromBasicAuth(ctx, username, password)
	return
}

// FromContext returns the Passenger value stored in ctx if any.
func FromContext(ctx context.Context) (p *Passenger, ok bool) {
	p, ok = ctx.Value(passengerKey).(*Passenger)
	return
}

// NewContext returns a new Context that carries value p.
func NewContext(ctx context.Context, p *Passenger) context.Context {
	return context.WithValue(ctx, passengerKey, p)
}

// NewContextFromRequest wraps FromRequest and NewContext.
func NewContextFromRequest(ctx context.Context, r *http.Request) (context.Context, error) {
	p, err := FromRequest(ctx, r)
	if err != nil {
		return ctx, err
	}
	return NewContext(ctx, p), nil
}

// DecodeToken will take a token as sent by a client and translate it into a
// key to look up the full token on the server side and the raw secret.
func decodeToken(ctx context.Context, token string) (*datastore.Key, []byte, error) {
	b, err := encoding.DecodeString(token)
	if err != nil {
		return nil, nil, err
	}

	if len(b) != 3*8+tokenLength {
		return nil, nil, errors.New("token length mismatch")
	}

	var intIDs [3]int64
	for i := range intIDs {
		intID, n := binary.Varint(b[i*8 : (i+1)*8])
		if n < 8 {
			return nil, nil, errors.New("varint read mismatch")
		}
		intIDs[len(intIDs)-1-i] = intID
	}

	kinds := [3]string{model.CompanyKind, model.UserKind, model.AccessTokenKind}

	var key *datastore.Key
	for i := range intIDs {
		if intIDs[i] == 0 {
			continue
		}
		key = datastore.NewKey(ctx, kinds[i], "", intIDs[i], key)
		if key == nil {
			return nil, nil, errors.New("could not construct key from token")
		}
	}

	return key, b[3*8:], nil
}

// EncodeToken translates the key and raw secret of a newly generated token to
// a form suitable for the client.
func encodeToken(key *datastore.Key, raw *[tokenLength]byte) (string, error) {
	// Buffer size will be 8 (size of an int64) plus the length of the raw
	// token itself.
	var b [3*8 + tokenLength]byte

	for i := 0; i < 3; i++ {
		if n := binary.PutVarint(b[i*8:(i+1)*8], key.IntID()); n < 8 {
			return "", errors.New("short write when encoding token")
		}
		if key != nil {
			key = key.Parent()
		}
	}
	for i := range raw {
		b[3*8+i] = raw[i]
	}
	return encoding.EncodeToString(b[:]), nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
)

type Input struct {
	Name string `json:"name"`
}

func main() {
	for _, e := range os.Environ() {
		fmt.Println(e)
	}
	input := &Input{}
	if err := json.NewDecoder(os.Stdin).Decode(input); err != nil {
		log.Fatalln("Error! Bad input. ", err)
	}
	if input.Name == "" {
		input.Name = "World"
	}
	fmt.Printf("Hello %v!\n", input.Name)
}
<commit_msg>Ignores bad payload and just prints hello world for iron/hello.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
)

type Input struct {
	Name string `json:"name"`
}

func main() {
	input := &Input{}
	if err := json.NewDecoder(os.Stdin).Decode(input); err != nil {
		// log.Println("Bad payload or no payload. Ignoring.", err)
	}
	if input.Name == "" {
		input.Name = "World"
	}
	fmt.Printf("Hello %v!\n", input.Name)
}
<|endoftext|>


<commit_before>package farm

import (
	"sync"

	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/common"
	"github.com/soundcloud/roshi/instrumentation"
	"github.com/tsenart/tb"
)

// RepairStrategy generates a core repair strategy for a specific set of
// Clusters.
type RepairStrategy func([]cluster.Cluster, instrumentation.RepairInstrumentation) coreRepairStrategy

// coreRepairStrategy encodes one way of performing repair requests.
type coreRepairStrategy func(kms []keyMember)

// keyMember is like a KeyScoreMember, just without score. It is used to pass
// "suspects" to the core repair strategy.
type keyMember struct {
	Key    string
	Member string
}

// AllRepairs is a RepairStrategy that does what you expect: actually issuing
// the repairs with 100% probability. You may want to use RateLimitedRepairs
// to cap the load to your clusters.
func AllRepairs(clusters []cluster.Cluster, instr instrumentation.RepairInstrumentation) coreRepairStrategy {
	return func(kms []keyMember) {
		go func() {
			instr.RepairCall()
			instr.RepairRequest(len(kms))
		}()

		type scoreTuple struct {
			score       float64
			wasInserted bool
			err         error
		}

		type responseTuple struct {
			clusterIndex int
			keyMember
			scoreTuple
		}

		responses := make(chan responseTuple, len(kms)*len(clusters))

		// Scatter Score requests.
		for _, km := range kms {
			for i, c := range clusters {
				go func(km keyMember, i int, c cluster.Cluster) {
					score, wasInserted, err := c.Score(km.Key, km.Member)
					t := scoreTuple{score: score, wasInserted: wasInserted, err: err}
					responses <- responseTuple{clusterIndex: i, keyMember: km, scoreTuple: t}
				}(km, i, c)
			}
		}

		gathered := map[keyMember][]scoreTuple{}

		// Gather Score responses.
		for i := 0; i < cap(responses); i++ {
			resp := <-responses
			if _, ok := gathered[resp.keyMember]; !ok {
				gathered[resp.keyMember] = make([]scoreTuple, len(clusters))
			}
			gathered[resp.keyMember][resp.clusterIndex] = resp.scoreTuple
		}

		type repair struct {
			common.KeyScoreMember
		}

		inserts := map[int][]common.KeyScoreMember{} // cluster index: KeyScoreMembers to Insert
		deletes := map[int][]common.KeyScoreMember{} // cluster index: KeyScoreMembers to Delete

		// Determine which clusters require repairs.
		for km, scoreTuples := range gathered {
			// Walk once, to determine the correct data.
			highestScore, wasInserted := 0., false
			for _, scoreTuple := range scoreTuples {
				if scoreTuple.score > highestScore {
					highestScore = scoreTuple.score
					wasInserted = scoreTuple.wasInserted
				}
			}

			// Choose the correct target write operation for this keyMember.
			target := inserts
			if !wasInserted {
				target = deletes
			}

			// Walk again, to schedule any necessary repairs.
			for clusterIndex, scoreTuple := range scoreTuples {
				if scoreTuple.score < highestScore || scoreTuple.wasInserted != wasInserted {
					target[clusterIndex] = append(target[clusterIndex], common.KeyScoreMember{
						Key:    km.Key,
						Score:  highestScore,
						Member: km.Member,
					})
				}
			}
		}

		var wg sync.WaitGroup
		wg.Add(len(inserts) + len(deletes))
		do := func(write func([]common.KeyScoreMember) error, tuples []common.KeyScoreMember) {
			if err := write(tuples); err == nil {
				go instr.RepairWriteSuccess(len(tuples))
			} else {
				go instr.RepairWriteFailure(len(tuples))
			}
			wg.Done()
		}

		// Launch repairs.
		for clusterIndex, tuples := range inserts {
			write := clusters[clusterIndex].Insert
			tuples := tuples
			go do(write, tuples)
		}
		for clusterIndex, tuples := range deletes {
			write := clusters[clusterIndex].Delete
			tuples := tuples
			go do(write, tuples)
		}

		wg.Wait()
	}
}

// NoRepairs is a no-op repair strategy.
func NoRepairs([]cluster.Cluster, instrumentation.RepairInstrumentation) coreRepairStrategy {
	return func([]keyMember) {}
}

// RateLimitedRepairs is a RepairStrategy generator, which limits the total
// number of repaired key-member tuples to maxKeysPerSecond. Pass a negative
// value to allow unlimited repair.
func RateLimitedRepairs(maxKeysPerSecond int) RepairStrategy {
	return func(clusters []cluster.Cluster, instr instrumentation.RepairInstrumentation) coreRepairStrategy {
		permits := permitter(allowAllPermitter{})
		if maxKeysPerSecond >= 0 {
			permits = permitter(tokenBucketPermitter{tb.NewBucket(int64(maxKeysPerSecond), 1)})
		}
		return func(kms []keyMember) {
			if n := len(kms); !permits.canHas(int64(n)) {
				instr.RepairDiscarded(n)
				return
			}
			AllRepairs(clusters, instr)(kms)
		}
	}
}
<commit_msg>farm: improve encapsulation of RateLimitedRepairs<commit_after>package farm

import (
	"sync"

	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/common"
	"github.com/soundcloud/roshi/instrumentation"
	"github.com/tsenart/tb"
)

// RepairStrategy generates a core repair strategy for a specific set of
// Clusters.
type RepairStrategy func([]cluster.Cluster, instrumentation.RepairInstrumentation) coreRepairStrategy

// coreRepairStrategy encodes one way of performing repair requests.
type coreRepairStrategy func(kms []keyMember)

// keyMember is like a KeyScoreMember, just without score. It is used to pass
// "suspects" to the core repair strategy.
type keyMember struct {
	Key    string
	Member string
}

// AllRepairs is a RepairStrategy that does what you expect: actually issuing
// the repairs with 100% probability. You may want to use RateLimitedRepairs
// to cap the load to your clusters.
func AllRepairs(clusters []cluster.Cluster, instr instrumentation.RepairInstrumentation) coreRepairStrategy {
	return func(kms []keyMember) {
		go func() {
			instr.RepairCall()
			instr.RepairRequest(len(kms))
		}()

		type scoreTuple struct {
			score       float64
			wasInserted bool
			err         error
		}

		type responseTuple struct {
			clusterIndex int
			keyMember
			scoreTuple
		}

		responses := make(chan responseTuple, len(kms)*len(clusters))

		// Scatter Score requests.
		for _, km := range kms {
			for i, c := range clusters {
				go func(km keyMember, i int, c cluster.Cluster) {
					score, wasInserted, err := c.Score(km.Key, km.Member)
					t := scoreTuple{score: score, wasInserted: wasInserted, err: err}
					responses <- responseTuple{clusterIndex: i, keyMember: km, scoreTuple: t}
				}(km, i, c)
			}
		}

		gathered := map[keyMember][]scoreTuple{}

		// Gather Score responses.
		for i := 0; i < cap(responses); i++ {
			resp := <-responses
			if _, ok := gathered[resp.keyMember]; !ok {
				gathered[resp.keyMember] = make([]scoreTuple, len(clusters))
			}
			gathered[resp.keyMember][resp.clusterIndex] = resp.scoreTuple
		}

		type repair struct {
			common.KeyScoreMember
		}

		inserts := map[int][]common.KeyScoreMember{} // cluster index: KeyScoreMembers to Insert
		deletes := map[int][]common.KeyScoreMember{} // cluster index: KeyScoreMembers to Delete

		// Determine which clusters require repairs.
		for km, scoreTuples := range gathered {
			// Walk once, to determine the correct data.
			highestScore, wasInserted := 0., false
			for _, scoreTuple := range scoreTuples {
				if scoreTuple.score > highestScore {
					highestScore = scoreTuple.score
					wasInserted = scoreTuple.wasInserted
				}
			}

			// Choose the correct target write operation for this keyMember.
			target := inserts
			if !wasInserted {
				target = deletes
			}

			// Walk again, to schedule any necessary repairs.
			for clusterIndex, scoreTuple := range scoreTuples {
				if scoreTuple.score < highestScore || scoreTuple.wasInserted != wasInserted {
					target[clusterIndex] = append(target[clusterIndex], common.KeyScoreMember{
						Key:    km.Key,
						Score:  highestScore,
						Member: km.Member,
					})
				}
			}
		}

		var wg sync.WaitGroup
		wg.Add(len(inserts) + len(deletes))
		do := func(write func([]common.KeyScoreMember) error, tuples []common.KeyScoreMember) {
			if err := write(tuples); err == nil {
				go instr.RepairWriteSuccess(len(tuples))
			} else {
				go instr.RepairWriteFailure(len(tuples))
			}
			wg.Done()
		}

		// Launch repairs.
		for clusterIndex, tuples := range inserts {
			write := clusters[clusterIndex].Insert
			tuples := tuples
			go do(write, tuples)
		}
		for clusterIndex, tuples := range deletes {
			write := clusters[clusterIndex].Delete
			tuples := tuples
			go do(write, tuples)
		}

		wg.Wait()
	}
}

// NoRepairs is a no-op repair strategy.
func NoRepairs([]cluster.Cluster, instrumentation.RepairInstrumentation) coreRepairStrategy {
	return func([]keyMember) {}
}

// RateLimitedRepairs is a RepairStrategy generator, which limits the total
// number of repaired key-member tuples to maxKeysPerSecond. Pass a negative
// value to allow unlimited repair.
func RateLimitedRepairs(maxKeysPerSecond int) RepairStrategy {
	permits := permitter(allowAllPermitter{})
	if maxKeysPerSecond >= 0 {
		permits = permitter(tokenBucketPermitter{tb.NewBucket(int64(maxKeysPerSecond), 1)})
	}
	return func(clusters []cluster.Cluster, instr instrumentation.RepairInstrumentation) coreRepairStrategy {
		return func(kms []keyMember) {
			if n := len(kms); !permits.canHas(int64(n)) {
				instr.RepairDiscarded(n)
				return
			}
			AllRepairs(clusters, instr)(kms)
		}
	}
}
<|endoftext|>


<commit_before>package buffalo

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/websocket"

	"github.com/gobuffalo/buffalo/render"
	"github.com/markbates/willie"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

func basicContext() DefaultContext {
	return DefaultContext{
		Context: context.Background(),
		logger:  NewLogger("debug"),
		data:    make(map[string]interface{}),
		flash:   &Flash{data: make(map[string][]string)},
	}
}

func Test_DefaultContext_Param(t *testing.T) {
	r := require.New(t)
	c := DefaultContext{
		params: url.Values{
			"name": []string{"Mark"},
		},
	}

	r.Equal("Mark", c.Param("name"))
}

func Test_DefaultContext_ParamInt(t *testing.T) {
	r := require.New(t)
	c := DefaultContext{
		params: url.Values{
			"name": []string{"Mark"},
			"id":   []string{"1"},
		},
	}

	id, err := c.ParamInt("id")
	r.NoError(err)
	r.Equal(1, id)

	_, err = c.ParamInt("badkey")
	r.Error(err)

	_, err = c.ParamInt("name")
	r.Error(err)
}

func Test_DefaultContext_GetSet(t *testing.T) {
	r := require.New(t)
	c := basicContext()
	r.Nil(c.Get("name"))

	c.Set("name", "Mark")
	r.NotNil(c.Get("name"))
	r.Equal("Mark", c.Get("name").(string))
}

func Test_DefaultContext_Value(t *testing.T) {
	r := require.New(t)
	c := basicContext()
	r.Nil(c.Value("name"))

	c.Set("name", "Mark")
	r.NotNil(c.Value("name"))
	r.Equal("Mark", c.Value("name").(string))
	r.Equal("Mark", c.Get("name").(string))
}

func Test_DefaultContext_Render(t *testing.T) {
	r := require.New(t)

	c := basicContext()
	res := httptest.NewRecorder()
	c.response = res
	c.params = url.Values{"name": []string{"Mark"}}
	c.Set("greet", "Hello")

	err := c.Render(123, render.String("{{greet}} {{params.name}}!"))
	r.NoError(err)

	r.Equal(123, res.Code)
	r.Equal("Hello Mark!", res.Body.String())
}

func Test_DefaultContext_Bind_Default(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	uv := url.Values{"first_name": []string{"Mark"}}
	res := w.Request("/").Post(uv)
	r.Equal(201, res.Code)

	r.Equal("Mark", user.FirstName)
}

func Test_DefaultContext_Bind_No_ContentType(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{
		FirstName: "Mark",
	}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return c.Error(422, err)
		}
		return c.Render(201, nil)
	})

	bb := &bytes.Buffer{}
	req, err := http.NewRequest("POST", "/", bb)
	r.NoError(err)
	req.Header.Del("Content-Type")
	res := httptest.NewRecorder()
	a.ServeHTTP(res, req)
	r.Equal(422, res.Code)
	r.Contains(res.Body.String(), "blank content type")
}

func Test_DefaultContext_Bind_Default_BlankFields(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{
		FirstName: "Mark",
	}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	uv := url.Values{"first_name": []string{""}}
	res := w.Request("/").Post(uv)
	r.Equal(201, res.Code)

	r.Equal("", user.FirstName)
}

func Test_DefaultContext_Bind_JSON(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `json:"first_name"`
	}{}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	res := w.JSON("/").Post(map[string]string{
		"first_name": "Mark",
	})
	r.Equal(201, res.Code)

	r.Equal("Mark", user.FirstName)
}

func Test_DefaultContext_Websocket(t *testing.T) {
	r := require.New(t)

	type Message struct {
		Original  string    `json:"original"`
		Formatted string    `json:"formatted"`
		Received  time.Time `json:"received"`
	}

	a := Automatic(Options{})
	a.GET("/socket", func(c Context) error {
		conn, err := c.Websocket()
		if err != nil {
			return err
		}
		for {

			_, m, err := conn.ReadMessage()
			if err != nil {
				return err
			}

			data := string(m)

			msg := Message{
				Original:  data,
				Formatted: strings.ToUpper(data),
				Received:  time.Now(),
			}

			if err := conn.WriteJSON(msg); err != nil {
				return err
			}
		}
	})

	ts := httptest.NewServer(a)
	defer ts.Close()

	wsURL := strings.Replace(ts.URL, "http", "ws", 1) + "/socket"

	ws, err := websocket.Dial(wsURL, "", "http://127.0.0.1")
	r.NoError(err)

	_, err = ws.Write([]byte("hello, world!"))
	r.NoError(err)

	msg := make([]byte, 512)
	read, err := ws.Read(msg)
	r.NoError(err)

	var message Message
	err = json.NewDecoder(bytes.NewReader(msg[:read])).Decode(&message)
	r.NoError(err)

	// Create a table of what we expect.
	tests := []struct {
		Got  string
		Want string
	}{
		{message.Formatted, "HELLO, WORLD!"},
		{message.Original, "hello, world!"},
	}

	// Check the different fields.
	for _, tt := range tests {
		r.Equal(tt.Want, tt.Got)
	}
}
<commit_msg>add test for empty value<commit_after>package buffalo

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/websocket"

	"github.com/gobuffalo/buffalo/render"
	"github.com/markbates/willie"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

func basicContext() DefaultContext {
	return DefaultContext{
		Context: context.Background(),
		logger:  NewLogger("debug"),
		data:    make(map[string]interface{}),
		flash:   &Flash{data: make(map[string][]string)},
	}
}

func Test_DefaultContext_Param(t *testing.T) {
	r := require.New(t)
	c := DefaultContext{
		params: url.Values{
			"name": []string{"Mark"},
		},
	}

	r.Equal("Mark", c.Param("name"))
}

func Test_DefaultContext_ParamInt(t *testing.T) {
	r := require.New(t)
	c := DefaultContext{
		params: url.Values{
			"name": []string{"Mark"},
			"id":   []string{"1"},
		},
	}

	id, err := c.ParamInt("id")
	r.NoError(err)
	r.Equal(1, id)

	_, err = c.ParamInt("badkey")
	r.Error(err)

	_, err = c.ParamInt("name")
	r.Error(err)
}

func Test_DefaultContext_GetSet(t *testing.T) {
	r := require.New(t)
	c := basicContext()
	r.Nil(c.Get("name"))

	c.Set("name", "Mark")
	r.NotNil(c.Get("name"))
	r.Equal("Mark", c.Get("name").(string))
}

func Test_DefaultContext_Value(t *testing.T) {
	r := require.New(t)
	c := basicContext()
	r.Nil(c.Value("name"))

	c.Set("name", "Mark")
	r.NotNil(c.Value("name"))
	r.Equal("Mark", c.Value("name").(string))
	r.Equal("Mark", c.Get("name").(string))
}

func Test_DefaultContext_Render(t *testing.T) {
	r := require.New(t)

	c := basicContext()
	res := httptest.NewRecorder()
	c.response = res
	c.params = url.Values{"name": []string{"Mark"}}
	c.Set("greet", "Hello")

	err := c.Render(123, render.String("{{greet}} {{params.name}}!"))
	r.NoError(err)

	r.Equal(123, res.Code)
	r.Equal("Hello Mark!", res.Body.String())
}

func Test_DefaultContext_Bind_Default(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	uv := url.Values{"first_name": []string{"Mark"}}
	res := w.Request("/").Post(uv)
	r.Equal(201, res.Code)

	r.Equal("Mark", user.FirstName)
}

func Test_DefaultContext_Bind_No_ContentType(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{
		FirstName: "Mark",
	}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return c.Error(422, err)
		}
		return c.Render(201, nil)
	})

	bb := &bytes.Buffer{}
	req, err := http.NewRequest("POST", "/", bb)
	r.NoError(err)
	req.Header.Del("Content-Type")
	res := httptest.NewRecorder()
	a.ServeHTTP(res, req)
	r.Equal(422, res.Code)
	r.Contains(res.Body.String(), "blank content type")
}

func Test_DefaultContext_Bind_Empty_ContentType(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{
		FirstName: "Mark",
	}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return c.Error(422, err)
		}
		return c.Render(201, nil)
	})

	bb := &bytes.Buffer{}
	req, err := http.NewRequest("POST", "/", bb)
	r.NoError(err)
	// Want to make sure that an empty string value does not cause an error on `split`
	req.Header.Set("Content-Type", "")
	res := httptest.NewRecorder()
	a.ServeHTTP(res, req)
	r.Equal(422, res.Code)
	r.Contains(res.Body.String(), "blank content type")
}

func Test_DefaultContext_Bind_Default_BlankFields(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `schema:"first_name"`
	}{
		FirstName: "Mark",
	}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	uv := url.Values{"first_name": []string{""}}
	res := w.Request("/").Post(uv)
	r.Equal(201, res.Code)

	r.Equal("", user.FirstName)
}

func Test_DefaultContext_Bind_JSON(t *testing.T) {
	r := require.New(t)

	user := struct {
		FirstName string `json:"first_name"`
	}{}

	a := New(Options{})
	a.POST("/", func(c Context) error {
		err := c.Bind(&user)
		if err != nil {
			return errors.WithStack(err)
		}
		return c.Render(201, nil)
	})

	w := willie.New(a)
	res := w.JSON("/").Post(map[string]string{
		"first_name": "Mark",
	})
	r.Equal(201, res.Code)

	r.Equal("Mark", user.FirstName)
}

func Test_DefaultContext_Websocket(t *testing.T) {
	r := require.New(t)

	type Message struct {
		Original  string    `json:"original"`
		Formatted string    `json:"formatted"`
		Received  time.Time `json:"received"`
	}

	a := Automatic(Options{})
	a.GET("/socket", func(c Context) error {
		conn, err := c.Websocket()
		if err != nil {
			return err
		}
		for {

			_, m, err := conn.ReadMessage()
			if err != nil {
				return err
			}

			data := string(m)

			msg := Message{
				Original:  data,
				Formatted: strings.ToUpper(data),
				Received:  time.Now(),
			}

			if err := conn.WriteJSON(msg); err != nil {
				return err
			}
		}
	})

	ts := httptest.NewServer(a)
	defer ts.Close()

	wsURL := strings.Replace(ts.URL, "http", "ws", 1) + "/socket"

	ws, err := websocket.Dial(wsURL, "", "http://127.0.0.1")
	r.NoError(err)

	_, err = ws.Write([]byte("hello, world!"))
	r.NoError(err)

	msg := make([]byte, 512)
	read, err := ws.Read(msg)
	r.NoError(err)

	var message Message
	err = json.NewDecoder(bytes.NewReader(msg[:read])).Decode(&message)
	r.NoError(err)

	// Create a table of what we expect.
	tests := []struct {
		Got  string
		Want string
	}{
		{message.Formatted, "HELLO, WORLD!"},
		{message.Original, "hello, world!"},
	}

	// Check the different fields.
	for _, tt := range tests {
		r.Equal(tt.Want, tt.Got)
	}
}
<|endoftext|>


<commit_before>package amazonamimanagement

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`

	Identifier   string `mapstructure:"identifier"`
	KeepReleases int    `mapstructure:"keep_releases"`

	ctx interpolate.Context
}

type PostProcessor struct {
	ec2conn ec2iface.EC2API
	config  Config
}

func (p *PostProcessor) Configure(raws ...interface{}) error {
	p.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)
	if err != nil {
		return err
	}

	return nil
}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	log.Println("Running Amazon AMI Management post-processor")

	ec2conn := p.ec2conn
	if ec2conn == nil {
		// If no ec2conn is set, then we use the real connection
		config, err := p.config.Config()
		if err != nil {
			return nil, true, err
		}
		log.Println("Creating AWS session")
		session := session.New(config)
		ec2conn = ec2.New(session)
	}

	log.Println("Describing images for generation management")
	output, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String("tag:Amazon_AMI_Management_Identifier"),
				Values: []*string{
					aws.String(p.config.Identifier),
				},
			},
		},
	})
	if err != nil {
		return nil, true, err
	}

	// AMIs are sorted in descending order by creation date
	sort(
		len(output.Images),
		func(i, j int) bool {
			iTime, _ := time.Parse("2006-01-02T15:04:05.000Z", *output.Images[i].CreationDate)
			jTime, _ := time.Parse("2006-01-02T15:04:05.000Z", *output.Images[j].CreationDate)
			return iTime.After(jTime)
		},
		func(i, j int) {
			output.Images[i], output.Images[j] = output.Images[j], output.Images[i]
		},
	)

	log.Println("Deleting old images...")
	for i, image := range output.Images {
		if i < p.config.KeepReleases {
			continue
		}
		ui.Message(fmt.Sprintf("Deleting image: %s", *image.ImageId))
		log.Printf("Deleting image AMI (%s)", *image.ImageId)
		if _, err := ec2conn.DeregisterImage(&ec2.DeregisterImageInput{
			ImageId: image.ImageId,
		}); err != nil {
			return nil, true, err
		}

		// DeregisterImage method only perform to AMI
		// Because it retain snapshots. Following operation is deleting snapshots.
		log.Printf("Deleting snapshot related to AMI (%s)", *image.ImageId)
		for _, device := range image.BlockDeviceMappings {
			// skip delete if use ephemeral devise
			if device.Ebs == nil {
				continue
			}
			log.Printf("Deleting snapshot (%s) related to AMI (%s)", *device.Ebs.SnapshotId, *image.ImageId)
			if _, err := ec2conn.DeleteSnapshot(&ec2.DeleteSnapshotInput{
				SnapshotId: device.Ebs.SnapshotId,
			}); err != nil {
				return nil, true, err
			}
		}
	}

	return artifact, true, nil
}

func sort(len int, lessFunc func(i, j int) bool, swapFunc func(i, j int)) error {
	for n := 0; n < len-1; n++ {
		for m := len - 1; m > n; m-- {
			if !lessFunc(m-1, m) {
				swapFunc(m-1, m)
			}
		}
	}
	return nil
}
<commit_msg>Use own credentials<commit_after>package amazonamimanagement

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`

	Identifier   string `mapstructure:"identifier"`
	KeepReleases int    `mapstructure:"keep_releases"`
	AccessKey    string `mapstructure:"access_key"`
	SecretKey    string `mapstructure:"secret_key"`
	Region       string `mapstructure:"region"`

	ctx interpolate.Context
}

type PostProcessor struct {
	ec2conn ec2iface.EC2API
	config  Config
}

func (p *PostProcessor) Configure(raws ...interface{}) error {
	p.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)
	if err != nil {
		return err
	}

	return nil
}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	log.Println("Running Amazon AMI Management post-processor")

	ec2conn := p.ec2conn
	if ec2conn == nil {
		// If no ec2conn is set, then we use the real connection
		config := aws.NewConfig().WithRegion(p.config.Region).WithMaxRetries(11)
		sess := session.New(config)
		creds := credentials.NewChainCredentials([]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     p.config.AccessKey,
				SecretAccessKey: p.config.SecretKey,
			}},
			&credentials.EnvProvider{},
			&credentials.SharedCredentialsProvider{Filename: "", Profile: ""},
			&ec2rolecreds.EC2RoleProvider{
				Client: ec2metadata.New(sess),
			},
		})

		log.Println("Creating AWS session")
		ec2Session := session.New(config.WithCredentials(creds))
		ec2conn = ec2.New(ec2Session)
	}

	log.Println("Describing images for generation management")
	output, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{
		Filters: []*ec2.Filter{
			{
				Name: aws.String("tag:Amazon_AMI_Management_Identifier"),
				Values: []*string{
					aws.String(p.config.Identifier),
				},
			},
		},
	})
	if err != nil {
		return nil, true, err
	}

	// AMIs are sorted in descending order by creation date
	sort(
		len(output.Images),
		func(i, j int) bool {
			iTime, _ := time.Parse("2006-01-02T15:04:05.000Z", *output.Images[i].CreationDate)
			jTime, _ := time.Parse("2006-01-02T15:04:05.000Z", *output.Images[j].CreationDate)
			return iTime.After(jTime)
		},
		func(i, j int) {
			output.Images[i], output.Images[j] = output.Images[j], output.Images[i]
		},
	)

	log.Println("Deleting old images...")
	for i, image := range output.Images {
		if i < p.config.KeepReleases {
			continue
		}
		ui.Message(fmt.Sprintf("Deleting image: %s", *image.ImageId))
		log.Printf("Deleting image AMI (%s)", *image.ImageId)
		if _, err := ec2conn.DeregisterImage(&ec2.DeregisterImageInput{
			ImageId: image.ImageId,
		}); err != nil {
			return nil, true, err
		}

		// DeregisterImage method only perform to AMI
		// Because it retain snapshots. Following operation is deleting snapshots.
		log.Printf("Deleting snapshot related to AMI (%s)", *image.ImageId)
		for _, device := range image.BlockDeviceMappings {
			// skip delete if use ephemeral devise
			if device.Ebs == nil {
				continue
			}
			log.Printf("Deleting snapshot (%s) related to AMI (%s)", *device.Ebs.SnapshotId, *image.ImageId)
			if _, err := ec2conn.DeleteSnapshot(&ec2.DeleteSnapshotInput{
				SnapshotId: device.Ebs.SnapshotId,
			}); err != nil {
				return nil, true, err
			}
		}
	}

	return artifact, true, nil
}

func sort(len int, lessFunc func(i, j int) bool, swapFunc func(i, j int)) error {
	for n := 0; n < len-1; n++ {
		for m := len - 1; m > n; m-- {
			if !lessFunc(m-1, m) {
				swapFunc(m-1, m)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2021 PingCAP, 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 placement

import (
	"fmt"
)

const (
	// BundleIDPrefix is the bundle prefix of all rule bundles from TiDB_DDL statements.
	BundleIDPrefix = "TiDB_DDL_"
	// PDBundleID is the bundle name of pd, the default bundle for all regions.
	PDBundleID = "pd"
)

// GroupID accepts a tableID or whatever integer, and encode the integer into a valid GroupID for PD.
func GroupID(id int64) string {
	return fmt.Sprintf("%s%d", BundleIDPrefix, id)
}

const (
	// RuleIndexDefault is the default index for a rule, check Rule.Index.
	RuleIndexDefault int = iota
	// RuleIndexDatabase is the index for a rule of database.
	RuleIndexDatabase
	// RuleIndexTable is the index for a rule of table.
	RuleIndexTable
	// RuleIndexPartition is the index for a rule of partition.
	RuleIndexPartition
	// RuleIndexIndex is the index for a rule of index.
	RuleIndexIndex
)

const (
	// DCLabelKey indicates the key of label which represents the dc for Store.
	// FIXME: currently we assumes "zone" is the dcLabel key in Store
	DCLabelKey = "zone"
	// EngineLabelKey is the label that indicates the backend of store instance:
	// tikv or tiflash. TiFlash instance will contain a label of 'engine: tiflash'.
	EngineLabelKey = "engine"
	// EngineLabelTiFlash is the label value, which a TiFlash instance will have with
	// a label key of EngineLabelKey.
	EngineLabelTiFlash = "tiflash"
	// EngineLabelTiKV is the label value used in some tests. And possibly TiKV will
	// set the engine label with a value of EngineLabelTiKV.
	EngineLabelTiKV = "tikv"
)
<commit_msg>*: Modify placement rule index to reserve some indexes for future work (#30737)<commit_after>// Copyright 2021 PingCAP, 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 placement

import (
	"fmt"
)

const (
	// BundleIDPrefix is the bundle prefix of all rule bundles from TiDB_DDL statements.
	BundleIDPrefix = "TiDB_DDL_"
	// PDBundleID is the bundle name of pd, the default bundle for all regions.
	PDBundleID = "pd"
)

// GroupID accepts a tableID or whatever integer, and encode the integer into a valid GroupID for PD.
func GroupID(id int64) string {
	return fmt.Sprintf("%s%d", BundleIDPrefix, id)
}

const (
	// RuleIndexTable is the index for a rule of table.
	RuleIndexTable = 40
	// RuleIndexPartition is the index for a rule of partition.
	RuleIndexPartition = 80
)

const (
	// DCLabelKey indicates the key of label which represents the dc for Store.
	// FIXME: currently we assumes "zone" is the dcLabel key in Store
	DCLabelKey = "zone"
	// EngineLabelKey is the label that indicates the backend of store instance:
	// tikv or tiflash. TiFlash instance will contain a label of 'engine: tiflash'.
	EngineLabelKey = "engine"
	// EngineLabelTiFlash is the label value, which a TiFlash instance will have with
	// a label key of EngineLabelKey.
	EngineLabelTiFlash = "tiflash"
	// EngineLabelTiKV is the label value used in some tests. And possibly TiKV will
	// set the engine label with a value of EngineLabelTiKV.
	EngineLabelTiKV = "tikv"
)
<|endoftext|>


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

package fslock_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	. "launchpad.net/gocheck"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/utils/fslock"
	"launchpad.net/tomb"
)

func Test(t *testing.T) {
	TestingT(t)
}

type fslockSuite struct {
	coretesting.LoggingSuite
	lockDelay time.Duration
}

var _ = Suite(&fslockSuite{})

func (s *fslockSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	s.lockDelay = fslock.SetLockWaitDelay(1 * time.Millisecond)
}

func (s *fslockSuite) TearDownSuite(c *C) {
	fslock.SetLockWaitDelay(s.lockDelay)
	s.LoggingSuite.TearDownSuite(c)
}

// This test also happens to test that locks can get created when the parent
// lock directory doesn't exist.
func (s *fslockSuite) TestValidNamesLockDir(c *C) {

	for _, name := range []string{
		"a",
		"longer",
		"longer-with.special-characters",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, IsNil)
	}
}

func (s *fslockSuite) TestInvalidNames(c *C) {

	for _, name := range []string{
		".start",
		"-start",
		"NoCapitals",
		"no+plus",
		"no/slash",
		"no\\backslash",
		"no$dollar",
		"no:colon",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, ErrorMatches, "Invalid lock name .*")
	}
}

func (s *fslockSuite) TestNewLockWithExistingDir(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	_, err = fslock.NewLock(dir, "special")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestNewLockWithExistingFileInPlace(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	path := path.Join(dir, "locks")
	err = ioutil.WriteFile(path, []byte("foo"), 0644)
	c.Assert(err, IsNil)

	_, err = fslock.NewLock(path, "special")
	c.Assert(err, ErrorMatches, `.* not a directory`)
}

func (s *fslockSuite) TestIsLockHeldBasics(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)

	err = lock.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, true)

	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestIsLockHeldTwoLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestLockBlocks(c *C) {

	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	acquired := make(chan struct{})
	err = lock1.Lock("")
	c.Assert(err, IsNil)

	go func() {
		lock2.Lock("")
		acquired <- struct{}{}
		close(acquired)
	}()

	// Waiting for something not to happen is inherently hard...
	select {
	case <-acquired:
		c.Fatalf("Unexpected lock acquisition")
	case <-time.After(coretesting.ShortWait):
		// all good
	}

	err = lock1.Unlock()
	c.Assert(err, IsNil)

	select {
	case <-acquired:
		// all good
	case <-time.After(coretesting.ShortWait):
		c.Fatalf("Expected lock acquisition")
	}

	c.Assert(lock2.IsLockHeld(), Equals, true)
}

func (s *fslockSuite) TestLockWithTimeoutUnlocked(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestLockWithTimeoutLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, Equals, fslock.ErrTimeout)
}

func (s *fslockSuite) TestUnlock(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)
}

func (s *fslockSuite) TestIsLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	c.Assert(lock1.IsLocked(), Equals, true)
	c.Assert(lock2.IsLocked(), Equals, true)
}

func (s *fslockSuite) TestBreakLock(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.BreakLock()
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLocked(), Equals, false)

	// Normally locks are broken due to client crashes, not duration.
	err = lock1.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	// Breaking a non-existant isn't an error
	err = lock2.BreakLock()
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestMessage(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")

	err = lock.Lock("my message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "my message")

	// Unlocking removes the message.
	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")
}

func (s *fslockSuite) TestMessageAcrossLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("very busy")
	c.Assert(err, IsNil)
	c.Assert(lock2.Message(), Equals, "very busy")
}

func (s *fslockSuite) TestInitialMessageWhenLocking(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Lock("initial message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial message")

	err = lock.Unlock()
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "initial timeout message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial timeout message")
}

func (s *fslockSuite) TestStress(c *C) {
	const lockAttempts = 200
	const concurrentLocks = 10

	var counter = new(int64)
	// Use atomics to update lockState to make sure the lock isn't held by
	// someone else. A value of 1 means locked, 0 means unlocked.
	var lockState = new(int32)
	var done = make(chan struct{})
	defer close(done)

	dir := c.MkDir()

	var stress = func(name string) {
		defer func() { done <- struct{}{} }()
		lock, err := fslock.NewLock(dir, "testing")
		if err != nil {
			c.Errorf("Failed to create a new lock")
			return
		}
		for i := 0; i < lockAttempts; i++ {
			err = lock.Lock(name)
			c.Assert(err, IsNil)
			state := atomic.AddInt32(lockState, 1)
			c.Assert(state, Equals, int32(1))
			// Tell the go routine scheduler to give a slice to someone else
			// while we have this locked.
			runtime.Gosched()
			// need to decrement prior to unlock to avoid the race of someone
			// else grabbing the lock before we decrement the state.
			atomic.AddInt32(lockState, -1)
			err = lock.Unlock()
			c.Assert(err, IsNil)
			// increment the general counter
			atomic.AddInt64(counter, 1)
		}
	}

	for i := 0; i < concurrentLocks; i++ {
		go stress(fmt.Sprintf("Lock %d", i))
	}
	for i := 0; i < concurrentLocks; i++ {
		<-done
	}
	c.Assert(*counter, Equals, int64(lockAttempts*concurrentLocks))
}

func (s *fslockSuite) TestTomb(c *C) {
	const timeToDie = 200 * time.Millisecond
	die := tomb.Tomb{}

	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	// Just use one lock, and try to lock it twice.
	err = lock.Lock("very busy")
	c.Assert(err, IsNil)

	checkTomb := func() error {
		select {
		case <-die.Dying():
			return tomb.ErrDying
		default:
			// no-op to fall through to return.
		}
		return nil
	}

	go func() {
		time.Sleep(timeToDie)
		die.Killf("time to die")
	}()

	err = lock.LockWithFunc("won't happen", checkTomb)
	c.Assert(err, Equals, tomb.ErrDying)
	c.Assert(lock.Message(), Equals, "very busy")

}
<commit_msg>utils/fslock: fix timeout in test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package fslock_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	. "launchpad.net/gocheck"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/utils/fslock"
	"launchpad.net/tomb"
)

func Test(t *testing.T) {
	TestingT(t)
}

type fslockSuite struct {
	coretesting.LoggingSuite
	lockDelay time.Duration
}

var _ = Suite(&fslockSuite{})

func (s *fslockSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	s.lockDelay = fslock.SetLockWaitDelay(1 * time.Millisecond)
}

func (s *fslockSuite) TearDownSuite(c *C) {
	fslock.SetLockWaitDelay(s.lockDelay)
	s.LoggingSuite.TearDownSuite(c)
}

// This test also happens to test that locks can get created when the parent
// lock directory doesn't exist.
func (s *fslockSuite) TestValidNamesLockDir(c *C) {

	for _, name := range []string{
		"a",
		"longer",
		"longer-with.special-characters",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, IsNil)
	}
}

func (s *fslockSuite) TestInvalidNames(c *C) {

	for _, name := range []string{
		".start",
		"-start",
		"NoCapitals",
		"no+plus",
		"no/slash",
		"no\\backslash",
		"no$dollar",
		"no:colon",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, ErrorMatches, "Invalid lock name .*")
	}
}

func (s *fslockSuite) TestNewLockWithExistingDir(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	_, err = fslock.NewLock(dir, "special")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestNewLockWithExistingFileInPlace(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	path := path.Join(dir, "locks")
	err = ioutil.WriteFile(path, []byte("foo"), 0644)
	c.Assert(err, IsNil)

	_, err = fslock.NewLock(path, "special")
	c.Assert(err, ErrorMatches, `.* not a directory`)
}

func (s *fslockSuite) TestIsLockHeldBasics(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)

	err = lock.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, true)

	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestIsLockHeldTwoLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestLockBlocks(c *C) {

	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	acquired := make(chan struct{})
	err = lock1.Lock("")
	c.Assert(err, IsNil)

	go func() {
		lock2.Lock("")
		acquired <- struct{}{}
		close(acquired)
	}()

	// Waiting for something not to happen is inherently hard...
	select {
	case <-acquired:
		c.Fatalf("Unexpected lock acquisition")
	case <-time.After(coretesting.ShortWait):
		// all good
	}

	err = lock1.Unlock()
	c.Assert(err, IsNil)

	select {
	case <-acquired:
		// all good
	case <-time.After(coretesting.LongWait):
		c.Fatalf("Expected lock acquisition")
	}

	c.Assert(lock2.IsLockHeld(), Equals, true)
}

func (s *fslockSuite) TestLockWithTimeoutUnlocked(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestLockWithTimeoutLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, Equals, fslock.ErrTimeout)
}

func (s *fslockSuite) TestUnlock(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)
}

func (s *fslockSuite) TestIsLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	c.Assert(lock1.IsLocked(), Equals, true)
	c.Assert(lock2.IsLocked(), Equals, true)
}

func (s *fslockSuite) TestBreakLock(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.BreakLock()
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLocked(), Equals, false)

	// Normally locks are broken due to client crashes, not duration.
	err = lock1.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	// Breaking a non-existant isn't an error
	err = lock2.BreakLock()
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestMessage(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")

	err = lock.Lock("my message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "my message")

	// Unlocking removes the message.
	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")
}

func (s *fslockSuite) TestMessageAcrossLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("very busy")
	c.Assert(err, IsNil)
	c.Assert(lock2.Message(), Equals, "very busy")
}

func (s *fslockSuite) TestInitialMessageWhenLocking(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Lock("initial message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial message")

	err = lock.Unlock()
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "initial timeout message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial timeout message")
}

func (s *fslockSuite) TestStress(c *C) {
	const lockAttempts = 200
	const concurrentLocks = 10

	var counter = new(int64)
	// Use atomics to update lockState to make sure the lock isn't held by
	// someone else. A value of 1 means locked, 0 means unlocked.
	var lockState = new(int32)
	var done = make(chan struct{})
	defer close(done)

	dir := c.MkDir()

	var stress = func(name string) {
		defer func() { done <- struct{}{} }()
		lock, err := fslock.NewLock(dir, "testing")
		if err != nil {
			c.Errorf("Failed to create a new lock")
			return
		}
		for i := 0; i < lockAttempts; i++ {
			err = lock.Lock(name)
			c.Assert(err, IsNil)
			state := atomic.AddInt32(lockState, 1)
			c.Assert(state, Equals, int32(1))
			// Tell the go routine scheduler to give a slice to someone else
			// while we have this locked.
			runtime.Gosched()
			// need to decrement prior to unlock to avoid the race of someone
			// else grabbing the lock before we decrement the state.
			atomic.AddInt32(lockState, -1)
			err = lock.Unlock()
			c.Assert(err, IsNil)
			// increment the general counter
			atomic.AddInt64(counter, 1)
		}
	}

	for i := 0; i < concurrentLocks; i++ {
		go stress(fmt.Sprintf("Lock %d", i))
	}
	for i := 0; i < concurrentLocks; i++ {
		<-done
	}
	c.Assert(*counter, Equals, int64(lockAttempts*concurrentLocks))
}

func (s *fslockSuite) TestTomb(c *C) {
	const timeToDie = 200 * time.Millisecond
	die := tomb.Tomb{}

	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	// Just use one lock, and try to lock it twice.
	err = lock.Lock("very busy")
	c.Assert(err, IsNil)

	checkTomb := func() error {
		select {
		case <-die.Dying():
			return tomb.ErrDying
		default:
			// no-op to fall through to return.
		}
		return nil
	}

	go func() {
		time.Sleep(timeToDie)
		die.Killf("time to die")
	}()

	err = lock.LockWithFunc("won't happen", checkTomb)
	c.Assert(err, Equals, tomb.ErrDying)
	c.Assert(lock.Message(), Equals, "very busy")

}
<|endoftext|>


<commit_before>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.com>

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/handlers"
	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
)

type Server struct {
	conf      *config
	udpClient *dns.Client
	tcpClient *dns.Client
	servemux  *http.ServeMux
}

type DNSRequest struct {
	request         *dns.Msg
	response        *dns.Msg
	transactionID   uint16
	currentUpstream string
	isTailored      bool
	errcode         int
	errtext         string
}

func NewServer(conf *config) (s *Server) {
	s = &Server{
		conf: conf,
		udpClient: &dns.Client{
			Net:     "udp",
			UDPSize: dns.DefaultMsgSize,
			Timeout: time.Duration(conf.Timeout) * time.Second,
		},
		tcpClient: &dns.Client{
			Net:     "tcp",
			Timeout: time.Duration(conf.Timeout) * time.Second,
		},
		servemux: http.NewServeMux(),
	}
	s.servemux.HandleFunc(conf.Path, s.handlerFunc)
	return
}

func (s *Server) Start() error {
	servemux := http.Handler(s.servemux)
	if s.conf.Verbose {
		servemux = handlers.CombinedLoggingHandler(os.Stdout, servemux)
	}
	results := make(chan error, len(s.conf.Listen))
	for _, addr := range s.conf.Listen {
		go func(addr string) {
			var err error
			if s.conf.Cert != "" || s.conf.Key != "" {
				err = http.ListenAndServeTLS(addr, s.conf.Cert, s.conf.Key, servemux)
			} else {
				err = http.ListenAndServe(addr, servemux)
			}
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(addr)
	}
	// wait for all handlers
	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (s *Server) handlerFunc(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	w.Header().Set("Server", USER_AGENT)
	w.Header().Set("X-Powered-By", USER_AGENT)

	if r.Form == nil {
		const maxMemory = 32 << 20 // 32 MB
		r.ParseMultipartForm(maxMemory)
	}

	for _, header := range s.conf.DebugHTTPHeaders {
		if value := r.Header.Get(header); value != "" {
			log.Printf("%s: %s\n", header, value)
		}
	}

	contentType := r.Header.Get("Content-Type")
	if ct := r.FormValue("ct"); ct != "" {
		contentType = ct
	}
	if contentType == "" {
		// Guess request Content-Type based on other parameters
		if r.FormValue("name") != "" {
			contentType = "application/dns-json"
		} else if r.FormValue("dns") != "" {
			contentType = "application/dns-message"
		}
	}
	var responseType string
	for _, responseCandidate := range strings.Split(r.Header.Get("Accept"), ",") {
		responseCandidate = strings.SplitN(responseCandidate, ";", 2)[0]
		if responseCandidate == "application/json" {
			responseType = "application/json"
			break
		} else if responseCandidate == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
			break
		} else if responseCandidate == "application/dns-message" {
			responseType = "application/dns-message"
			break
		}
	}
	if responseType == "" {
		// Guess response Content-Type based on request Content-Type
		if contentType == "application/dns-json" {
			responseType = "application/json"
		} else if contentType == "application/dns-message" {
			responseType = "application/dns-message"
		} else if contentType == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if contentType == "application/dns-json" {
		req = s.parseRequestGoogle(ctx, w, r)
	} else if contentType == "application/dns-message" {
		req = s.parseRequestIETF(ctx, w, r)
	} else if contentType == "application/dns-udpwireformat" {
		req = s.parseRequestIETF(ctx, w, r)
	} else {
		jsonDNS.FormatError(w, fmt.Sprintf("Invalid argument value: \"ct\" = %q", contentType), 415)
		return
	}
	if req.errcode == 444 {
		return
	}
	if req.errcode != 0 {
		jsonDNS.FormatError(w, req.errtext, req.errcode)
		return
	}

	req = s.patchRootRD(req)

	var err error
	req, err = s.doDNSQuery(ctx, req)
	if err != nil {
		jsonDNS.FormatError(w, fmt.Sprintf("DNS query failure (%s)", err.Error()), 503)
		return
	}

	if responseType == "application/json" {
		s.generateResponseGoogle(ctx, w, r, req)
	} else if responseType == "application/dns-message" {
		s.generateResponseIETF(ctx, w, r, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (s *Server) findClientIP(r *http.Request) net.IP {
	XForwardedFor := r.Header.Get("X-Forwarded-For")
	if XForwardedFor != "" {
		for _, addr := range strings.Split(XForwardedFor, ",") {
			addr = strings.TrimSpace(addr)
			ip := net.ParseIP(addr)
			if jsonDNS.IsGlobalIP(ip) {
				return ip
			}
		}
	}
	XRealIP := r.Header.Get("X-Real-IP")
	if XRealIP != "" {
		addr := strings.TrimSpace(XRealIP)
		ip := net.ParseIP(addr)
		if jsonDNS.IsGlobalIP(ip) {
			return ip
		}
	}
	remoteAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
	if err != nil {
		return nil
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		return ip
	}
	return nil
}

// Workaround a bug causing Unbound to refuse returning anything about the root
func (s *Server) patchRootRD(req *DNSRequest) *DNSRequest {
	for _, question := range req.request.Question {
		if question.Name == "." {
			req.request.RecursionDesired = true
		}
	}
	return req
}

func (s *Server) doDNSQuery(ctx context.Context, req *DNSRequest) (resp *DNSRequest, err error) {
	// TODO(m13253): Make ctx work. Waiting for a patch for ExchangeContext from miekg/dns.
	numServers := len(s.conf.Upstream)
	for i := uint(0); i < s.conf.Tries; i++ {
		req.currentUpstream = s.conf.Upstream[rand.Intn(numServers)]
		if !s.conf.TCPOnly {
			req.response, _, err = s.udpClient.Exchange(req.request, req.currentUpstream)
			if err == dns.ErrTruncated {
				log.Println(err)
				req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
			}
		} else {
			req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
		}
		if err == nil || err == dns.ErrTruncated {
			return req, nil
		}
		log.Printf("DNS error from upstream %s: %s\n", req.currentUpstream, err.Error())
	}
	return req, err
}
<commit_msg>Enable CORS by default, which is necessary for AJAX resolver to run<commit_after>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.com>

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/handlers"
	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
)

type Server struct {
	conf      *config
	udpClient *dns.Client
	tcpClient *dns.Client
	servemux  *http.ServeMux
}

type DNSRequest struct {
	request         *dns.Msg
	response        *dns.Msg
	transactionID   uint16
	currentUpstream string
	isTailored      bool
	errcode         int
	errtext         string
}

func NewServer(conf *config) (s *Server) {
	s = &Server{
		conf: conf,
		udpClient: &dns.Client{
			Net:     "udp",
			UDPSize: dns.DefaultMsgSize,
			Timeout: time.Duration(conf.Timeout) * time.Second,
		},
		tcpClient: &dns.Client{
			Net:     "tcp",
			Timeout: time.Duration(conf.Timeout) * time.Second,
		},
		servemux: http.NewServeMux(),
	}
	s.servemux.HandleFunc(conf.Path, s.handlerFunc)
	return
}

func (s *Server) Start() error {
	servemux := http.Handler(s.servemux)
	if s.conf.Verbose {
		servemux = handlers.CombinedLoggingHandler(os.Stdout, servemux)
	}
	results := make(chan error, len(s.conf.Listen))
	for _, addr := range s.conf.Listen {
		go func(addr string) {
			var err error
			if s.conf.Cert != "" || s.conf.Key != "" {
				err = http.ListenAndServeTLS(addr, s.conf.Cert, s.conf.Key, servemux)
			} else {
				err = http.ListenAndServe(addr, servemux)
			}
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(addr)
	}
	// wait for all handlers
	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (s *Server) handlerFunc(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
	w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS, POST")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Max-Age", "3600")
	w.Header().Set("Server", USER_AGENT)
	w.Header().Set("X-Powered-By", USER_AGENT)

	if r.Method == "OPTIONS" {
		w.Header().Set("Content-Length", "0")
		return
	}

	if r.Form == nil {
		const maxMemory = 32 << 20 // 32 MB
		r.ParseMultipartForm(maxMemory)
	}

	for _, header := range s.conf.DebugHTTPHeaders {
		if value := r.Header.Get(header); value != "" {
			log.Printf("%s: %s\n", header, value)
		}
	}

	contentType := r.Header.Get("Content-Type")
	if ct := r.FormValue("ct"); ct != "" {
		contentType = ct
	}
	if contentType == "" {
		// Guess request Content-Type based on other parameters
		if r.FormValue("name") != "" {
			contentType = "application/dns-json"
		} else if r.FormValue("dns") != "" {
			contentType = "application/dns-message"
		}
	}
	var responseType string
	for _, responseCandidate := range strings.Split(r.Header.Get("Accept"), ",") {
		responseCandidate = strings.SplitN(responseCandidate, ";", 2)[0]
		if responseCandidate == "application/json" {
			responseType = "application/json"
			break
		} else if responseCandidate == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
			break
		} else if responseCandidate == "application/dns-message" {
			responseType = "application/dns-message"
			break
		}
	}
	if responseType == "" {
		// Guess response Content-Type based on request Content-Type
		if contentType == "application/dns-json" {
			responseType = "application/json"
		} else if contentType == "application/dns-message" {
			responseType = "application/dns-message"
		} else if contentType == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if contentType == "application/dns-json" {
		req = s.parseRequestGoogle(ctx, w, r)
	} else if contentType == "application/dns-message" {
		req = s.parseRequestIETF(ctx, w, r)
	} else if contentType == "application/dns-udpwireformat" {
		req = s.parseRequestIETF(ctx, w, r)
	} else {
		jsonDNS.FormatError(w, fmt.Sprintf("Invalid argument value: \"ct\" = %q", contentType), 415)
		return
	}
	if req.errcode == 444 {
		return
	}
	if req.errcode != 0 {
		jsonDNS.FormatError(w, req.errtext, req.errcode)
		return
	}

	req = s.patchRootRD(req)

	var err error
	req, err = s.doDNSQuery(ctx, req)
	if err != nil {
		jsonDNS.FormatError(w, fmt.Sprintf("DNS query failure (%s)", err.Error()), 503)
		return
	}

	if responseType == "application/json" {
		s.generateResponseGoogle(ctx, w, r, req)
	} else if responseType == "application/dns-message" {
		s.generateResponseIETF(ctx, w, r, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (s *Server) findClientIP(r *http.Request) net.IP {
	XForwardedFor := r.Header.Get("X-Forwarded-For")
	if XForwardedFor != "" {
		for _, addr := range strings.Split(XForwardedFor, ",") {
			addr = strings.TrimSpace(addr)
			ip := net.ParseIP(addr)
			if jsonDNS.IsGlobalIP(ip) {
				return ip
			}
		}
	}
	XRealIP := r.Header.Get("X-Real-IP")
	if XRealIP != "" {
		addr := strings.TrimSpace(XRealIP)
		ip := net.ParseIP(addr)
		if jsonDNS.IsGlobalIP(ip) {
			return ip
		}
	}
	remoteAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
	if err != nil {
		return nil
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		return ip
	}
	return nil
}

// Workaround a bug causing Unbound to refuse returning anything about the root
func (s *Server) patchRootRD(req *DNSRequest) *DNSRequest {
	for _, question := range req.request.Question {
		if question.Name == "." {
			req.request.RecursionDesired = true
		}
	}
	return req
}

func (s *Server) doDNSQuery(ctx context.Context, req *DNSRequest) (resp *DNSRequest, err error) {
	// TODO(m13253): Make ctx work. Waiting for a patch for ExchangeContext from miekg/dns.
	numServers := len(s.conf.Upstream)
	for i := uint(0); i < s.conf.Tries; i++ {
		req.currentUpstream = s.conf.Upstream[rand.Intn(numServers)]
		if !s.conf.TCPOnly {
			req.response, _, err = s.udpClient.Exchange(req.request, req.currentUpstream)
			if err == dns.ErrTruncated {
				log.Println(err)
				req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
			}
		} else {
			req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
		}
		if err == nil || err == dns.ErrTruncated {
			return req, nil
		}
		log.Printf("DNS error from upstream %s: %s\n", req.currentUpstream, err.Error())
	}
	return req, err
}
<|endoftext|>