<commit_before>package slack

import (
	"context"
	"errors"
	"fmt"
	"net/url"
)

type adminResponse struct {
	OK    bool   `json:"ok"`
	Error string `json:"error"`
}

func adminRequest(ctx context.Context, client HTTPRequester, method string, teamName string, values url.Values, debug bool) (*adminResponse, error) {
	adminResponse := &adminResponse{}
	err := parseAdminResponse(ctx, client, method, teamName, values, adminResponse, debug)
	if err != nil {
		return nil, err
	}

	if !adminResponse.OK {
		return nil, errors.New(adminResponse.Error)
	}

	return adminResponse, nil
}

// DisableUser disabled a user account, given a user ID
func (api *Client) DisableUser(teamName string, uid string) error {
	return api.DisableUserContext(context.Background(), teamName, uid)
}

// DisableUserContext disabled a user account, given a user ID with a custom context
func (api *Client) DisableUserContext(ctx context.Context, teamName string, uid string) error {
	values := url.Values{
		"user":       {uid},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setInactive", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to disable user with id '%s': %s", uid, err)
	}

	return nil
}

// InviteGuest invites a user to Slack as a single-channel guest
func (api *Client) InviteGuest(teamName, channel, firstName, lastName, emailAddress string) error {
	return api.InviteGuestContext(context.Background(), teamName, channel, firstName, lastName, emailAddress)
}

// InviteGuestContext invites a user to Slack as a single-channel guest with a custom context
func (api *Client) InviteGuestContext(ctx context.Context, teamName, channel, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":            {emailAddress},
		"channels":         {channel},
		"first_name":       {firstName},
		"last_name":        {lastName},
		"ultra_restricted": {"1"},
		"token":            {api.token},
		"set_active":       {"true"},
		"_attempts":        {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to invite single-channel guest: %s", err)
	}

	return nil
}

// InviteRestricted invites a user to Slack as a restricted account
func (api *Client) InviteRestricted(teamName, channel, firstName, lastName, emailAddress string) error {
	return api.InviteRestrictedContext(context.Background(), teamName, channel, firstName, lastName, emailAddress)
}

// InviteRestrictedContext invites a user to Slack as a restricted account with a custom context
func (api *Client) InviteRestrictedContext(ctx context.Context, teamName, channel, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":      {emailAddress},
		"channels":   {channel},
		"first_name": {firstName},
		"last_name":  {lastName},
		"restricted": {"1"},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to restricted account: %s", err)
	}

	return nil
}

// InviteToTeam invites a user to a Slack team
func (api *Client) InviteToTeam(teamName, firstName, lastName, emailAddress string) error {
	return api.InviteToTeamContext(context.Background(), teamName, firstName, lastName, emailAddress)
}

// InviteToTeamContext invites a user to a Slack team with a custom context
func (api *Client) InviteToTeamContext(ctx context.Context, teamName, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":      {emailAddress},
		"first_name": {firstName},
		"last_name":  {lastName},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to invite to team: %s", err)
	}

	return nil
}

// SetRegular enables the specified user
func (api *Client) SetRegular(teamName, user string) error {
	return api.SetRegularContext(context.Background(), teamName, user)
}

// SetRegularContext enables the specified user with a custom context
func (api *Client) SetRegularContext(ctx context.Context, teamName, user string) error {
	values := url.Values{
		"user":       {user},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setRegular", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to change the user (%s) to a regular user: %s", user, err)
	}

	return nil
}

// SendSSOBindingEmail sends an SSO binding email to the specified user
func (api *Client) SendSSOBindingEmail(teamName, user string) error {
	return api.SendSSOBindingEmailContext(context.Background(), teamName, user)
}

// SendSSOBindingEmailContext sends an SSO binding email to the specified user with a custom context
func (api *Client) SendSSOBindingEmailContext(ctx context.Context, teamName, user string) error {
	values := url.Values{
		"user":       {user},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "sendSSOBind", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to send SSO binding email for user (%s): %s", user, err)
	}

	return nil
}

// SetUltraRestricted converts a user into a single-channel guest
func (api *Client) SetUltraRestricted(teamName, uid, channel string) error {
	return api.SetUltraRestrictedContext(context.Background(), teamName, uid, channel)
}

// SetUltraRestrictedContext converts a user into a single-channel guest with a custom context
func (api *Client) SetUltraRestrictedContext(ctx context.Context, teamName, uid, channel string) error {
	values := url.Values{
		"user":       {uid},
		"channel":    {channel},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setUltraRestricted", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to ultra-restrict account: %s", err)
	}

	return nil
}

// SetRestricted converts a user into a restricted account
func (api *Client) SetRestricted(teamName, uid string) error {
	return api.SetRestrictedContext(context.Background(), teamName, uid)
}

// SetRestrictedContext converts a user into a restricted account with a custom context
func (api *Client) SetRestrictedContext(ctx context.Context, teamName, uid string) error {
	values := url.Values{
		"user":       {uid},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setRestricted", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to restrict account: %s", err)
	}

	return nil
}
<commit_msg>Include the resend param when sending guest invites<commit_after>package slack

import (
	"context"
	"errors"
	"fmt"
	"net/url"
)

type adminResponse struct {
	OK    bool   `json:"ok"`
	Error string `json:"error"`
}

func adminRequest(ctx context.Context, client HTTPRequester, method string, teamName string, values url.Values, debug bool) (*adminResponse, error) {
	adminResponse := &adminResponse{}
	err := parseAdminResponse(ctx, client, method, teamName, values, adminResponse, debug)
	if err != nil {
		return nil, err
	}

	if !adminResponse.OK {
		return nil, errors.New(adminResponse.Error)
	}

	return adminResponse, nil
}

// DisableUser disabled a user account, given a user ID
func (api *Client) DisableUser(teamName string, uid string) error {
	return api.DisableUserContext(context.Background(), teamName, uid)
}

// DisableUserContext disabled a user account, given a user ID with a custom context
func (api *Client) DisableUserContext(ctx context.Context, teamName string, uid string) error {
	values := url.Values{
		"user":       {uid},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setInactive", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to disable user with id '%s': %s", uid, err)
	}

	return nil
}

// InviteGuest invites a user to Slack as a single-channel guest
func (api *Client) InviteGuest(teamName, channel, firstName, lastName, emailAddress string) error {
	return api.InviteGuestContext(context.Background(), teamName, channel, firstName, lastName, emailAddress)
}

// InviteGuestContext invites a user to Slack as a single-channel guest with a custom context
func (api *Client) InviteGuestContext(ctx context.Context, teamName, channel, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":            {emailAddress},
		"channels":         {channel},
		"first_name":       {firstName},
		"last_name":        {lastName},
		"ultra_restricted": {"1"},
		"token":            {api.token},
		"resend":           {"true"},
		"set_active":       {"true"},
		"_attempts":        {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to invite single-channel guest: %s", err)
	}

	return nil
}

// InviteRestricted invites a user to Slack as a restricted account
func (api *Client) InviteRestricted(teamName, channel, firstName, lastName, emailAddress string) error {
	return api.InviteRestrictedContext(context.Background(), teamName, channel, firstName, lastName, emailAddress)
}

// InviteRestrictedContext invites a user to Slack as a restricted account with a custom context
func (api *Client) InviteRestrictedContext(ctx context.Context, teamName, channel, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":      {emailAddress},
		"channels":   {channel},
		"first_name": {firstName},
		"last_name":  {lastName},
		"restricted": {"1"},
		"token":      {api.token},
		"resend":     {"true"},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to restricted account: %s", err)
	}

	return nil
}

// InviteToTeam invites a user to a Slack team
func (api *Client) InviteToTeam(teamName, firstName, lastName, emailAddress string) error {
	return api.InviteToTeamContext(context.Background(), teamName, firstName, lastName, emailAddress)
}

// InviteToTeamContext invites a user to a Slack team with a custom context
func (api *Client) InviteToTeamContext(ctx context.Context, teamName, firstName, lastName, emailAddress string) error {
	values := url.Values{
		"email":      {emailAddress},
		"first_name": {firstName},
		"last_name":  {lastName},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "invite", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to invite to team: %s", err)
	}

	return nil
}

// SetRegular enables the specified user
func (api *Client) SetRegular(teamName, user string) error {
	return api.SetRegularContext(context.Background(), teamName, user)
}

// SetRegularContext enables the specified user with a custom context
func (api *Client) SetRegularContext(ctx context.Context, teamName, user string) error {
	values := url.Values{
		"user":       {user},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setRegular", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to change the user (%s) to a regular user: %s", user, err)
	}

	return nil
}

// SendSSOBindingEmail sends an SSO binding email to the specified user
func (api *Client) SendSSOBindingEmail(teamName, user string) error {
	return api.SendSSOBindingEmailContext(context.Background(), teamName, user)
}

// SendSSOBindingEmailContext sends an SSO binding email to the specified user with a custom context
func (api *Client) SendSSOBindingEmailContext(ctx context.Context, teamName, user string) error {
	values := url.Values{
		"user":       {user},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "sendSSOBind", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to send SSO binding email for user (%s): %s", user, err)
	}

	return nil
}

// SetUltraRestricted converts a user into a single-channel guest
func (api *Client) SetUltraRestricted(teamName, uid, channel string) error {
	return api.SetUltraRestrictedContext(context.Background(), teamName, uid, channel)
}

// SetUltraRestrictedContext converts a user into a single-channel guest with a custom context
func (api *Client) SetUltraRestrictedContext(ctx context.Context, teamName, uid, channel string) error {
	values := url.Values{
		"user":       {uid},
		"channel":    {channel},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setUltraRestricted", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to ultra-restrict account: %s", err)
	}

	return nil
}

// SetRestricted converts a user into a restricted account
func (api *Client) SetRestricted(teamName, uid string) error {
	return api.SetRestrictedContext(context.Background(), teamName, uid)
}

// SetRestrictedContext converts a user into a restricted account with a custom context
func (api *Client) SetRestrictedContext(ctx context.Context, teamName, uid string) error {
	values := url.Values{
		"user":       {uid},
		"token":      {api.token},
		"set_active": {"true"},
		"_attempts":  {"1"},
	}

	_, err := adminRequest(ctx, api.httpclient, "setRestricted", teamName, values, api.debug)
	if err != nil {
		return fmt.Errorf("Failed to restrict account: %s", err)
	}

	return nil
}
<|endoftext|>


<commit_before>/**
 * keys
 * This module handles all the key generation logic for NaCl and AES.
 * It passes off reading from the rand reader to the entropy module to
 * limit code duplication. It also provides helper methods to turn byte
 * array data into the proper key types so they can be used in the crypto
 * methods where they are needed, and read and write key files to disk.
 */
package msg

import (
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/pbkdf2"
	"io/ioutil"
	"log"
	"os"
)

// Consts used in this module
const PBKDF2_SALT_BYTES = 32
const PBKDF2_PASSWORD_BYTES = 32
const PBKDF2_KEY_BYTES = 32
const PBKDF2_CYCLES = 5000

// Include getters to make it clear other modules within cryptoUtils are using
// these constants
const NACL_FINGERPRINT_BYTES = 20

func getNACLFingerprintLength() int {
	return NACL_FINGERPRINT_BYTES
}

const NACL_KEY_BYTES = 32

/**
 * GenAES256Key
 * Generate a completley random 32 byte key using PBKDF2 with a completely random
 * "passphrase" and salt. sha256 is being used as the hashing function.
 */
func GenAES256Key() []byte {
	password := ReadFromRand(PBKDF2_PASSWORD_BYTES)
	salt := ReadFromRand(PBKDF2_SALT_BYTES)

	return pbkdf2.Key(password, salt, PBKDF2_CYCLES, PBKDF2_KEY_BYTES, sha256.New)
}

/**
 * GenerateNACLKeyPair
 * Helper method that wraps NaCl's built in key generation function.
 */
func GenerateNACLKeyPair() (pubKey, privKey *[32]byte) {
	pubKey, privKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatal(err)
	}
	return
}

/**
 * WriteNACLKeyFile
 * Helper method that encodes NACL keys to base64 then writes them to disk.
 */
func WriteNACLKeyFile(keyFilePath string, key *[32]byte, perms os.FileMode) {
	encodedKey := base64.StdEncoding.EncodeToString(key[:])
	err := ioutil.WriteFile(keyFilePath, []byte(encodedKey), perms)
	if err != nil {
		log.Fatal(err)
	}
}

/**
 * ReadNACLKeyFile
 * Helper method to read NACL key files that were written to disk with WriteNACLKeyFile.
 * Because this is a companion method, we assume the file being read has been encoded
 * as base64.
 */
func ReadNACLKeyFile(keyFilePath string) *[32]byte {
	keyFile, err := ioutil.ReadFile(keyFilePath)
	if err != nil {
		log.Fatal(err)
	}

	keyBytes, err := base64.StdEncoding.DecodeString(string(keyFile))
	if err != nil {
		log.Fatal(err)
	}

	return ConvertBytesToNACLKey(keyBytes)
}

/**
 * GetNACLKeyFingerprint
 * Generates 20 byte fingerprints for NaCl Keys using sha1
 * This is only meant to be used with PUBLIC keys.
 */
func GetNACLKeyFingerprint(key *[32]byte) []byte {
	h := sha1.New()
	h.Write(key[:])
	return h.Sum(nil)
}

/**
 * ConvertBytesToNACLKey
 * Helper that takes in a byte array and casts it to *[32]byte so it can be used
 * as an NACL key. This is used by ReadNACLKeyFile.
 */
func ConvertBytesToNACLKey(bytes []byte) *[32]byte {
	if len(bytes) < 32 {
		log.Fatal("Too few bytes! This does not appear to be a valid NACL Key")
	}
	var naclKey [32]byte
	copy(naclKey[:], bytes[:])
	return &naclKey
}
<commit_msg>Support passing encoded strings and returning keys<commit_after>/**
 * keys
 * This module handles all the key generation logic for NaCl and AES.
 * It passes off reading from the rand reader to the entropy module to
 * limit code duplication. It also provides helper methods to turn byte
 * array data into the proper key types so they can be used in the crypto
 * methods where they are needed, and read and write key files to disk.
 */
package msg

import (
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/pbkdf2"
	"io/ioutil"
	"log"
	"os"
)

// Consts used in this module
const PBKDF2_SALT_BYTES = 32
const PBKDF2_PASSWORD_BYTES = 32
const PBKDF2_KEY_BYTES = 32
const PBKDF2_CYCLES = 5000

// Include getters to make it clear other modules within cryptoUtils are using
// these constants
const NACL_FINGERPRINT_BYTES = 20

func getNACLFingerprintLength() int {
	return NACL_FINGERPRINT_BYTES
}

const NACL_KEY_BYTES = 32

/**
 * GenAES256Key
 * Generate a completley random 32 byte key using PBKDF2 with a completely random
 * "passphrase" and salt. sha256 is being used as the hashing function.
 */
func GenAES256Key() []byte {
	password := ReadFromRand(PBKDF2_PASSWORD_BYTES)
	salt := ReadFromRand(PBKDF2_SALT_BYTES)

	return pbkdf2.Key(password, salt, PBKDF2_CYCLES, PBKDF2_KEY_BYTES, sha256.New)
}

/**
 * GenerateNACLKeyPair
 * Helper method that wraps NaCl's built in key generation function.
 */
func GenerateNACLKeyPair() (pubKey, privKey *[32]byte) {
	pubKey, privKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatal(err)
	}
	return
}

/**
 * WriteNACLKeyFile
 * Helper method that encodes NACL keys to base64 then writes them to disk.
 */
func WriteNACLKeyFile(keyFilePath string, key *[32]byte, perms os.FileMode) {
	encodedKey := base64.StdEncoding.EncodeToString(key[:])
	err := ioutil.WriteFile(keyFilePath, []byte(encodedKey), perms)
	if err != nil {
		log.Fatal(err)
	}
}

/**
 * ReadNACLKeyFile
 * Helper method to read NACL key files that were written to disk with WriteNACLKeyFile.
 * Because this is a companion method, we assume the file being read has been encoded
 * as base64.
 */
func ReadNACLKeyFile(keyFilePath string) *[32]byte {
	keyFile, err := ioutil.ReadFile(keyFilePath)
	if err != nil {
		log.Fatal(err)
	}

	return ReadNACLKeyString(string(keyFile))
}

/**
 * ReadNACLKeyString
 * Helper method to read encoded NACL key data strings and return an NACL key.
 * Because this is a helper method, we assume the string passed in is encoded
 * as base64.
 */
func ReadNACLKeyString(encodedKey string) *[32]byte {
	keyBytes, err := base64.StdEncoding.DecodeString(encodedKey)
	if err != nil {
		log.Fatal(err)
	}

	naclKey := ConvertBytesToNACLKey(keyBytes)
	return naclKey
}

/**
 * GetNACLKeyFingerprint
 * Generates 20 byte fingerprints for NaCl Keys using sha1
 * This is only meant to be used with PUBLIC keys.
 */
func GetNACLKeyFingerprint(key *[32]byte) []byte {
	h := sha1.New()
	h.Write(key[:])
	return h.Sum(nil)
}

/**
 * ConvertBytesToNACLKey
 * Helper that takes in a byte array and casts it to *[32]byte so it can be used
 * as an NACL key. This is used by ReadNACLKeyFile.
 */
func ConvertBytesToNACLKey(bytes []byte) *[32]byte {
	if len(bytes) < 32 {
		log.Fatal("Too few bytes! This does not appear to be a valid NACL Key")
	}
	var naclKey [32]byte
	copy(naclKey[:], bytes[:])
	return &naclKey
}
<|endoftext|>


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

package gabi

import (
	"crypto/rand"
	"encoding/xml"
	"io"
	"io/ioutil"
	"math/big"
	"os"
	"strconv"
	"time"

	"github.com/credentials/safeprime"
)

const (
	//XMLHeader can be a used as the XML header when writing keys in XML format.
	XMLHeader = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n"
	// DefaultEpochLength is the default epoch length for public keys.
	DefaultEpochLength = 432000
)

// PrivateKey represents an issuer's private key.
type PrivateKey struct {
	XMLName xml.Name `xml:"http://www.zurich.ibm.com/security/idemix IssuerPrivateKey"`
	Counter uint     `xml:"Counter"`
	Expiry  int64    `xml:"ExpiryDate"`
	P       *big.Int `xml:"Elements>p"`
	Q       *big.Int `xml:"Elements>q"`
	PPrime  *big.Int `xml:"Elements>pPrime"`
	QPrime  *big.Int `xml:"Elements>qPrime"`
}

// NewPrivateKey creates a new issuer private key using the provided parameters.
func NewPrivateKey(p, q *big.Int, counter uint, expiry time.Time) *PrivateKey {
	sk := PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int), Counter: counter, Expiry: expiry.Unix()}

	sk.PPrime.Sub(p, bigONE)
	sk.PPrime.Rsh(sk.PPrime, 1)

	sk.QPrime.Sub(q, bigONE)
	sk.QPrime.Rsh(sk.QPrime, 1)

	return &sk
}

// NewPrivateKeyFromXML creates a new issuer private key using the xml data
// provided.
func NewPrivateKeyFromXML(xmlInput string) (*PrivateKey, error) {
	privk := &PrivateKey{}
	err := xml.Unmarshal([]byte(xmlInput), privk)
	if err != nil {
		return nil, err
	}
	return privk, nil
}

// NewPrivateKeyFromFile create a new issuer private key from an xml file.
func NewPrivateKeyFromFile(filename string) (*PrivateKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	privk := &PrivateKey{}

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(b, privk)
	if err != nil {
		return nil, err
	}
	return privk, nil
}

// Print prints the key to stdout.
func (privk *PrivateKey) Print() error {
	return privk.WriteTo(os.Stdout)
}

// WriteTo writes the XML-serialized public key to the given writer.
func (privk *PrivateKey) WriteTo(writer io.Writer) error {
	// Write the standard XML header
	_, err := writer.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(privk, "", "   ")
	if err != nil {
		return err
	}
	_, err = writer.Write(b)
	return err
}

// WriteToFile writes the private key to an xml file.
func (privk *PrivateKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	return privk.WriteTo(f)
}

// xmlBases is an auxiliary struct to encode/decode the odd way bases are
// represented in the xml representation of public keys
type xmlBases struct {
	Num   int        `xml:"num,attr"`
	Bases []*xmlBase `xml:",any"`
}

type xmlBase struct {
	XMLName xml.Name
	Bigint  string `xml:",innerxml"` // Has to be a string for ",innerxml" to work
}

// xmlFeatures is an auxiliary struct to make the XML encoding/decoding a bit
// easier while keeping the struct for PublicKey somewhat simple.
type xmlFeatures struct {
	Epoch struct {
		Length int `xml:"length,attr"`
	}
}

// Bases is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type Bases []*big.Int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (bl *Bases) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlBases

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}

	arr := make([]*big.Int, t.Num)
	for i := range arr {
		arr[i], _ = new(big.Int).SetString(t.Bases[i].Bigint, 10)
	}

	*bl = Bases(arr)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (bl *Bases) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	l := len(*bl)
	bases := make([]*xmlBase, l)

	for i := range bases {
		bases[i] = &xmlBase{
			XMLName: xml.Name{Local: "Base_" + strconv.Itoa(i)},
			Bigint:  (*bl)[i].String(),
		}
	}

	t := xmlBases{
		Num:   l,
		Bases: bases,
	}
	return e.EncodeElement(t, start)
}

// EpochLength is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type EpochLength int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (el *EpochLength) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlFeatures

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*el = EpochLength(t.Epoch.Length)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (el *EpochLength) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	var t xmlFeatures
	t.Epoch.Length = int(*el)
	return e.EncodeElement(t, start)
}

// PublicKey represents an issuer's public key.
type PublicKey struct {
	XMLName     xml.Name          `xml:"http://www.zurich.ibm.com/security/idemix IssuerPublicKey"`
	Counter     uint              `xml:"Counter"`
	Expiry      int64             `xml:"ExpiryDate"`
	N           *big.Int          `xml:"Elements>n"` // Modulus n
	Z           *big.Int          `xml:"Elements>Z"` // Generator Z
	S           *big.Int          `xml:"Elements>S"` // Generator S
	R           Bases             `xml:"Elements>Bases"`
	EpochLength EpochLength       `xml:"Features"`
	Params      *SystemParameters `xml:"-"`
}

// NewPublicKey creates and returns a new public key based on the provided parameters.
func NewPublicKey(N, Z, S *big.Int, R []*big.Int, counter uint, expiry time.Time) *PublicKey {
	return &PublicKey{
		Counter:     counter,
		Expiry:      expiry.Unix(),
		N:           N,
		Z:           Z,
		S:           S,
		R:           R,
		EpochLength: DefaultEpochLength,
		Params:      &DefaultSystemParameters,
	}
}

// NewPublicKeyFromXML creates a new issuer public key using the xml data
// provided.
func NewPublicKeyFromXML(xmlInput string) (*PublicKey, error) {
	// TODO: this might fail in the future. The DefaultSystemParameters and the
	// public key might not match!
	pubk := &PublicKey{Params: &DefaultSystemParameters}
	err := xml.Unmarshal([]byte(xmlInput), pubk)
	if err != nil {
		return nil, err
	}
	return pubk, nil
}

// NewPublicKeyFromFile create a new issuer public key from an xml file.
func NewPublicKeyFromFile(filename string) (*PublicKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	pubk := &PublicKey{}

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(b, pubk)
	if err != nil {
		return nil, err
	}
	return pubk, nil
}

// Print prints the key to stdout.
func (pubk *PublicKey) Print() error {
	return pubk.WriteTo(os.Stdout)
}

// WriteTo writes the XML-serialized public key to the given writer.
func (pubk *PublicKey) WriteTo(writer io.Writer) error {
	// Write the standard XML header
	_, err := writer.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(pubk, "", "   ")
	if err != nil {
		return err
	}
	_, err = writer.Write(b)
	return err
}

// WriteToFile writes the public key to an xml file.
func (pubk *PublicKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	return pubk.WriteTo(f)
}

// randomSafePrime produces a safe prime of the requested number of bits
func randomSafePrime(bits int) (*big.Int, error) {
	p2 := new(big.Int)
	for {
		p, err := rand.Prime(rand.Reader, bits)
		if err != nil {
			return nil, err
		}
		p2.Rsh(p, 1) // p2 = (p - 1)/2
		if p2.ProbablyPrime(20) {
			return p, nil
		}
	}
}

// GenerateKeyPair generates a private/public keypair for an Issuer
func GenerateKeyPair(param *SystemParameters, attrsAmount int, counter uint, expiry time.Time) (*PrivateKey, *PublicKey, error) {
	primeSize := param.Ln / 2

	// p and q need to be safe primes
	p, err := safeprime.Generate(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	q, err := safeprime.Generate(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	priv := &PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int), Counter: counter, Expiry: expiry.Unix()}

	// compute p' and q'
	priv.PPrime.Sub(priv.P, bigONE)
	priv.PPrime.Rsh(priv.PPrime, 1)

	priv.QPrime.Sub(priv.Q, bigONE)
	priv.QPrime.Rsh(priv.QPrime, 1)

	// compute n
	pubk := &PublicKey{Params: param, EpochLength: DefaultEpochLength, Counter: counter, Expiry: expiry.Unix()}
	pubk.N = new(big.Int).Mul(priv.P, priv.Q)

	// Find an acceptable value for S; we follow lead of the Silvia code here:
	// Pick a random l_n value and check whether it is a quadratic residue modulo n

	var s *big.Int
	for {
		s, err = randomBigInt(param.Ln)
		if err != nil {
			return nil, nil, err
		}
		// check if S \elem Z_n
		if s.Cmp(pubk.N) > 0 {
			continue
		}
		if legendreSymbol(s, priv.P) == 1 && legendreSymbol(s, priv.Q) == 1 {
			break
		}
	}

	pubk.S = s

	// Derive Z from S
	var x *big.Int
	for {
		x, _ = randomBigInt(primeSize)
		if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
			break
		}
	}

	// Compute Z = S^x mod n
	pubk.Z = new(big.Int).Exp(pubk.S, x, pubk.N)

	// Derive R_i for i = 0...attrsAmount from S
	pubk.R = make([]*big.Int, attrsAmount)
	for i := 0; i < attrsAmount; i++ {
		pubk.R[i] = new(big.Int)

		var x *big.Int
		for {
			x, _ = randomBigInt(primeSize)
			if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
				break
			}
		}
		// Compute R_i = S^x mod n
		pubk.R[i].Exp(pubk.S, x, pubk.N)
	}

	return priv, pubk, nil
}
<commit_msg>Rename Expiry to ExpiryDate in keys.<commit_after>// Copyright 2016 Maarten Everts. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gabi

import (
	"crypto/rand"
	"encoding/xml"
	"io"
	"io/ioutil"
	"math/big"
	"os"
	"strconv"
	"time"

	"github.com/credentials/safeprime"
)

const (
	//XMLHeader can be a used as the XML header when writing keys in XML format.
	XMLHeader = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n"
	// DefaultEpochLength is the default epoch length for public keys.
	DefaultEpochLength = 432000
)

// PrivateKey represents an issuer's private key.
type PrivateKey struct {
	XMLName    xml.Name `xml:"http://www.zurich.ibm.com/security/idemix IssuerPrivateKey"`
	Counter    uint     `xml:"Counter"`
	ExpiryDate int64    `xml:"ExpiryDate"`
	P          *big.Int `xml:"Elements>p"`
	Q          *big.Int `xml:"Elements>q"`
	PPrime     *big.Int `xml:"Elements>pPrime"`
	QPrime     *big.Int `xml:"Elements>qPrime"`
}

// NewPrivateKey creates a new issuer private key using the provided parameters.
func NewPrivateKey(p, q *big.Int, counter uint, expiryDate time.Time) *PrivateKey {
	sk := PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int), Counter: counter, ExpiryDate: expiryDate.Unix()}

	sk.PPrime.Sub(p, bigONE)
	sk.PPrime.Rsh(sk.PPrime, 1)

	sk.QPrime.Sub(q, bigONE)
	sk.QPrime.Rsh(sk.QPrime, 1)

	return &sk
}

// NewPrivateKeyFromXML creates a new issuer private key using the xml data
// provided.
func NewPrivateKeyFromXML(xmlInput string) (*PrivateKey, error) {
	privk := &PrivateKey{}
	err := xml.Unmarshal([]byte(xmlInput), privk)
	if err != nil {
		return nil, err
	}
	return privk, nil
}

// NewPrivateKeyFromFile create a new issuer private key from an xml file.
func NewPrivateKeyFromFile(filename string) (*PrivateKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	privk := &PrivateKey{}

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(b, privk)
	if err != nil {
		return nil, err
	}
	return privk, nil
}

// Print prints the key to stdout.
func (privk *PrivateKey) Print() error {
	return privk.WriteTo(os.Stdout)
}

// WriteTo writes the XML-serialized public key to the given writer.
func (privk *PrivateKey) WriteTo(writer io.Writer) error {
	// Write the standard XML header
	_, err := writer.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(privk, "", "   ")
	if err != nil {
		return err
	}
	_, err = writer.Write(b)
	return err
}

// WriteToFile writes the private key to an xml file.
func (privk *PrivateKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	return privk.WriteTo(f)
}

// xmlBases is an auxiliary struct to encode/decode the odd way bases are
// represented in the xml representation of public keys
type xmlBases struct {
	Num   int        `xml:"num,attr"`
	Bases []*xmlBase `xml:",any"`
}

type xmlBase struct {
	XMLName xml.Name
	Bigint  string `xml:",innerxml"` // Has to be a string for ",innerxml" to work
}

// xmlFeatures is an auxiliary struct to make the XML encoding/decoding a bit
// easier while keeping the struct for PublicKey somewhat simple.
type xmlFeatures struct {
	Epoch struct {
		Length int `xml:"length,attr"`
	}
}

// Bases is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type Bases []*big.Int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (bl *Bases) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlBases

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}

	arr := make([]*big.Int, t.Num)
	for i := range arr {
		arr[i], _ = new(big.Int).SetString(t.Bases[i].Bigint, 10)
	}

	*bl = Bases(arr)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (bl *Bases) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	l := len(*bl)
	bases := make([]*xmlBase, l)

	for i := range bases {
		bases[i] = &xmlBase{
			XMLName: xml.Name{Local: "Base_" + strconv.Itoa(i)},
			Bigint:  (*bl)[i].String(),
		}
	}

	t := xmlBases{
		Num:   l,
		Bases: bases,
	}
	return e.EncodeElement(t, start)
}

// EpochLength is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type EpochLength int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (el *EpochLength) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlFeatures

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*el = EpochLength(t.Epoch.Length)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (el *EpochLength) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	var t xmlFeatures
	t.Epoch.Length = int(*el)
	return e.EncodeElement(t, start)
}

// PublicKey represents an issuer's public key.
type PublicKey struct {
	XMLName     xml.Name          `xml:"http://www.zurich.ibm.com/security/idemix IssuerPublicKey"`
	Counter     uint              `xml:"Counter"`
	ExpiryDate  int64             `xml:"ExpiryDate"`
	N           *big.Int          `xml:"Elements>n"` // Modulus n
	Z           *big.Int          `xml:"Elements>Z"` // Generator Z
	S           *big.Int          `xml:"Elements>S"` // Generator S
	R           Bases             `xml:"Elements>Bases"`
	EpochLength EpochLength       `xml:"Features"`
	Params      *SystemParameters `xml:"-"`
}

// NewPublicKey creates and returns a new public key based on the provided parameters.
func NewPublicKey(N, Z, S *big.Int, R []*big.Int, counter uint, expiry time.Time) *PublicKey {
	return &PublicKey{
		Counter:     counter,
		ExpiryDate:  expiry.Unix(),
		N:           N,
		Z:           Z,
		S:           S,
		R:           R,
		EpochLength: DefaultEpochLength,
		Params:      &DefaultSystemParameters,
	}
}

// NewPublicKeyFromXML creates a new issuer public key using the xml data
// provided.
func NewPublicKeyFromXML(xmlInput string) (*PublicKey, error) {
	// TODO: this might fail in the future. The DefaultSystemParameters and the
	// public key might not match!
	pubk := &PublicKey{Params: &DefaultSystemParameters}
	err := xml.Unmarshal([]byte(xmlInput), pubk)
	if err != nil {
		return nil, err
	}
	return pubk, nil
}

// NewPublicKeyFromFile create a new issuer public key from an xml file.
func NewPublicKeyFromFile(filename string) (*PublicKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	pubk := &PublicKey{}

	b, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(b, pubk)
	if err != nil {
		return nil, err
	}
	return pubk, nil
}

// Print prints the key to stdout.
func (pubk *PublicKey) Print() error {
	return pubk.WriteTo(os.Stdout)
}

// WriteTo writes the XML-serialized public key to the given writer.
func (pubk *PublicKey) WriteTo(writer io.Writer) error {
	// Write the standard XML header
	_, err := writer.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(pubk, "", "   ")
	if err != nil {
		return err
	}
	_, err = writer.Write(b)
	return err
}

// WriteToFile writes the public key to an xml file.
func (pubk *PublicKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	return pubk.WriteTo(f)
}

// randomSafePrime produces a safe prime of the requested number of bits
func randomSafePrime(bits int) (*big.Int, error) {
	p2 := new(big.Int)
	for {
		p, err := rand.Prime(rand.Reader, bits)
		if err != nil {
			return nil, err
		}
		p2.Rsh(p, 1) // p2 = (p - 1)/2
		if p2.ProbablyPrime(20) {
			return p, nil
		}
	}
}

// GenerateKeyPair generates a private/public keypair for an Issuer
func GenerateKeyPair(param *SystemParameters, attrsAmount int, counter uint, expiryDate time.Time) (*PrivateKey, *PublicKey, error) {
	primeSize := param.Ln / 2

	// p and q need to be safe primes
	p, err := safeprime.Generate(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	q, err := safeprime.Generate(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	priv := &PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int), Counter: counter, ExpiryDate: expiryDate.Unix()}

	// compute p' and q'
	priv.PPrime.Sub(priv.P, bigONE)
	priv.PPrime.Rsh(priv.PPrime, 1)

	priv.QPrime.Sub(priv.Q, bigONE)
	priv.QPrime.Rsh(priv.QPrime, 1)

	// compute n
	pubk := &PublicKey{Params: param, EpochLength: DefaultEpochLength, Counter: counter, ExpiryDate: expiryDate.Unix()}
	pubk.N = new(big.Int).Mul(priv.P, priv.Q)

	// Find an acceptable value for S; we follow lead of the Silvia code here:
	// Pick a random l_n value and check whether it is a quadratic residue modulo n

	var s *big.Int
	for {
		s, err = randomBigInt(param.Ln)
		if err != nil {
			return nil, nil, err
		}
		// check if S \elem Z_n
		if s.Cmp(pubk.N) > 0 {
			continue
		}
		if legendreSymbol(s, priv.P) == 1 && legendreSymbol(s, priv.Q) == 1 {
			break
		}
	}

	pubk.S = s

	// Derive Z from S
	var x *big.Int
	for {
		x, _ = randomBigInt(primeSize)
		if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
			break
		}
	}

	// Compute Z = S^x mod n
	pubk.Z = new(big.Int).Exp(pubk.S, x, pubk.N)

	// Derive R_i for i = 0...attrsAmount from S
	pubk.R = make([]*big.Int, attrsAmount)
	for i := 0; i < attrsAmount; i++ {
		pubk.R[i] = new(big.Int)

		var x *big.Int
		for {
			x, _ = randomBigInt(primeSize)
			if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
				break
			}
		}
		// Compute R_i = S^x mod n
		pubk.R[i].Exp(pubk.S, x, pubk.N)
	}

	return priv, pubk, nil
}
<|endoftext|>


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

package monmq

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/jroimartin/rpcmq"
)

// The type CommandFunction declares the signature of the methods that can be
// registered by an Agent. The data parameter contains auxiliary data.
type CommandFunction func(data []byte) ([]byte, error)

// An Agent is responsible for sending its status when a supervisor on the same
// exchange asks for it. It also executes the control operations requested by
// the supervisors.
type Agent struct {
	s      *rpcmq.Server
	status Status

	// TLSConfig allows to configure the TLS parameters used to connect to
	// the broker via amqps.
	TLSConfig *tls.Config

	// SoftShutdownFunc will be called when a supervisor invokes the
	// command SoftShutdown.
	SoftShutdownFunc CommandFunction

	// HardShutdownFunc will be called when a supervisor invokes the
	// command HardShutdown.
	HardShutdownFunc CommandFunction

	// PauseFunc will be called when a supervisor invokes the command
	// Pause.
	PauseFunc CommandFunction

	// ResumeFunc will be called when a supervisor invokes the command
	// Resume.
	ResumeFunc CommandFunction

	// CustomFunc will be called when a supervisor invokes the command
	// CustomCmd.
	CustomFunc CommandFunction

	// KillTaskFunc will be called when a supervisor invokes the command
	// KillTask.
	KillTaskFunc CommandFunction
}

// NewAgent returns a reference to an Agent object. The paremeter uri is the
// network address of the broker and exchange is the name of exchange that will
// be created.
func NewAgent(uri, exchange, name string) *Agent {
	a := &Agent{}
	a.s = rpcmq.NewServer(uri, "", exchange, "fanout")
	a.s.Parallel = 1
	a.status.Name = name
	return a
}

// Init initializes the Agent object. It establishes the connection with the
// broker, creating a channel and the exchange that will be used under the hood.
func (a *Agent) Init() error {
	a.s.TLSConfig = a.TLSConfig
	if err := a.s.Register("invoke", a.invoke); err != nil {
		return err
	}
	if err := a.s.Init(); err != nil {
		return err
	}
	a.status.Running = true
	return nil
}

// Shutdown shuts down the agent gracefully. Using this method will ensure that
// all requests sent by the supervisors to the agent will be received by the
// latter.
func (a *Agent) Shutdown() {
	a.s.Shutdown()
}

// RegisterTask adds a task to the list of tasks handled by the agent.
func (a *Agent) RegisterTask(id string) {
	a.status.Tasks = append(a.status.Tasks, id)
}

// RemoveTask removes a task from the list of tasks handled by the agent.
func (a *Agent) RemoveTask(id string) error {
	idx := -1
	for i, t := range a.status.Tasks {
		if t == id {
			idx = i
			break
		}
	}
	if idx == -1 {
		return errors.New("task not found")
	}
	a.status.Tasks = append(a.status.Tasks[:idx], a.status.Tasks[idx+1:]...)
	return nil
}

func (a *Agent) invoke(id string, data []byte) ([]byte, error) {
	var f CommandFunction
	running := a.status.Running
	cmd, aux := Command(data[0]), data[1:]
	switch {
	case cmd == GetStatus:
		f = a.getStatus
	case a.status.Name == string(aux) && cmd == SoftShutdown:
		f = a.SoftShutdownFunc
		running = false
	case a.status.Name == string(aux) && cmd == HardShutdown:
		f = a.HardShutdownFunc
		running = false
	case a.status.Name == string(aux) && cmd == Pause:
		f = a.PauseFunc
		running = false
	case a.status.Name == string(aux) && cmd == Resume:
		f = a.ResumeFunc
		running = true
	case a.status.Name == string(aux) && cmd == CustomCmd:
		f = a.CustomFunc
	case a.ownsTask(string(aux)) && cmd == KillTask:
		f = a.KillTaskFunc
	}
	if f == nil {
		// The command is not for this agent
		return nil, nil
	}
	b, err := f(data)
	if err != nil {
		return nil, err
	}
	a.status.Running = running
	out := fmt.Sprintf("%c%s", cmd, b)
	return []byte(out), nil
}

func (a *Agent) getStatus(data []byte) ([]byte, error) {
	info, err := getSystemInfo()
	if err != nil {
		return nil, err
	}
	a.status.Info = info
	return json.Marshal(a.status)
}

func (a *Agent) ownsTask(id string) bool {
	for _, t := range a.status.Tasks {
		if id == t {
			return true
		}
	}
	return false
}
<commit_msg>Make agent concurrent safe<commit_after>// Copyright 2015 The monmq Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package monmq

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"sync"

	"github.com/jroimartin/rpcmq"
)

// The type CommandFunction declares the signature of the methods that can be
// registered by an Agent. The data parameter contains auxiliary data.
type CommandFunction func(data []byte) ([]byte, error)

// An Agent is responsible for sending its status when a supervisor on the same
// exchange asks for it. It also executes the control operations requested by
// the supervisors.
type Agent struct {
	s      *rpcmq.Server
	status Status

	mu sync.Mutex

	// TLSConfig allows to configure the TLS parameters used to connect to
	// the broker via amqps.
	TLSConfig *tls.Config

	// SoftShutdownFunc will be called when a supervisor invokes the
	// command SoftShutdown.
	SoftShutdownFunc CommandFunction

	// HardShutdownFunc will be called when a supervisor invokes the
	// command HardShutdown.
	HardShutdownFunc CommandFunction

	// PauseFunc will be called when a supervisor invokes the command
	// Pause.
	PauseFunc CommandFunction

	// ResumeFunc will be called when a supervisor invokes the command
	// Resume.
	ResumeFunc CommandFunction

	// CustomFunc will be called when a supervisor invokes the command
	// CustomCmd.
	CustomFunc CommandFunction

	// KillTaskFunc will be called when a supervisor invokes the command
	// KillTask.
	KillTaskFunc CommandFunction
}

// NewAgent returns a reference to an Agent object. The paremeter uri is the
// network address of the broker and exchange is the name of exchange that will
// be created.
func NewAgent(uri, exchange, name string) *Agent {
	a := &Agent{}
	a.s = rpcmq.NewServer(uri, "", exchange, "fanout")
	a.s.Parallel = 1
	a.status.Name = name
	return a
}

// Init initializes the Agent object. It establishes the connection with the
// broker, creating a channel and the exchange that will be used under the hood.
func (a *Agent) Init() error {
	a.s.TLSConfig = a.TLSConfig
	if err := a.s.Register("invoke", a.invoke); err != nil {
		return err
	}
	if err := a.s.Init(); err != nil {
		return err
	}
	a.status.Running = true
	return nil
}

// Shutdown shuts down the agent gracefully. Using this method will ensure that
// all requests sent by the supervisors to the agent will be received by the
// latter.
func (a *Agent) Shutdown() {
	a.s.Shutdown()
}

// RegisterTask adds a task to the list of tasks handled by the agent.
func (a *Agent) RegisterTask(id string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	a.status.Tasks = append(a.status.Tasks, id)
}

// RemoveTask removes a task from the list of tasks handled by the agent.
func (a *Agent) RemoveTask(id string) error {
	a.mu.Lock()
	defer a.mu.Unlock()

	idx := -1
	for i, t := range a.status.Tasks {
		if t == id {
			idx = i
			break
		}
	}
	if idx == -1 {
		return errors.New("task not found")
	}
	a.status.Tasks = append(a.status.Tasks[:idx], a.status.Tasks[idx+1:]...)
	return nil
}

func (a *Agent) invoke(id string, data []byte) ([]byte, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	var f CommandFunction
	running := a.status.Running
	cmd, aux := Command(data[0]), data[1:]
	switch {
	case cmd == GetStatus:
		f = a.getStatus
	case a.status.Name == string(aux) && cmd == SoftShutdown:
		f = a.SoftShutdownFunc
		running = false
	case a.status.Name == string(aux) && cmd == HardShutdown:
		f = a.HardShutdownFunc
		running = false
	case a.status.Name == string(aux) && cmd == Pause:
		f = a.PauseFunc
		running = false
	case a.status.Name == string(aux) && cmd == Resume:
		f = a.ResumeFunc
		running = true
	case a.status.Name == string(aux) && cmd == CustomCmd:
		f = a.CustomFunc
	case a.ownsTask(string(aux)) && cmd == KillTask:
		f = a.KillTaskFunc
	}
	if f == nil {
		// The command is not for this agent
		return nil, nil
	}
	b, err := f(data)
	if err != nil {
		return nil, err
	}
	a.status.Running = running
	out := fmt.Sprintf("%c%s", cmd, b)
	return []byte(out), nil
}

func (a *Agent) getStatus(data []byte) ([]byte, error) {
	info, err := getSystemInfo()
	if err != nil {
		return nil, err
	}
	a.status.Info = info
	return json.Marshal(a.status)
}

func (a *Agent) ownsTask(id string) bool {
	for _, t := range a.status.Tasks {
		if id == t {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package ingress

import (
	"context"
	"errors"
	"fmt"

	"knative.dev/net-kourier/pkg/config"
	envoy "knative.dev/net-kourier/pkg/envoy/server"
	"knative.dev/net-kourier/pkg/generator"
	"knative.dev/networking/pkg/apis/networking/v1alpha1"
	"knative.dev/networking/pkg/client/injection/reconciler/networking/v1alpha1/ingress"
	"knative.dev/networking/pkg/status"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/reconciler"
)

const conflictReason = "DomainConflict"

type Reconciler struct {
	xdsServer         *envoy.XdsServer
	caches            *generator.Caches
	statusManager     *status.Prober
	ingressTranslator *generator.IngressTranslator
	extAuthz          bool
}

var _ ingress.Interface = (*Reconciler)(nil)
var _ ingress.ReadOnlyInterface = (*Reconciler)(nil)
var _ ingress.Finalizer = (*Reconciler)(nil)
var _ ingress.ReadOnlyFinalizer = (*Reconciler)(nil)

func (r *Reconciler) ReconcileKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	ing.SetDefaults(ctx)
	before := ing.DeepCopy()

	if err := r.updateIngress(ctx, ing); errors.Is(err, generator.ErrDomainConflict) {
		// If we had an error due to a duplicated domain, we must mark the ingress as failed with a
		// custom status. We don't want to return an error in this case as we want to update its status.
		logging.FromContext(ctx).Info(err.Error())
		ing.Status.MarkLoadBalancerFailed(conflictReason, "Ingress rejected: "+err.Error())
		return nil
	} else if err != nil {
		return fmt.Errorf("failed to update ingress: %w", err)
	}

	if !ing.IsReady() || !isExpectedLoadBalancer(ing) {
		ready, err := r.statusManager.IsReady(ctx, before)
		if err != nil {
			return fmt.Errorf("failed to probe Ingress: %w", err)
		}
		if ready {
			external, internal := config.ServiceHostnames()
			ing.Status.MarkLoadBalancerReady(
				[]v1alpha1.LoadBalancerIngressStatus{{DomainInternal: external}},
				[]v1alpha1.LoadBalancerIngressStatus{{DomainInternal: internal}},
			)
			ing.Status.MarkNetworkConfigured()
		} else {
			ing.Status.MarkLoadBalancerNotReady()
		}
	}

	return nil
}

// isExpectedLoadBalancer verifies if expected Loadbalancer is set in status field.
func isExpectedLoadBalancer(ing *v1alpha1.Ingress) bool {
	external, internal := config.ServiceHostnames()
	if ing.Status.PublicLoadBalancer == nil || len(ing.Status.PublicLoadBalancer.Ingress) < 1 ||
		ing.Status.PublicLoadBalancer.Ingress[0].DomainInternal != external {
		return false
	}
	if ing.Status.PrivateLoadBalancer == nil || len(ing.Status.PrivateLoadBalancer.Ingress) < 1 ||
		ing.Status.PrivateLoadBalancer.Ingress[0].DomainInternal != internal {
		return false
	}
	return true
}

func (r *Reconciler) ObserveKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	ing.SetDefaults(ctx)

	if err := r.updateIngress(ctx, ing); errors.Is(err, generator.ErrDomainConflict) {
		// If we had an error due to a duplicated domain, just abort.
		logging.FromContext(ctx).Info(err.Error())
		return nil
	} else if err != nil {
		return fmt.Errorf("failed to update ingress: %w", err)
	}

	return nil
}

func (r *Reconciler) FinalizeKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	return r.ObserveFinalizeKind(ctx, ing)
}

func (r *Reconciler) ObserveFinalizeKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	logger := logging.FromContext(ctx)
	logger.Infof("Deleting Ingress")

	r.statusManager.CancelIngressProbing(ing)

	if err := r.caches.DeleteIngressInfo(ctx, ing.Name, ing.Namespace); err != nil {
		return err
	}

	return r.updateEnvoyConfig(ctx)
}

func (r *Reconciler) updateIngress(ctx context.Context, ingress *v1alpha1.Ingress) error {
	logger := logging.FromContext(ctx)
	logger.Infof("Updating Ingress")

	if err := generator.UpdateInfoForIngress(
		ctx, r.caches, ingress, r.ingressTranslator, r.extAuthz); err != nil {
		return err
	}

	return r.updateEnvoyConfig(ctx)
}

func (r *Reconciler) updateEnvoyConfig(ctx context.Context) error {
	logger := logging.FromContext(ctx)
	logger.Debugf("Preparing Envoy Snapshot")

	newSnapshot, err := r.caches.ToEnvoySnapshot(ctx)
	if err != nil {
		return err
	}

	return r.xdsServer.SetSnapshot(nodeID, newSnapshot)
}
<commit_msg>Mark network configured as soon as it is (#493)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package ingress

import (
	"context"
	"errors"
	"fmt"

	"knative.dev/net-kourier/pkg/config"
	envoy "knative.dev/net-kourier/pkg/envoy/server"
	"knative.dev/net-kourier/pkg/generator"
	"knative.dev/networking/pkg/apis/networking/v1alpha1"
	"knative.dev/networking/pkg/client/injection/reconciler/networking/v1alpha1/ingress"
	"knative.dev/networking/pkg/status"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/reconciler"
)

const conflictReason = "DomainConflict"

type Reconciler struct {
	xdsServer         *envoy.XdsServer
	caches            *generator.Caches
	statusManager     *status.Prober
	ingressTranslator *generator.IngressTranslator
	extAuthz          bool
}

var _ ingress.Interface = (*Reconciler)(nil)
var _ ingress.ReadOnlyInterface = (*Reconciler)(nil)
var _ ingress.Finalizer = (*Reconciler)(nil)
var _ ingress.ReadOnlyFinalizer = (*Reconciler)(nil)

func (r *Reconciler) ReconcileKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	ing.SetDefaults(ctx)
	before := ing.DeepCopy()

	if err := r.updateIngress(ctx, ing); errors.Is(err, generator.ErrDomainConflict) {
		// If we had an error due to a duplicated domain, we must mark the ingress as failed with a
		// custom status. We don't want to return an error in this case as we want to update its status.
		logging.FromContext(ctx).Info(err.Error())
		ing.Status.MarkLoadBalancerFailed(conflictReason, "Ingress rejected: "+err.Error())
		return nil
	} else if err != nil {
		return fmt.Errorf("failed to update ingress: %w", err)
	}

	ing.Status.MarkNetworkConfigured()
	if !ing.IsReady() || !isExpectedLoadBalancer(ing) {
		ready, err := r.statusManager.IsReady(ctx, before)
		if err != nil {
			return fmt.Errorf("failed to probe Ingress: %w", err)
		}
		if ready {
			external, internal := config.ServiceHostnames()
			ing.Status.MarkLoadBalancerReady(
				[]v1alpha1.LoadBalancerIngressStatus{{DomainInternal: external}},
				[]v1alpha1.LoadBalancerIngressStatus{{DomainInternal: internal}},
			)
		} else {
			ing.Status.MarkLoadBalancerNotReady()
		}
	}

	return nil
}

// isExpectedLoadBalancer verifies if expected Loadbalancer is set in status field.
func isExpectedLoadBalancer(ing *v1alpha1.Ingress) bool {
	external, internal := config.ServiceHostnames()
	if ing.Status.PublicLoadBalancer == nil || len(ing.Status.PublicLoadBalancer.Ingress) < 1 ||
		ing.Status.PublicLoadBalancer.Ingress[0].DomainInternal != external {
		return false
	}
	if ing.Status.PrivateLoadBalancer == nil || len(ing.Status.PrivateLoadBalancer.Ingress) < 1 ||
		ing.Status.PrivateLoadBalancer.Ingress[0].DomainInternal != internal {
		return false
	}
	return true
}

func (r *Reconciler) ObserveKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	ing.SetDefaults(ctx)

	if err := r.updateIngress(ctx, ing); errors.Is(err, generator.ErrDomainConflict) {
		// If we had an error due to a duplicated domain, just abort.
		logging.FromContext(ctx).Info(err.Error())
		return nil
	} else if err != nil {
		return fmt.Errorf("failed to update ingress: %w", err)
	}

	return nil
}

func (r *Reconciler) FinalizeKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	return r.ObserveFinalizeKind(ctx, ing)
}

func (r *Reconciler) ObserveFinalizeKind(ctx context.Context, ing *v1alpha1.Ingress) reconciler.Event {
	logger := logging.FromContext(ctx)
	logger.Infof("Deleting Ingress")

	r.statusManager.CancelIngressProbing(ing)

	if err := r.caches.DeleteIngressInfo(ctx, ing.Name, ing.Namespace); err != nil {
		return err
	}

	return r.updateEnvoyConfig(ctx)
}

func (r *Reconciler) updateIngress(ctx context.Context, ingress *v1alpha1.Ingress) error {
	logger := logging.FromContext(ctx)
	logger.Infof("Updating Ingress")

	if err := generator.UpdateInfoForIngress(
		ctx, r.caches, ingress, r.ingressTranslator, r.extAuthz); err != nil {
		return err
	}

	return r.updateEnvoyConfig(ctx)
}

func (r *Reconciler) updateEnvoyConfig(ctx context.Context) error {
	logger := logging.FromContext(ctx)
	logger.Debugf("Preparing Envoy Snapshot")

	newSnapshot, err := r.caches.ToEnvoySnapshot(ctx)
	if err != nil {
		return err
	}

	return r.xdsServer.SetSnapshot(nodeID, newSnapshot)
}
<|endoftext|>


<commit_before>// Package form provides an easy to use way to parse form values from an HTTP
// request into a struct
package form // import "vimagination.zapto.org/form"

import (
	"net/http"
	"reflect"
	"strings"
	"sync"
)

var interType = reflect.TypeOf((*formParser)(nil)).Elem()

type processorDetails struct {
	processor
	Post, Required bool
	Index          []int
}

type typeMap map[string]processorDetails

var (
	tmMu     sync.RWMutex
	typeMaps = make(map[reflect.Type]typeMap)
)

func getTypeMap(t reflect.Type) typeMap {
	tmMu.RLock()
	tm, ok := typeMaps[t]
	tmMu.RUnlock()
	if ok {
		return tm
	}
	tmMu.Lock()
	tm = createTypeMap(t)
	tmMu.Unlock()
	return tm
}

func basicTypeProcessor(t reflect.Type, tag reflect.StructTag) processor {
	switch t.Kind() {
	case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Int:
		return newInum(tag, t.Bits())
	case reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uint:
		return newUnum(tag, t.Bits())
	case reflect.Float32, reflect.Float64:
		return newFloat(tag, t.Bits())
	case reflect.String:
		return str{}
	case reflect.Bool:
		return boolean{}
	}
	return nil
}

func createTypeMap(t reflect.Type) typeMap {
	tm, ok := typeMaps[t]
	if ok {
		return tm
	}
	tm = make(typeMap)
	for i := 0; i < t.Len(); i++ {
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		name := f.Name
		var required, post bool
		if n := f.Tag.Get("form"); n == "-" {
			continue
		} else if n != "" {
			p := strings.IndexByte(n, ',')
			if p >= 0 {
				if p > 0 {
					name = n[:p]
				}
				rest := n[p:]
				required = strings.Contains(rest, ",required,") || strings.HasPrefix(rest, ",required")
				post = strings.Contains(rest, ",post,") || strings.HasPrefix(rest, ",post")
			} else {
				name = n
			}
		}
		var p processor
		if f.Type.Implements(interType) {
			p = inter(false)
		} else if reflect.PtrTo(f.Type).Implements(interType) {
			p = inter(true)
		} else if k := f.Type.Kind(); k == reflect.Slice || k == reflect.Ptr {
			et := f.Type.Elem()
			s := basicTypeProcessor(et, f.Tag)
			if s == nil {
				continue
			}
			if k == reflect.Slice {
				p = slice{
					processor: s,
					typ:       et,
				}
			} else {
				p = pointer{
					processor: s,
					typ:       et,
				}
			}
		} else if k == reflect.Struct && f.Anonymous {
			for n, p := range createTypeMap(f.Type) {
				if _, ok := tm[n]; !ok {
					tm[n] = processorDetails{
						processor: p.processor,
						Required:  p.Required,
						Post:      p.Post,
						Index:     append(append(make([]int, 0, len(p.Index)+1), i), p.Index...),
					}
				}
			}
			continue
		} else {
			p = basicTypeProcessor(f.Type, f.Tag)
			if p == nil {
				continue
			}
		}
		tm[name] = processorDetails{
			processor: p,
			Required:  required,
			Post:      post,
			Index:     []int{i},
		}
	}
	typeMaps[t] = tm
	return tm
}

// Process parses the form data from the request into the passed value, which
// must be a pointer to a struct.
//
// Form keys are assumed to be the field names unless a 'form' tag is provided
// with an alternate name, for example, in the following struct, the int is
// parse with key 'A' and the bool is parsed with key 'C'.
//
// type Example struct {
//	A int
//	B bool `form:"C"`
// }
//
// Two options can be added to the form tag to modify the processing. The
// 'post' option forces the processer to parse a value from the PostForm field
// of the Request, and the 'required' option will have an error thrown if the
// key in not set.
//
// Number types can also have minimums and maximums checked during processing
// by setting the min and max tags accordingly.
//
// Lastly, a custom data processor can be specified by attaching a method to
// the field type with the following specification:
//
// ParseForm([]string) error
func Process(r *http.Request, fv interface{}) error {
	v := reflect.ValueOf(fv)
	if v.Kind() != reflect.Ptr {
		return ErrNeedPointer
	}
	for v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	if v.Kind() != reflect.Struct {
		return ErrNeedStruct
	}
	tm := getTypeMap(v.Type())
	if err := r.ParseForm(); err != nil {
		return err
	}
	var errors Errors
	for key, pd := range tm {
		var (
			val []string
			ok  bool
		)
		if pd.Post {
			val, ok = r.PostForm[key]
		} else {
			val, ok = r.Form[key]
		}
		if ok {
			if err := pd.processor.process(v.FieldByIndex(pd.Index), val); err != nil {
				errors = append(errors, ErrProcessingFailed{
					Key: key,
					Err: err,
				})
			}
		} else if pd.Required {
			errors = append(errors, ErrRequiredMissing(key))
		}
	}
	if len(errors) > 0 {
		return errors
	}
	return nil
}
<commit_msg>corrected call to get number of struct fields<commit_after>// Package form provides an easy to use way to parse form values from an HTTP
// request into a struct
package form // import "vimagination.zapto.org/form"

import (
	"net/http"
	"reflect"
	"strings"
	"sync"
)

var interType = reflect.TypeOf((*formParser)(nil)).Elem()

type processorDetails struct {
	processor
	Post, Required bool
	Index          []int
}

type typeMap map[string]processorDetails

var (
	tmMu     sync.RWMutex
	typeMaps = make(map[reflect.Type]typeMap)
)

func getTypeMap(t reflect.Type) typeMap {
	tmMu.RLock()
	tm, ok := typeMaps[t]
	tmMu.RUnlock()
	if ok {
		return tm
	}
	tmMu.Lock()
	tm = createTypeMap(t)
	tmMu.Unlock()
	return tm
}

func basicTypeProcessor(t reflect.Type, tag reflect.StructTag) processor {
	switch t.Kind() {
	case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Int:
		return newInum(tag, t.Bits())
	case reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uint:
		return newUnum(tag, t.Bits())
	case reflect.Float32, reflect.Float64:
		return newFloat(tag, t.Bits())
	case reflect.String:
		return str{}
	case reflect.Bool:
		return boolean{}
	}
	return nil
}

func createTypeMap(t reflect.Type) typeMap {
	tm, ok := typeMaps[t]
	if ok {
		return tm
	}
	tm = make(typeMap)
	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		name := f.Name
		var required, post bool
		if n := f.Tag.Get("form"); n == "-" {
			continue
		} else if n != "" {
			p := strings.IndexByte(n, ',')
			if p >= 0 {
				if p > 0 {
					name = n[:p]
				}
				rest := n[p:]
				required = strings.Contains(rest, ",required,") || strings.HasPrefix(rest, ",required")
				post = strings.Contains(rest, ",post,") || strings.HasPrefix(rest, ",post")
			} else {
				name = n
			}
		}
		var p processor
		if f.Type.Implements(interType) {
			p = inter(false)
		} else if reflect.PtrTo(f.Type).Implements(interType) {
			p = inter(true)
		} else if k := f.Type.Kind(); k == reflect.Slice || k == reflect.Ptr {
			et := f.Type.Elem()
			s := basicTypeProcessor(et, f.Tag)
			if s == nil {
				continue
			}
			if k == reflect.Slice {
				p = slice{
					processor: s,
					typ:       et,
				}
			} else {
				p = pointer{
					processor: s,
					typ:       et,
				}
			}
		} else if k == reflect.Struct && f.Anonymous {
			for n, p := range createTypeMap(f.Type) {
				if _, ok := tm[n]; !ok {
					tm[n] = processorDetails{
						processor: p.processor,
						Required:  p.Required,
						Post:      p.Post,
						Index:     append(append(make([]int, 0, len(p.Index)+1), i), p.Index...),
					}
				}
			}
			continue
		} else {
			p = basicTypeProcessor(f.Type, f.Tag)
			if p == nil {
				continue
			}
		}
		tm[name] = processorDetails{
			processor: p,
			Required:  required,
			Post:      post,
			Index:     []int{i},
		}
	}
	typeMaps[t] = tm
	return tm
}

// Process parses the form data from the request into the passed value, which
// must be a pointer to a struct.
//
// Form keys are assumed to be the field names unless a 'form' tag is provided
// with an alternate name, for example, in the following struct, the int is
// parse with key 'A' and the bool is parsed with key 'C'.
//
// type Example struct {
//	A int
//	B bool `form:"C"`
// }
//
// Two options can be added to the form tag to modify the processing. The
// 'post' option forces the processer to parse a value from the PostForm field
// of the Request, and the 'required' option will have an error thrown if the
// key in not set.
//
// Number types can also have minimums and maximums checked during processing
// by setting the min and max tags accordingly.
//
// Lastly, a custom data processor can be specified by attaching a method to
// the field type with the following specification:
//
// ParseForm([]string) error
func Process(r *http.Request, fv interface{}) error {
	v := reflect.ValueOf(fv)
	if v.Kind() != reflect.Ptr {
		return ErrNeedPointer
	}
	for v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	if v.Kind() != reflect.Struct {
		return ErrNeedStruct
	}
	tm := getTypeMap(v.Type())
	if err := r.ParseForm(); err != nil {
		return err
	}
	var errors Errors
	for key, pd := range tm {
		var (
			val []string
			ok  bool
		)
		if pd.Post {
			val, ok = r.PostForm[key]
		} else {
			val, ok = r.Form[key]
		}
		if ok {
			if err := pd.processor.process(v.FieldByIndex(pd.Index), val); err != nil {
				errors = append(errors, ErrProcessingFailed{
					Key: key,
					Err: err,
				})
			}
		} else if pd.Required {
			errors = append(errors, ErrRequiredMissing(key))
		}
	}
	if len(errors) > 0 {
		return errors
	}
	return nil
}
<|endoftext|>


<commit_before>package extensions

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/url"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"github.com/telepresenceio/telepresence/rpc/v2/common"
	"github.com/telepresenceio/telepresence/rpc/v2/systema"
	"github.com/telepresenceio/telepresence/v2/pkg/client"
	"github.com/telepresenceio/telepresence/v2/pkg/client/cli/cliutil"
)

type systemaCredentials string

// GetRequestMetadata implements credentials.PerRPCCredentials.
func (c systemaCredentials) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) {
	md := map[string]string{
		"Authorization": "Bearer " + string(c),
	}
	return md, nil
}

// RequireTransportSecurity implements credentials.PerRPCCredentials.
func (c systemaCredentials) RequireTransportSecurity() bool {
	return true
}

func systemaGetPreferredAgentImageName(ctx context.Context, urlStr string) (string, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", err
	}

	accessToken, err := cliutil.GetCloudAccessToken(ctx, true)
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: login error: %w", err)
	}
	creds := systemaCredentials(accessToken)

	conn, err := grpc.DialContext(ctx,
		(&url.URL{Scheme: "dns", Path: "/" + u.Host}).String(), // https://github.com/grpc/grpc/blob/master/doc/naming.md
		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ServerName: u.Hostname()})),
		grpc.WithPerRPCCredentials(creds))
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: dial error: %w", err)
	}
	defer conn.Close()

	systemaClient := systema.NewSystemACRUDClient(conn)

	resp, err := systemaClient.PreferredAgent(ctx, &common.VersionInfo{
		ApiVersion: client.APIVersion,
		Version:    client.Version(),
	})
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: gRPC error: %w", err)
	}

	return resp.GetImageName(), nil
}
<commit_msg>cli: Switch over to using API keys instead of access tokens<commit_after>package extensions

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/url"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"github.com/telepresenceio/telepresence/rpc/v2/common"
	"github.com/telepresenceio/telepresence/rpc/v2/systema"
	"github.com/telepresenceio/telepresence/v2/pkg/client"
	"github.com/telepresenceio/telepresence/v2/pkg/client/cli/cliutil"
)

type systemaCredentials string

// GetRequestMetadata implements credentials.PerRPCCredentials.
func (c systemaCredentials) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) {
	md := map[string]string{
		"X-Ambassador-Api-Key": string(c),
	}
	return md, nil
}

// RequireTransportSecurity implements credentials.PerRPCCredentials.
func (c systemaCredentials) RequireTransportSecurity() bool {
	return true
}

func systemaGetPreferredAgentImageName(ctx context.Context, urlStr string) (string, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", err
	}

	apikey, err := cliutil.GetCloudAPIKey(ctx, "laptop", true)
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: login error: %w", err)
	}
	creds := systemaCredentials(apikey)

	conn, err := grpc.DialContext(ctx,
		(&url.URL{Scheme: "dns", Path: "/" + u.Host}).String(), // https://github.com/grpc/grpc/blob/master/doc/naming.md
		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ServerName: u.Hostname()})),
		grpc.WithPerRPCCredentials(creds))
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: dial error: %w", err)
	}
	defer conn.Close()

	systemaClient := systema.NewSystemACRUDClient(conn)

	resp, err := systemaClient.PreferredAgent(ctx, &common.VersionInfo{
		ApiVersion: client.APIVersion,
		Version:    client.Version(),
	})
	if err != nil {
		return "", fmt.Errorf("getting Ambassador Cloud preferred agent image: gRPC error: %w", err)
	}

	return resp.GetImageName(), nil
}
<|endoftext|>


<commit_before>package convertutil

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/grokify/gotilla/image/imageutil"
	"github.com/grokify/gotilla/io/ioutilmore"
)

type CopyType int

const (
	PDFFormat    CopyType = iota // 0 convert cli value
	KindleFormat                 // 1 convert cli value
)

// ReformatImages converts images in one dir to another using default
// formats for Kindle and PDF.
func ReformatImages(baseSrcDir, baseOutDir string, copyType CopyType) error {
	var err error
	baseSrcDir, err = filepath.Abs(strings.TrimSpace(baseSrcDir))
	if err != nil {
		return err
	}
	baseOutDir, err = filepath.Abs(strings.TrimSpace(baseOutDir))
	if err != nil {
		return err
	}
	return reformatImagesSubdir(baseSrcDir, baseOutDir, "", copyType)
}

func reformatImagesSubdir(baseSrcDir, baseOutDir, dirPart string, copyType CopyType) error {
	thisSrcDir := baseSrcDir
	thisOutDir := baseOutDir
	dirPart = strings.TrimSpace(dirPart)
	if len(dirPart) > 0 {
		thisSrcDir = filepath.Join(thisSrcDir, dirPart)
		thisOutDir = filepath.Join(thisOutDir, dirPart)

		isDir, err := ioutilmore.IsDir(thisSrcDir)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("Need Dir [%s]", thisSrcDir)
		}
	}

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

	sdirs, files, err := ioutilmore.ReadDirSplit(thisSrcDir, false)
	if err != nil {
		return err
	}
	for _, file := range files {
		thisSrcFile := filepath.Join(thisSrcDir, file.Name())
		thisOutFile := filepath.Join(thisOutDir, file.Name())
		if !imageutil.IsImageExt(thisSrcFile) {
			continue
		}
		switch copyType {
		case PDFFormat:
			_, stderr, err := ConvertToPDF(thisSrcFile, thisOutFile)
			err = CheckError(err, stderr)
			if err != nil {
				return err
			}
		case KindleFormat:
			_, stderr, err := ConvertToKindle(thisSrcFile, thisOutFile)
			err = CheckError(err, stderr)
			if err != nil {
				return err
			}
		}
	}
	for _, sdir := range sdirs {
		subDir := sdir.Name()
		if len(dirPart) > 0 {
			subDir = filepath.Join(dirPart, subDir)
		}
		err := reformatImagesSubdir(baseSrcDir, baseOutDir, subDir, copyType)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>add image/convertutil constants<commit_after>package convertutil

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/grokify/gotilla/image/imageutil"
	"github.com/grokify/gotilla/io/ioutilmore"
)

const (
	PDFSpecs    = "1950pxw300dpi"
	KindleSpecs = "600pxw300dpi"
)

type CopyType int

const (
	PDFFormat    CopyType = iota // 0 convert cli value
	KindleFormat                 // 1 convert cli value
)

// ReformatImages converts images in one dir to another using default
// formats for Kindle and PDF.
func ReformatImages(baseSrcDir, baseOutDir string, copyType CopyType) error {
	var err error
	baseSrcDir, err = filepath.Abs(strings.TrimSpace(baseSrcDir))
	if err != nil {
		return err
	}
	baseOutDir, err = filepath.Abs(strings.TrimSpace(baseOutDir))
	if err != nil {
		return err
	}
	return reformatImagesSubdir(baseSrcDir, baseOutDir, "", copyType)
}

func reformatImagesSubdir(baseSrcDir, baseOutDir, dirPart string, copyType CopyType) error {
	thisSrcDir := baseSrcDir
	thisOutDir := baseOutDir
	dirPart = strings.TrimSpace(dirPart)
	if len(dirPart) > 0 {
		thisSrcDir = filepath.Join(thisSrcDir, dirPart)
		thisOutDir = filepath.Join(thisOutDir, dirPart)

		isDir, err := ioutilmore.IsDir(thisSrcDir)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("Need Dir [%s]", thisSrcDir)
		}
	}

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

	sdirs, files, err := ioutilmore.ReadDirSplit(thisSrcDir, false)
	if err != nil {
		return err
	}
	for _, file := range files {
		thisSrcFile := filepath.Join(thisSrcDir, file.Name())
		thisOutFile := filepath.Join(thisOutDir, file.Name())
		if !imageutil.IsImageExt(thisSrcFile) {
			continue
		}
		switch copyType {
		case PDFFormat:
			_, stderr, err := ConvertToPDF(thisSrcFile, thisOutFile)
			err = CheckError(err, stderr)
			if err != nil {
				return err
			}
		case KindleFormat:
			_, stderr, err := ConvertToKindle(thisSrcFile, thisOutFile)
			err = CheckError(err, stderr)
			if err != nil {
				return err
			}
		}
	}
	for _, sdir := range sdirs {
		subDir := sdir.Name()
		if len(dirPart) > 0 {
			subDir = filepath.Join(dirPart, subDir)
		}
		err := reformatImagesSubdir(baseSrcDir, baseOutDir, subDir, copyType)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package connector

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/datawire/dlib/dexec"
	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	rpc "github.com/datawire/telepresence2/rpc/v2/connector"
	"github.com/datawire/telepresence2/rpc/v2/manager"
	"github.com/datawire/telepresence2/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/datawire/telepresence2/v2/pkg/client/actions"
)

func (s *service) interceptStatus() (rpc.InterceptError, string) {
	ie := rpc.InterceptError_UNSPECIFIED
	msg := ""
	switch {
	case s.cluster == nil:
		ie = rpc.InterceptError_NO_CONNECTION
	case s.trafficMgr == nil:
		ie = rpc.InterceptError_NO_TRAFFIC_MANAGER
	case s.trafficMgr.managerClient == nil:
		if s.trafficMgr.managerErr != nil {
			ie = rpc.InterceptError_TRAFFIC_MANAGER_ERROR
			msg = s.trafficMgr.managerErr.Error()
		} else {
			ie = rpc.InterceptError_TRAFFIC_MANAGER_CONNECTING
		}
	}
	return ie, msg
}

type portForward struct {
	ManagerPort int32
	TargetHost  string
	TargetPort  int32
}

func (tm *trafficManager) workerPortForwardIntercepts(ctx context.Context) error {
	// Don't use a dgroup.Group because:
	//  1. we don't actually care about tracking errors (we just always retry) or any of
	//     dgroup's other functionality
	//  2. because goroutines may churn as intercepts are created and deleted, tracking all of
	//     their exit statuses is just a memory leak
	//  3. because we want a per-worker cancel, we'd have to implement our own Context
	//     management on top anyway, so dgroup wouldn't actually save us any complexity.
	var wg sync.WaitGroup

	livePortForwards := make(map[portForward]context.CancelFunc)

	backoff := 100 * time.Millisecond
	for ctx.Err() == nil {
		stream, err := tm.managerClient.WatchIntercepts(ctx, tm.session())
		for err == nil {
			var snapshot *manager.InterceptInfoSnapshot
			snapshot, err = stream.Recv()
			if err != nil {
				break
			}
			snapshotPortForwards := make(map[portForward]struct{})
			for _, intercept := range snapshot.Intercepts {
				pf := portForward{
					ManagerPort: intercept.ManagerPort,
					TargetHost:  intercept.Spec.TargetHost,
					TargetPort:  intercept.Spec.TargetPort,
				}
				snapshotPortForwards[pf] = struct{}{}
				if _, isLive := livePortForwards[pf]; !isLive {
					pfCtx, pfCancel := context.WithCancel(ctx)
					pfCtx = dgroup.WithGoroutineName(pfCtx,
						fmt.Sprintf("/%d:%s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort))
					livePortForwards[pf] = pfCancel

					wg.Add(1)
					go tm.workerPortForwardIntercept(pfCtx, pf)
				}
			}
			for pf, cancel := range livePortForwards {
				if _, isWanted := snapshotPortForwards[pf]; !isWanted {
					dlog.Infof(ctx, "Terminating port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)
					cancel()
					delete(livePortForwards, pf)
				}
			}
		}

		if ctx.Err() == nil {
			dlog.Errorf(ctx, "communicating with manager: %v", err)
			dtime.SleepWithContext(ctx, backoff)
			backoff *= 2
			if backoff > 3*time.Second {
				backoff = 3 * time.Second
			}
		}
	}

	wg.Wait()

	return nil
}

// addIntercept adds one intercept
func (tm *trafficManager) addIntercept(c, longLived context.Context, ir *manager.CreateInterceptRequest) (*rpc.InterceptResult, error) {
	spec := ir.InterceptSpec
	intercepts, err := actions.ListMyIntercepts(c, tm.managerClient, tm.session().SessionId)
	if err != nil {
		return nil, err
	}
	for _, iCept := range intercepts {
		if iCept.Spec.Name == spec.Name {
			return &rpc.InterceptResult{
				Error:     rpc.InterceptError_ALREADY_EXISTS,
				ErrorText: iCept.Spec.Name,
			}, nil
		}
		if iCept.Spec.TargetPort == spec.TargetPort && iCept.Spec.TargetHost == spec.TargetHost {
			return &rpc.InterceptResult{
				InterceptInfo: &manager.InterceptInfo{Spec: spec},
				Error:         rpc.InterceptError_LOCAL_TARGET_IN_USE,
				ErrorText:     iCept.Spec.Name,
			}, nil
		}
	}

	spec.Client = tm.userAndHost
	if spec.Mechanism == "" {
		spec.Mechanism = "tcp"
	}

	agentName := spec.Agent
	if spec.Name == "" {
		spec.Name = agentName
	}

	hasSpecMechanism := func(a *manager.AgentInfo) bool {
		for _, mech := range a.Mechanisms {
			if spec.Mechanism == mech.Name {
				return true
			}
		}
		return false
	}

	var found *manager.AgentInfo
	if ags, _ := actions.ListAllAgents(c, tm.managerClient, tm.session().SessionId); ags != nil {
		for _, ag := range ags {
			if !(ag.Name == spec.Agent && hasSpecMechanism(ag)) {
				continue
			}
			if found == nil {
				found = ag
				continue
			}
			if ag.Version == found.Version && ag.Product == found.Product {
				// Just hostname that differs, this is a replica
				continue
			}
			txt, _ := json.Marshal([]*manager.AgentInfo{ag, found})
			return &rpc.InterceptResult{
				InterceptInfo: nil,
				Error:         rpc.InterceptError_AMBIGUOUS_MATCH,
				ErrorText:     string(txt),
			}, nil
		}
	}

	var result *rpc.InterceptResult
	if found == nil {
		if result = tm.addAgent(c, agentName, agentImageName(c, tm.env)); result.Error != rpc.InterceptError_UNSPECIFIED {
			return result, nil
		}
	} else {
		result = &rpc.InterceptResult{
			Environment: found.Environment,
		}
	}

	ir.Session = tm.session()
	dlog.Debugf(c, "creating intercept %s", ir.InterceptSpec.Name)
	ii, err := tm.managerClient.CreateIntercept(c, ir)
	if err != nil {
		dlog.Debugf(c, "manager responded to CreateIntercept with error %v", err)
		return &rpc.InterceptResult{Error: rpc.InterceptError_TRAFFIC_MANAGER_ERROR, ErrorText: err.Error()}, nil
	}
	dlog.Debugf(c, "created intercept %s", ii.Spec.Name)

	ii, err = tm.waitForActiveIntercept(c, ii.Id)
	if err != nil {
		_ = tm.removeIntercept(c, spec.Name)
		return &rpc.InterceptResult{Error: rpc.InterceptError_FAILED_TO_ESTABLISH, ErrorText: err.Error()}, nil
	}
	result.InterceptInfo = ii

	return result, nil
}

func (tm *trafficManager) addAgent(c context.Context, agentName, agentImageName string) *rpc.InterceptResult {
	if err := tm.installer.ensureAgent(c, agentName, "", agentImageName); err != nil {
		if err == agentNotFound {
			return &rpc.InterceptResult{
				Error:     rpc.InterceptError_NOT_FOUND,
				ErrorText: agentName,
			}
		}
		dlog.Error(c, err)
		return &rpc.InterceptResult{
			Error:     rpc.InterceptError_FAILED_TO_ESTABLISH,
			ErrorText: err.Error(),
		}
	}

	dlog.Infof(c, "waiting for agent for deployment %q", agentName)
	agent, err := tm.waitForAgent(c, agentName)
	if err != nil {
		dlog.Error(c, err)
		return &rpc.InterceptResult{
			Error:     rpc.InterceptError_FAILED_TO_ESTABLISH,
			ErrorText: err.Error(),
		}
	}
	dlog.Infof(c, "agent created for deployment %q", agentName)
	return &rpc.InterceptResult{
		Error:       rpc.InterceptError_UNSPECIFIED,
		Environment: agent.Environment,
	}
}

func (tm *trafficManager) waitForActiveIntercept(ctx context.Context, id string) (*manager.InterceptInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	dlog.Debugf(ctx, "waiting for intercept id=%q to become active", id)
	stream, err := tm.managerClient.WatchIntercepts(ctx, tm.session())
	for err != nil {
		return nil, fmt.Errorf("waiting for intercept id=%q to become active: %w", id, err)
	}
	for {
		snapshot, err := stream.Recv()
		if err != nil {
			return nil, fmt.Errorf("waiting for intercept id=%q to become active: %w", id, err)
		}

		var intercept *manager.InterceptInfo
		for _, straw := range snapshot.Intercepts {
			if straw.Id == id {
				intercept = straw
				break
			}
		}

		switch {
		case intercept == nil:
			dlog.Debugf(ctx, "wait status: intercept id=%q does not yet exist", id)
		case intercept.Disposition == manager.InterceptDispositionType_WAITING:
			dlog.Debugf(ctx, "wait status: intercept id=%q is still WAITING", id)
		default:
			dlog.Debugf(ctx, "wait status: intercept id=%q is no longer WAITING; is now %v", id, intercept.Disposition)
			if intercept.Disposition != manager.InterceptDispositionType_ACTIVE {
				return nil, errors.Errorf("intercept in error state %v: %v", intercept.Disposition, intercept.Message)
			}
			return intercept, nil
		}
	}
}

func (tm *trafficManager) waitForAgent(ctx context.Context, name string) (*manager.AgentInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 120*time.Second) // installing a new agent can take some time
	defer cancel()

	stream, err := tm.managerClient.WatchAgents(ctx, tm.session())
	if err != nil {
		return nil, fmt.Errorf("waiting for agent %q to be present: %q", name, err)
	}
	for {
		snapshot, err := stream.Recv()
		if err != nil {
			return nil, fmt.Errorf("waiting for agent %q to be present: %q", name, err)
		}

		var agentList []*manager.AgentInfo
		for _, agent := range snapshot.Agents {
			if agent.Name == name {
				agentList = append(agentList, agent)
			}
		}
		if managerutil.AgentsAreCompatible(agentList) {
			return agentList[0], nil
		}
	}
}

func (tm *trafficManager) workerPortForwardIntercept(ctx context.Context, pf portForward) {
	dlog.Infof(ctx, "Initiating port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)

	sshArgs := []string{
		"ssh",

		"-F", "none", // don't load the user's config file

		// connection settings
		"-C", // compression
		"-oConnectTimeout=10",
		"-oStrictHostKeyChecking=no",     // don't bother checking the host key...
		"-oUserKnownHostsFile=/dev/null", // and since we're not checking it, don't bother remembering it either

		// port-forward settings
		"-N", // no remote command; just connect and forward ports
		"-oExitOnForwardFailure=yes",
		"-R", fmt.Sprintf("%d:%s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort), // port to forward

		// where to connect to
		"-p", strconv.Itoa(int(tm.sshPort)),
		"telepresence@localhost",
	}

	backoff := 100 * time.Millisecond
	for ctx.Err() == nil {
		start := time.Now()
		err := dexec.CommandContext(ctx, sshArgs[0], sshArgs[1:]...).Run()
		lifetime := time.Since(start)

		if ctx.Err() == nil {
			if err == nil {
				err = errors.New("process terminated unexpectedly")
			}
			dlog.Errorf(ctx, "communicating with manager: %v", err)
			if lifetime >= 20*time.Second {
				backoff = 100 * time.Millisecond
				dtime.SleepWithContext(ctx, backoff)
			} else {
				dtime.SleepWithContext(ctx, backoff)
				backoff *= 2
				if backoff > 3*time.Second {
					backoff = 3 * time.Second
				}
			}
		}
	}

	dlog.Infof(ctx, "Terminated port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)
}

// removeIntercept removes one intercept by name
func (tm *trafficManager) removeIntercept(c context.Context, name string) error {
	dlog.Debugf(c, "telling manager to remove intercept %s", name)
	_, err := tm.managerClient.RemoveIntercept(c, &manager.RemoveInterceptRequest2{
		Session: tm.session(),
		Name:    name,
	})
	return err
}

// clearIntercepts removes all intercepts
func (tm *trafficManager) clearIntercepts(c context.Context) error {
	intercepts, _ := actions.ListMyIntercepts(c, tm.managerClient, tm.session().SessionId)
	for _, cept := range intercepts {
		err := tm.removeIntercept(c, cept.Spec.Name)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Don't create port forwards for intercepts that aren't active yet.<commit_after>package connector

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/datawire/dlib/dexec"
	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	rpc "github.com/datawire/telepresence2/rpc/v2/connector"
	"github.com/datawire/telepresence2/rpc/v2/manager"
	"github.com/datawire/telepresence2/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/datawire/telepresence2/v2/pkg/client/actions"
)

func (s *service) interceptStatus() (rpc.InterceptError, string) {
	ie := rpc.InterceptError_UNSPECIFIED
	msg := ""
	switch {
	case s.cluster == nil:
		ie = rpc.InterceptError_NO_CONNECTION
	case s.trafficMgr == nil:
		ie = rpc.InterceptError_NO_TRAFFIC_MANAGER
	case s.trafficMgr.managerClient == nil:
		if s.trafficMgr.managerErr != nil {
			ie = rpc.InterceptError_TRAFFIC_MANAGER_ERROR
			msg = s.trafficMgr.managerErr.Error()
		} else {
			ie = rpc.InterceptError_TRAFFIC_MANAGER_CONNECTING
		}
	}
	return ie, msg
}

type portForward struct {
	ManagerPort int32
	TargetHost  string
	TargetPort  int32
}

func (tm *trafficManager) workerPortForwardIntercepts(ctx context.Context) error {
	// Don't use a dgroup.Group because:
	//  1. we don't actually care about tracking errors (we just always retry) or any of
	//     dgroup's other functionality
	//  2. because goroutines may churn as intercepts are created and deleted, tracking all of
	//     their exit statuses is just a memory leak
	//  3. because we want a per-worker cancel, we'd have to implement our own Context
	//     management on top anyway, so dgroup wouldn't actually save us any complexity.
	var wg sync.WaitGroup

	livePortForwards := make(map[portForward]context.CancelFunc)

	backoff := 100 * time.Millisecond
	for ctx.Err() == nil {
		stream, err := tm.managerClient.WatchIntercepts(ctx, tm.session())
		for err == nil {
			var snapshot *manager.InterceptInfoSnapshot
			snapshot, err = stream.Recv()
			if err != nil {
				break
			}
			snapshotPortForwards := make(map[portForward]struct{})
			for _, intercept := range snapshot.Intercepts {
				if intercept.Disposition != manager.InterceptDispositionType_ACTIVE {
					continue
				}
				pf := portForward{
					ManagerPort: intercept.ManagerPort,
					TargetHost:  intercept.Spec.TargetHost,
					TargetPort:  intercept.Spec.TargetPort,
				}
				snapshotPortForwards[pf] = struct{}{}
				if _, isLive := livePortForwards[pf]; !isLive {
					pfCtx, pfCancel := context.WithCancel(ctx)
					pfCtx = dgroup.WithGoroutineName(pfCtx,
						fmt.Sprintf("/%d:%s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort))
					livePortForwards[pf] = pfCancel

					wg.Add(1)
					go tm.workerPortForwardIntercept(pfCtx, pf)
				}
			}
			for pf, cancel := range livePortForwards {
				if _, isWanted := snapshotPortForwards[pf]; !isWanted {
					dlog.Infof(ctx, "Terminating port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)
					cancel()
					delete(livePortForwards, pf)
				}
			}
		}

		if ctx.Err() == nil {
			dlog.Errorf(ctx, "communicating with manager: %v", err)
			dtime.SleepWithContext(ctx, backoff)
			backoff *= 2
			if backoff > 3*time.Second {
				backoff = 3 * time.Second
			}
		}
	}

	wg.Wait()

	return nil
}

// addIntercept adds one intercept
func (tm *trafficManager) addIntercept(c, longLived context.Context, ir *manager.CreateInterceptRequest) (*rpc.InterceptResult, error) {
	spec := ir.InterceptSpec
	intercepts, err := actions.ListMyIntercepts(c, tm.managerClient, tm.session().SessionId)
	if err != nil {
		return nil, err
	}
	for _, iCept := range intercepts {
		if iCept.Spec.Name == spec.Name {
			return &rpc.InterceptResult{
				Error:     rpc.InterceptError_ALREADY_EXISTS,
				ErrorText: iCept.Spec.Name,
			}, nil
		}
		if iCept.Spec.TargetPort == spec.TargetPort && iCept.Spec.TargetHost == spec.TargetHost {
			return &rpc.InterceptResult{
				InterceptInfo: &manager.InterceptInfo{Spec: spec},
				Error:         rpc.InterceptError_LOCAL_TARGET_IN_USE,
				ErrorText:     iCept.Spec.Name,
			}, nil
		}
	}

	spec.Client = tm.userAndHost
	if spec.Mechanism == "" {
		spec.Mechanism = "tcp"
	}

	agentName := spec.Agent
	if spec.Name == "" {
		spec.Name = agentName
	}

	hasSpecMechanism := func(a *manager.AgentInfo) bool {
		for _, mech := range a.Mechanisms {
			if spec.Mechanism == mech.Name {
				return true
			}
		}
		return false
	}

	var found *manager.AgentInfo
	if ags, _ := actions.ListAllAgents(c, tm.managerClient, tm.session().SessionId); ags != nil {
		for _, ag := range ags {
			if !(ag.Name == spec.Agent && hasSpecMechanism(ag)) {
				continue
			}
			if found == nil {
				found = ag
				continue
			}
			if ag.Version == found.Version && ag.Product == found.Product {
				// Just hostname that differs, this is a replica
				continue
			}
			txt, _ := json.Marshal([]*manager.AgentInfo{ag, found})
			return &rpc.InterceptResult{
				InterceptInfo: nil,
				Error:         rpc.InterceptError_AMBIGUOUS_MATCH,
				ErrorText:     string(txt),
			}, nil
		}
	}

	var result *rpc.InterceptResult
	if found == nil {
		if result = tm.addAgent(c, agentName, agentImageName(c, tm.env)); result.Error != rpc.InterceptError_UNSPECIFIED {
			return result, nil
		}
	} else {
		result = &rpc.InterceptResult{
			Environment: found.Environment,
		}
	}

	ir.Session = tm.session()
	dlog.Debugf(c, "creating intercept %s", ir.InterceptSpec.Name)
	ii, err := tm.managerClient.CreateIntercept(c, ir)
	if err != nil {
		dlog.Debugf(c, "manager responded to CreateIntercept with error %v", err)
		return &rpc.InterceptResult{Error: rpc.InterceptError_TRAFFIC_MANAGER_ERROR, ErrorText: err.Error()}, nil
	}
	dlog.Debugf(c, "created intercept %s", ii.Spec.Name)

	ii, err = tm.waitForActiveIntercept(c, ii.Id)
	if err != nil {
		_ = tm.removeIntercept(c, spec.Name)
		return &rpc.InterceptResult{Error: rpc.InterceptError_FAILED_TO_ESTABLISH, ErrorText: err.Error()}, nil
	}
	result.InterceptInfo = ii

	return result, nil
}

func (tm *trafficManager) addAgent(c context.Context, agentName, agentImageName string) *rpc.InterceptResult {
	if err := tm.installer.ensureAgent(c, agentName, "", agentImageName); err != nil {
		if err == agentNotFound {
			return &rpc.InterceptResult{
				Error:     rpc.InterceptError_NOT_FOUND,
				ErrorText: agentName,
			}
		}
		dlog.Error(c, err)
		return &rpc.InterceptResult{
			Error:     rpc.InterceptError_FAILED_TO_ESTABLISH,
			ErrorText: err.Error(),
		}
	}

	dlog.Infof(c, "waiting for agent for deployment %q", agentName)
	agent, err := tm.waitForAgent(c, agentName)
	if err != nil {
		dlog.Error(c, err)
		return &rpc.InterceptResult{
			Error:     rpc.InterceptError_FAILED_TO_ESTABLISH,
			ErrorText: err.Error(),
		}
	}
	dlog.Infof(c, "agent created for deployment %q", agentName)
	return &rpc.InterceptResult{
		Error:       rpc.InterceptError_UNSPECIFIED,
		Environment: agent.Environment,
	}
}

func (tm *trafficManager) waitForActiveIntercept(ctx context.Context, id string) (*manager.InterceptInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	dlog.Debugf(ctx, "waiting for intercept id=%q to become active", id)
	stream, err := tm.managerClient.WatchIntercepts(ctx, tm.session())
	for err != nil {
		return nil, fmt.Errorf("waiting for intercept id=%q to become active: %w", id, err)
	}
	for {
		snapshot, err := stream.Recv()
		if err != nil {
			return nil, fmt.Errorf("waiting for intercept id=%q to become active: %w", id, err)
		}

		var intercept *manager.InterceptInfo
		for _, straw := range snapshot.Intercepts {
			if straw.Id == id {
				intercept = straw
				break
			}
		}

		switch {
		case intercept == nil:
			dlog.Debugf(ctx, "wait status: intercept id=%q does not yet exist", id)
		case intercept.Disposition == manager.InterceptDispositionType_WAITING:
			dlog.Debugf(ctx, "wait status: intercept id=%q is still WAITING", id)
		default:
			dlog.Debugf(ctx, "wait status: intercept id=%q is no longer WAITING; is now %v", id, intercept.Disposition)
			if intercept.Disposition != manager.InterceptDispositionType_ACTIVE {
				return nil, errors.Errorf("intercept in error state %v: %v", intercept.Disposition, intercept.Message)
			}
			return intercept, nil
		}
	}
}

func (tm *trafficManager) waitForAgent(ctx context.Context, name string) (*manager.AgentInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, 120*time.Second) // installing a new agent can take some time
	defer cancel()

	stream, err := tm.managerClient.WatchAgents(ctx, tm.session())
	if err != nil {
		return nil, fmt.Errorf("waiting for agent %q to be present: %q", name, err)
	}
	for {
		snapshot, err := stream.Recv()
		if err != nil {
			return nil, fmt.Errorf("waiting for agent %q to be present: %q", name, err)
		}

		var agentList []*manager.AgentInfo
		for _, agent := range snapshot.Agents {
			if agent.Name == name {
				agentList = append(agentList, agent)
			}
		}
		if managerutil.AgentsAreCompatible(agentList) {
			return agentList[0], nil
		}
	}
}

func (tm *trafficManager) workerPortForwardIntercept(ctx context.Context, pf portForward) {
	dlog.Infof(ctx, "Initiating port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)

	sshArgs := []string{
		"ssh",

		"-F", "none", // don't load the user's config file

		// connection settings
		"-C", // compression
		"-oConnectTimeout=10",
		"-oStrictHostKeyChecking=no",     // don't bother checking the host key...
		"-oUserKnownHostsFile=/dev/null", // and since we're not checking it, don't bother remembering it either

		// port-forward settings
		"-N", // no remote command; just connect and forward ports
		"-oExitOnForwardFailure=yes",
		"-R", fmt.Sprintf("%d:%s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort), // port to forward

		// where to connect to
		"-p", strconv.Itoa(int(tm.sshPort)),
		"telepresence@localhost",
	}

	backoff := 100 * time.Millisecond
	for ctx.Err() == nil {
		start := time.Now()
		err := dexec.CommandContext(ctx, sshArgs[0], sshArgs[1:]...).Run()
		lifetime := time.Since(start)

		if ctx.Err() == nil {
			if err == nil {
				err = errors.New("process terminated unexpectedly")
			}
			dlog.Errorf(ctx, "communicating with manager: %v", err)
			if lifetime >= 20*time.Second {
				backoff = 100 * time.Millisecond
				dtime.SleepWithContext(ctx, backoff)
			} else {
				dtime.SleepWithContext(ctx, backoff)
				backoff *= 2
				if backoff > 3*time.Second {
					backoff = 3 * time.Second
				}
			}
		}
	}

	dlog.Infof(ctx, "Terminated port-forward manager:%d -> %s:%d", pf.ManagerPort, pf.TargetHost, pf.TargetPort)
}

// removeIntercept removes one intercept by name
func (tm *trafficManager) removeIntercept(c context.Context, name string) error {
	dlog.Debugf(c, "telling manager to remove intercept %s", name)
	_, err := tm.managerClient.RemoveIntercept(c, &manager.RemoveInterceptRequest2{
		Session: tm.session(),
		Name:    name,
	})
	return err
}

// clearIntercepts removes all intercepts
func (tm *trafficManager) clearIntercepts(c context.Context) error {
	intercepts, _ := actions.ListMyIntercepts(c, tm.managerClient, tm.session().SessionId)
	for _, cept := range intercepts {
		err := tm.removeIntercept(c, cept.Spec.Name)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>