<commit_before>package ephemeral

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"time"

	"github.com/keybase/client/go/kbcrypto"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/teambot"
	"github.com/keybase/client/go/teams"
)

type TeambotEKSeed keybase1.Bytes32

func (s *TeambotEKSeed) DeriveDHKey() *libkb.NaclDHKeyPair {
	return deriveDHKey(keybase1.Bytes32(*s), libkb.DeriveReasonTeambotEKEncryption)
}

func deriveTeambotEKFromTeamEK(mctx libkb.MetaContext, teamEK keybase1.TeamEk, botUID keybase1.UID) TeambotEKSeed {
	hasher := hmac.New(sha256.New, teamEK.Seed[:])
	_, _ = hasher.Write(botUID.ToBytes())
	_, _ = hasher.Write([]byte(libkb.EncryptionReasonTeambotEphemeralKey))
	return TeambotEKSeed(libkb.MakeByte32(hasher.Sum(nil)))
}

func newTeambotEKSeedFromBytes(b []byte) (s TeambotEKSeed, err error) {
	seed, err := newEKSeedFromBytes(b)
	if err != nil {
		return s, err
	}
	return TeambotEKSeed(seed), nil
}

type TeambotEphemeralKeyer struct{}

var _ EphemeralKeyer = (*TeambotEphemeralKeyer)(nil)

func NewTeambotEphemeralKeyer() *TeambotEphemeralKeyer {
	return &TeambotEphemeralKeyer{}
}

func (k *TeambotEphemeralKeyer) Type() keybase1.TeamEphemeralKeyType {
	return keybase1.TeamEphemeralKeyType_TEAMBOT
}

func postNewTeambotEK(mctx libkb.MetaContext, teamID keybase1.TeamID, sig, box string) (err error) {
	defer mctx.TraceTimed("postNewTeambotEK", func() error { return err })()

	apiArg := libkb.APIArg{
		Endpoint:    "teambot/key",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"team_id":      libkb.S{Val: string(teamID)},
			"sig":          libkb.S{Val: sig},
			"box":          libkb.S{Val: box},
			"is_ephemeral": libkb.B{Val: true},
		},
		AppStatusCodes: []int{libkb.SCOk, libkb.SCTeambotKeyGenerationExists},
	}
	_, err = mctx.G().GetAPI().Post(mctx, apiArg)
	return err
}

func prepareNewTeambotEK(mctx libkb.MetaContext, team *teams.Team, botUID keybase1.UID,
	seed TeambotEKSeed, metadata *keybase1.TeambotEkMetadata,
	merkleRoot libkb.MerkleRoot) (sig string, box *keybase1.TeambotEkBoxed, err error) {
	defer mctx.TraceTimed("prepareNewTeambotEK", func() error { return err })()

	statement, _, _, err := fetchUserEKStatement(mctx, botUID)
	if err != nil {
		return "", nil, err
	}
	activeMetadataMap, err := activeUserEKMetadata(mctx, map[keybase1.UID]*keybase1.UserEkStatement{botUID: statement}, merkleRoot)
	if err != nil {
		return "", nil, err
	} else if len(activeMetadataMap) == 0 {
		mctx.Debug("unable to make teambot key, bot has no active user EKs")
		return "", nil, nil
	}
	activeMetadata := activeMetadataMap[botUID]
	metadata.UserEkGeneration = activeMetadata.Generation

	// Encrypting with a nil sender means we'll generate a random sender
	// private key.
	recipientKey, err := libkb.ImportKeypairFromKID(activeMetadata.Kid)
	if err != nil {
		return "", nil, err
	}
	boxedSeed, err := recipientKey.EncryptToString(seed[:], nil)
	if err != nil {
		return "", nil, err
	}

	boxed := keybase1.TeambotEkBoxed{
		Box:      boxedSeed,
		Metadata: *metadata,
	}

	metadataJSON, err := json.Marshal(metadata)
	if err != nil {
		return "", nil, err
	}

	signingKey, err := team.SigningKey(mctx.Ctx())
	if err != nil {
		return "", nil, err
	}
	sig, _, err = signingKey.SignToString(metadataJSON)
	if err != nil {
		return "", nil, err
	}
	return sig, &boxed, nil
}

func fetchLatestTeambotEK(mctx libkb.MetaContext, teamID keybase1.TeamID) (metadata *keybase1.TeambotEkMetadata, wrongKID bool, err error) {
	defer mctx.TraceTimed("fetchLatestTeambotEK", func() error { return err })()

	var parsedResponse teambot.TeambotKeyResponse
	gotResult := false
	for i := 0; i < 10; i++ {
		apiArg := libkb.APIArg{
			Endpoint:    "teambot/key",
			SessionType: libkb.APISessionTypeREQUIRED,
			Args: libkb.HTTPArgs{
				"team_id":      libkb.S{Val: string(teamID)},
				"is_ephemeral": libkb.B{Val: true},
			},
		}
		res, err := mctx.G().GetAPI().Get(mctx, apiArg)
		if err != nil {
			return nil, false, err
		}

		if err = res.Body.UnmarshalAgain(&parsedResponse); err != nil {
			return nil, false, err
		}
		if parsedResponse.Result == nil {
			mctx.Debug("no teambot ek response, trying again, attempt: %d", i)
			time.Sleep(time.Second)
			continue
		}
		gotResult = true
		break
	}
	if !gotResult {
		return nil, false, nil
	}
	return verifyTeambotSigWithLatestPTK(mctx, teamID, parsedResponse.Result.Sig)
}

func extractTeambotEKMetadataFromSig(sig string) (*kbcrypto.NaclSigningKeyPublic, *keybase1.TeambotEkMetadata, error) {
	signerKey, payload, _, err := kbcrypto.NaclVerifyAndExtract(sig)
	if err != nil {
		return signerKey, nil, err
	}

	parsedMetadata := keybase1.TeambotEkMetadata{}
	if err = json.Unmarshal(payload, &parsedMetadata); err != nil {
		return signerKey, nil, err
	}
	return signerKey, &parsedMetadata, nil
}

// Verify that the blob is validly signed, and that the signing key is the
// given team's latest PTK, then parse its contents.
func verifyTeambotSigWithLatestPTK(mctx libkb.MetaContext, teamID keybase1.TeamID, sig string) (
	metadata *keybase1.TeambotEkMetadata, wrongKID bool, err error) {
	defer mctx.TraceTimed("verifyTeambotSigWithLatestPTK", func() error { return err })()

	signerKey, metadata, err := extractTeambotEKMetadataFromSig(sig)
	if err != nil {
		return nil, false, err
	}

	team, err := teams.Load(mctx.Ctx(), mctx.G(), keybase1.LoadTeamArg{
		ID: teamID,
	})
	if err != nil {
		return nil, false, err
	}

	// Verify the signing key corresponds to the latest PTK. We load the team's
	// from cache, but if the KID doesn't match, we try a forced reload to see
	// if the cache might've been stale. Only if the KID still doesn't match
	// after the reload do we complain.
	teamSigningKID, err := team.SigningKID(mctx.Ctx())
	if err != nil {
		return nil, false, err
	}
	if !teamSigningKID.Equal(signerKey.GetKID()) {
		// The latest PTK might be stale. Force a reload, then check this over again.
		team, err := teams.Load(mctx.Ctx(), mctx.G(), keybase1.LoadTeamArg{
			ID:          team.ID,
			ForceRepoll: true,
		})
		if err != nil {
			return nil, false, err
		}
		teamSigningKID, err = team.SigningKID(mctx.Ctx())
		if err != nil {
			return nil, false, err
		}
		// return the metdata with wrongKID=true
		if !teamSigningKID.Equal(signerKey.GetKID()) {
			return metadata, true, fmt.Errorf("teambotEK returned for PTK signing KID %s, but latest is %s",
				signerKey.GetKID(), teamSigningKID)
		}
	}

	// If we didn't short circuit above, then the signing key is correct.
	return metadata, false, nil
}

func (k *TeambotEphemeralKeyer) Unbox(mctx libkb.MetaContext, boxed keybase1.TeamEphemeralKeyBoxed,
	contentCtime *gregor1.Time) (ek keybase1.TeamEphemeralKey, err error) {
	defer mctx.TraceTimed(fmt.Sprintf("TeambotEphemeralKeyer#Unbox: teambotEKGeneration: %v", boxed.Generation()), func() error { return err })()

	typ, err := boxed.KeyType()
	if err != nil {
		return ek, err
	}
	if !typ.IsTeambot() {
		return ek, NewIncorrectTeamEphemeralKeyTypeError(typ, keybase1.TeamEphemeralKeyType_TEAMBOT)
	}

	teambotEKBoxed := boxed.Teambot()
	teambotEKGeneration := teambotEKBoxed.Metadata.Generation
	userEKBoxStorage := mctx.G().GetUserEKBoxStorage()
	userEK, err := userEKBoxStorage.Get(mctx, teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
	if err != nil {
		mctx.Debug("unable to get from userEKStorage %v", err)
		if _, ok := err.(EphemeralKeyError); ok {
			return ek, newEKUnboxErr(mctx, TeambotEKKind, teambotEKGeneration, UserEKKind,
				teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
		}
		return ek, err
	}

	userSeed := UserEKSeed(userEK.Seed)
	userKeypair := userSeed.DeriveDHKey()

	msg, _, err := userKeypair.DecryptFromString(teambotEKBoxed.Box)
	if err != nil {
		mctx.Debug("unable to decrypt teambotEKBoxed %v", err)
		return ek, newEKUnboxErr(mctx, TeambotEKKind, teambotEKGeneration, UserEKKind,
			teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
	}

	seed, err := newTeambotEKSeedFromBytes(msg)
	if err != nil {
		return ek, err
	}

	keypair := seed.DeriveDHKey()
	if !keypair.GetKID().Equal(teambotEKBoxed.Metadata.Kid) {
		return ek, fmt.Errorf("Failed to verify server given seed [%s] against signed KID [%s]. Box: %+v",
			teambotEKBoxed.Metadata.Kid, keypair.GetKID(), teambotEKBoxed)
	}

	return keybase1.NewTeamEphemeralKeyWithTeambot(keybase1.TeambotEk{
		Seed:     keybase1.Bytes32(seed),
		Metadata: teambotEKBoxed.Metadata,
	}), nil
}

func (k *TeambotEphemeralKeyer) Fetch(mctx libkb.MetaContext, teamID keybase1.TeamID, generation keybase1.EkGeneration,
	contentCtime *gregor1.Time) (teambotEK keybase1.TeamEphemeralKeyBoxed, err error) {
	apiArg := libkb.APIArg{
		Endpoint:    "teambot/box",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"team_id":      libkb.S{Val: string(teamID)},
			"generation":   libkb.U{Val: uint64(generation)},
			"is_ephemeral": libkb.B{Val: true},
		},
	}

	var resp teambot.TeambotKeyBoxedResponse
	res, err := mctx.G().GetAPI().Get(mctx, apiArg)
	if err != nil {
		err = errFromAppStatus(err)
		return teambotEK, err
	}

	if err = res.Body.UnmarshalAgain(&resp); err != nil {
		return teambotEK, err
	}

	if resp.Result == nil {
		err = newEKMissingBoxErr(mctx, TeambotEKKind, generation)
		return teambotEK, err
	}

	// Although we verify the signature is valid, it's possible that this key
	// was signed with a PTK that is not our latest and greatest. We allow this
	// when we are using this ek for *decryption*. When getting a key for
	// *encryption* callers are responsible for verifying the signature is
	// signed by the latest PTK or requesting the generation of a new EK. This
	// logic currently lives in ephemeral/lib.go#getLatestTeambotEK
	_, metadata, err := extractTeambotEKMetadataFromSig(resp.Result.Sig)
	if err != nil {
		return teambotEK, err
	} else if metadata == nil { // shouldn't happen
		return teambotEK, fmt.Errorf("unable to fetch valid teambotEKMetadata")
	}

	if generation != metadata.Generation {
		// sanity check that we got the right generation
		return teambotEK, newEKCorruptedErr(mctx, TeambotEKKind, generation, metadata.Generation)
	}
	teambotEKBoxed := keybase1.TeambotEkBoxed{
		Box:      resp.Result.Box,
		Metadata: *metadata,
	}
	return keybase1.NewTeamEphemeralKeyBoxedWithTeambot(teambotEKBoxed), nil
}
<commit_msg>add application arg for teambot endpoints (#22449)<commit_after>package ephemeral

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"time"

	"github.com/keybase/client/go/kbcrypto"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/teambot"
	"github.com/keybase/client/go/teams"
)

type TeambotEKSeed keybase1.Bytes32

func (s *TeambotEKSeed) DeriveDHKey() *libkb.NaclDHKeyPair {
	return deriveDHKey(keybase1.Bytes32(*s), libkb.DeriveReasonTeambotEKEncryption)
}

func deriveTeambotEKFromTeamEK(mctx libkb.MetaContext, teamEK keybase1.TeamEk, botUID keybase1.UID) TeambotEKSeed {
	hasher := hmac.New(sha256.New, teamEK.Seed[:])
	_, _ = hasher.Write(botUID.ToBytes())
	_, _ = hasher.Write([]byte(libkb.EncryptionReasonTeambotEphemeralKey))
	return TeambotEKSeed(libkb.MakeByte32(hasher.Sum(nil)))
}

func newTeambotEKSeedFromBytes(b []byte) (s TeambotEKSeed, err error) {
	seed, err := newEKSeedFromBytes(b)
	if err != nil {
		return s, err
	}
	return TeambotEKSeed(seed), nil
}

type TeambotEphemeralKeyer struct{}

var _ EphemeralKeyer = (*TeambotEphemeralKeyer)(nil)

func NewTeambotEphemeralKeyer() *TeambotEphemeralKeyer {
	return &TeambotEphemeralKeyer{}
}

func (k *TeambotEphemeralKeyer) Type() keybase1.TeamEphemeralKeyType {
	return keybase1.TeamEphemeralKeyType_TEAMBOT
}

func postNewTeambotEK(mctx libkb.MetaContext, teamID keybase1.TeamID, sig, box string) (err error) {
	defer mctx.TraceTimed("postNewTeambotEK", func() error { return err })()

	apiArg := libkb.APIArg{
		Endpoint:    "teambot/key",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"team_id":      libkb.S{Val: string(teamID)},
			"sig":          libkb.S{Val: sig},
			"box":          libkb.S{Val: box},
			"is_ephemeral": libkb.B{Val: true},
			"application":  libkb.I{Val: int(keybase1.TeamApplication_CHAT)},
		},
		AppStatusCodes: []int{libkb.SCOk, libkb.SCTeambotKeyGenerationExists},
	}
	_, err = mctx.G().GetAPI().Post(mctx, apiArg)
	return err
}

func prepareNewTeambotEK(mctx libkb.MetaContext, team *teams.Team, botUID keybase1.UID,
	seed TeambotEKSeed, metadata *keybase1.TeambotEkMetadata,
	merkleRoot libkb.MerkleRoot) (sig string, box *keybase1.TeambotEkBoxed, err error) {
	defer mctx.TraceTimed("prepareNewTeambotEK", func() error { return err })()

	statement, _, _, err := fetchUserEKStatement(mctx, botUID)
	if err != nil {
		return "", nil, err
	}
	activeMetadataMap, err := activeUserEKMetadata(mctx, map[keybase1.UID]*keybase1.UserEkStatement{botUID: statement}, merkleRoot)
	if err != nil {
		return "", nil, err
	} else if len(activeMetadataMap) == 0 {
		mctx.Debug("unable to make teambot key, bot has no active user EKs")
		return "", nil, nil
	}
	activeMetadata := activeMetadataMap[botUID]
	metadata.UserEkGeneration = activeMetadata.Generation

	// Encrypting with a nil sender means we'll generate a random sender
	// private key.
	recipientKey, err := libkb.ImportKeypairFromKID(activeMetadata.Kid)
	if err != nil {
		return "", nil, err
	}
	boxedSeed, err := recipientKey.EncryptToString(seed[:], nil)
	if err != nil {
		return "", nil, err
	}

	boxed := keybase1.TeambotEkBoxed{
		Box:      boxedSeed,
		Metadata: *metadata,
	}

	metadataJSON, err := json.Marshal(metadata)
	if err != nil {
		return "", nil, err
	}

	signingKey, err := team.SigningKey(mctx.Ctx())
	if err != nil {
		return "", nil, err
	}
	sig, _, err = signingKey.SignToString(metadataJSON)
	if err != nil {
		return "", nil, err
	}
	return sig, &boxed, nil
}

func fetchLatestTeambotEK(mctx libkb.MetaContext, teamID keybase1.TeamID) (metadata *keybase1.TeambotEkMetadata, wrongKID bool, err error) {
	defer mctx.TraceTimed("fetchLatestTeambotEK", func() error { return err })()

	var parsedResponse teambot.TeambotKeyResponse
	gotResult := false
	for i := 0; i < 10; i++ {
		apiArg := libkb.APIArg{
			Endpoint:    "teambot/key",
			SessionType: libkb.APISessionTypeREQUIRED,
			Args: libkb.HTTPArgs{
				"team_id":      libkb.S{Val: string(teamID)},
				"is_ephemeral": libkb.B{Val: true},
				"application":  libkb.I{Val: int(keybase1.TeamApplication_CHAT)},
			},
		}
		res, err := mctx.G().GetAPI().Get(mctx, apiArg)
		if err != nil {
			return nil, false, err
		}

		if err = res.Body.UnmarshalAgain(&parsedResponse); err != nil {
			return nil, false, err
		}
		if parsedResponse.Result == nil {
			mctx.Debug("no teambot ek response, trying again, attempt: %d", i)
			time.Sleep(time.Second)
			continue
		}
		gotResult = true
		break
	}
	if !gotResult {
		return nil, false, nil
	}
	return verifyTeambotSigWithLatestPTK(mctx, teamID, parsedResponse.Result.Sig)
}

func extractTeambotEKMetadataFromSig(sig string) (*kbcrypto.NaclSigningKeyPublic, *keybase1.TeambotEkMetadata, error) {
	signerKey, payload, _, err := kbcrypto.NaclVerifyAndExtract(sig)
	if err != nil {
		return signerKey, nil, err
	}

	parsedMetadata := keybase1.TeambotEkMetadata{}
	if err = json.Unmarshal(payload, &parsedMetadata); err != nil {
		return signerKey, nil, err
	}
	return signerKey, &parsedMetadata, nil
}

// Verify that the blob is validly signed, and that the signing key is the
// given team's latest PTK, then parse its contents.
func verifyTeambotSigWithLatestPTK(mctx libkb.MetaContext, teamID keybase1.TeamID, sig string) (
	metadata *keybase1.TeambotEkMetadata, wrongKID bool, err error) {
	defer mctx.TraceTimed("verifyTeambotSigWithLatestPTK", func() error { return err })()

	signerKey, metadata, err := extractTeambotEKMetadataFromSig(sig)
	if err != nil {
		return nil, false, err
	}

	team, err := teams.Load(mctx.Ctx(), mctx.G(), keybase1.LoadTeamArg{
		ID: teamID,
	})
	if err != nil {
		return nil, false, err
	}

	// Verify the signing key corresponds to the latest PTK. We load the team's
	// from cache, but if the KID doesn't match, we try a forced reload to see
	// if the cache might've been stale. Only if the KID still doesn't match
	// after the reload do we complain.
	teamSigningKID, err := team.SigningKID(mctx.Ctx())
	if err != nil {
		return nil, false, err
	}
	if !teamSigningKID.Equal(signerKey.GetKID()) {
		// The latest PTK might be stale. Force a reload, then check this over again.
		team, err := teams.Load(mctx.Ctx(), mctx.G(), keybase1.LoadTeamArg{
			ID:          team.ID,
			ForceRepoll: true,
		})
		if err != nil {
			return nil, false, err
		}
		teamSigningKID, err = team.SigningKID(mctx.Ctx())
		if err != nil {
			return nil, false, err
		}
		// return the metdata with wrongKID=true
		if !teamSigningKID.Equal(signerKey.GetKID()) {
			return metadata, true, fmt.Errorf("teambotEK returned for PTK signing KID %s, but latest is %s",
				signerKey.GetKID(), teamSigningKID)
		}
	}

	// If we didn't short circuit above, then the signing key is correct.
	return metadata, false, nil
}

func (k *TeambotEphemeralKeyer) Unbox(mctx libkb.MetaContext, boxed keybase1.TeamEphemeralKeyBoxed,
	contentCtime *gregor1.Time) (ek keybase1.TeamEphemeralKey, err error) {
	defer mctx.TraceTimed(fmt.Sprintf("TeambotEphemeralKeyer#Unbox: teambotEKGeneration: %v", boxed.Generation()), func() error { return err })()

	typ, err := boxed.KeyType()
	if err != nil {
		return ek, err
	}
	if !typ.IsTeambot() {
		return ek, NewIncorrectTeamEphemeralKeyTypeError(typ, keybase1.TeamEphemeralKeyType_TEAMBOT)
	}

	teambotEKBoxed := boxed.Teambot()
	teambotEKGeneration := teambotEKBoxed.Metadata.Generation
	userEKBoxStorage := mctx.G().GetUserEKBoxStorage()
	userEK, err := userEKBoxStorage.Get(mctx, teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
	if err != nil {
		mctx.Debug("unable to get from userEKStorage %v", err)
		if _, ok := err.(EphemeralKeyError); ok {
			return ek, newEKUnboxErr(mctx, TeambotEKKind, teambotEKGeneration, UserEKKind,
				teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
		}
		return ek, err
	}

	userSeed := UserEKSeed(userEK.Seed)
	userKeypair := userSeed.DeriveDHKey()

	msg, _, err := userKeypair.DecryptFromString(teambotEKBoxed.Box)
	if err != nil {
		mctx.Debug("unable to decrypt teambotEKBoxed %v", err)
		return ek, newEKUnboxErr(mctx, TeambotEKKind, teambotEKGeneration, UserEKKind,
			teambotEKBoxed.Metadata.UserEkGeneration, contentCtime)
	}

	seed, err := newTeambotEKSeedFromBytes(msg)
	if err != nil {
		return ek, err
	}

	keypair := seed.DeriveDHKey()
	if !keypair.GetKID().Equal(teambotEKBoxed.Metadata.Kid) {
		return ek, fmt.Errorf("Failed to verify server given seed [%s] against signed KID [%s]. Box: %+v",
			teambotEKBoxed.Metadata.Kid, keypair.GetKID(), teambotEKBoxed)
	}

	return keybase1.NewTeamEphemeralKeyWithTeambot(keybase1.TeambotEk{
		Seed:     keybase1.Bytes32(seed),
		Metadata: teambotEKBoxed.Metadata,
	}), nil
}

func (k *TeambotEphemeralKeyer) Fetch(mctx libkb.MetaContext, teamID keybase1.TeamID, generation keybase1.EkGeneration,
	contentCtime *gregor1.Time) (teambotEK keybase1.TeamEphemeralKeyBoxed, err error) {
	apiArg := libkb.APIArg{
		Endpoint:    "teambot/box",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"team_id":      libkb.S{Val: string(teamID)},
			"generation":   libkb.U{Val: uint64(generation)},
			"is_ephemeral": libkb.B{Val: true},
			"application":  libkb.I{Val: int(keybase1.TeamApplication_CHAT)},
		},
	}

	var resp teambot.TeambotKeyBoxedResponse
	res, err := mctx.G().GetAPI().Get(mctx, apiArg)
	if err != nil {
		err = errFromAppStatus(err)
		return teambotEK, err
	}

	if err = res.Body.UnmarshalAgain(&resp); err != nil {
		return teambotEK, err
	}

	if resp.Result == nil {
		err = newEKMissingBoxErr(mctx, TeambotEKKind, generation)
		return teambotEK, err
	}

	// Although we verify the signature is valid, it's possible that this key
	// was signed with a PTK that is not our latest and greatest. We allow this
	// when we are using this ek for *decryption*. When getting a key for
	// *encryption* callers are responsible for verifying the signature is
	// signed by the latest PTK or requesting the generation of a new EK. This
	// logic currently lives in ephemeral/lib.go#getLatestTeambotEK
	_, metadata, err := extractTeambotEKMetadataFromSig(resp.Result.Sig)
	if err != nil {
		return teambotEK, err
	} else if metadata == nil { // shouldn't happen
		return teambotEK, fmt.Errorf("unable to fetch valid teambotEKMetadata")
	}

	if generation != metadata.Generation {
		// sanity check that we got the right generation
		return teambotEK, newEKCorruptedErr(mctx, TeambotEKKind, generation, metadata.Generation)
	}
	teambotEKBoxed := keybase1.TeambotEkBoxed{
		Box:      resp.Result.Box,
		Metadata: *metadata,
	}
	return keybase1.NewTeamEphemeralKeyBoxedWithTeambot(teambotEKBoxed), nil
}
<|endoftext|>


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

// +build cgo,!arm,!arm64,!ios

package x509

/*
#cgo CFLAGS: -mmacosx-version-min=10.6 -D__MAC_OS_X_VERSION_MAX_ALLOWED=1060
#cgo LDFLAGS: -framework CoreFoundation -framework Security

#include <errno.h>
#include <sys/sysctl.h>

#include <CoreFoundation/CoreFoundation.h>
#include <Security/Security.h>

// FetchPEMRoots_MountainLion is the version of FetchPEMRoots from Go 1.6
// which still works on OS X 10.8 (Mountain Lion).
// It lacks support for admin & user cert domains.
// See golang.org/issue/16473
int FetchPEMRoots_MountainLion(CFDataRef *pemRoots) {
	if (pemRoots == NULL) {
		return -1;
	}
	CFArrayRef certs = NULL;
	OSStatus err = SecTrustCopyAnchorCertificates(&certs);
	if (err != noErr) {
		return -1;
	}
	CFMutableDataRef combinedData = CFDataCreateMutable(kCFAllocatorDefault, 0);
	int i, ncerts = CFArrayGetCount(certs);
	for (i = 0; i < ncerts; i++) {
		CFDataRef data = NULL;
		SecCertificateRef cert = (SecCertificateRef)CFArrayGetValueAtIndex(certs, i);
		if (cert == NULL) {
			continue;
		}
		// Note: SecKeychainItemExport is deprecated as of 10.7 in favor of SecItemExport.
		// Once we support weak imports via cgo we should prefer that, and fall back to this
		// for older systems.
		err = SecKeychainItemExport(cert, kSecFormatX509Cert, kSecItemPemArmour, NULL, &data);
		if (err != noErr) {
			continue;
		}
		if (data != NULL) {
			CFDataAppendBytes(combinedData, CFDataGetBytePtr(data), CFDataGetLength(data));
			CFRelease(data);
		}
	}
	CFRelease(certs);
	*pemRoots = combinedData;
	return 0;
}

// useOldCode reports whether the running machine is OS X 10.8 Mountain Lion
// or older. We only support Mountain Lion and higher, but we'll at least try our
// best on older machines and continue to use the old code path.
//
// See golang.org/issue/16473
int useOldCode() {
	char str[256];
	size_t size = sizeof(str);
	memset(str, 0, size);
	sysctlbyname("kern.osrelease", str, &size, NULL, 0);
	// OS X 10.8 is osrelease "12.*", 10.7 is 11.*, 10.6 is 10.*.
	// We never supported things before that.
	return memcmp(str, "12.", 3) == 0 || memcmp(str, "11.", 3) == 0 || memcmp(str, "10.", 3) == 0;
}

// FetchPEMRootsCTX509 fetches the system's list of trusted X.509 root certificates.
//
// On success it returns 0 and fills pemRoots with a CFDataRef that contains the extracted root
// certificates of the system. On failure, the function returns -1.
//
// Note: The CFDataRef returned in pemRoots must be released (using CFRelease) after
// we've consumed its content.
int FetchPEMRootsCTX509(CFDataRef *pemRoots) {
	if (useOldCode()) {
		return FetchPEMRoots_MountainLion(pemRoots);
	}

	// Get certificates from all domains, not just System, this lets
	// the user add CAs to their "login" keychain, and Admins to add
	// to the "System" keychain
	SecTrustSettingsDomain domains[] = { kSecTrustSettingsDomainSystem,
					     kSecTrustSettingsDomainAdmin,
					     kSecTrustSettingsDomainUser };

	int numDomains = sizeof(domains)/sizeof(SecTrustSettingsDomain);
	if (pemRoots == NULL) {
		return -1;
	}

	CFMutableDataRef combinedData = CFDataCreateMutable(kCFAllocatorDefault, 0);
	for (int i = 0; i < numDomains; i++) {
		CFArrayRef certs = NULL;
		// Only get certificates from domain that are trusted
		OSStatus err = SecTrustSettingsCopyCertificates(domains[i], &certs);
		if (err != noErr) {
			continue;
		}

		int numCerts = CFArrayGetCount(certs);
		for (int j = 0; j < numCerts; j++) {
			CFDataRef data = NULL;
			CFErrorRef errRef = NULL;
			SecCertificateRef cert = (SecCertificateRef)CFArrayGetValueAtIndex(certs, j);
			if (cert == NULL) {
				continue;
			}
			// We only want to add Root CAs, so make sure Subject and Issuer Name match
			CFDataRef subjectName = SecCertificateCopyNormalizedSubjectContent(cert, &errRef);
			if (errRef != NULL) {
				CFRelease(errRef);
				continue;
			}
			CFDataRef issuerName = SecCertificateCopyNormalizedIssuerContent(cert, &errRef);
			if (errRef != NULL) {
				CFRelease(subjectName);
				CFRelease(errRef);
				continue;
			}
			Boolean equal = CFEqual(subjectName, issuerName);
			CFRelease(subjectName);
			CFRelease(issuerName);
			if (!equal) {
				continue;
			}

			// Note: SecKeychainItemExport is deprecated as of 10.7 in favor of SecItemExport.
			// Once we support weak imports via cgo we should prefer that, and fall back to this
			// for older systems.
			err = SecKeychainItemExport(cert, kSecFormatX509Cert, kSecItemPemArmour, NULL, &data);
			if (err != noErr) {
				continue;
			}

			if (data != NULL) {
				CFDataAppendBytes(combinedData, CFDataGetBytePtr(data), CFDataGetLength(data));
				CFRelease(data);
			}
		}
		CFRelease(certs);
	}
	*pemRoots = combinedData;
	return 0;
}
*/
import "C"
import (
	"errors"
	"unsafe"
)

func loadSystemRoots() (*CertPool, error) {
	roots := NewCertPool()

	var data C.CFDataRef = nil
	err := C.FetchPEMRootsCTX509(&data)
	if err == -1 {
		// TODO: better error message
		return nil, errors.New("crypto/x509: failed to load darwin system roots with cgo")
	}

	defer C.CFRelease(C.CFTypeRef(data))
	buf := C.GoBytes(unsafe.Pointer(C.CFDataGetBytePtr(data)), C.int(C.CFDataGetLength(data)))
	roots.AppendCertsFromPEM(buf)
	return roots, nil
}
<commit_msg>[go] Rename Darwin C functions to not clash with golang<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build cgo,!arm,!arm64,!ios

package x509

/*
#cgo CFLAGS: -mmacosx-version-min=10.6 -D__MAC_OS_X_VERSION_MAX_ALLOWED=1060
#cgo LDFLAGS: -framework CoreFoundation -framework Security

#include <errno.h>
#include <sys/sysctl.h>

#include <CoreFoundation/CoreFoundation.h>
#include <Security/Security.h>

// FetchPEMRootsCTX509_MountainLion is the version of FetchPEMRoots from Go 1.6
// which still works on OS X 10.8 (Mountain Lion).
// It lacks support for admin & user cert domains.
// See golang.org/issue/16473
int FetchPEMRootsCTX509_MountainLion(CFDataRef *pemRoots) {
	if (pemRoots == NULL) {
		return -1;
	}
	CFArrayRef certs = NULL;
	OSStatus err = SecTrustCopyAnchorCertificates(&certs);
	if (err != noErr) {
		return -1;
	}
	CFMutableDataRef combinedData = CFDataCreateMutable(kCFAllocatorDefault, 0);
	int i, ncerts = CFArrayGetCount(certs);
	for (i = 0; i < ncerts; i++) {
		CFDataRef data = NULL;
		SecCertificateRef cert = (SecCertificateRef)CFArrayGetValueAtIndex(certs, i);
		if (cert == NULL) {
			continue;
		}
		// Note: SecKeychainItemExport is deprecated as of 10.7 in favor of SecItemExport.
		// Once we support weak imports via cgo we should prefer that, and fall back to this
		// for older systems.
		err = SecKeychainItemExport(cert, kSecFormatX509Cert, kSecItemPemArmour, NULL, &data);
		if (err != noErr) {
			continue;
		}
		if (data != NULL) {
			CFDataAppendBytes(combinedData, CFDataGetBytePtr(data), CFDataGetLength(data));
			CFRelease(data);
		}
	}
	CFRelease(certs);
	*pemRoots = combinedData;
	return 0;
}

// useOldCodeCTX509 reports whether the running machine is OS X 10.8 Mountain Lion
// or older. We only support Mountain Lion and higher, but we'll at least try our
// best on older machines and continue to use the old code path.
//
// See golang.org/issue/16473
int useOldCodeCTX509() {
	char str[256];
	size_t size = sizeof(str);
	memset(str, 0, size);
	sysctlbyname("kern.osrelease", str, &size, NULL, 0);
	// OS X 10.8 is osrelease "12.*", 10.7 is 11.*, 10.6 is 10.*.
	// We never supported things before that.
	return memcmp(str, "12.", 3) == 0 || memcmp(str, "11.", 3) == 0 || memcmp(str, "10.", 3) == 0;
}

// FetchPEMRootsCTX509 fetches the system's list of trusted X.509 root certificates.
//
// On success it returns 0 and fills pemRoots with a CFDataRef that contains the extracted root
// certificates of the system. On failure, the function returns -1.
//
// Note: The CFDataRef returned in pemRoots must be released (using CFRelease) after
// we've consumed its content.
int FetchPEMRootsCTX509(CFDataRef *pemRoots) {
	if (useOldCodeCTX509()) {
		return FetchPEMRootsCTX509_MountainLion(pemRoots);
	}

	// Get certificates from all domains, not just System, this lets
	// the user add CAs to their "login" keychain, and Admins to add
	// to the "System" keychain
	SecTrustSettingsDomain domains[] = { kSecTrustSettingsDomainSystem,
					     kSecTrustSettingsDomainAdmin,
					     kSecTrustSettingsDomainUser };

	int numDomains = sizeof(domains)/sizeof(SecTrustSettingsDomain);
	if (pemRoots == NULL) {
		return -1;
	}

	CFMutableDataRef combinedData = CFDataCreateMutable(kCFAllocatorDefault, 0);
	for (int i = 0; i < numDomains; i++) {
		CFArrayRef certs = NULL;
		// Only get certificates from domain that are trusted
		OSStatus err = SecTrustSettingsCopyCertificates(domains[i], &certs);
		if (err != noErr) {
			continue;
		}

		int numCerts = CFArrayGetCount(certs);
		for (int j = 0; j < numCerts; j++) {
			CFDataRef data = NULL;
			CFErrorRef errRef = NULL;
			SecCertificateRef cert = (SecCertificateRef)CFArrayGetValueAtIndex(certs, j);
			if (cert == NULL) {
				continue;
			}
			// We only want to add Root CAs, so make sure Subject and Issuer Name match
			CFDataRef subjectName = SecCertificateCopyNormalizedSubjectContent(cert, &errRef);
			if (errRef != NULL) {
				CFRelease(errRef);
				continue;
			}
			CFDataRef issuerName = SecCertificateCopyNormalizedIssuerContent(cert, &errRef);
			if (errRef != NULL) {
				CFRelease(subjectName);
				CFRelease(errRef);
				continue;
			}
			Boolean equal = CFEqual(subjectName, issuerName);
			CFRelease(subjectName);
			CFRelease(issuerName);
			if (!equal) {
				continue;
			}

			// Note: SecKeychainItemExport is deprecated as of 10.7 in favor of SecItemExport.
			// Once we support weak imports via cgo we should prefer that, and fall back to this
			// for older systems.
			err = SecKeychainItemExport(cert, kSecFormatX509Cert, kSecItemPemArmour, NULL, &data);
			if (err != noErr) {
				continue;
			}

			if (data != NULL) {
				CFDataAppendBytes(combinedData, CFDataGetBytePtr(data), CFDataGetLength(data));
				CFRelease(data);
			}
		}
		CFRelease(certs);
	}
	*pemRoots = combinedData;
	return 0;
}
*/
import "C"
import (
	"errors"
	"unsafe"
)

func loadSystemRoots() (*CertPool, error) {
	roots := NewCertPool()

	var data C.CFDataRef = nil
	err := C.FetchPEMRootsCTX509(&data)
	if err == -1 {
		// TODO: better error message
		return nil, errors.New("crypto/x509: failed to load darwin system roots with cgo")
	}

	defer C.CFRelease(C.CFTypeRef(data))
	buf := C.GoBytes(unsafe.Pointer(C.CFDataGetBytePtr(data)), C.int(C.CFDataGetLength(data)))
	roots.AppendCertsFromPEM(buf)
	return roots, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	docopt "github.com/docopt/docopt-go"
	"github.com/jurriaan/kafkatools"
)

var (
	version     = "0.1"
	gitrev      = "unknown"
	versionInfo = `reset_consumer_group %s (git rev %s)`
	usage       = `reset_consumer_group - a tool to reset the consumer group offset for a specific topic

usage:
  reset_consumer_group [options] <group> <topic>

options:
  -h --help             show this screen.
  --version             show version.
  --broker [broker]     the kafka bootstrap broker
  --to-time [timestamp]   set offsets to a specific timestamp
  --partition [partition] only update a specific partition
  --offset [offset] update to a certain offset
`
)

func main() {
	docOpts, err := docopt.Parse(usage, nil, true, fmt.Sprintf(versionInfo, version, gitrev), false)

	if err != nil {
		log.Panicf("[PANIC] We couldn't parse doc opts params: %v", err)
	}

	if docOpts["--broker"] == nil {
		log.Fatal("You have to provide a broker")
	}
	broker := docOpts["--broker"].(string)
	topics := strings.Split(docOpts["<topic>"].(string), ",")
	consumerGroup := docOpts["<group>"].(string)
	partition := int32(-1)
	if docOpts["--partition"] != nil {
		r, err := strconv.Atoi(docOpts["--partition"].(string))
		if err != nil {
			log.Fatal("Couldn't parse partition", err)
		}
		partition = int32(r)
	}

	client := kafkatools.GetSaramaClient(broker)
	consumer := kafkatools.GetSaramaConsumer(broker, consumerGroup, topics)
	defer func() {
		err := consumer.Close()
		if err != nil {
			log.Fatal("Could not properly close the consumer")
		}
		err = client.Close()
		if err != nil {
			log.Fatal("Could not properly close the client")
		}
		log.Println("Connection closed. Bye.")
	}()

	offset := sarama.OffsetNewest
	if docOpts["--to-time"] != nil {
		atTime, err := time.Parse(time.RFC3339, docOpts["--to-time"].(string))
		if err != nil {
			log.Fatal("Invalid time format specified (RFC3339 required): ", err)
		}

		// Compute time in milliseconds
		offset = atTime.UnixNano() / (int64(time.Millisecond) / int64(time.Nanosecond))
	} else if docOpts["--offset"] != nil {
		r, err := strconv.Atoi(docOpts["--offset"].(string))
		if err != nil {
			log.Fatal("Couldn't parse offset", err)
		}
		offset = int64(r)
	}
	groupOffsets, topicOffsets := kafkatools.FetchOffsets(client, offset)

	go func() {
		for err := range consumer.Errors() {
			log.Printf("Error: %s\n", err.Error())
		}
	}()

	groupOffsetMap := generateGroupOffsetMap(groupOffsets, topics, consumerGroup)
	setConsumerOffsets(consumer, topics, topicOffsets, groupOffsetMap, partition, offset)
}

func generateGroupOffsetMap(groupOffsets kafkatools.GroupOffsetSlice, topics []string, consumerGroup string) (groupOffsetMap map[string]map[int32]int64) {
	groupOffsetMap = make(map[string]map[int32]int64)
	for _, groupOffset := range groupOffsets {
		if groupOffset.Group == consumerGroup {
			for _, topicOffset := range groupOffset.GroupTopicOffsets {
				for _, topic := range topics {
					if topicOffset.Topic == topic {
						for _, partitionOffset := range topicOffset.TopicPartitionOffsets {
							if _, ok := groupOffsetMap[partitionOffset.Topic]; !ok {
								groupOffsetMap[partitionOffset.Topic] = make(map[int32]int64)
							}
							groupOffsetMap[partitionOffset.Topic][partitionOffset.Partition] = partitionOffset.Offset
						}
					}
				}
			}
		}
	}
	return
}
func setConsumerOffsets(consumer *cluster.Consumer, topics []string, topicOffsets map[string]map[int32]kafkatools.TopicPartitionOffset, groupOffsetMap map[string]map[int32]int64, providedPartition int32, providedOffset int64) {
	log.Println("Waiting for consumer to join all partitions")
	log.Println("Make sure there are no other consumers listening for this to work")

	for note := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", note)
		if len(note.Current) != len(topics) {
			continue
		} else {
			for _, topic := range topics {
				if len(note.Current[topic]) != len(topicOffsets[topic]) {
					continue
				}
			}
		}

		for _, topic := range topics {
			topicPartitionOffset := topicOffsets[topic]
			if providedPartition != -1 {
				if providedOffset == -1 {
					providedOffset = topicPartitionOffset[providedPartition].Offset
				}

				log.Printf("Setting %s:%d's offset from %d to %d", topic, providedPartition, groupOffsetMap[topic][providedPartition], providedOffset)
				consumer.MarkPartitionOffset(topic, providedPartition, providedOffset-1, "")
			} else {
				for partition, offset := range topicPartitionOffset {
					log.Printf("Setting %s:%d's offset from %d to %d", topic, partition, groupOffsetMap[topic][partition], offset.Offset)
					consumer.MarkPartitionOffset(topic, partition, offset.Offset-1, "")
				}
			}
		}

		err := consumer.CommitOffsets()
		if err != nil {
			log.Fatal("Error committing offsets", err)
		}

		log.Println("Quitting")
		return
	}

}
<commit_msg>Refactor main to avoid gocyclo lint errors<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	docopt "github.com/docopt/docopt-go"
	"github.com/jurriaan/kafkatools"
)

var (
	version     = "0.1"
	gitrev      = "unknown"
	versionInfo = `reset_consumer_group %s (git rev %s)`
	usage       = `reset_consumer_group - a tool to reset the consumer group offset for a specific topic

usage:
  reset_consumer_group [options] <group> <topic>

options:
  -h --help             show this screen.
  --version             show version.
  --broker [broker]     the kafka bootstrap broker
  --to-time [timestamp]   set offsets to a specific timestamp
  --partition [partition] only update a specific partition
  --offset [offset] update to a certain offset
`
)

func main() {
	docOpts, err := docopt.Parse(usage, nil, true, fmt.Sprintf(versionInfo, version, gitrev), false)

	if err != nil {
		log.Panicf("[PANIC] We couldn't parse doc opts params: %v", err)
	}

	if docOpts["--broker"] == nil {
		log.Fatal("You have to provide a broker")
	}
	broker := docOpts["--broker"].(string)
	topics := strings.Split(docOpts["<topic>"].(string), ",")
	consumerGroup := docOpts["<group>"].(string)
	partition := getPartition(docOpts)

	client := kafkatools.GetSaramaClient(broker)
	consumer := kafkatools.GetSaramaConsumer(broker, consumerGroup, topics)
	defer func() {
		err := consumer.Close()
		if err != nil {
			log.Fatal("Could not properly close the consumer")
		}
		err = client.Close()
		if err != nil {
			log.Fatal("Could not properly close the client")
		}
		log.Println("Connection closed. Bye.")
	}()

	offset := getOffset(docOpts)
	groupOffsets, topicOffsets := kafkatools.FetchOffsets(client, offset)

	go func() {
		for err := range consumer.Errors() {
			log.Printf("Error: %s\n", err.Error())
		}
	}()

	groupOffsetMap := generateGroupOffsetMap(groupOffsets, topics, consumerGroup)
	setConsumerOffsets(consumer, topics, topicOffsets, groupOffsetMap, partition, offset)
}

func getPartition(docOpts map[string]interface{}) int32 {
	if docOpts["--partition"] == nil {
		return int32(-1)
	}

	r, err := strconv.Atoi(docOpts["--partition"].(string))
	if err != nil {
		log.Fatal("Couldn't parse partition", err)
	}
	return int32(r)
}

func getOffset(docOpts map[string]interface{}) int64 {
	if docOpts["--to-time"] != nil {
		atTime, err := time.Parse(time.RFC3339, docOpts["--to-time"].(string))
		if err != nil {
			log.Fatal("Invalid time format specified (RFC3339 required): ", err)
		}

		// Compute time in milliseconds
		return atTime.UnixNano() / (int64(time.Millisecond) / int64(time.Nanosecond))
	} else if docOpts["--offset"] != nil {
		r, err := strconv.Atoi(docOpts["--offset"].(string))
		if err != nil {
			log.Fatal("Couldn't parse offset", err)
		}
		return int64(r)
	}

	return sarama.OffsetNewest
}

func generateGroupOffsetMap(groupOffsets kafkatools.GroupOffsetSlice, topics []string, consumerGroup string) (groupOffsetMap map[string]map[int32]int64) {
	groupOffsetMap = make(map[string]map[int32]int64)
	for _, groupOffset := range groupOffsets {
		if groupOffset.Group == consumerGroup {
			for _, topicOffset := range groupOffset.GroupTopicOffsets {
				for _, topic := range topics {
					if topicOffset.Topic == topic {
						for _, partitionOffset := range topicOffset.TopicPartitionOffsets {
							if _, ok := groupOffsetMap[partitionOffset.Topic]; !ok {
								groupOffsetMap[partitionOffset.Topic] = make(map[int32]int64)
							}
							groupOffsetMap[partitionOffset.Topic][partitionOffset.Partition] = partitionOffset.Offset
						}
					}
				}
			}
		}
	}
	return
}

func setConsumerOffsets(consumer *cluster.Consumer, topics []string, topicOffsets map[string]map[int32]kafkatools.TopicPartitionOffset, groupOffsetMap map[string]map[int32]int64, providedPartition int32, providedOffset int64) {
	log.Println("Waiting for consumer to join all partitions")
	log.Println("Make sure there are no other consumers listening for this to work")

	for note := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", note)
		if len(note.Current) != len(topics) {
			continue
		} else {
			for _, topic := range topics {
				if len(note.Current[topic]) != len(topicOffsets[topic]) {
					continue
				}
			}
		}

		for _, topic := range topics {
			topicPartitionOffset := topicOffsets[topic]
			if providedPartition != -1 {
				if providedOffset == -1 {
					providedOffset = topicPartitionOffset[providedPartition].Offset
				}

				log.Printf("Setting %s:%d's offset from %d to %d", topic, providedPartition, groupOffsetMap[topic][providedPartition], providedOffset)
				consumer.MarkPartitionOffset(topic, providedPartition, providedOffset-1, "")
			} else {
				for partition, offset := range topicPartitionOffset {
					log.Printf("Setting %s:%d's offset from %d to %d", topic, partition, groupOffsetMap[topic][partition], offset.Offset)
					consumer.MarkPartitionOffset(topic, partition, offset.Offset-1, "")
				}
			}
		}

		err := consumer.CommitOffsets()
		if err != nil {
			log.Fatal("Error committing offsets", err)
		}

		log.Println("Quitting")
		return
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	docopt "github.com/docopt/docopt-go"
	"github.com/jurriaan/kafkatools"
)

var (
	version     = "0.1"
	gitrev      = "unknown"
	versionInfo = `reset_consumer_group %s (git rev %s)`
	usage       = `reset_consumer_group - a tool to reset the consumer group offset for a specific topic

usage:
  reset_consumer_group [options] <group> <topic>

options:
  -h --help             show this screen.
  --version             show version.
  --broker [broker]     the kafka bootstrap broker
  --to-time [timestamp] set offsets to a specific timestamp
`
)

func main() {
	docOpts, err := docopt.Parse(usage, nil, true, fmt.Sprintf(versionInfo, version, gitrev), false)

	if err != nil {
		log.Panicf("[PANIC] We couldn't parse doc opts params: %v", err)
	}

	if docOpts["--broker"] == nil {
		log.Fatal("You have to provide a broker")
	}
	broker := docOpts["--broker"].(string)
	topics := strings.Split(docOpts["<topic>"].(string), ",")
	consumerGroup := docOpts["<group>"].(string)

	client := kafkatools.GetSaramaClient(broker)
	consumer := kafkatools.GetSaramaConsumer(broker, consumerGroup, topics)
	defer func() {
		err := consumer.Close()
		if err != nil {
			log.Fatal("Could not properly close the consumer")
		}
		err = client.Close()
		if err != nil {
			log.Fatal("Could not properly close the client")
		}
		log.Println("Connection closed. Bye.")
	}()

	offset := sarama.OffsetNewest
	if docOpts["--to-time"] != nil {
		atTime, err := time.Parse(time.RFC3339, docOpts["--to-time"].(string))
		if err != nil {
			log.Fatal("Invalid time format specified (RFC3339 required): ", err)
		}

		// Compute time in milliseconds
		offset = atTime.UnixNano() / (int64(time.Millisecond) / int64(time.Nanosecond))
	}
	groupOffsets, topicOffsets := kafkatools.FetchOffsets(client, offset)

	go func() {
		for err := range consumer.Errors() {
			log.Printf("Error: %s\n", err.Error())
		}
	}()

	groupOffsetMap := generateGroupOffsetMap(groupOffsets, topics, consumerGroup)
	setConsumerOffsets(consumer, topics, topicOffsets, groupOffsetMap)
}

func generateGroupOffsetMap(groupOffsets kafkatools.GroupOffsetSlice, topics []string, consumerGroup string) (groupOffsetMap map[string]map[int32]int64) {
	groupOffsetMap = make(map[string]map[int32]int64)
	for _, groupOffset := range groupOffsets {
		if groupOffset.Group == consumerGroup {
			for _, topicOffset := range groupOffset.GroupTopicOffsets {
				for _, topic := range topics {
					if topicOffset.Topic == topic {
						for _, partitionOffset := range topicOffset.TopicPartitionOffsets {
							if _, ok := groupOffsetMap[partitionOffset.Topic]; !ok {
								groupOffsetMap[partitionOffset.Topic] = make(map[int32]int64)
							}
							groupOffsetMap[partitionOffset.Topic][partitionOffset.Partition] = partitionOffset.Offset
						}
					}
				}
			}
		}
	}
	return
}
func setConsumerOffsets(consumer *cluster.Consumer, topics []string, topicOffsets map[string]map[int32]kafkatools.TopicPartitionOffset, groupOffsetMap map[string]map[int32]int64) {
	log.Println("Waiting for consumer to join all partitions")
	log.Println("Make sure there are no other consumers listening for this to work")

	for note := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", note)
		if len(note.Current) != len(topics) {
			continue
		} else {
			for _, topic := range topics {
				if len(note.Current[topic]) != len(topicOffsets[topic]) {
					continue
				}
			}
		}

		for _, topic := range topics {
			topicPartitionOffset := topicOffsets[topic]
			for partition, offset := range topicPartitionOffset {
				log.Printf("Setting %s:%d's offset from %d to %d", topic, partition, groupOffsetMap[topic][partition], offset.Offset)
				consumer.MarkPartitionOffset(topic, partition, offset.Offset-1, "")
			}
		}

		err := consumer.CommitOffsets()
		if err != nil {
			log.Fatal("Error committing offsets", err)
		}

		log.Println("Quitting")
		return
	}

}
<commit_msg>Add updating of specific partition to specific offset<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	docopt "github.com/docopt/docopt-go"
	"github.com/jurriaan/kafkatools"
)

var (
	version     = "0.1"
	gitrev      = "unknown"
	versionInfo = `reset_consumer_group %s (git rev %s)`
	usage       = `reset_consumer_group - a tool to reset the consumer group offset for a specific topic

usage:
  reset_consumer_group [options] <group> <topic>

options:
  -h --help             show this screen.
  --version             show version.
  --broker [broker]     the kafka bootstrap broker
  --to-time [timestamp]   set offsets to a specific timestamp
  --partition [partition] only update a specific partition
  --offset [offset] update to a certain offset
`
)

func main() {
	docOpts, err := docopt.Parse(usage, nil, true, fmt.Sprintf(versionInfo, version, gitrev), false)

	if err != nil {
		log.Panicf("[PANIC] We couldn't parse doc opts params: %v", err)
	}

	if docOpts["--broker"] == nil {
		log.Fatal("You have to provide a broker")
	}
	broker := docOpts["--broker"].(string)
	topics := strings.Split(docOpts["<topic>"].(string), ",")
	consumerGroup := docOpts["<group>"].(string)
	partition := int32(-1)
	if docOpts["--partition"] != nil {
		r, err := strconv.Atoi(docOpts["--partition"].(string))
		if err != nil {
			log.Fatal("Couldn't parse partition", err)
		}
		partition = int32(r)
	}

	client := kafkatools.GetSaramaClient(broker)
	consumer := kafkatools.GetSaramaConsumer(broker, consumerGroup, topics)
	defer func() {
		err := consumer.Close()
		if err != nil {
			log.Fatal("Could not properly close the consumer")
		}
		err = client.Close()
		if err != nil {
			log.Fatal("Could not properly close the client")
		}
		log.Println("Connection closed. Bye.")
	}()

	offset := sarama.OffsetNewest
	if docOpts["--to-time"] != nil {
		atTime, err := time.Parse(time.RFC3339, docOpts["--to-time"].(string))
		if err != nil {
			log.Fatal("Invalid time format specified (RFC3339 required): ", err)
		}

		// Compute time in milliseconds
		offset = atTime.UnixNano() / (int64(time.Millisecond) / int64(time.Nanosecond))
	} else if docOpts["--offset"] != nil {
		r, err := strconv.Atoi(docOpts["--offset"].(string))
		if err != nil {
			log.Fatal("Couldn't parse offset", err)
		}
		offset = int64(r)
	}
	groupOffsets, topicOffsets := kafkatools.FetchOffsets(client, offset)

	go func() {
		for err := range consumer.Errors() {
			log.Printf("Error: %s\n", err.Error())
		}
	}()

	groupOffsetMap := generateGroupOffsetMap(groupOffsets, topics, consumerGroup)
	setConsumerOffsets(consumer, topics, topicOffsets, groupOffsetMap, partition, offset)
}

func generateGroupOffsetMap(groupOffsets kafkatools.GroupOffsetSlice, topics []string, consumerGroup string) (groupOffsetMap map[string]map[int32]int64) {
	groupOffsetMap = make(map[string]map[int32]int64)
	for _, groupOffset := range groupOffsets {
		if groupOffset.Group == consumerGroup {
			for _, topicOffset := range groupOffset.GroupTopicOffsets {
				for _, topic := range topics {
					if topicOffset.Topic == topic {
						for _, partitionOffset := range topicOffset.TopicPartitionOffsets {
							if _, ok := groupOffsetMap[partitionOffset.Topic]; !ok {
								groupOffsetMap[partitionOffset.Topic] = make(map[int32]int64)
							}
							groupOffsetMap[partitionOffset.Topic][partitionOffset.Partition] = partitionOffset.Offset
						}
					}
				}
			}
		}
	}
	return
}
func setConsumerOffsets(consumer *cluster.Consumer, topics []string, topicOffsets map[string]map[int32]kafkatools.TopicPartitionOffset, groupOffsetMap map[string]map[int32]int64, providedPartition int32, providedOffset int64) {
	log.Println("Waiting for consumer to join all partitions")
	log.Println("Make sure there are no other consumers listening for this to work")

	for note := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", note)
		if len(note.Current) != len(topics) {
			continue
		} else {
			for _, topic := range topics {
				if len(note.Current[topic]) != len(topicOffsets[topic]) {
					continue
				}
			}
		}

		for _, topic := range topics {
			topicPartitionOffset := topicOffsets[topic]
			if providedPartition != -1 {
				if providedOffset == -1 {
					providedOffset = topicPartitionOffset[providedPartition].Offset
				}

				consumer.MarkPartitionOffset(topic, providedPartition, providedOffset-1, "")
			} else {
				for partition, offset := range topicPartitionOffset {
					log.Printf("Setting %s:%d's offset from %d to %d", topic, partition, groupOffsetMap[topic][partition], offset.Offset)
					consumer.MarkPartitionOffset(topic, partition, offset.Offset-1, "")
				}
			}
		}

		err := consumer.CommitOffsets()
		if err != nil {
			log.Fatal("Error committing offsets", err)
		}

		log.Println("Quitting")
		return
	}

}
<|endoftext|>


<commit_before>package soac

import (
	"fmt"
)

type Color byte
type Attribute byte
type BackGround byte

type AttrFactor interface {
	Get() int
}

const (
	Black Color = iota + 30
	Red
	Green
	Yellow
	Blue
	Magenda
	Cyan
	White
)

func (c Color) Get() int {
	return int(c)
}

const (
	Reset Attribute = iota
	Bold
	Faint
	Italic
	Underline
	Blink1
	Blink2
	Reverse
	Concealed
	Crossedout
)

func (a Attribute) Get() int {
	return int(a)
}

const (
	BgBlack BackGround = iota + 40
	BgRed
	BgGreen
	BgYellow
	BgBlue
	BgMagenda
	BgCyan
	BgWhite
)

func (b BackGround) Get() int {
	return int(b)
}

type Changer struct {
	fg   Color
	attr Attribute
	bg   BackGround
}

func NewChanger() *Changer {
	return &Changer{39, 0, 49}
}

func (self *Changer) Apply(val interface{}) string {
	return fmt.Sprintf("\x1b[%d;%d;%dm%v\x1b[0m", self.attr, self.fg, self.bg, val)
}

func (self *Changer) Set(factors ...AttrFactor) {
	for _, factor := range factors {
		i := factor.Get()
		if 0 <= i && i <= 9 {
			self.attr = Attribute(i)
		} else if 30 <= i && i <= 37 {
			self.fg = Color(i)
		} else if 40 <= i && i <= 47 {
			self.bg = BackGround(i)
		} else {
			// return error or emit warnning
		}
	}
}

func (self *Changer) Reset() {
	self.fg = 39
	self.attr = 0
	self.bg = 49
}

func (self *Changer) Black() *Changer {
	self.fg = Black
	return self
}

func (self *Changer) Red() *Changer {
	self.fg = Red
	return self
}

func (self *Changer) Green() *Changer {
	self.fg = Green
	return self
}

func (self *Changer) Yellow() *Changer {
	self.fg = Yellow
	return self
}

func (self *Changer) Blue() *Changer {
	self.fg = Blue
	return self
}

func (self *Changer) Magenda() *Changer {
	self.fg = Magenda
	return self
}

func (self *Changer) Cyan() *Changer {
	self.fg = Cyan
	return self
}

func (self *Changer) White() *Changer {
	self.fg = White
	return self
}

func (self *Changer) Bold() *Changer {
	self.attr = Bold
	return self
}

func (self *Changer) Faint() *Changer {
	self.attr = Faint
	return self
}

func (self *Changer) Italic() *Changer {
	self.attr = Italic
	return self
}

func (self *Changer) Underline() *Changer {
	self.attr = Underline
	return self
}

func (self *Changer) Blink1() *Changer {
	self.attr = Blink1
	return self
}

func (self *Changer) Blink2() *Changer {
	self.attr = Blink2
	return self
}

func (self *Changer) Reverse() *Changer {
	self.attr = Reverse
	return self
}

func (self *Changer) Concealed() *Changer {
	self.attr = Concealed
	return self
}

func (self *Changer) CrossedOut() *Changer {
	self.attr = Crossedout
	return self
}

func (self *Changer) Bblack() *Changer {
	self.bg = BgBlack
	return self
}

func (self *Changer) Bred() *Changer {
	self.bg = BgRed
	return self
}

func (self *Changer) Bgreen() *Changer {
	self.bg = BgGreen
	return self
}

func (self *Changer) Byellow() *Changer {
	self.bg = BgYellow
	return self
}

func (self *Changer) Bblue() *Changer {
	self.bg = BgBlue
	return self
}

func (self *Changer) Bmagenda() *Changer {
	self.bg = BgMagenda
	return self
}

func (self *Changer) Bcyan() *Changer {
	self.bg = BgCyan
	return self
}

func (self *Changer) Bwhite() *Changer {
	self.bg = BgWhite
	return self
}
<commit_msg>implement 256 color<commit_after>package soac

import (
	"fmt"
)

type Color byte
type Attribute byte
type BackGround byte

type AttrFactor interface {
	Get() int
}

const (
	Black Color = iota + 30
	Red
	Green
	Yellow
	Blue
	Magenda
	Cyan
	White
)

func (c Color) Get() int {
	return int(c)
}

const (
	Reset Attribute = iota
	Bold
	Faint
	Italic
	Underline
	Blink1
	Blink2
	Reverse
	Concealed
	Crossedout
)

func (a Attribute) Get() int {
	return int(a)
}

const (
	BgBlack BackGround = iota + 40
	BgRed
	BgGreen
	BgYellow
	BgBlue
	BgMagenda
	BgCyan
	BgWhite
)

func (b BackGround) Get() int {
	return int(b)
}

type Changer struct {
	fg    Color
	attr  Attribute
	bg    BackGround
	is256 bool
}

func NewChanger() *Changer {
	return &Changer{39, 0, 49, false}
}

func (self *Changer) Apply(val interface{}) string {
	if !self.is256 {
		return fmt.Sprintf("\x1b[%d;%d;%dm%v\x1b[0m", self.attr, self.fg, self.bg, val)
	}
	return fmt.Sprintf("\x1b[%d;38;5;%d;%dm%v\x1b[0m", self.attr, self.fg, self.bg, val)
}

func (self *Changer) Set(factors ...AttrFactor) {
	for _, factor := range factors {
		i := factor.Get()
		if 0 <= i && i <= 9 {
			self.attr = Attribute(i)
		} else if 30 <= i && i <= 37 {
			self.fg = Color(i)
			self.is256 = false
		} else if 40 <= i && i <= 47 {
			self.bg = BackGround(i)
		} else {
			// return error or emit warnning
		}
	}
}

func (self *Changer) Set256(color byte) {
	self.fg = Color(color)
	self.is256 = true
}

func (self *Changer) Reset() {
	self.fg = 39
	self.attr = 0
	self.bg = 49
	self.is256 = false
}

func (self *Changer) Black() *Changer {
	self.fg = Black
	self.is256 = false
	return self
}

func (self *Changer) Red() *Changer {
	self.fg = Red
	self.is256 = false
	return self
}

func (self *Changer) Green() *Changer {
	self.fg = Green
	self.is256 = false
	return self
}

func (self *Changer) Yellow() *Changer {
	self.fg = Yellow
	self.is256 = false
	return self
}

func (self *Changer) Blue() *Changer {
	self.fg = Blue
	self.is256 = false
	return self
}

func (self *Changer) Magenda() *Changer {
	self.fg = Magenda
	self.is256 = false
	return self
}

func (self *Changer) Cyan() *Changer {
	self.fg = Cyan
	self.is256 = false
	return self
}

func (self *Changer) White() *Changer {
	self.fg = White
	self.is256 = false
	return self
}

func (self *Changer) Bold() *Changer {
	self.attr = Bold
	return self
}

func (self *Changer) Faint() *Changer {
	self.attr = Faint
	return self
}

func (self *Changer) Italic() *Changer {
	self.attr = Italic
	return self
}

func (self *Changer) Underline() *Changer {
	self.attr = Underline
	return self
}

func (self *Changer) Blink1() *Changer {
	self.attr = Blink1
	return self
}

func (self *Changer) Blink2() *Changer {
	self.attr = Blink2
	return self
}

func (self *Changer) Reverse() *Changer {
	self.attr = Reverse
	return self
}

func (self *Changer) Concealed() *Changer {
	self.attr = Concealed
	return self
}

func (self *Changer) CrossedOut() *Changer {
	self.attr = Crossedout
	return self
}

func (self *Changer) Bblack() *Changer {
	self.bg = BgBlack
	return self
}

func (self *Changer) Bred() *Changer {
	self.bg = BgRed
	return self
}

func (self *Changer) Bgreen() *Changer {
	self.bg = BgGreen
	return self
}

func (self *Changer) Byellow() *Changer {
	self.bg = BgYellow
	return self
}

func (self *Changer) Bblue() *Changer {
	self.bg = BgBlue
	return self
}

func (self *Changer) Bmagenda() *Changer {
	self.bg = BgMagenda
	return self
}

func (self *Changer) Bcyan() *Changer {
	self.bg = BgCyan
	return self
}

func (self *Changer) Bwhite() *Changer {
	self.bg = BgWhite
	return self
}
<|endoftext|>


<commit_before>package goczmq

/*
#cgo !windows pkg-config: libczmq
#cgo windows CFLAGS: -I/usr/local/include
#cgo windows LDFLAGS: -L/usr/local/lib -lczmq
#include "czmq.h"
#include <stdlib.h>
#include <string.h>

int Sock_connect(zsock_t *self, const char *format) {return zsock_connect(self, format, NULL);}
int Sock_disconnect(zsock_t *self, const char *format) {return zsock_disconnect(self, format, NULL);}
int Sock_bind(zsock_t *self, const char *format) {return zsock_bind(self, format, NULL);}
int Sock_unbind(zsock_t *self, const char *format) {return zsock_unbind(self, format, NULL);}
*/
import "C"

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"
	"unsafe"
)

var (
	// ErrSliceFull is returned if a []byte passed to Read was not
	// large enough to hold the contents of a message
	ErrSliceFull = errors.New("goczmq: slice full")
)

// Sock wraps the zsock_t class in CZMQ.
type Sock struct {
	zsockT       *C.struct__zsock_t
	file         string
	line         int
	zType        int
	lastClientID string
}

func init() {
	if err := os.Setenv("ZSYS_SIGHANDLER", "false"); err != nil {
		panic(err)
	}
}

// GetLastClientID returns the id of the last client you received
// a message from if the underlying socket is a ROUTER or SERVER
// socket
func (s *Sock) GetLastClientID() []byte {
	return []byte(s.lastClientID)
}

// SetLastClientID sets the client id that will be used when
// sending from a ROUTER or SERVER socket
func (s *Sock) SetLastClientID(id []byte) {
	s.lastClientID = string(id)
}

// NewSock creates a new socket.  The caller source and
// line number are passed so CZMQ can report socket leaks
// intelligently.
func NewSock(t int) *Sock {
	var s *Sock
	_, file, line, ok := runtime.Caller(1)

	if ok {
		s = &Sock{
			file:  file,
			line:  line,
			zType: t,
		}
	} else {
		s = &Sock{
			file:  "",
			line:  0,
			zType: t,
		}
	}

	s.zsockT = C.zsock_new_(C.int(s.zType), C.CString(s.file), C.size_t(s.line))
	return s
}

// Connect connects a socket to an endpoint
// returns an error if the connect failed.
func (s *Sock) Connect(endpoint string) error {
	rc := C.Sock_connect(s.zsockT, C.CString(endpoint))
	if rc != C.int(0) {
		return errors.New("failed")
	}
	return nil
}

// Disconnect disconnects a socket from an endpoint.  If returns
// an error if the endpoint was not found
func (s *Sock) Disconnect(endpoint string) error {
	rc := C.Sock_disconnect(s.zsockT, C.CString(endpoint))
	if int(rc) == -1 {
		return fmt.Errorf("endopint was not bound")
	}
	return nil
}

// Bind binds a socket to an endpoint.  On success returns
// the port number used for tcp transports, or 0 for other
// transports.  On failure returns a -1 for port, and an error.
func (s *Sock) Bind(endpoint string) (int, error) {
	port := C.Sock_bind(s.zsockT, C.CString(endpoint))
	if port == C.int(-1) {
		return -1, errors.New("failed")
	}
	return int(port), nil
}

// Unbind unbinds a socket from an endpoint.  If returns
// an error if the endpoint was not found
func (s *Sock) Unbind(endpoint string) error {
	rc := C.Sock_unbind(s.zsockT, C.CString(endpoint))
	if int(rc) == -1 {
		return fmt.Errorf("endopint was not bound")
	}
	return nil
}

// Attach attaches a socket to zero or more endpoints. If endpoints is not null,
// parses as list of ZeroMQ endpoints, separated by commas, and prefixed by
// '@' (to bind the socket) or '>' (to attach the socket). If the endpoint
// does not start with '@' or '>', the serverish argument determines whether
// it is used to bind (serverish = true) or connect (serverish = false)
func (s *Sock) Attach(endpoints string, serverish bool) error {
	rc := C.zsock_attach(s.zsockT, C.CString(endpoints), C._Bool(serverish))
	if rc == -1 {
		return ErrSockAttach
	}
	return nil
}

// NewPUB creates a PUB socket and calls Attach.
// The socket will Bind by default.
func NewPUB(endpoints string) (*Sock, error) {
	s := NewSock(PUB)
	return s, s.Attach(endpoints, true)
}

// NewSUB creates a SUB socket and calls Attach.
// 'subscribe' is a comma delimited list of topics to subscribe to.
// The socket will Connect by default.
func NewSUB(endpoints string, subscribe string) (*Sock, error) {
	s := NewSock(SUB)
	subscriptions := strings.Split(subscribe, ",")

	for _, topic := range subscriptions {
		s.SetSubscribe(topic)
	}

	return s, s.Attach(endpoints, false)
}

// NewREP creates a REP socket and calls Attach.
// The socket will Bind by default.
func NewREP(endpoints string) (*Sock, error) {
	s := NewSock(REP)
	return s, s.Attach(endpoints, true)
}

// NewREQ creates a REQ socket and calls Attach.
// The socket will Connect by default.
func NewREQ(endpoints string) (*Sock, error) {
	s := NewSock(REQ)
	return s, s.Attach(endpoints, false)
}

// NewPULL creates a PULL socket and calls Attach.
// The socket will Bind by default.
func NewPULL(endpoints string) (*Sock, error) {
	s := NewSock(PULL)
	return s, s.Attach(endpoints, true)
}

// NewPUSH creates a PUSH socket and calls Attach.
// The socket will Connect by default.
func NewPUSH(endpoints string) (*Sock, error) {
	s := NewSock(PUSH)
	return s, s.Attach(endpoints, false)
}

// NewROUTER creates a ROUTER socket and calls Attach.
// The socket will Bind by default.
func NewROUTER(endpoints string) (*Sock, error) {
	s := NewSock(ROUTER)
	return s, s.Attach(endpoints, true)
}

// NewDEALER creates a DEALER socket and calls Attach.
// The socket will Connect by default.
func NewDEALER(endpoints string) (*Sock, error) {
	s := NewSock(DEALER)
	return s, s.Attach(endpoints, false)
}

// NewXPUB creates an XPUB socket and calls Attach.
// The socket will Bind by default.
func NewXPUB(endpoints string) (*Sock, error) {
	s := NewSock(XPUB)
	return s, s.Attach(endpoints, true)
}

// NewXSUB creates an XSUB socket and calls Attach.
// The socket will Connect by default.
func NewXSUB(endpoints string) (*Sock, error) {
	s := NewSock(XSUB)
	return s, s.Attach(endpoints, false)
}

// NewPAIR creates a PAIR socket and calls Attach.
// The socket will Connect by default.
func NewPAIR(endpoints string) (*Sock, error) {
	s := NewSock(PAIR)
	return s, s.Attach(endpoints, false)
}

// NewSTREAM creates a STREAM socket and calls Attach.
// The socket will Connect by default.
func NewSTREAM(endpoints string) (*Sock, error) {
	s := NewSock(STREAM)
	return s, s.Attach(endpoints, false)
}

// Pollin returns true if there is a POLLIN
// event on the socket
func (s *Sock) Pollin() bool {
	return s.Events() == POLLIN
	// return C.zsock_events(unsafe.Pointer(s.zsockT)) == C.ZMQ_POLLIN
}

// Pollout returns true if there is a POLLOUT
// event on the socket
func (s *Sock) Pollout() bool {
	return s.Events() == POLLOUT
}

// SendFrame sends a byte array via the socket.  For the flags
// value, use 0 for a single message, or SNDMORE if it is
// a multi-part message
func (s *Sock) SendFrame(data []byte, flags int) error {
	var rc C.int

	if len(data) == 0 {
		frame := C.zframe_new(unsafe.Pointer(C.CString("")), C.size_t(len(data)))
		rc = C.zframe_send(&frame, unsafe.Pointer(s.zsockT), C.int(flags))
	} else {
		frame := C.zframe_new(unsafe.Pointer(&data[0]), C.size_t(len(data)))
		rc = C.zframe_send(&frame, unsafe.Pointer(s.zsockT), C.int(flags))
	}

	if rc == C.int(-1) {
		return errors.New("failed")
	}

	return nil
}

// RecvFrame reads a frame from the socket and returns it
// as a byte array, along with a more flag and and error
// (if there is an error)
func (s *Sock) RecvFrame() ([]byte, int, error) {
	frame := C.zframe_recv(unsafe.Pointer(s.zsockT))
	if frame == nil {
		return []byte{0}, 0, errors.New("failed")
	}
	dataSize := C.zframe_size(frame)
	dataPtr := C.zframe_data(frame)
	b := C.GoBytes(unsafe.Pointer(dataPtr), C.int(dataSize))
	more := C.zframe_more(frame)
	C.zframe_destroy(&frame)
	return b, int(more), nil
}

// RecvFrameNoWait receives a frame from the socket
// and returns it as a byte array if one is waiting.
// Returns an empty frame, a 0 more flag and an error
// if one is not immediately available
func (s *Sock) RecvFrameNoWait() ([]byte, int, error) {
	if !s.Pollin() {
		return []byte{0}, 0, fmt.Errorf("no frame")
	}

	return s.RecvFrame()
}

// SendMessage accepts an array of byte arrays and
// sends it as a multi-part message.
func (s *Sock) SendMessage(parts [][]byte) error {
	var f int
	numParts := len(parts)
	for i, val := range parts {
		if i == numParts-1 {
			f = 0
		} else {
			f = MORE
		}

		err := s.SendFrame(val, f)
		if err != nil {
			return err
		}
	}
	return nil
}

// RecvMessage receives a full message from the socket
// and returns it as an array of byte arrays.
func (s *Sock) RecvMessage() ([][]byte, error) {
	var msg [][]byte

	for {
		frame, flag, err := s.RecvFrame()
		if err != nil {
			return msg, err
		}
		msg = append(msg, frame)
		if flag != MORE {
			break
		}
	}
	return msg, nil
}

// Read provides an io.Reader interface to a zeromq socket
func (s *Sock) Read(p []byte) (int, error) {
	var total int
	frame, flag, err := s.RecvFrame()
	if err != nil {
		return total, err
	}

	if s.GetType() == ROUTER {
		s.lastClientID = string(frame)
	} else {
		copy(p[:], frame[:])
		total += len(frame)
	}

	for flag == MORE {
		frame, flag, err = s.RecvFrame()
		if err != nil {
			return total, err
		}
		copy(p[total:], frame[:])
		total += len(frame)
	}

	if total > len(p) {
		err = ErrSliceFull
	} else {
		err = io.EOF
	}

	return total, err
}

// Write provides an io.Writer interface to a zeromq socket
func (s *Sock) Write(p []byte) (int, error) {
	var total int
	if s.GetType() == ROUTER {
		err := s.SendFrame(s.GetLastClientID(), MORE)
		if err != nil {
			return total, err
		}
	}
	err := s.SendFrame(p, 0)
	if err != nil {
		return total, err
	}

	return len(p), nil
}

// RecvMessageNoWait receives a full message from the socket
// and returns it as an array of byte arrays if one is waiting.
// Returns an empty message and an error if one is not immediately
// available
func (s *Sock) RecvMessageNoWait() ([][]byte, error) {
	var msg [][]byte
	if !s.Pollin() {
		return msg, fmt.Errorf("no message")
	}

	for {
		frame, flag, err := s.RecvFrame()
		if err != nil {
			return msg, err
		}
		msg = append(msg, frame)
		if flag != MORE {
			break
		}
	}
	return msg, nil
}

// GetType returns the socket's type
func (s *Sock) GetType() int {
	return s.zType
}

// Destroy destroys the underlying zsockT.
func (s *Sock) Destroy() {
	C.zsock_destroy_(&s.zsockT, C.CString(s.file), C.size_t(s.line))
}
<commit_msg>Problem: goczmq no longer compiles due to czmq function rename<commit_after>package goczmq

/*
#cgo !windows pkg-config: libczmq
#cgo windows CFLAGS: -I/usr/local/include
#cgo windows LDFLAGS: -L/usr/local/lib -lczmq
#include "czmq.h"
#include <stdlib.h>
#include <string.h>

int Sock_connect(zsock_t *self, const char *format) {return zsock_connect(self, format, NULL);}
int Sock_disconnect(zsock_t *self, const char *format) {return zsock_disconnect(self, format, NULL);}
int Sock_bind(zsock_t *self, const char *format) {return zsock_bind(self, format, NULL);}
int Sock_unbind(zsock_t *self, const char *format) {return zsock_unbind(self, format, NULL);}
*/
import "C"

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strings"
	"unsafe"
)

var (
	// ErrSliceFull is returned if a []byte passed to Read was not
	// large enough to hold the contents of a message
	ErrSliceFull = errors.New("goczmq: slice full")
)

// Sock wraps the zsock_t class in CZMQ.
type Sock struct {
	zsockT       *C.struct__zsock_t
	file         string
	line         int
	zType        int
	lastClientID string
}

func init() {
	if err := os.Setenv("ZSYS_SIGHANDLER", "false"); err != nil {
		panic(err)
	}
}

// GetLastClientID returns the id of the last client you received
// a message from if the underlying socket is a ROUTER or SERVER
// socket
func (s *Sock) GetLastClientID() []byte {
	return []byte(s.lastClientID)
}

// SetLastClientID sets the client id that will be used when
// sending from a ROUTER or SERVER socket
func (s *Sock) SetLastClientID(id []byte) {
	s.lastClientID = string(id)
}

// NewSock creates a new socket.  The caller source and
// line number are passed so CZMQ can report socket leaks
// intelligently.
func NewSock(t int) *Sock {
	var s *Sock
	_, file, line, ok := runtime.Caller(1)

	if ok {
		s = &Sock{
			file:  file,
			line:  line,
			zType: t,
		}
	} else {
		s = &Sock{
			file:  "",
			line:  0,
			zType: t,
		}
	}

	s.zsockT = C.zsock_new_checked(C.int(s.zType), C.CString(s.file), C.size_t(s.line))
	return s
}

// Connect connects a socket to an endpoint
// returns an error if the connect failed.
func (s *Sock) Connect(endpoint string) error {
	rc := C.Sock_connect(s.zsockT, C.CString(endpoint))
	if rc != C.int(0) {
		return errors.New("failed")
	}
	return nil
}

// Disconnect disconnects a socket from an endpoint.  If returns
// an error if the endpoint was not found
func (s *Sock) Disconnect(endpoint string) error {
	rc := C.Sock_disconnect(s.zsockT, C.CString(endpoint))
	if int(rc) == -1 {
		return fmt.Errorf("endopint was not bound")
	}
	return nil
}

// Bind binds a socket to an endpoint.  On success returns
// the port number used for tcp transports, or 0 for other
// transports.  On failure returns a -1 for port, and an error.
func (s *Sock) Bind(endpoint string) (int, error) {
	port := C.Sock_bind(s.zsockT, C.CString(endpoint))
	if port == C.int(-1) {
		return -1, errors.New("failed")
	}
	return int(port), nil
}

// Unbind unbinds a socket from an endpoint.  If returns
// an error if the endpoint was not found
func (s *Sock) Unbind(endpoint string) error {
	rc := C.Sock_unbind(s.zsockT, C.CString(endpoint))
	if int(rc) == -1 {
		return fmt.Errorf("endopint was not bound")
	}
	return nil
}

// Attach attaches a socket to zero or more endpoints. If endpoints is not null,
// parses as list of ZeroMQ endpoints, separated by commas, and prefixed by
// '@' (to bind the socket) or '>' (to attach the socket). If the endpoint
// does not start with '@' or '>', the serverish argument determines whether
// it is used to bind (serverish = true) or connect (serverish = false)
func (s *Sock) Attach(endpoints string, serverish bool) error {
	rc := C.zsock_attach(s.zsockT, C.CString(endpoints), C._Bool(serverish))
	if rc == -1 {
		return ErrSockAttach
	}
	return nil
}

// NewPUB creates a PUB socket and calls Attach.
// The socket will Bind by default.
func NewPUB(endpoints string) (*Sock, error) {
	s := NewSock(PUB)
	return s, s.Attach(endpoints, true)
}

// NewSUB creates a SUB socket and calls Attach.
// 'subscribe' is a comma delimited list of topics to subscribe to.
// The socket will Connect by default.
func NewSUB(endpoints string, subscribe string) (*Sock, error) {
	s := NewSock(SUB)
	subscriptions := strings.Split(subscribe, ",")

	for _, topic := range subscriptions {
		s.SetSubscribe(topic)
	}

	return s, s.Attach(endpoints, false)
}

// NewREP creates a REP socket and calls Attach.
// The socket will Bind by default.
func NewREP(endpoints string) (*Sock, error) {
	s := NewSock(REP)
	return s, s.Attach(endpoints, true)
}

// NewREQ creates a REQ socket and calls Attach.
// The socket will Connect by default.
func NewREQ(endpoints string) (*Sock, error) {
	s := NewSock(REQ)
	return s, s.Attach(endpoints, false)
}

// NewPULL creates a PULL socket and calls Attach.
// The socket will Bind by default.
func NewPULL(endpoints string) (*Sock, error) {
	s := NewSock(PULL)
	return s, s.Attach(endpoints, true)
}

// NewPUSH creates a PUSH socket and calls Attach.
// The socket will Connect by default.
func NewPUSH(endpoints string) (*Sock, error) {
	s := NewSock(PUSH)
	return s, s.Attach(endpoints, false)
}

// NewROUTER creates a ROUTER socket and calls Attach.
// The socket will Bind by default.
func NewROUTER(endpoints string) (*Sock, error) {
	s := NewSock(ROUTER)
	return s, s.Attach(endpoints, true)
}

// NewDEALER creates a DEALER socket and calls Attach.
// The socket will Connect by default.
func NewDEALER(endpoints string) (*Sock, error) {
	s := NewSock(DEALER)
	return s, s.Attach(endpoints, false)
}

// NewXPUB creates an XPUB socket and calls Attach.
// The socket will Bind by default.
func NewXPUB(endpoints string) (*Sock, error) {
	s := NewSock(XPUB)
	return s, s.Attach(endpoints, true)
}

// NewXSUB creates an XSUB socket and calls Attach.
// The socket will Connect by default.
func NewXSUB(endpoints string) (*Sock, error) {
	s := NewSock(XSUB)
	return s, s.Attach(endpoints, false)
}

// NewPAIR creates a PAIR socket and calls Attach.
// The socket will Connect by default.
func NewPAIR(endpoints string) (*Sock, error) {
	s := NewSock(PAIR)
	return s, s.Attach(endpoints, false)
}

// NewSTREAM creates a STREAM socket and calls Attach.
// The socket will Connect by default.
func NewSTREAM(endpoints string) (*Sock, error) {
	s := NewSock(STREAM)
	return s, s.Attach(endpoints, false)
}

// Pollin returns true if there is a POLLIN
// event on the socket
func (s *Sock) Pollin() bool {
	return s.Events() == POLLIN
	// return C.zsock_events(unsafe.Pointer(s.zsockT)) == C.ZMQ_POLLIN
}

// Pollout returns true if there is a POLLOUT
// event on the socket
func (s *Sock) Pollout() bool {
	return s.Events() == POLLOUT
}

// SendFrame sends a byte array via the socket.  For the flags
// value, use 0 for a single message, or SNDMORE if it is
// a multi-part message
func (s *Sock) SendFrame(data []byte, flags int) error {
	var rc C.int

	if len(data) == 0 {
		frame := C.zframe_new(unsafe.Pointer(C.CString("")), C.size_t(len(data)))
		rc = C.zframe_send(&frame, unsafe.Pointer(s.zsockT), C.int(flags))
	} else {
		frame := C.zframe_new(unsafe.Pointer(&data[0]), C.size_t(len(data)))
		rc = C.zframe_send(&frame, unsafe.Pointer(s.zsockT), C.int(flags))
	}

	if rc == C.int(-1) {
		return errors.New("failed")
	}

	return nil
}

// RecvFrame reads a frame from the socket and returns it
// as a byte array, along with a more flag and and error
// (if there is an error)
func (s *Sock) RecvFrame() ([]byte, int, error) {
	frame := C.zframe_recv(unsafe.Pointer(s.zsockT))
	if frame == nil {
		return []byte{0}, 0, errors.New("failed")
	}
	dataSize := C.zframe_size(frame)
	dataPtr := C.zframe_data(frame)
	b := C.GoBytes(unsafe.Pointer(dataPtr), C.int(dataSize))
	more := C.zframe_more(frame)
	C.zframe_destroy(&frame)
	return b, int(more), nil
}

// RecvFrameNoWait receives a frame from the socket
// and returns it as a byte array if one is waiting.
// Returns an empty frame, a 0 more flag and an error
// if one is not immediately available
func (s *Sock) RecvFrameNoWait() ([]byte, int, error) {
	if !s.Pollin() {
		return []byte{0}, 0, fmt.Errorf("no frame")
	}

	return s.RecvFrame()
}

// SendMessage accepts an array of byte arrays and
// sends it as a multi-part message.
func (s *Sock) SendMessage(parts [][]byte) error {
	var f int
	numParts := len(parts)
	for i, val := range parts {
		if i == numParts-1 {
			f = 0
		} else {
			f = MORE
		}

		err := s.SendFrame(val, f)
		if err != nil {
			return err
		}
	}
	return nil
}

// RecvMessage receives a full message from the socket
// and returns it as an array of byte arrays.
func (s *Sock) RecvMessage() ([][]byte, error) {
	var msg [][]byte

	for {
		frame, flag, err := s.RecvFrame()
		if err != nil {
			return msg, err
		}
		msg = append(msg, frame)
		if flag != MORE {
			break
		}
	}
	return msg, nil
}

// Read provides an io.Reader interface to a zeromq socket
func (s *Sock) Read(p []byte) (int, error) {
	var total int
	frame, flag, err := s.RecvFrame()
	if err != nil {
		return total, err
	}

	if s.GetType() == ROUTER {
		s.lastClientID = string(frame)
	} else {
		copy(p[:], frame[:])
		total += len(frame)
	}

	for flag == MORE {
		frame, flag, err = s.RecvFrame()
		if err != nil {
			return total, err
		}
		copy(p[total:], frame[:])
		total += len(frame)
	}

	if total > len(p) {
		err = ErrSliceFull
	} else {
		err = io.EOF
	}

	return total, err
}

// Write provides an io.Writer interface to a zeromq socket
func (s *Sock) Write(p []byte) (int, error) {
	var total int
	if s.GetType() == ROUTER {
		err := s.SendFrame(s.GetLastClientID(), MORE)
		if err != nil {
			return total, err
		}
	}
	err := s.SendFrame(p, 0)
	if err != nil {
		return total, err
	}

	return len(p), nil
}

// RecvMessageNoWait receives a full message from the socket
// and returns it as an array of byte arrays if one is waiting.
// Returns an empty message and an error if one is not immediately
// available
func (s *Sock) RecvMessageNoWait() ([][]byte, error) {
	var msg [][]byte
	if !s.Pollin() {
		return msg, fmt.Errorf("no message")
	}

	for {
		frame, flag, err := s.RecvFrame()
		if err != nil {
			return msg, err
		}
		msg = append(msg, frame)
		if flag != MORE {
			break
		}
	}
	return msg, nil
}

// GetType returns the socket's type
func (s *Sock) GetType() int {
	return s.zType
}

// Destroy destroys the underlying zsockT.
func (s *Sock) Destroy() {
	C.zsock_destroy_checked(&s.zsockT, C.CString(s.file), C.size_t(s.line))
}
<|endoftext|>


<commit_before>package binary

import (
	"testing"
)

// You must implement the function,
//
//    func ParseBinary(string) (int, error)
//
// It is standard for Go functions to return error values to report error conditions.
// The test cases have some inputs that are invalid.
// For invalid inputs, return an error that signals to the user why the error happened.
// The test cases can only check that you return *some* error,
// but it's still good practice to return useful errors.
//
// Also define a testVersion with a value that matches
// the targetTestVersion here.

const targetTestVersion = 1

var testCases = []struct {
	binary   string
	expected int
	ok bool
}{
	{"1", 1, true},
	{"10", 2, true},
	{"11", 3, true},
	{"100", 4, true},
	{"1001", 9, true},
	{"11010", 26, true},
	{"10001101000", 1128, true},
	{"0", 0, true},
	{"foo101", 0, false},
	{"101bar", 0, false},
	{"101baz010", 0, false},
	{"22", 0, false},
}

func TestParseBinary(t *testing.T) {
	for _, tt := range testCases {
		actual, err := ParseBinary(tt.binary)
		if tt.ok {
			if err != nil {
				t.Fatalf("ParseBinary(%v) returned error %q.  Error not expected.",
					tt.binary, err)
			}
			if actual != tt.expected {
				t.Fatalf("ParseBinary(%v): actual %d, expected %v",
					tt.binary, actual, tt.expected)
			}
		} else if err == nil {
			t.Fatalf("ParseBinary(%v) returned %d and no error.  Expected an error.",
				tt.binary, actual)
		}
	}
}

// Benchmark combined time for all tests
func BenchmarkBinary(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, tt := range testCases {
			ParseBinary(tt.binary)
		}
	}
}
<commit_msg>binary: go fmt<commit_after>package binary

import (
	"testing"
)

// You must implement the function,
//
//    func ParseBinary(string) (int, error)
//
// It is standard for Go functions to return error values to report error conditions.
// The test cases have some inputs that are invalid.
// For invalid inputs, return an error that signals to the user why the error happened.
// The test cases can only check that you return *some* error,
// but it's still good practice to return useful errors.
//
// Also define a testVersion with a value that matches
// the targetTestVersion here.

const targetTestVersion = 1

var testCases = []struct {
	binary   string
	expected int
	ok       bool
}{
	{"1", 1, true},
	{"10", 2, true},
	{"11", 3, true},
	{"100", 4, true},
	{"1001", 9, true},
	{"11010", 26, true},
	{"10001101000", 1128, true},
	{"0", 0, true},
	{"foo101", 0, false},
	{"101bar", 0, false},
	{"101baz010", 0, false},
	{"22", 0, false},
}

func TestParseBinary(t *testing.T) {
	for _, tt := range testCases {
		actual, err := ParseBinary(tt.binary)
		if tt.ok {
			if err != nil {
				t.Fatalf("ParseBinary(%v) returned error %q.  Error not expected.",
					tt.binary, err)
			}
			if actual != tt.expected {
				t.Fatalf("ParseBinary(%v): actual %d, expected %v",
					tt.binary, actual, tt.expected)
			}
		} else if err == nil {
			t.Fatalf("ParseBinary(%v) returned %d and no error.  Expected an error.",
				tt.binary, actual)
		}
	}
}

// Benchmark combined time for all tests
func BenchmarkBinary(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, tt := range testCases {
			ParseBinary(tt.binary)
		}
	}
}
<|endoftext|>


<commit_before>//package soda provides HTTP GET tools for SODA (Socrata Open Data API) webservices, see http://dev.socrata.com/
package soda

import (
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
)

var ErrDone = fmt.Errorf("Done")

//A wrapper/container for SODA requests.
//This is NOT safe for use by multiple goroutines as Format, Filters and Query will be overwritten.
//Create a new GetRequest in each goroutine you use or use an OffsetGetRequest
type GetRequest struct {
	apptoken string
	endpoint string //endpoint without format (not .json etc at the end)
	Format   string //json, csv etc
	Filters  SimpleFilters
	Query    SoSQL
}

//Create a new GET request, the endpoint must be specified without the format.
//For example https://data.ct.gov/resource/hma6-9xbg
func NewGetRequest(endpoint, apptoken string) *GetRequest {
	return &GetRequest{
		apptoken: apptoken,
		endpoint: endpoint,
		Filters:  make(SimpleFilters),
	}
}

//Execute Get request
func (r *GetRequest) Get() (*http.Response, error) {
	return get(r, r.URLValues().Encode())
}

func (r *GetRequest) GetEndpoint() string {
	if r.Format == "" {
		r.Format = "json"
	}
	return fmt.Sprintf("%s.%s", r.endpoint, r.Format)
}

func (r *GetRequest) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range r.Filters.URLValues() {
		uv[key] = val
	}
	for key, val := range r.Query.URLValues() {
		uv[key] = val
	}
	return uv
}

//Counts the total number of records in the dataset
func (r *GetRequest) Count() (uint, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Order = oldorder
		r.Query.Select = oldselect
	}()

	r.Format = "json"
	r.Query.Select = []string{"count(*)"}
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	count := []struct {
		Count string
	}{}
	err = json.NewDecoder(resp.Body).Decode(&count)
	if err != nil {
		return 0, err
	}
	if len(count) == 0 {
		return 0, fmt.Errorf("Empty count response")
	}
	icount, err := strconv.Atoi(count[0].Count)
	if err != nil {
		return 0, err
	}
	return uint(icount), nil
}

//Gets all the fields present in the dataset (ignores select fields).
//Spaces in fieldnames are replaced by underscores.
func (r *GetRequest) Fields() ([]string, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	fields := []string{}

	r.Format = "csv"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return fields, err
	}
	defer resp.Body.Close()

	csvreader := csv.NewReader(resp.Body)
	record, err := csvreader.Read()
	if err != nil && err != io.EOF {
		return fields, err
	}
	for _, field := range record {
		fields = append(fields, strings.Replace(field, " ", "_", -1))
	}
	return fields, nil
}

//http://dev.socrata.com/docs/filtering.html
type SimpleFilters map[string]string

func (sf SimpleFilters) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range sf {
		uv.Add(key, val)
	}
	return uv
}

//http://dev.socrata.com/docs/queries.html
type SoSQL struct {
	Select []string //The set of columns to be returned. Default: All columns, equivalent to $select=*
	Where  string   //Filters the rows to be returned. Default: No filter, and returning a max of $limit values
	Order  []struct {
		Column string //Column name
		Desc   bool   //Descending. Default: false = Ascending
	} //Specifies the order of results. Default: Unspecified order, but it will be consistent across paging
	Group  string //Column to group results on, similar to SQL Grouping. Default: No grouping
	Limit  uint   //Maximum number of results to return. Default: 1000 (with a maximum of 50,000)
	Offset uint   //Offset count into the results to start at, used for paging. Default: 0
	Q      string //Performs a full text search for a value. Default: No search

}

func (sq *SoSQL) AddOrder(column string, descending bool) {
	sq.Order = append(sq.Order, struct {
		Column string
		Desc   bool
	}{column, descending})
}

func (sq *SoSQL) ClearOrder() {
	sq.Order = []struct {
		Column string
		Desc   bool
	}{}
}

func (sq *SoSQL) URLValues() url.Values {
	uv := make(url.Values)
	if len(sq.Select) > 0 {
		uv.Add("$select", strings.Join(sq.Select, ","))
	}
	if len(sq.Where) > 0 {
		uv.Add("$where", sq.Where)
	}
	if len(sq.Order) > 0 {
		order := []string{}
		for _, o := range sq.Order {
			if o.Desc {
				order = append(order, o.Column+" DESC")
			} else {
				order = append(order, o.Column+" ASC")
			}
		}
		uv.Add("$order", strings.Join(order, ","))
	}
	if len(sq.Q) > 0 {
		uv.Add("$q", sq.Q)
	}
	if len(sq.Group) > 0 {
		uv.Add("$group", sq.Group)
	}
	if sq.Limit > 0 {
		uv.Add("$limit", fmt.Sprintf("%d", sq.Limit))
	}
	if sq.Offset > 0 {
		uv.Add("$offset", fmt.Sprintf("%d", sq.Offset))
	}
	return uv
}

//An OffsetGetRequest is a request getter that gets all the records using the filters and limits from gr and
//is safe to use by multiple goroutines, use Next(number) to get the next number of records.
//A sync.WaitGroup is embedded for easy parallelism.
type OffsetGetRequest struct {
	gr     *GetRequest
	m      sync.Mutex
	offset uint
	count  uint
	sync.WaitGroup
}

//Gets the next number of records
func (o *OffsetGetRequest) Next(number uint) (*http.Response, error) {
	if o.offset >= o.count {
		return nil, ErrDone
	}
	o.m.Lock() //lock to protect offset
	o.gr.Query.Offset = o.offset
	o.gr.Query.Limit = number
	rawquery := o.gr.URLValues().Encode()
	o.offset += number
	o.m.Unlock() //unlock before the request is done
	return get(o.gr, rawquery)
}

func (o *OffsetGetRequest) Count() uint {
	return o.count
}

//returns if we have gotten all records
func (o *OffsetGetRequest) Done() bool {
	return o.offset >= o.count
}

//Creates a new OffsetGetRequest from gr and does a count request to determine the number of records to get
func NewOffsetGetRequest(gr *GetRequest) (*OffsetGetRequest, error) {
	count, err := gr.Count()
	if err != nil {
		return nil, err
	}
	return &OffsetGetRequest{gr: gr, offset: 0, count: count}, nil
}

func get(r *GetRequest, rawquery string) (*http.Response, error) {

	//If offset is used we must specify an order
	if r.Query.Offset > 0 && len(r.Query.Order) == 0 {
		return nil, fmt.Errorf("Cannot use an offset without setting the order")
	}

	req, err := http.NewRequest("GET", r.GetEndpoint(), nil)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = rawquery
	req.Header.Set("X-App-Token", r.apptoken)

	//Execute
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		errormessage, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("SODA error %d:\nURL: GET %s\nResponse: %s", resp.StatusCode, req.URL.String(), errormessage)
	}

	return resp, nil
}
<commit_msg>fix waitgroup<commit_after>//package soda provides HTTP GET tools for SODA (Socrata Open Data API) webservices, see http://dev.socrata.com/
package soda

import (
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
)

var ErrDone = fmt.Errorf("Done")

//A wrapper/container for SODA requests.
//This is NOT safe for use by multiple goroutines as Format, Filters and Query will be overwritten.
//Create a new GetRequest in each goroutine you use or use an OffsetGetRequest
type GetRequest struct {
	apptoken string
	endpoint string //endpoint without format (not .json etc at the end)
	Format   string //json, csv etc
	Filters  SimpleFilters
	Query    SoSQL
}

//Create a new GET request, the endpoint must be specified without the format.
//For example https://data.ct.gov/resource/hma6-9xbg
func NewGetRequest(endpoint, apptoken string) *GetRequest {
	return &GetRequest{
		apptoken: apptoken,
		endpoint: endpoint,
		Filters:  make(SimpleFilters),
	}
}

//Execute Get request
func (r *GetRequest) Get() (*http.Response, error) {
	return get(r, r.URLValues().Encode())
}

func (r *GetRequest) GetEndpoint() string {
	if r.Format == "" {
		r.Format = "json"
	}
	return fmt.Sprintf("%s.%s", r.endpoint, r.Format)
}

func (r *GetRequest) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range r.Filters.URLValues() {
		uv[key] = val
	}
	for key, val := range r.Query.URLValues() {
		uv[key] = val
	}
	return uv
}

//Counts the total number of records in the dataset
func (r *GetRequest) Count() (uint, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Order = oldorder
		r.Query.Select = oldselect
	}()

	r.Format = "json"
	r.Query.Select = []string{"count(*)"}
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	count := []struct {
		Count string
	}{}
	err = json.NewDecoder(resp.Body).Decode(&count)
	if err != nil {
		return 0, err
	}
	if len(count) == 0 {
		return 0, fmt.Errorf("Empty count response")
	}
	icount, err := strconv.Atoi(count[0].Count)
	if err != nil {
		return 0, err
	}
	return uint(icount), nil
}

//Gets all the fields present in the dataset (ignores select fields).
//Spaces in fieldnames are replaced by underscores.
func (r *GetRequest) Fields() ([]string, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	fields := []string{}

	r.Format = "csv"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return fields, err
	}
	defer resp.Body.Close()

	csvreader := csv.NewReader(resp.Body)
	record, err := csvreader.Read()
	if err != nil && err != io.EOF {
		return fields, err
	}
	for _, field := range record {
		fields = append(fields, strings.Replace(field, " ", "_", -1))
	}
	return fields, nil
}

//http://dev.socrata.com/docs/filtering.html
type SimpleFilters map[string]string

func (sf SimpleFilters) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range sf {
		uv.Add(key, val)
	}
	return uv
}

//http://dev.socrata.com/docs/queries.html
type SoSQL struct {
	Select []string //The set of columns to be returned. Default: All columns, equivalent to $select=*
	Where  string   //Filters the rows to be returned. Default: No filter, and returning a max of $limit values
	Order  []struct {
		Column string //Column name
		Desc   bool   //Descending. Default: false = Ascending
	} //Specifies the order of results. Default: Unspecified order, but it will be consistent across paging
	Group  string //Column to group results on, similar to SQL Grouping. Default: No grouping
	Limit  uint   //Maximum number of results to return. Default: 1000 (with a maximum of 50,000)
	Offset uint   //Offset count into the results to start at, used for paging. Default: 0
	Q      string //Performs a full text search for a value. Default: No search

}

func (sq *SoSQL) AddOrder(column string, descending bool) {
	sq.Order = append(sq.Order, struct {
		Column string
		Desc   bool
	}{column, descending})
}

func (sq *SoSQL) ClearOrder() {
	sq.Order = []struct {
		Column string
		Desc   bool
	}{}
}

func (sq *SoSQL) URLValues() url.Values {
	uv := make(url.Values)
	if len(sq.Select) > 0 {
		uv.Add("$select", strings.Join(sq.Select, ","))
	}
	if len(sq.Where) > 0 {
		uv.Add("$where", sq.Where)
	}
	if len(sq.Order) > 0 {
		order := []string{}
		for _, o := range sq.Order {
			if o.Desc {
				order = append(order, o.Column+" DESC")
			} else {
				order = append(order, o.Column+" ASC")
			}
		}
		uv.Add("$order", strings.Join(order, ","))
	}
	if len(sq.Q) > 0 {
		uv.Add("$q", sq.Q)
	}
	if len(sq.Group) > 0 {
		uv.Add("$group", sq.Group)
	}
	if sq.Limit > 0 {
		uv.Add("$limit", fmt.Sprintf("%d", sq.Limit))
	}
	if sq.Offset > 0 {
		uv.Add("$offset", fmt.Sprintf("%d", sq.Offset))
	}
	return uv
}

//An OffsetGetRequest is a request getter that gets all the records using the filters and limits from gr and
//is safe to use by multiple goroutines, use Next(number) to get the next number of records.
//A sync.WaitGroup is embedded for easy concurrency.
type OffsetGetRequest struct {
	gr     *GetRequest
	m      sync.Mutex
	offset uint
	count  uint
	sync.WaitGroup
}

//Gets the next number of records
func (o *OffsetGetRequest) Next(number uint) (*http.Response, error) {
	if o.IsDone() {
		return nil, ErrDone
	}
	o.m.Lock() //lock to protect offset
	o.gr.Query.Offset = o.offset
	o.gr.Query.Limit = number
	rawquery := o.gr.URLValues().Encode()
	o.offset += number
	o.m.Unlock() //unlock before the request is done
	return get(o.gr, rawquery)
}

//Returns the number of records
func (o *OffsetGetRequest) Count() uint {
	return o.count
}

//returns if we have gotten all records
func (o *OffsetGetRequest) IsDone() bool {
	return o.offset >= o.count
}

//Creates a new OffsetGetRequest from gr and does a count request to determine the number of records to get
func NewOffsetGetRequest(gr *GetRequest) (*OffsetGetRequest, error) {
	count, err := gr.Count()
	if err != nil {
		return nil, err
	}
	return &OffsetGetRequest{gr: gr, offset: 0, count: count}, nil
}

func get(r *GetRequest, rawquery string) (*http.Response, error) {

	//If offset is used we must specify an order
	if r.Query.Offset > 0 && len(r.Query.Order) == 0 {
		return nil, fmt.Errorf("Cannot use an offset without setting the order")
	}

	req, err := http.NewRequest("GET", r.GetEndpoint(), nil)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = rawquery
	req.Header.Set("X-App-Token", r.apptoken)

	//Execute
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		errormessage, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("SODA error %d:\nURL: GET %s\nResponse: %s", resp.StatusCode, req.URL.String(), errormessage)
	}

	return resp, nil
}
<|endoftext|>