<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/kbfscrypto"
	"golang.org/x/net/context"
)

// keybaseServiceOwner is a wrapper around a KeybaseService, to allow
// switching the underlying service at runtime. It is usually
// implemented by Config.
type keybaseServiceOwner interface {
	KeybaseService() KeybaseService
}

// KBPKIClient uses a KeybaseService.
type KBPKIClient struct {
	serviceOwner keybaseServiceOwner
	log          logger.Logger
}

var _ KBPKI = (*KBPKIClient)(nil)

// NewKBPKIClient returns a new KBPKIClient with the given service.
func NewKBPKIClient(
	serviceOwner keybaseServiceOwner, log logger.Logger) *KBPKIClient {
	return &KBPKIClient{serviceOwner, log}
}

// GetCurrentSession implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCurrentSession(ctx context.Context) (
	SessionInfo, error) {
	const sessionID = 0
	return k.serviceOwner.KeybaseService().CurrentSession(ctx, sessionID)
}

// Resolve implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Resolve(ctx context.Context, assertion string) (
	libkb.NormalizedUsername, keybase1.UserOrTeamID, error) {
	return k.serviceOwner.KeybaseService().Resolve(ctx, assertion)
}

// Identify implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Identify(ctx context.Context, assertion, reason string) (
	libkb.NormalizedUsername, keybase1.UserOrTeamID, error) {
	return k.serviceOwner.KeybaseService().Identify(ctx, assertion, reason)
}

// GetNormalizedUsername implements the KBPKI interface for
// KBPKIClient.
func (k *KBPKIClient) GetNormalizedUsername(
	ctx context.Context, id keybase1.UserOrTeamID) (
	libkb.NormalizedUsername, error) {
	var assertion string
	if id.IsUser() {
		assertion = fmt.Sprintf("uid:%s", id)
	} else {
		assertion = fmt.Sprintf("tid:%s", id)
	}
	username, _, err := k.Resolve(ctx, assertion)
	if err != nil {
		return libkb.NormalizedUsername(""), err
	}
	return username, nil
}

func (k *KBPKIClient) hasVerifyingKey(ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey, atServerTime time.Time) (bool, error) {
	userInfo, err := k.loadUserPlusKeys(ctx, uid, verifyingKey.KID())
	if err != nil {
		return false, err
	}

	for _, key := range userInfo.VerifyingKeys {
		if verifyingKey.KID().Equal(key.KID()) {
			return true, nil
		}
	}

	for key, t := range userInfo.RevokedVerifyingKeys {
		if !verifyingKey.KID().Equal(key.KID()) {
			continue
		}
		// We add some slack to the revoke time, because the MD server
		// won't instanteneously find out about the revoke -- it might
		// keep accepting writes from the revoked device for a short
		// period of time until it learns about the revoke.
		const revokeSlack = 1 * time.Minute
		revokedTime := keybase1.FromTime(t.Unix)
		// Trust the server times -- if the key was valid at the given
		// time, we are good to go.  TODO: use Merkle data to check
		// the server timestamps, to prove the server isn't lying.
		if atServerTime.Before(revokedTime.Add(revokeSlack)) {
			k.log.CDebugf(ctx, "Trusting revoked verifying key %s for user %s "+
				"(revoked time: %v vs. server time %v, slack=%s)",
				verifyingKey.KID(), uid, revokedTime, atServerTime, revokeSlack)
			return true, nil
		}
		k.log.CDebugf(ctx, "Not trusting revoked verifying key %s for "+
			"user %s (revoked time: %v vs. server time %v, slack=%s)",
			verifyingKey.KID(), uid, revokedTime, atServerTime, revokeSlack)
		return false, nil
	}

	return false, nil
}

func (k *KBPKIClient) hasUnverifiedVerifyingKey(
	ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) (bool, error) {
	keys, err := k.loadUnverifiedKeys(ctx, uid)
	if err != nil {
		return false, err
	}

	for _, key := range keys {
		if !verifyingKey.KID().Equal(key.KID) {
			continue
		}
		k.log.CDebugf(ctx, "Trusting potentially unverified key %s for user %s",
			verifyingKey.KID(), uid)
		return true, nil
	}

	return false, nil
}

// HasVerifyingKey implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) HasVerifyingKey(ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey, atServerTime time.Time) error {
	ok, err := k.hasVerifyingKey(ctx, uid, verifyingKey, atServerTime)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}

	// If the first attempt couldn't find the key, try again after
	// clearing our local cache.  We might have stale info if the
	// service hasn't learned of the users' new key yet.
	k.serviceOwner.KeybaseService().FlushUserFromLocalCache(ctx, uid)

	ok, err = k.hasVerifyingKey(ctx, uid, verifyingKey, atServerTime)
	if err != nil {
		return err
	}
	if !ok {
		return VerifyingKeyNotFoundError{verifyingKey}
	}
	return nil
}

// HasUnverifiedVerifyingKey implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) HasUnverifiedVerifyingKey(
	ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) error {
	ok, err := k.hasUnverifiedVerifyingKey(ctx, uid, verifyingKey)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}
	k.serviceOwner.KeybaseService().FlushUserUnverifiedKeysFromLocalCache(ctx, uid)
	ok, err = k.hasUnverifiedVerifyingKey(ctx, uid, verifyingKey)
	if err != nil {
		return err
	}
	if !ok {
		return VerifyingKeyNotFoundError{verifyingKey}
	}
	return nil
}

func (k *KBPKIClient) loadUserPlusKeys(ctx context.Context,
	uid keybase1.UID, pollForKID keybase1.KID) (UserInfo, error) {
	return k.serviceOwner.KeybaseService().LoadUserPlusKeys(ctx, uid, pollForKID)
}

// GetCryptPublicKeys implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCryptPublicKeys(ctx context.Context,
	uid keybase1.UID) (keys []kbfscrypto.CryptPublicKey, err error) {
	userInfo, err := k.loadUserPlusKeys(ctx, uid, "")
	if err != nil {
		return nil, err
	}
	return userInfo.CryptPublicKeys, nil
}

// GetTeamTLFCryptKeys implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetTeamTLFCryptKeys(
	ctx context.Context, tid keybase1.TeamID, desiredKeyGen KeyGen) (
	map[KeyGen]kbfscrypto.TLFCryptKey, KeyGen, error) {
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, desiredKeyGen, keybase1.UserVersion{}, keybase1.TeamRole_NONE)
	if err != nil {
		return nil, 0, err
	}
	return teamInfo.CryptKeys, teamInfo.LatestKeyGen, nil
}

// GetCurrentMerkleSeqNo implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCurrentMerkleSeqNo(ctx context.Context) (
	MerkleSeqNo, error) {
	return k.serviceOwner.KeybaseService().GetCurrentMerkleSeqNo(ctx)
}

// IsTeamWriter implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) IsTeamWriter(
	ctx context.Context, tid keybase1.TeamID, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) (bool, error) {
	// Use the verifying key to find out the eldest seqno of the user.
	userInfo, err := k.loadUserPlusKeys(ctx, uid, verifyingKey.KID())
	if err != nil {
		return false, err
	}

	found := false
	for _, key := range userInfo.VerifyingKeys {
		if verifyingKey.KID().Equal(key.KID()) {
			found = true
			break
		}
	}
	if !found {
		// The user doesn't currently have this KID, therefore they
		// shouldn't be treated as a writer.  The caller should check
		// historical device records and team membership.
		k.log.CDebugf(ctx, "User %s doesn't currently have verifying key %s",
			uid, verifyingKey.KID())
		return false, nil
	}

	desiredUser := keybase1.UserVersion{
		Uid:         uid,
		EldestSeqno: userInfo.EldestSeqno,
	}
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, UnspecifiedKeyGen, desiredUser, keybase1.TeamRole_WRITER)
	if err != nil {
		return false, err
	}
	return teamInfo.Writers[uid], nil
}

// IsTeamReader implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) IsTeamReader(
	ctx context.Context, tid keybase1.TeamID, uid keybase1.UID) (bool, error) {
	desiredUser := keybase1.UserVersion{Uid: uid}
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, UnspecifiedKeyGen, desiredUser, keybase1.TeamRole_READER)
	if err != nil {
		return false, err
	}
	return teamInfo.Writers[uid] || teamInfo.Readers[uid], nil
}

// FavoriteAdd implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteAdd(ctx context.Context, folder keybase1.Folder) error {
	return k.serviceOwner.KeybaseService().FavoriteAdd(ctx, folder)
}

// FavoriteDelete implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteDelete(ctx context.Context, folder keybase1.Folder) error {
	return k.serviceOwner.KeybaseService().FavoriteDelete(ctx, folder)
}

// FavoriteList implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteList(ctx context.Context) ([]keybase1.Folder, error) {
	const sessionID = 0
	return k.serviceOwner.KeybaseService().FavoriteList(ctx, sessionID)
}

// Notify implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Notify(ctx context.Context, notification *keybase1.FSNotification) error {
	return k.serviceOwner.KeybaseService().Notify(ctx, notification)
}

func (k *KBPKIClient) loadUnverifiedKeys(ctx context.Context, uid keybase1.UID) (
	[]keybase1.PublicKey, error) {
	return k.serviceOwner.KeybaseService().LoadUnverifiedKeys(ctx, uid)
}

// GetCurrentSessionIfPossible returns the current username and UID
// from kbpki.GetCurrentSession.  If sessionNotRequired is true
// NoCurrentSessionError is ignored and empty username and uid will be
// returned. If it is false all errors are returned.
func GetCurrentSessionIfPossible(
	ctx context.Context, kbpki KBPKI, sessionNotRequired bool) (
	SessionInfo, error) {
	session, err := kbpki.GetCurrentSession(ctx)
	if err == nil {
		return session, nil
	}
	// Return all errors if a session is required.
	if !sessionNotRequired {
		return SessionInfo{}, err
	}

	// If not logged in, return empty session.
	if _, notLoggedIn := err.(NoCurrentSessionError); notLoggedIn {
		return SessionInfo{}, nil
	}

	// Otherwise, just return the error.
	return SessionInfo{}, err
}
<commit_msg>kbpki_client: use revoked keys when getting client's eldest seqno<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/kbfs/kbfscrypto"
	"golang.org/x/net/context"
)

// keybaseServiceOwner is a wrapper around a KeybaseService, to allow
// switching the underlying service at runtime. It is usually
// implemented by Config.
type keybaseServiceOwner interface {
	KeybaseService() KeybaseService
}

// KBPKIClient uses a KeybaseService.
type KBPKIClient struct {
	serviceOwner keybaseServiceOwner
	log          logger.Logger
}

var _ KBPKI = (*KBPKIClient)(nil)

// NewKBPKIClient returns a new KBPKIClient with the given service.
func NewKBPKIClient(
	serviceOwner keybaseServiceOwner, log logger.Logger) *KBPKIClient {
	return &KBPKIClient{serviceOwner, log}
}

// GetCurrentSession implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCurrentSession(ctx context.Context) (
	SessionInfo, error) {
	const sessionID = 0
	return k.serviceOwner.KeybaseService().CurrentSession(ctx, sessionID)
}

// Resolve implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Resolve(ctx context.Context, assertion string) (
	libkb.NormalizedUsername, keybase1.UserOrTeamID, error) {
	return k.serviceOwner.KeybaseService().Resolve(ctx, assertion)
}

// Identify implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Identify(ctx context.Context, assertion, reason string) (
	libkb.NormalizedUsername, keybase1.UserOrTeamID, error) {
	return k.serviceOwner.KeybaseService().Identify(ctx, assertion, reason)
}

// GetNormalizedUsername implements the KBPKI interface for
// KBPKIClient.
func (k *KBPKIClient) GetNormalizedUsername(
	ctx context.Context, id keybase1.UserOrTeamID) (
	libkb.NormalizedUsername, error) {
	var assertion string
	if id.IsUser() {
		assertion = fmt.Sprintf("uid:%s", id)
	} else {
		assertion = fmt.Sprintf("tid:%s", id)
	}
	username, _, err := k.Resolve(ctx, assertion)
	if err != nil {
		return libkb.NormalizedUsername(""), err
	}
	return username, nil
}

func (k *KBPKIClient) hasVerifyingKey(ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey, atServerTime time.Time) (bool, error) {
	userInfo, err := k.loadUserPlusKeys(ctx, uid, verifyingKey.KID())
	if err != nil {
		return false, err
	}

	for _, key := range userInfo.VerifyingKeys {
		if verifyingKey.KID().Equal(key.KID()) {
			return true, nil
		}
	}

	for key, t := range userInfo.RevokedVerifyingKeys {
		if !verifyingKey.KID().Equal(key.KID()) {
			continue
		}
		// We add some slack to the revoke time, because the MD server
		// won't instanteneously find out about the revoke -- it might
		// keep accepting writes from the revoked device for a short
		// period of time until it learns about the revoke.
		const revokeSlack = 1 * time.Minute
		revokedTime := keybase1.FromTime(t.Unix)
		// Trust the server times -- if the key was valid at the given
		// time, we are good to go.  TODO: use Merkle data to check
		// the server timestamps, to prove the server isn't lying.
		if atServerTime.Before(revokedTime.Add(revokeSlack)) {
			k.log.CDebugf(ctx, "Trusting revoked verifying key %s for user %s "+
				"(revoked time: %v vs. server time %v, slack=%s)",
				verifyingKey.KID(), uid, revokedTime, atServerTime, revokeSlack)
			return true, nil
		}
		k.log.CDebugf(ctx, "Not trusting revoked verifying key %s for "+
			"user %s (revoked time: %v vs. server time %v, slack=%s)",
			verifyingKey.KID(), uid, revokedTime, atServerTime, revokeSlack)
		return false, nil
	}

	return false, nil
}

func (k *KBPKIClient) hasUnverifiedVerifyingKey(
	ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) (bool, error) {
	keys, err := k.loadUnverifiedKeys(ctx, uid)
	if err != nil {
		return false, err
	}

	for _, key := range keys {
		if !verifyingKey.KID().Equal(key.KID) {
			continue
		}
		k.log.CDebugf(ctx, "Trusting potentially unverified key %s for user %s",
			verifyingKey.KID(), uid)
		return true, nil
	}

	return false, nil
}

// HasVerifyingKey implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) HasVerifyingKey(ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey, atServerTime time.Time) error {
	ok, err := k.hasVerifyingKey(ctx, uid, verifyingKey, atServerTime)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}

	// If the first attempt couldn't find the key, try again after
	// clearing our local cache.  We might have stale info if the
	// service hasn't learned of the users' new key yet.
	k.serviceOwner.KeybaseService().FlushUserFromLocalCache(ctx, uid)

	ok, err = k.hasVerifyingKey(ctx, uid, verifyingKey, atServerTime)
	if err != nil {
		return err
	}
	if !ok {
		return VerifyingKeyNotFoundError{verifyingKey}
	}
	return nil
}

// HasUnverifiedVerifyingKey implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) HasUnverifiedVerifyingKey(
	ctx context.Context, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) error {
	ok, err := k.hasUnverifiedVerifyingKey(ctx, uid, verifyingKey)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}
	k.serviceOwner.KeybaseService().FlushUserUnverifiedKeysFromLocalCache(ctx, uid)
	ok, err = k.hasUnverifiedVerifyingKey(ctx, uid, verifyingKey)
	if err != nil {
		return err
	}
	if !ok {
		return VerifyingKeyNotFoundError{verifyingKey}
	}
	return nil
}

func (k *KBPKIClient) loadUserPlusKeys(ctx context.Context,
	uid keybase1.UID, pollForKID keybase1.KID) (UserInfo, error) {
	return k.serviceOwner.KeybaseService().LoadUserPlusKeys(ctx, uid, pollForKID)
}

// GetCryptPublicKeys implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCryptPublicKeys(ctx context.Context,
	uid keybase1.UID) (keys []kbfscrypto.CryptPublicKey, err error) {
	userInfo, err := k.loadUserPlusKeys(ctx, uid, "")
	if err != nil {
		return nil, err
	}
	return userInfo.CryptPublicKeys, nil
}

// GetTeamTLFCryptKeys implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetTeamTLFCryptKeys(
	ctx context.Context, tid keybase1.TeamID, desiredKeyGen KeyGen) (
	map[KeyGen]kbfscrypto.TLFCryptKey, KeyGen, error) {
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, desiredKeyGen, keybase1.UserVersion{}, keybase1.TeamRole_NONE)
	if err != nil {
		return nil, 0, err
	}
	return teamInfo.CryptKeys, teamInfo.LatestKeyGen, nil
}

// GetCurrentMerkleSeqNo implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) GetCurrentMerkleSeqNo(ctx context.Context) (
	MerkleSeqNo, error) {
	return k.serviceOwner.KeybaseService().GetCurrentMerkleSeqNo(ctx)
}

// IsTeamWriter implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) IsTeamWriter(
	ctx context.Context, tid keybase1.TeamID, uid keybase1.UID,
	verifyingKey kbfscrypto.VerifyingKey) (bool, error) {
	// Use the verifying key to find out the eldest seqno of the user.
	userInfo, err := k.loadUserPlusKeys(ctx, uid, verifyingKey.KID())
	if err != nil {
		return false, err
	}

	found := false
	for _, key := range userInfo.VerifyingKeys {
		if verifyingKey.KID().Equal(key.KID()) {
			found = true
			break
		}
	}
	if !found {
		// For the purposes of finding the eldest seqno, we need to
		// check the verified key against the list of revoked keys as
		// well.  (The caller should use `HasVerifyingKey` later to
		// check whether the revoked key was valid at the time of the
		// update or not.)
		for key := range userInfo.RevokedVerifyingKeys {
			if verifyingKey.KID().Equal(key.KID()) {
				found = true
				break
			}
		}
	}
	if !found {
		// The user doesn't currently have this KID, therefore they
		// shouldn't be treated as a writer.  The caller should check
		// historical device records and team membership.
		k.log.CDebugf(ctx, "User %s doesn't currently have verifying key %s",
			uid, verifyingKey.KID())
		return false, nil
	}

	desiredUser := keybase1.UserVersion{
		Uid:         uid,
		EldestSeqno: userInfo.EldestSeqno,
	}
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, UnspecifiedKeyGen, desiredUser, keybase1.TeamRole_WRITER)
	if err != nil {
		return false, err
	}
	return teamInfo.Writers[uid], nil
}

// IsTeamReader implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) IsTeamReader(
	ctx context.Context, tid keybase1.TeamID, uid keybase1.UID) (bool, error) {
	desiredUser := keybase1.UserVersion{Uid: uid}
	teamInfo, err := k.serviceOwner.KeybaseService().LoadTeamPlusKeys(
		ctx, tid, UnspecifiedKeyGen, desiredUser, keybase1.TeamRole_READER)
	if err != nil {
		return false, err
	}
	return teamInfo.Writers[uid] || teamInfo.Readers[uid], nil
}

// FavoriteAdd implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteAdd(ctx context.Context, folder keybase1.Folder) error {
	return k.serviceOwner.KeybaseService().FavoriteAdd(ctx, folder)
}

// FavoriteDelete implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteDelete(ctx context.Context, folder keybase1.Folder) error {
	return k.serviceOwner.KeybaseService().FavoriteDelete(ctx, folder)
}

// FavoriteList implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) FavoriteList(ctx context.Context) ([]keybase1.Folder, error) {
	const sessionID = 0
	return k.serviceOwner.KeybaseService().FavoriteList(ctx, sessionID)
}

// Notify implements the KBPKI interface for KBPKIClient.
func (k *KBPKIClient) Notify(ctx context.Context, notification *keybase1.FSNotification) error {
	return k.serviceOwner.KeybaseService().Notify(ctx, notification)
}

func (k *KBPKIClient) loadUnverifiedKeys(ctx context.Context, uid keybase1.UID) (
	[]keybase1.PublicKey, error) {
	return k.serviceOwner.KeybaseService().LoadUnverifiedKeys(ctx, uid)
}

// GetCurrentSessionIfPossible returns the current username and UID
// from kbpki.GetCurrentSession.  If sessionNotRequired is true
// NoCurrentSessionError is ignored and empty username and uid will be
// returned. If it is false all errors are returned.
func GetCurrentSessionIfPossible(
	ctx context.Context, kbpki KBPKI, sessionNotRequired bool) (
	SessionInfo, error) {
	session, err := kbpki.GetCurrentSession(ctx)
	if err == nil {
		return session, nil
	}
	// Return all errors if a session is required.
	if !sessionNotRequired {
		return SessionInfo{}, err
	}

	// If not logged in, return empty session.
	if _, notLoggedIn := err.(NoCurrentSessionError); notLoggedIn {
		return SessionInfo{}, nil
	}

	// Otherwise, just return the error.
	return SessionInfo{}, err
}
<|endoftext|>


<commit_before>package authkit

import (
	"html/template"
	"strings"
)

var funcMap = template.FuncMap{
	"join": strings.Join,
}

var loginTemplate = template.Must(template.New("login").Funcs(funcMap).Parse(loginJS))
var loginJS = `
var __authkit__state__ = {
  authCallback : "__authkit__"+parseInt(Math.random()*1e12,10).toString(36),
  providers : {}
};

{{ $base := .Base }}
{{ range $key, $value := .Providers }}
__authkit__state__.providers["{{ $value.Network }}"] = {
  network : "{{ $value.Network }}",
  authurl : "{{ $value.AuthURL }}",
  client_id : "{{ $value.ClientID }}",
  scopes : "{{ join $value.Scopes " " }}",
  redirect: "{{ $base }}redirect"
};
{{ end }}

var authkit = {
  providers : function () {
    return __authkit__state__.providers;
  },
  provider : function (p) {
    return __authkit__state__.providers[p];
  },
  login : function (provider) {
    var p = authkit.provider(provider);
    if (!p) return;
    var scopes = encodeURIComponent(p.scopes);
    var redir = encodeURIComponent(window.location.origin+p.redirect);
    var state = encodeURIComponent(JSON.stringify({network:p.network,redirect_uri:redir,cbid:__authkit__state__.authCallback}));
    var authU = p.authurl+"?redirect_uri="+redir+"&response_type=code&client_id="+p.client_id+"&state="+state;
    if (p.scopes) {
      authU = authU + "&scope="+scopes;
    }
    var self = this;
    self.user = function (cb) { this.authCB = cb; };
    window[__authkit__state__.authCallback] = function (token, usr) {
      console.log(usr);
      self.authCB(usr, token)
    }
    popup(authU, 500, 600);
    return self;
  }
};

function popup (u, w, h) {
  var documentElement = document.documentElement;

	// Multi Screen Popup Positioning (http://stackoverflow.com/a/16861050)
	//   Credit: http://www.xtf.dk/2011/08/center-new-popup-window-even-on.html
	// Fixes dual-screen position                         Most browsers      Firefox
	var dualScreenLeft = window.screenLeft !== undefined ? window.screenLeft : screen.left;
	var dualScreenTop = window.screenTop !== undefined ? window.screenTop : screen.top;

	var width = window.innerWidth || documentElement.clientWidth || screen.width;
	var height = window.innerHeight || documentElement.clientHeight || screen.height;

	var left = ((width - w) / 2) + dualScreenLeft;
	var top  = ((height - h) / 2) + dualScreenTop;
	var feat = "resizeable=true,height=" + h + ",width=" + w + ",left=" + left + ",top="  + top
	window.open(u, "_blank", feat);
}
`
var redirectTemplate = template.Must(template.New("redirect").Funcs(funcMap).Parse(redirect))
var redirect = `
<html>
<body>
  <script>
	  // First, parse the query string
    var params = {}, queryString = location.search.substring(1),
        regex = /([^&=]+)=([^&]*)/g, m;
    while (m = regex.exec(queryString)) {
      params[decodeURIComponent(m[1])] = decodeURIComponent(m[2]);
    }

    var req = new XMLHttpRequest();
    var cbid = JSON.parse(params.state).cbid;
    req.open('GET', '{{ .Base }}auth?code='+params.code+"&state="+params.state, true);

    req.onreadystatechange = function (e) {
      if (req.readyState == 4) {
        if(req.status == 200){
           var tok = req.getResponseHeader("Authorization");
           var usr = JSON.parse(req.responseText);
           window.opener[cbid](tok, usr);
           window.close();
        }
        else if(req.status == 400) {
            console.log('There was an error processing the access code:',req.responseText)
        }
        else {
          console.log('something other than 200 was returned:',req.responseText)
        }
      }
    };
    req.send(null);
	</script>
</body>
</html>
`
<commit_msg>option to set another base url for the redirect url<commit_after>package authkit

import (
	"html/template"
	"strings"
)

var funcMap = template.FuncMap{
	"join": strings.Join,
}

var loginTemplate = template.Must(template.New("login").Funcs(funcMap).Parse(loginJS))
var loginJS = `
var __authkit__state__ = {
  authCallback : "__authkit__"+parseInt(Math.random()*1e12,10).toString(36),
  providers : {}
};

{{ $base := .Base }}
{{ range $key, $value := .Providers }}
__authkit__state__.providers["{{ $value.Network }}"] = {
  network : "{{ $value.Network }}",
  authurl : "{{ $value.AuthURL }}",
  client_id : "{{ $value.ClientID }}",
  scopes : "{{ join $value.Scopes " " }}",
  redirect: "{{ $base }}redirect"
};
{{ end }}

var authkit = {
  providers : function () {
    return __authkit__state__.providers;
  },
  provider : function (p) {
    return __authkit__state__.providers[p];
  },
  login : function (provider,base) {
    var p = authkit.provider(provider);
    if (!p) return;
    if (!base) base = window.location.origin;
    var scopes = encodeURIComponent(p.scopes);
    var redir = encodeURIComponent(base+p.redirect);
    var state = encodeURIComponent(JSON.stringify({network:p.network,redirect_uri:redir,cbid:__authkit__state__.authCallback}));
    var authU = p.authurl+"?redirect_uri="+redir+"&response_type=code&client_id="+p.client_id+"&state="+state;
    if (p.scopes) {
      authU = authU + "&scope="+scopes;
    }
    var self = this;
    self.user = function (cb) { this.authCB = cb; };
    window[__authkit__state__.authCallback] = function (token, usr) {
      console.log(usr);
      self.authCB(usr, token)
    }
    popup(authU, 500, 600);
    return self;
  }
};

function popup (u, w, h) {
  var documentElement = document.documentElement;

	// Multi Screen Popup Positioning (http://stackoverflow.com/a/16861050)
	//   Credit: http://www.xtf.dk/2011/08/center-new-popup-window-even-on.html
	// Fixes dual-screen position                         Most browsers      Firefox
	var dualScreenLeft = window.screenLeft !== undefined ? window.screenLeft : screen.left;
	var dualScreenTop = window.screenTop !== undefined ? window.screenTop : screen.top;

	var width = window.innerWidth || documentElement.clientWidth || screen.width;
	var height = window.innerHeight || documentElement.clientHeight || screen.height;

	var left = ((width - w) / 2) + dualScreenLeft;
	var top  = ((height - h) / 2) + dualScreenTop;
	var feat = "resizeable=true,height=" + h + ",width=" + w + ",left=" + left + ",top="  + top
	window.open(u, "_blank", feat);
}
`
var redirectTemplate = template.Must(template.New("redirect").Funcs(funcMap).Parse(redirect))
var redirect = `
<html>
<body>
  <script>
	  // First, parse the query string
    var params = {}, queryString = location.search.substring(1),
        regex = /([^&=]+)=([^&]*)/g, m;
    while (m = regex.exec(queryString)) {
      params[decodeURIComponent(m[1])] = decodeURIComponent(m[2]);
    }

    var req = new XMLHttpRequest();
    var cbid = JSON.parse(params.state).cbid;
    req.open('GET', '{{ .Base }}auth?code='+params.code+"&state="+params.state, true);

    req.onreadystatechange = function (e) {
      if (req.readyState == 4) {
        if(req.status == 200){
           var tok = req.getResponseHeader("Authorization");
           var usr = JSON.parse(req.responseText);
           window.opener[cbid](tok, usr);
           window.close();
        }
        else if(req.status == 400) {
            console.log('There was an error processing the access code:',req.responseText)
        }
        else {
          console.log('something other than 200 was returned:',req.responseText)
        }
      }
    };
    req.send(null);
	</script>
</body>
</html>
`
<|endoftext|>


<commit_before>package networkevent

import (
	"encoding/json"
	"time"

	"code.uber.internal/infra/kraken/torlib"
	"code.uber.internal/infra/kraken/utils/log"

	"github.com/willf/bitset"
)

// Name defines event names.
type Name string

// Possible event names.
const (
	AddTorrent       Name = "add_torrent"
	AddPendingConn   Name = "add_pending_conn"
	DropPendingConn  Name = "drop_pending_conn"
	AddActiveConn    Name = "add_active_conn"
	DropActiveConn   Name = "drop_active_conn"
	BlacklistConn    Name = "blacklist_conn"
	ReceivePiece     Name = "receive_piece"
	TorrentComplete  Name = "torrent_complete"
	TorrentCancelled Name = "torrent_cancelled"

	// Piece request timing:

	// Leecher
	DispatcherSentPieceRequest Name = "dispatcher_sent_piece_request"
	ConnSenderGotPieceRequest  Name = "conn_sender_got_piece_request"
	ConnSenderSentPieceRequest Name = "conn_sender_sent_piece_request"

	// Seeder
	ConnReaderGotPieceRequest  Name = "conn_reader_got_piece_request"
	ConnReaderSentPieceRequest Name = "conn_reader_sent_piece_request"
	DispatcherGotPieceRequest  Name = "dispatcher_got_piece_request"
	DispatcherReadPiece        Name = "dispatcher_read_piece"
	DispatcherSentPiecePayload Name = "dispatcher_sent_piece_payload"
	ConnSenderGotPiecePayload  Name = "conn_sender_got_piece_payload"
	ConnSenderSentPiecePayload Name = "conn_sender_sent_piece_payload"

	// Leecher
	ConnReaderGotPiecePayload  Name = "conn_reader_got_piece_payload"
	ConnReaderSentPiecePayload Name = "conn_reader_sent_piece_payload"
	DispatcherGotPiecePayload  Name = "dispatcher_got_piece_payload"
	DispatcherWrotePiece       Name = "dispatcher_wrote_piece"

	// Errors
	ConnSendDroppedPieceRequest Name = "conn_send_dropped_piece_request"
	ConnSendDroppedPiecePayload Name = "conn_send_dropped_piece_payload"
)

// Event consolidates all possible event fields.
type Event struct {
	Name    Name      `json:"event"`
	Torrent string    `json:"torrent"`
	Self    string    `json:"self"`
	Time    time.Time `json:"ts"`

	// Optional fields.
	Peer         string         `json:"peer,omitempty"`
	Piece        int            `json:"piece,omitempty"`
	Bitfield     *bitset.BitSet `json:"bitfield,omitempty"`
	DurationMS   int64          `json:"duration_ms,omitempty"`
	ConnCapacity int            `json:"conn_capacity,omitempty"`
}

func baseEvent(name Name, h torlib.InfoHash, self torlib.PeerID) *Event {
	return &Event{
		Name:    name,
		Torrent: h.String(),
		Self:    self.String(),
		Time:    time.Now(),
	}
}

// JSON converts event into a json string primarely for logging purposes
func (e *Event) JSON() string {
	b, err := json.Marshal(e)
	if err != nil {
		log.Errorf("json marshal error %s", err)
		return ""
	}
	return string(b)
}

// AddTorrentEvent returns an event for an added torrent with initial bitfield.
func AddTorrentEvent(h torlib.InfoHash, self torlib.PeerID, b *bitset.BitSet, connCapacity int) *Event {
	e := baseEvent(AddTorrent, h, self)
	e.Bitfield = bitset.New(b.Len())
	b.Copy(e.Bitfield)
	e.ConnCapacity = connCapacity
	return e
}

// AddPendingConnEvent returns an event for an added pending conn from self to peer.
func AddPendingConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(AddPendingConn, h, self)
	e.Peer = peer.String()
	return e
}

// DropPendingConnEvent returns an event for a dropped pending conn from self to peer.
func DropPendingConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(DropPendingConn, h, self)
	e.Peer = peer.String()
	return e
}

// AddActiveConnEvent returns an event for an added active conn from self to peer.
func AddActiveConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(AddActiveConn, h, self)
	e.Peer = peer.String()
	return e
}

// DropActiveConnEvent returns an event for a dropped active conn from self to peer.
func DropActiveConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(DropActiveConn, h, self)
	e.Peer = peer.String()
	return e
}

// BlacklistConnEvent returns an event for a blacklisted connection.
func BlacklistConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, dur time.Duration) *Event {
	e := baseEvent(BlacklistConn, h, self)
	e.Peer = peer.String()
	e.DurationMS = int64(dur.Seconds() * 1000)
	return e
}

// ReceivePieceEvent returns an event for a piece received from peer.
func ReceivePieceEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {
	e := baseEvent(ReceivePiece, h, self)
	e.Peer = peer.String()
	e.Piece = piece
	return e
}

// TorrentCompleteEvent returns an event for a completed torrent.
func TorrentCompleteEvent(h torlib.InfoHash, self torlib.PeerID) *Event {
	return baseEvent(TorrentComplete, h, self)
}

// TorrentCancelledEvent returns an event for a cancelled torrent.
func TorrentCancelledEvent(h torlib.InfoHash, self torlib.PeerID) *Event {
	return baseEvent(TorrentCancelled, h, self)
}

func pieceEvent(
	name Name, h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	e := baseEvent(name, h, self)
	e.Peer = peer.String()
	e.Piece = piece
	return e
}

// DispatcherSentPieceRequestEvent ...
func DispatcherSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherSentPieceRequest, h, self, peer, piece)
}

// ConnSenderSentPieceRequestEvent ...
func ConnSenderSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderSentPieceRequest, h, self, peer, piece)
}

// ConnSenderGotPieceRequestEvent ...
func ConnSenderGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderGotPieceRequest, h, self, peer, piece)
}

// ConnReaderGotPieceRequestEvent ...
func ConnReaderGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderGotPieceRequest, h, self, peer, piece)
}

// ConnReaderSentPieceRequestEvent ...
func ConnReaderSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderSentPieceRequest, h, self, peer, piece)
}

// DispatcherGotPieceRequestEvent ...
func DispatcherGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherGotPieceRequest, h, self, peer, piece)
}

// DispatcherReadPieceEvent ...
func DispatcherReadPieceEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherReadPiece, h, self, peer, piece)
}

// DispatcherSentPiecePayloadEvent ...
func DispatcherSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherSentPiecePayload, h, self, peer, piece)
}

// ConnSenderGotPiecePayloadEvent ...
func ConnSenderGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderGotPiecePayload, h, self, peer, piece)
}

// ConnSenderSentPiecePayloadEvent ...
func ConnSenderSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderSentPiecePayload, h, self, peer, piece)
}

// ConnReaderGotPiecePayloadEvent ...
func ConnReaderGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderGotPiecePayload, h, self, peer, piece)
}

// ConnReaderSentPiecePayloadEvent ...
func ConnReaderSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderSentPiecePayload, h, self, peer, piece)
}

// DispatcherGotPiecePayloadEvent ...
func DispatcherGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherGotPiecePayload, h, self, peer, piece)
}

// DispatcherWrotePieceEvent ...
func DispatcherWrotePieceEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherWrotePiece, h, self, peer, piece)
}

// ConnSendDroppedPieceRequestEvent ...
func ConnSendDroppedPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSendDroppedPieceRequest, h, self, peer, piece)
}

// ConnSendDroppedPiecePayloadEvent ...
func ConnSendDroppedPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSendDroppedPiecePayload, h, self, peer, piece)
}
<commit_msg>Change networkevent bitfield to []bool<commit_after>package networkevent

import (
	"encoding/json"
	"time"

	"code.uber.internal/infra/kraken/torlib"
	"code.uber.internal/infra/kraken/utils/log"

	"github.com/willf/bitset"
)

// Name defines event names.
type Name string

// Possible event names.
const (
	AddTorrent       Name = "add_torrent"
	AddPendingConn   Name = "add_pending_conn"
	DropPendingConn  Name = "drop_pending_conn"
	AddActiveConn    Name = "add_active_conn"
	DropActiveConn   Name = "drop_active_conn"
	BlacklistConn    Name = "blacklist_conn"
	ReceivePiece     Name = "receive_piece"
	TorrentComplete  Name = "torrent_complete"
	TorrentCancelled Name = "torrent_cancelled"

	// Piece request timing:

	// Leecher
	DispatcherSentPieceRequest Name = "dispatcher_sent_piece_request"
	ConnSenderGotPieceRequest  Name = "conn_sender_got_piece_request"
	ConnSenderSentPieceRequest Name = "conn_sender_sent_piece_request"

	// Seeder
	ConnReaderGotPieceRequest  Name = "conn_reader_got_piece_request"
	ConnReaderSentPieceRequest Name = "conn_reader_sent_piece_request"
	DispatcherGotPieceRequest  Name = "dispatcher_got_piece_request"
	DispatcherReadPiece        Name = "dispatcher_read_piece"
	DispatcherSentPiecePayload Name = "dispatcher_sent_piece_payload"
	ConnSenderGotPiecePayload  Name = "conn_sender_got_piece_payload"
	ConnSenderSentPiecePayload Name = "conn_sender_sent_piece_payload"

	// Leecher
	ConnReaderGotPiecePayload  Name = "conn_reader_got_piece_payload"
	ConnReaderSentPiecePayload Name = "conn_reader_sent_piece_payload"
	DispatcherGotPiecePayload  Name = "dispatcher_got_piece_payload"
	DispatcherWrotePiece       Name = "dispatcher_wrote_piece"

	// Errors
	ConnSendDroppedPieceRequest Name = "conn_send_dropped_piece_request"
	ConnSendDroppedPiecePayload Name = "conn_send_dropped_piece_payload"
)

// Event consolidates all possible event fields.
type Event struct {
	Name    Name      `json:"event"`
	Torrent string    `json:"torrent"`
	Self    string    `json:"self"`
	Time    time.Time `json:"ts"`

	// Optional fields.
	Peer         string `json:"peer,omitempty"`
	Piece        int    `json:"piece,omitempty"`
	Bitfield     []bool `json:"bitfield,omitempty"`
	DurationMS   int64  `json:"duration_ms,omitempty"`
	ConnCapacity int    `json:"conn_capacity,omitempty"`
}

func baseEvent(name Name, h torlib.InfoHash, self torlib.PeerID) *Event {
	return &Event{
		Name:    name,
		Torrent: h.String(),
		Self:    self.String(),
		Time:    time.Now(),
	}
}

// JSON converts event into a json string primarely for logging purposes
func (e *Event) JSON() string {
	b, err := json.Marshal(e)
	if err != nil {
		log.Errorf("json marshal error %s", err)
		return ""
	}
	return string(b)
}

// AddTorrentEvent returns an event for an added torrent with initial bitfield.
func AddTorrentEvent(h torlib.InfoHash, self torlib.PeerID, b *bitset.BitSet, connCapacity int) *Event {
	e := baseEvent(AddTorrent, h, self)
	bools := make([]bool, b.Len())
	for i := uint(0); i < b.Len(); i++ {
		bools[i] = b.Test(i)
	}
	e.Bitfield = bools
	e.ConnCapacity = connCapacity
	return e
}

// AddPendingConnEvent returns an event for an added pending conn from self to peer.
func AddPendingConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(AddPendingConn, h, self)
	e.Peer = peer.String()
	return e
}

// DropPendingConnEvent returns an event for a dropped pending conn from self to peer.
func DropPendingConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(DropPendingConn, h, self)
	e.Peer = peer.String()
	return e
}

// AddActiveConnEvent returns an event for an added active conn from self to peer.
func AddActiveConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(AddActiveConn, h, self)
	e.Peer = peer.String()
	return e
}

// DropActiveConnEvent returns an event for a dropped active conn from self to peer.
func DropActiveConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID) *Event {
	e := baseEvent(DropActiveConn, h, self)
	e.Peer = peer.String()
	return e
}

// BlacklistConnEvent returns an event for a blacklisted connection.
func BlacklistConnEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, dur time.Duration) *Event {
	e := baseEvent(BlacklistConn, h, self)
	e.Peer = peer.String()
	e.DurationMS = int64(dur.Seconds() * 1000)
	return e
}

// ReceivePieceEvent returns an event for a piece received from peer.
func ReceivePieceEvent(h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {
	e := baseEvent(ReceivePiece, h, self)
	e.Peer = peer.String()
	e.Piece = piece
	return e
}

// TorrentCompleteEvent returns an event for a completed torrent.
func TorrentCompleteEvent(h torlib.InfoHash, self torlib.PeerID) *Event {
	return baseEvent(TorrentComplete, h, self)
}

// TorrentCancelledEvent returns an event for a cancelled torrent.
func TorrentCancelledEvent(h torlib.InfoHash, self torlib.PeerID) *Event {
	return baseEvent(TorrentCancelled, h, self)
}

func pieceEvent(
	name Name, h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	e := baseEvent(name, h, self)
	e.Peer = peer.String()
	e.Piece = piece
	return e
}

// DispatcherSentPieceRequestEvent ...
func DispatcherSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherSentPieceRequest, h, self, peer, piece)
}

// ConnSenderSentPieceRequestEvent ...
func ConnSenderSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderSentPieceRequest, h, self, peer, piece)
}

// ConnSenderGotPieceRequestEvent ...
func ConnSenderGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderGotPieceRequest, h, self, peer, piece)
}

// ConnReaderGotPieceRequestEvent ...
func ConnReaderGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderGotPieceRequest, h, self, peer, piece)
}

// ConnReaderSentPieceRequestEvent ...
func ConnReaderSentPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderSentPieceRequest, h, self, peer, piece)
}

// DispatcherGotPieceRequestEvent ...
func DispatcherGotPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherGotPieceRequest, h, self, peer, piece)
}

// DispatcherReadPieceEvent ...
func DispatcherReadPieceEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherReadPiece, h, self, peer, piece)
}

// DispatcherSentPiecePayloadEvent ...
func DispatcherSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherSentPiecePayload, h, self, peer, piece)
}

// ConnSenderGotPiecePayloadEvent ...
func ConnSenderGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderGotPiecePayload, h, self, peer, piece)
}

// ConnSenderSentPiecePayloadEvent ...
func ConnSenderSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSenderSentPiecePayload, h, self, peer, piece)
}

// ConnReaderGotPiecePayloadEvent ...
func ConnReaderGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderGotPiecePayload, h, self, peer, piece)
}

// ConnReaderSentPiecePayloadEvent ...
func ConnReaderSentPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnReaderSentPiecePayload, h, self, peer, piece)
}

// DispatcherGotPiecePayloadEvent ...
func DispatcherGotPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherGotPiecePayload, h, self, peer, piece)
}

// DispatcherWrotePieceEvent ...
func DispatcherWrotePieceEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(DispatcherWrotePiece, h, self, peer, piece)
}

// ConnSendDroppedPieceRequestEvent ...
func ConnSendDroppedPieceRequestEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSendDroppedPieceRequest, h, self, peer, piece)
}

// ConnSendDroppedPiecePayloadEvent ...
func ConnSendDroppedPiecePayloadEvent(
	h torlib.InfoHash, self torlib.PeerID, peer torlib.PeerID, piece int) *Event {

	return pieceEvent(ConnSendDroppedPiecePayload, h, self, peer, piece)
}
<|endoftext|>


<commit_before>package servo

import (
	"github.com/fgrosse/goldi"
	"github.com/fgrosse/servo/configuration"
	"github.com/mgutz/logxi/v1"
	"fmt"
	"github.com/fgrosse/goldi/validation"
)

var KernelVersion = "unknown"

// The Kernel is basically a goldi.TypeRegistry on which all necessary types are registered.
// Once type registration is done it is started with its Run function.
type Kernel struct {
	goldi.TypeRegistry
	log    Logger
	Config ConfigurationLoader
}

// NewKernel creates a new kernel.
// It will initialize the goldi.TypeRegistry and use the given configuration loader.
// The actual loading of the configuration is deferred until kernel.Run is called.
func NewKernel(config ConfigurationLoader) *Kernel {
	kernel := &Kernel{
		TypeRegistry: goldi.NewTypeRegistry(),
		Config:       config,
		log:          new(NullLogger),
	}

	registerInternalTypes(kernel.TypeRegistry)
	return kernel
}

// NewDebugKernel creates a new kernel just as NewKernel but immediately injects a logxi logger.
func NewDebugKernel(config ConfigurationLoader) *Kernel {
	kernel := NewKernel(config)
	kernel.log = log.New("kernel")
	kernel.log.SetLevel(log.LevelDebug)
	return kernel
}

func (k *Kernel) Register(bundle Bundle) {
	k.log.Info("Loading bundle", "bundle", fmt.Sprintf("%T", bundle))
	bundle.Boot(k)
}

// Run creates a goldi.Container based on the TypeRegistry of the kernel and used the configuration loader.
// It does then instantiate the "kernel.server" type and calls Run on the resulting Server implementation.
// This method blocks until the server returns from Run.
func (k *Kernel) Run() error {
	k.log.Info("Starting servo kernel", "version", KernelVersion)
	container, err := k.createContainer()
	if err != nil {
		return err
	}

	server, err := container.Get("kernel.http.server")
	if err != nil {
		k.log.Fatal(err.Error())
	}

	return server.(Server).Run()
}

func (k *Kernel) createContainer() (*goldi.Container, error) {
	k.log.Debug("Loading the configuration..")
	config, err := k.Config.Load()
	if err != nil {
		return nil, err
	}

	k.log.Trace("Flattening the configuration..")
	flattenedConfig := new(configuration.Flattener).Flatten(config)
	k.log.Debug("Configuration has been loaded", "config", flattenedConfig)

	k.log.Debug("Creating goldi container")
	container := goldi.NewContainer(k.TypeRegistry, flattenedConfig)
	container.InjectInstance("container", container)

	err = k.validateContainer(container)
	if err != nil {
		k.log.Error("Container is invalid", "error", err)
	} else {
		k.log.Debug("Container passed validation")
	}

	return container, err
}

const TypeContainerValidator = "container.validator"

func (k *Kernel) validateContainer(container *goldi.Container) error {
	k.log.Trace("Retrieving validator from container", "service_name", TypeContainerValidator)
	validator, err := container.Get(TypeContainerValidator)
	if err != nil {
		k.log.Fatal(err.Error())
	}

	k.log.Debug("Validating container")

	return validator.(*validation.ContainerValidator).Validate(container)
}
<commit_msg>Improve kernel documentation<commit_after>package servo

import (
	"fmt"
	"os"

	"github.com/fgrosse/goldi"
	"github.com/fgrosse/goldi/validation"
	"github.com/fgrosse/servo/configuration"
	"github.com/mgutz/logxi/v1"
)

var KernelVersion = "unknown"

// The Kernel is basically a goldi.TypeRegistry on which all necessary types are registered.
// Once type registration is done it is started with its Run function.
type Kernel struct {
	goldi.TypeRegistry

	// Config is a ConfigurationLoader that loads the configuration upon Kernel.Run
	Config ConfigurationLoader

	// Env identifies the environment the kernel runs in.
	// If Env is empty a production environment is assumed. The value is usually set via
	// the environment variable `SERVO_ENV`. Setting this variable to `dev` will configure
	// the kernel for local development (i.e. print debug messages).
	Env string

	log Logger
}

// NewKernel creates a new kernel.
// It will initialize the goldi.TypeRegistry and use the given configuration loader.
// The actual loading of the configuration is deferred until kernel.Run is called.
//
// The kernel environment identifier is read from the os environment variable
// `SERVO_ENV`. If SERVO_ENV is set to `dev` the internal log level will be set to debug.
func NewKernel(config ConfigurationLoader) *Kernel {
	kernel := &Kernel{
		TypeRegistry: goldi.NewTypeRegistry(),
		Env:          os.Getenv("SERVO_ENV"),
		Config:       config,
		log:          new(NullLogger),
	}

	if kernel.Env == "dev" {
		kernel.log = log.New("kernel")
		kernel.log.SetLevel(log.LevelDebug)
	}

	registerInternalTypes(kernel.TypeRegistry)
	return kernel
}

// Register is used to boot servo.Bundle with this kernel instance.
func (k *Kernel) Register(bundle Bundle) {
	k.log.Info("Loading bundle", "bundle", fmt.Sprintf("%T", bundle))
	bundle.Boot(k)
}

// Run load the configuration and creates the DI container.
// Afterwards the `kernel.http.server` type is created by the container and started.
func (k *Kernel) Run() error {
	k.log.Info("Starting servo kernel", "version", KernelVersion)
	container, err := k.createContainer()
	if err != nil {
		return err
	}

	server, err := container.Get("kernel.http.server")
	if err != nil {
		k.log.Fatal(err.Error())
	}

	return server.(Server).Run()
}

func (k *Kernel) createContainer() (*goldi.Container, error) {
	k.log.Debug("Loading the configuration..")
	config, err := k.Config.Load()
	if err != nil {
		return nil, err
	}

	k.log.Trace("Flattening the configuration..")
	flattenedConfig := new(configuration.Flattener).Flatten(config)
	k.log.Debug("Configuration has been loaded", "config", flattenedConfig)

	k.log.Debug("Creating goldi container")
	container := goldi.NewContainer(k.TypeRegistry, flattenedConfig)
	container.InjectInstance("container", container)

	err = k.validateContainer(container)
	if err != nil {
		k.log.Error("Container is invalid", "error", err)
	} else {
		k.log.Debug("Container passed validation")
	}

	return container, err
}

// TypeContainerValidator is the type ID that is used to get the container validator from the DI container.
const TypeContainerValidator = "container.validator"

func (k *Kernel) validateContainer(container *goldi.Container) error {
	k.log.Trace("Retrieving validator from container", "service_name", TypeContainerValidator)
	validator, err := container.Get(TypeContainerValidator)
	if err != nil {
		k.log.Fatal(err.Error())
	}

	k.log.Debug("Validating container")

	return validator.(*validation.ContainerValidator).Validate(container)
}
<|endoftext|>


<commit_before>package libvirt

import (
	"encoding/xml"
	"fmt"
	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/schema"
	"io"
	"log"
	"net/http"
	"strconv"
)

func volumeCommonSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
			ForceNew: true,
		},
		"pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			Default:  "default",
			ForceNew: true,
		},
		"source": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"size": &schema.Schema{
			Type:     schema.TypeInt,
			Optional: true,
			Computed: true,
			ForceNew: true,
		},
		"base_volume_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
	}
}

func resourceLibvirtVolume() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtVolumeCreate,
		Read:   resourceLibvirtVolumeRead,
		Delete: resourceLibvirtVolumeDelete,
		Schema: volumeCommonSchema(),
	}
}

func remoteImageSize(url string) (int, error) {
	response, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return length, nil
}

func resourceLibvirtVolumeCreate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	poolName := "default"
	if _, ok := d.GetOk("pool"); ok {
		poolName = d.Get("pool").(string)
	}

	pool, err := virConn.LookupStoragePoolByName(poolName)
	if err != nil {
		return fmt.Errorf("can't find storage pool '%s'", poolName)
	}
	defer pool.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	err = pool.Refresh(0)
	if err != nil {
		return fmt.Errorf("Error refreshing pool for volume: %s", err)
	}

	volumeDef := newDefVolume()

	if name, ok := d.GetOk("name"); ok {
		volumeDef.Name = name.(string)
	}

	// an existing image was given, this mean we can't choose size
	if url, ok := d.GetOk("source"); ok {
		// source and size conflict
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'source' is given (the size will be set to the size of the source image.")
		}

		size, err := remoteImageSize(url.(string))
		if err != nil {
			return err
		}
		log.Printf("Remote image is: %d bytes", size)
		volumeDef.Capacity.Unit = "B"
		volumeDef.Capacity.Amount = size
	} else {
		_, noSize := d.GetOk("size")
		_, noBaseVol := d.GetOk("base_volume_id")

		if noSize && noBaseVol {
			return fmt.Errorf("'size' needs to be specified if no 'source' or 'base_volume_id' is given.")
		}
		volumeDef.Capacity.Amount = d.Get("size").(int)
	}

	if baseVolumeId, ok := d.GetOk("base_volume_id"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_id' is given (the size will be set to the size of the backing image.")
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := virConn.LookupStorageVolByKey(baseVolumeId.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeId.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeId)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	volumeDefXml, err := xml.Marshal(volumeDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt volume: %s", err)
	}

	// create the volume
	volume, err := pool.StorageVolCreateXML(string(volumeDefXml), 0)
	if err != nil {
		return fmt.Errorf("Error creating libvirt volume: %s", err)
	}
	defer volume.Free()

	// we use the key as the id
	key, err := volume.GetKey()
	if err != nil {
		return fmt.Errorf("Error retrieving volume key: %s", err)
	}
	d.SetId(key)
	log.Printf("[INFO] Volume ID: %s", d.Id())

	// upload source if present
	if url, ok := d.GetOk("source"); ok {
		stream, err := libvirt.NewVirStream(virConn, 0)
		defer stream.Close()

		volume.Upload(stream, 0, uint64(volumeDef.Capacity.Amount), 0)
		response, err := http.Get(url.(string))
		defer response.Body.Close()
		if err != nil {
			return fmt.Errorf("Error while downloading %s: %s", url.(string), err)
		}

		n, err := io.Copy(stream, response.Body)
		if err != nil {
			return fmt.Errorf("Error while downloading %s: %s", url.(string), err)
		}
		log.Printf("%d bytes uploaded\n", n)
	}

	return resourceLibvirtVolumeRead(d, meta)
}

func resourceLibvirtVolumeRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s", d.Id())
	}
	defer volume.Free()

	volName, err := volume.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}

	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	volPoolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving pool name: %s", err)
	}

	d.Set("pool", volPoolName)
	d.Set("name", volName)

	info, err := volume.GetInfo()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}
	d.Set("size", info.GetCapacityInBytes())

	return nil
}

func resourceLibvirtVolumeDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s", d.Id())
	}
	defer volume.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	err = volPool.Refresh(0)
	if err != nil {
		return fmt.Errorf("Error refreshing pool for volume: %s", err)
	}

	// Workaround for redhat#1293804
	// https://bugzilla.redhat.com/show_bug.cgi?id=1293804#c12
	// Does not solve the problem but it makes it happen less often.
	_, err = volume.GetXMLDesc(0)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s XML desc: %s", d.Id(), err)
	}

	err = volume.Delete(0)
	if err != nil {
		return fmt.Errorf("Can't delete volume %s: %s", d.Id(), err)
	}

	return nil
}
<commit_msg>fail if base_volume_id and source are specified together<commit_after>package libvirt

import (
	"encoding/xml"
	"fmt"
	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/schema"
	"io"
	"log"
	"net/http"
	"strconv"
)

func volumeCommonSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
			ForceNew: true,
		},
		"pool": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			Default:  "default",
			ForceNew: true,
		},
		"source": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"size": &schema.Schema{
			Type:     schema.TypeInt,
			Optional: true,
			Computed: true,
			ForceNew: true,
		},
		"base_volume_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
	}
}

func resourceLibvirtVolume() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtVolumeCreate,
		Read:   resourceLibvirtVolumeRead,
		Delete: resourceLibvirtVolumeDelete,
		Schema: volumeCommonSchema(),
	}
}

func remoteImageSize(url string) (int, error) {
	response, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		return 0, err
	}
	return length, nil
}

func resourceLibvirtVolumeCreate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	poolName := "default"
	if _, ok := d.GetOk("pool"); ok {
		poolName = d.Get("pool").(string)
	}

	pool, err := virConn.LookupStoragePoolByName(poolName)
	if err != nil {
		return fmt.Errorf("can't find storage pool '%s'", poolName)
	}
	defer pool.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	err = pool.Refresh(0)
	if err != nil {
		return fmt.Errorf("Error refreshing pool for volume: %s", err)
	}

	volumeDef := newDefVolume()

	if name, ok := d.GetOk("name"); ok {
		volumeDef.Name = name.(string)
	}

	// an existing image was given, this mean we can't choose size
	if url, ok := d.GetOk("source"); ok {
		// source and size conflict
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'source' is given (the size will be set to the size of the source image.")
		}
		if _, ok := d.GetOk("base_volume_id"); ok {
			return fmt.Errorf("'base_volume_id' can't be specified when also 'source' is given (the size will be set to the size of the base image.")
		}

		size, err := remoteImageSize(url.(string))
		if err != nil {
			return err
		}
		log.Printf("Remote image is: %d bytes", size)
		volumeDef.Capacity.Unit = "B"
		volumeDef.Capacity.Amount = size
	} else {
		_, noSize := d.GetOk("size")
		_, noBaseVol := d.GetOk("base_volume_id")

		if noSize && noBaseVol {
			return fmt.Errorf("'size' needs to be specified if no 'source' or 'base_volume_id' is given.")
		}
		volumeDef.Capacity.Amount = d.Get("size").(int)
	}

	if baseVolumeId, ok := d.GetOk("base_volume_id"); ok {
		if _, ok := d.GetOk("size"); ok {
			return fmt.Errorf("'size' can't be specified when also 'base_volume_id' is given (the size will be set to the size of the backing image.")
		}

		volumeDef.BackingStore = new(defBackingStore)
		volumeDef.BackingStore.Format.Type = "qcow2"
		baseVolume, err := virConn.LookupStorageVolByKey(baseVolumeId.(string))
		if err != nil {
			return fmt.Errorf("Can't retrieve volume %s", baseVolumeId.(string))
		}
		baseVolPath, err := baseVolume.GetPath()
		if err != nil {
			return fmt.Errorf("can't get name for base image '%s'", baseVolumeId)
		}
		volumeDef.BackingStore.Path = baseVolPath
	}

	volumeDefXml, err := xml.Marshal(volumeDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt volume: %s", err)
	}

	// create the volume
	volume, err := pool.StorageVolCreateXML(string(volumeDefXml), 0)
	if err != nil {
		return fmt.Errorf("Error creating libvirt volume: %s", err)
	}
	defer volume.Free()

	// we use the key as the id
	key, err := volume.GetKey()
	if err != nil {
		return fmt.Errorf("Error retrieving volume key: %s", err)
	}
	d.SetId(key)
	log.Printf("[INFO] Volume ID: %s", d.Id())

	// upload source if present
	if url, ok := d.GetOk("source"); ok {
		stream, err := libvirt.NewVirStream(virConn, 0)
		defer stream.Close()

		volume.Upload(stream, 0, uint64(volumeDef.Capacity.Amount), 0)
		response, err := http.Get(url.(string))
		defer response.Body.Close()
		if err != nil {
			return fmt.Errorf("Error while downloading %s: %s", url.(string), err)
		}

		n, err := io.Copy(stream, response.Body)
		if err != nil {
			return fmt.Errorf("Error while downloading %s: %s", url.(string), err)
		}
		log.Printf("%d bytes uploaded\n", n)
	}

	return resourceLibvirtVolumeRead(d, meta)
}

func resourceLibvirtVolumeRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s", d.Id())
	}
	defer volume.Free()

	volName, err := volume.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}

	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	volPoolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving pool name: %s", err)
	}

	d.Set("pool", volPoolName)
	d.Set("name", volName)

	info, err := volume.GetInfo()
	if err != nil {
		return fmt.Errorf("Error retrieving volume name: %s", err)
	}
	d.Set("size", info.GetCapacityInBytes())

	return nil
}

func resourceLibvirtVolumeDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	volume, err := virConn.LookupStorageVolByKey(d.Id())
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s", d.Id())
	}
	defer volume.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	err = volPool.Refresh(0)
	if err != nil {
		return fmt.Errorf("Error refreshing pool for volume: %s", err)
	}

	// Workaround for redhat#1293804
	// https://bugzilla.redhat.com/show_bug.cgi?id=1293804#c12
	// Does not solve the problem but it makes it happen less often.
	_, err = volume.GetXMLDesc(0)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s XML desc: %s", d.Id(), err)
	}

	err = volume.Delete(0)
	if err != nil {
		return fmt.Errorf("Can't delete volume %s: %s", d.Id(), err)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"../nsq"
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

const maxTimeout = time.Hour
const maxReadyCount = 2500

type ProtocolV2 struct {
	nsq.Protocol
}

func init() {
	// BigEndian client byte sequence "  V2"
	var magicInt int32
	buf := bytes.NewBuffer([]byte(nsq.MagicV2))
	binary.Read(buf, binary.BigEndian, &magicInt)
	protocols[magicInt] = &ProtocolV2{}
}

func (p *ProtocolV2) IOLoop(conn net.Conn) error {
	var err error
	var line string

	client := NewClientV2(conn)
	atomic.StoreInt32(&client.State, nsq.StateInit)

	err = nil
	reader := bufio.NewReader(client)
	for {
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		if *verbose {
			log.Printf("PROTOCOL(V2): [%s] %#v", client.RemoteAddr().String(), params)
		}

		response, err := p.Exec(client, params)
		if err != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeError, []byte(err.Error()))
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
			continue
		}

		if response != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeResponse, response)
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
		}
	}

	log.Printf("PROTOCOL(V2): [%s] exiting ioloop", client.RemoteAddr().String())
	// TODO: gracefully send clients the close signal
	close(client.ExitChan)

	return err
}

func (p *ProtocolV2) Write(client *ClientV2, data []byte) error {
	attempts := 0
	for {
		client.Lock()
		client.SetWriteDeadline(time.Now().Add(time.Second))
		_, err := nsq.SendResponse(client, data)
		client.Unlock()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				attempts++
				if attempts == 3 {
					return errors.New("E_WRITE_TIMEOUT")
				}
				continue
			}
			return err
		}
		break
	}
	return nil
}

func (p *ProtocolV2) Exec(client *ClientV2, params []string) ([]byte, error) {
	switch params[0] {
	case "SUB":
		return p.SUB(client, params)
	case "RDY":
		return p.RDY(client, params)
	case "FIN":
		return p.FIN(client, params)
	case "REQ":
		return p.REQ(client, params)
	case "CLS":
		return p.CLS(client, params)
	}
	return nil, nsq.ClientErrInvalid
}

func (p *ProtocolV2) messagePump(client *ClientV2) {
	var err error

	// ReadyStateChan has a buffer of 1 to guarantee that in the event
	// there is a race before we enter either of the select loops where
	// ReadyStateChan is read from that the update is not lost
	for {
		if !client.IsReadyForMessages() {
			// wait for a change in state
			select {
			case <-client.ReadyStateChan:
			case <-client.ExitChan:
				goto exit
			}
		} else {
			select {
			case <-client.ReadyStateChan:
			case msg := <-client.Channel.clientMsgChan:
				if *verbose {
					log.Printf("PROTOCOL(V2): writing msg(%s) to client(%s) - %s",
						msg.Id, client, msg.Body)
				}

				data, err := msg.Encode()
				if err != nil {
					goto exit
				}

				clientData, err := nsq.Frame(nsq.FrameTypeMessage, data)
				if err != nil {
					goto exit
				}

				client.Channel.StartInFlightTimeout(msg, client)
				client.SendingMessage()

				err = p.Write(client, clientData)
				if err != nil {
					goto exit
				}
			case <-client.ExitChan:
				goto exit
			}
		}
	}

exit:
	log.Printf("PROTOCOL(V2): [%s] exiting messagePump", client.RemoteAddr().String())
	client.Channel.RemoveClient(client)
	if err != nil {
		log.Printf("PROTOCOL(V2): messagePump error - %s", err.Error())
	}
}

func (p *ProtocolV2) SUB(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateInit {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	topicName := params[1]
	if !nsq.IsValidName(topicName) {
		return nil, nsq.ClientErrBadTopic
	}

	channelName := params[2]
	if !nsq.IsValidName(channelName) {
		return nil, nsq.ClientErrBadChannel
	}

	if len(params) == 5 {
		client.ShortIdentifier = params[3]
		client.LongIdentifier = params[4]
	}

	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel(channelName)
	channel.AddClient(client)

	client.Channel = channel
	atomic.StoreInt32(&client.State, nsq.StateSubscribed)

	go p.messagePump(client)

	return nil, nil
}

func (p *ProtocolV2) RDY(client *ClientV2, params []string) ([]byte, error) {
	var err error

	state := atomic.LoadInt32(&client.State)

	if state == nsq.StateClosing {
		// just ignore ready changes on a closing channel
		log.Printf("PROTOCOL(V2): ignoring RDY after CLS in state ClientStateV2Closing")
		return nil, nil
	}

	if state != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	count := 1
	if len(params) > 1 {
		count, err = strconv.Atoi(params[1])
		if err != nil {
			return nil, nsq.ClientErrInvalid
		}
	}

	if count > maxReadyCount {
		log.Printf("ERROR: client(%s) sent ready count %d > %d", client, count, maxReadyCount)
		count = maxReadyCount
	}

	client.SetReadyCount(int64(count))

	return nil, nil
}

func (p *ProtocolV2) FIN(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 2 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	err := client.Channel.FinishMessage(client, []byte(idStr))
	if err != nil {
		return nil, nsq.ClientErrFinishFailed
	}

	client.FinishedMessage()

	return nil, nil
}

func (p *ProtocolV2) REQ(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	timeoutMs, err := strconv.Atoi(params[2])
	if err != nil {
		return nil, nsq.ClientErrInvalid
	}
	timeoutDuration := time.Duration(timeoutMs) * time.Millisecond

	if timeoutDuration < 0 || timeoutDuration > maxTimeout {
		return nil, nsq.ClientErrInvalid
	}

	err = client.Channel.RequeueMessage(client, []byte(idStr), timeoutDuration)
	if err != nil {
		return nil, nsq.ClientErrRequeueFailed
	}

	client.RequeuedMessage()

	return nil, nil
}

func (p *ProtocolV2) CLS(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) > 1 {
		return nil, nsq.ClientErrInvalid
	}

	client.StartClose()

	return []byte("CLOSE_WAIT"), nil
}
<commit_msg>nqsd: close the connection when exiting read loop (assumed idempotent)<commit_after>package main

import (
	"../nsq"
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

const maxTimeout = time.Hour
const maxReadyCount = 2500

type ProtocolV2 struct {
	nsq.Protocol
}

func init() {
	// BigEndian client byte sequence "  V2"
	var magicInt int32
	buf := bytes.NewBuffer([]byte(nsq.MagicV2))
	binary.Read(buf, binary.BigEndian, &magicInt)
	protocols[magicInt] = &ProtocolV2{}
}

func (p *ProtocolV2) IOLoop(conn net.Conn) error {
	var err error
	var line string

	client := NewClientV2(conn)
	atomic.StoreInt32(&client.State, nsq.StateInit)

	err = nil
	reader := bufio.NewReader(client)
	for {
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		if *verbose {
			log.Printf("PROTOCOL(V2): [%s] %#v", client.RemoteAddr().String(), params)
		}

		response, err := p.Exec(client, params)
		if err != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeError, []byte(err.Error()))
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
			continue
		}

		if response != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeResponse, response)
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
		}
	}

	log.Printf("PROTOCOL(V2): [%s] exiting ioloop", client.RemoteAddr().String())
	// TODO: gracefully send clients the close signal
	conn.Close()
	close(client.ExitChan)

	return err
}

func (p *ProtocolV2) Write(client *ClientV2, data []byte) error {
	attempts := 0
	for {
		client.Lock()
		client.SetWriteDeadline(time.Now().Add(time.Second))
		_, err := nsq.SendResponse(client, data)
		client.Unlock()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				attempts++
				if attempts == 3 {
					return errors.New("E_WRITE_TIMEOUT")
				}
				continue
			}
			return err
		}
		break
	}
	return nil
}

func (p *ProtocolV2) Exec(client *ClientV2, params []string) ([]byte, error) {
	switch params[0] {
	case "SUB":
		return p.SUB(client, params)
	case "RDY":
		return p.RDY(client, params)
	case "FIN":
		return p.FIN(client, params)
	case "REQ":
		return p.REQ(client, params)
	case "CLS":
		return p.CLS(client, params)
	}
	return nil, nsq.ClientErrInvalid
}

func (p *ProtocolV2) messagePump(client *ClientV2) {
	var err error

	// ReadyStateChan has a buffer of 1 to guarantee that in the event
	// there is a race before we enter either of the select loops where
	// ReadyStateChan is read from that the update is not lost
	for {
		if !client.IsReadyForMessages() {
			// wait for a change in state
			select {
			case <-client.ReadyStateChan:
			case <-client.ExitChan:
				goto exit
			}
		} else {
			select {
			case <-client.ReadyStateChan:
			case msg := <-client.Channel.clientMsgChan:
				if *verbose {
					log.Printf("PROTOCOL(V2): writing msg(%s) to client(%s) - %s",
						msg.Id, client, msg.Body)
				}

				data, err := msg.Encode()
				if err != nil {
					goto exit
				}

				clientData, err := nsq.Frame(nsq.FrameTypeMessage, data)
				if err != nil {
					goto exit
				}

				client.Channel.StartInFlightTimeout(msg, client)
				client.SendingMessage()

				err = p.Write(client, clientData)
				if err != nil {
					goto exit
				}
			case <-client.ExitChan:
				goto exit
			}
		}
	}

exit:
	log.Printf("PROTOCOL(V2): [%s] exiting messagePump", client.RemoteAddr().String())
	client.Channel.RemoveClient(client)
	if err != nil {
		log.Printf("PROTOCOL(V2): messagePump error - %s", err.Error())
	}
}

func (p *ProtocolV2) SUB(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateInit {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	topicName := params[1]
	if !nsq.IsValidName(topicName) {
		return nil, nsq.ClientErrBadTopic
	}

	channelName := params[2]
	if !nsq.IsValidName(channelName) {
		return nil, nsq.ClientErrBadChannel
	}

	if len(params) == 5 {
		client.ShortIdentifier = params[3]
		client.LongIdentifier = params[4]
	}

	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel(channelName)
	channel.AddClient(client)

	client.Channel = channel
	atomic.StoreInt32(&client.State, nsq.StateSubscribed)

	go p.messagePump(client)

	return nil, nil
}

func (p *ProtocolV2) RDY(client *ClientV2, params []string) ([]byte, error) {
	var err error

	state := atomic.LoadInt32(&client.State)

	if state == nsq.StateClosing {
		// just ignore ready changes on a closing channel
		log.Printf("PROTOCOL(V2): ignoring RDY after CLS in state ClientStateV2Closing")
		return nil, nil
	}

	if state != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	count := 1
	if len(params) > 1 {
		count, err = strconv.Atoi(params[1])
		if err != nil {
			return nil, nsq.ClientErrInvalid
		}
	}

	if count > maxReadyCount {
		log.Printf("ERROR: client(%s) sent ready count %d > %d", client, count, maxReadyCount)
		count = maxReadyCount
	}

	client.SetReadyCount(int64(count))

	return nil, nil
}

func (p *ProtocolV2) FIN(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 2 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	err := client.Channel.FinishMessage(client, []byte(idStr))
	if err != nil {
		return nil, nsq.ClientErrFinishFailed
	}

	client.FinishedMessage()

	return nil, nil
}

func (p *ProtocolV2) REQ(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	timeoutMs, err := strconv.Atoi(params[2])
	if err != nil {
		return nil, nsq.ClientErrInvalid
	}
	timeoutDuration := time.Duration(timeoutMs) * time.Millisecond

	if timeoutDuration < 0 || timeoutDuration > maxTimeout {
		return nil, nsq.ClientErrInvalid
	}

	err = client.Channel.RequeueMessage(client, []byte(idStr), timeoutDuration)
	if err != nil {
		return nil, nsq.ClientErrRequeueFailed
	}

	client.RequeuedMessage()

	return nil, nil
}

func (p *ProtocolV2) CLS(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) > 1 {
		return nil, nsq.ClientErrInvalid
	}

	client.StartClose()

	return []byte("CLOSE_WAIT"), nil
}
<|endoftext|>


<commit_before>package worker

import (
	"net/http"

	gconn "code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/garden/routes"
	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/worker/transport"
	"github.com/concourse/retryhttp"
	"github.com/tedsuo/rata"
)

//go:generate counterfeiter . GardenConnectionFactory
type GardenConnectionFactory interface {
	BuildConnection() gconn.Connection
}

type gardenConnectionFactory struct {
	db                  transport.TransportDB
	logger              lager.Logger
	workerName          string
	workerHost          string
	retryBackOffFactory retryhttp.BackOffFactory
}

func NewGardenConnectionFactory(
	db transport.TransportDB,
	logger lager.Logger,
	workerName string,
	workerHost string,
	retryBackOffFactory retryhttp.BackOffFactory,
) GardenConnectionFactory {
	return &gardenConnectionFactory{
		db:                  db,
		logger:              logger,
		workerName:          workerName,
		workerHost:          workerHost,
		retryBackOffFactory: retryBackOffFactory,
	}
}

func (gcf *gardenConnectionFactory) BuildConnection() gconn.Connection {
	httpClient := &http.Client{
		Transport: &retryhttp.RetryRoundTripper{
			Logger:         gcf.logger.Session("retryable-http-client"),
			BackOffFactory: gcf.retryBackOffFactory,
			RoundTripper:   transport.NewRoundTripper(gcf.workerName, gcf.workerHost, gcf.db, &http.Transport{DisableKeepAlives: true}),
		},
	}

	hijackableClient := &retryhttp.RetryHijackableClient{
		Logger:           gcf.logger.Session("retry-hijackable-client"),
		BackOffFactory:   gcf.retryBackOffFactory,
		HijackableClient: transport.NewHijackableClient(gcf.workerName, gcf.db, retryhttp.DefaultHijackableClient),
	}

	// the request generator's address doesn't matter because it's overwritten by the worker lookup clients
	hijackStreamer := &transport.WorkerHijackStreamer{
		HttpClient:       httpClient,
		HijackableClient: hijackableClient,
		Req:              rata.NewRequestGenerator("http://127.0.0.1:8080", routes.Routes),
	}

	return gconn.NewWithHijacker(hijackStreamer, gcf.logger)
}
<commit_msg>set default Retryer for RoundTrippper<commit_after>package worker

import (
	"net/http"

	gconn "code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/garden/routes"
	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/worker/transport"
	"github.com/concourse/retryhttp"
	"github.com/tedsuo/rata"
)

//go:generate counterfeiter . GardenConnectionFactory
type GardenConnectionFactory interface {
	BuildConnection() gconn.Connection
}

type gardenConnectionFactory struct {
	db                  transport.TransportDB
	logger              lager.Logger
	workerName          string
	workerHost          string
	retryBackOffFactory retryhttp.BackOffFactory
}

func NewGardenConnectionFactory(
	db transport.TransportDB,
	logger lager.Logger,
	workerName string,
	workerHost string,
	retryBackOffFactory retryhttp.BackOffFactory,
) GardenConnectionFactory {
	return &gardenConnectionFactory{
		db:                  db,
		logger:              logger,
		workerName:          workerName,
		workerHost:          workerHost,
		retryBackOffFactory: retryBackOffFactory,
	}
}

func (gcf *gardenConnectionFactory) BuildConnection() gconn.Connection {
	httpClient := &http.Client{
		Transport: &retryhttp.RetryRoundTripper{
			Logger:         gcf.logger.Session("retryable-http-client"),
			BackOffFactory: gcf.retryBackOffFactory,
			RoundTripper:   transport.NewRoundTripper(gcf.workerName, gcf.workerHost, gcf.db, &http.Transport{DisableKeepAlives: true}),
			Retryer:        &retryhttp.DefaultRetryer{},
		},
	}

	hijackableClient := &retryhttp.RetryHijackableClient{
		Logger:           gcf.logger.Session("retry-hijackable-client"),
		BackOffFactory:   gcf.retryBackOffFactory,
		HijackableClient: transport.NewHijackableClient(gcf.workerName, gcf.db, retryhttp.DefaultHijackableClient),
		Retryer:          &retryhttp.DefaultRetryer{},
	}

	// the request generator's address doesn't matter because it's overwritten by the worker lookup clients
	hijackStreamer := &transport.WorkerHijackStreamer{
		HttpClient:       httpClient,
		HijackableClient: hijackableClient,
		Req:              rata.NewRequestGenerator("http://127.0.0.1:8080", routes.Routes),
	}

	return gconn.NewWithHijacker(hijackStreamer, gcf.logger)
}
<|endoftext|>


<commit_before>package tchannel_test

// This file contains functions for tests to access internal tchannel state.
// Since it has a _test.go suffix, it is only compiled with tests in this package.

// Copyright (c) 2015 Uber Technologies, Inc.

// 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.

import (
	"bytes"
	"math/rand"
	"strings"
	"sync"
	"testing"
	"time"

	. "github.com/uber/tchannel/golang"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/raw"
	"github.com/uber/tchannel/golang/testutils"
	"golang.org/x/net/context"
)

type swapper struct {
	t *testing.T
}

func (s *swapper) OnError(ctx context.Context, err error) {
	s.t.Errorf("OnError: %v", err)
}

func (*swapper) Handle(ctx context.Context, args *raw.Args) (*raw.Res, error) {
	return &raw.Res{
		Arg2: args.Arg3,
		Arg3: args.Arg2,
	}, nil
}

func checkEmptyExchangesConns(connections []*Connection) string {
	var errors []string
	for _, c := range connections {
		if v := CheckEmptyExchanges(c); v != "" {
			errors = append(errors, v)
		}
	}
	return strings.Join(errors, "\n")
}

func TestFramesReleased(t *testing.T) {
	if testing.Short() {
		return
	}

	defer testutils.SetTimeout(t, 10*time.Second)()
	const (
		requestsPerGoroutine = 10
		numGoroutines        = 10
		maxRandArg           = 512 * 1024
	)

	// Generate random bytes used to create arguments.
	randBytes := make([]byte, maxRandArg)
	for i := 0; i < len(randBytes); i += 8 {
		n := rand.Int63()
		for j := 0; j < 8; j++ {
			randBytes[i+j] = byte(n & 0xff)
			n = n << 1
		}
	}

	var connections []*Connection
	pool := NewRecordingFramePool()
	require.NoError(t, testutils.WithServer(&testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: pool,
		},
	}, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		clientCh, err := NewChannel("swap-client", nil)
		require.NoError(t, err)
		defer clientCh.Close()

		// Create an active connection that can be shared by the goroutines by calling Ping.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, clientCh.Ping(ctx, hostPort))

		generateArg := func(n int) []byte {
			from := rand.Intn(maxRandArg - n)
			return randBytes[from : from+n]
		}

		var wg sync.WaitGroup
		worker := func() {
			for i := 0; i < requestsPerGoroutine; i++ {
				ctx, cancel := NewContext(time.Second * 5)
				defer cancel()

				require.NoError(t, clientCh.Ping(ctx, hostPort))

				argSize := rand.Intn(maxRandArg)
				arg2 := generateArg(argSize)
				arg3 := generateArg(argSize)
				resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
				if !assert.NoError(t, err, "error during sendRecv") {
					continue
				}

				// We expect the arguments to be swapped.
				if bytes.Compare(arg3, resArg2) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
				}
				if bytes.Compare(arg2, resArg3) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
				}
			}
			wg.Done()
		}

		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go worker()
		}

		wg.Wait()

		connections = append(connections, GetConnections(serverCh)...)
		connections = append(connections, GetConnections(clientCh)...)
	}))

	// Wait a few milliseconds for the closing of channels to take effect.
	time.Sleep(10 * time.Millisecond)

	if unreleasedCount, isEmpty := pool.CheckEmpty(); isEmpty != "" || unreleasedCount > 0 {
		t.Errorf("Frame pool has %v unreleased frames, errors:\n%v", unreleasedCount, isEmpty)
	}

	// Check the message exchanges and make sure they are all empty.
	if exchangesLeft := checkEmptyExchangesConns(connections); exchangesLeft != "" {
		t.Errorf("Found uncleared message exchanges:\n%v", exchangesLeft)
	}
}
<commit_msg>Update test to use testutils.RandBytes<commit_after>package tchannel_test

// This file contains functions for tests to access internal tchannel state.
// Since it has a _test.go suffix, it is only compiled with tests in this package.

// Copyright (c) 2015 Uber Technologies, Inc.

// 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.

import (
	"bytes"
	"math/rand"
	"strings"
	"sync"
	"testing"
	"time"

	. "github.com/uber/tchannel/golang"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/raw"
	"github.com/uber/tchannel/golang/testutils"
	"golang.org/x/net/context"
)

type swapper struct {
	t *testing.T
}

func (s *swapper) OnError(ctx context.Context, err error) {
	s.t.Errorf("OnError: %v", err)
}

func (*swapper) Handle(ctx context.Context, args *raw.Args) (*raw.Res, error) {
	return &raw.Res{
		Arg2: args.Arg3,
		Arg3: args.Arg2,
	}, nil
}

func checkEmptyExchangesConns(connections []*Connection) string {
	var errors []string
	for _, c := range connections {
		if v := CheckEmptyExchanges(c); v != "" {
			errors = append(errors, v)
		}
	}
	return strings.Join(errors, "\n")
}

func TestFramesReleased(t *testing.T) {
	if testing.Short() {
		return
	}

	defer testutils.SetTimeout(t, 10*time.Second)()
	const (
		requestsPerGoroutine = 10
		numGoroutines        = 10
		maxRandArg           = 512 * 1024
	)

	var connections []*Connection
	pool := NewRecordingFramePool()
	require.NoError(t, testutils.WithServer(&testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: pool,
		},
	}, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		clientCh, err := NewChannel("swap-client", nil)
		require.NoError(t, err)
		defer clientCh.Close()

		// Create an active connection that can be shared by the goroutines by calling Ping.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, clientCh.Ping(ctx, hostPort))

		var wg sync.WaitGroup
		worker := func() {
			for i := 0; i < requestsPerGoroutine; i++ {
				ctx, cancel := NewContext(time.Second * 5)
				defer cancel()

				require.NoError(t, clientCh.Ping(ctx, hostPort))

				arg2 := testutils.RandBytes(rand.Intn(maxRandArg))
				arg3 := testutils.RandBytes(rand.Intn(maxRandArg))
				resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
				if !assert.NoError(t, err, "error during sendRecv") {
					continue
				}

				// We expect the arguments to be swapped.
				if bytes.Compare(arg3, resArg2) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
				}
				if bytes.Compare(arg2, resArg3) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
				}
			}
			wg.Done()
		}

		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go worker()
		}

		wg.Wait()

		connections = append(connections, GetConnections(serverCh)...)
		connections = append(connections, GetConnections(clientCh)...)
	}))

	// Wait a few milliseconds for the closing of channels to take effect.
	time.Sleep(10 * time.Millisecond)

	if unreleasedCount, isEmpty := pool.CheckEmpty(); isEmpty != "" || unreleasedCount > 0 {
		t.Errorf("Frame pool has %v unreleased frames, errors:\n%v", unreleasedCount, isEmpty)
	}

	// Check the message exchanges and make sure they are all empty.
	if exchangesLeft := checkEmptyExchangesConns(connections); exchangesLeft != "" {
		t.Errorf("Found uncleared message exchanges:\n%v", exchangesLeft)
	}
}
<|endoftext|>