<commit_before>package isolated

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"

	"code.cloudfoundry.org/cli/integration/helpers"
	"code.cloudfoundry.org/cli/resources"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("update-security-group command", func() {
	When("--help flag is set", func() {
		It("Displays command usage to output", func() {
			session := helpers.CF("update-security-group", "--help")
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(0))
		})
	})

	var (
		securityGroup             resources.SecurityGroup
		securityGroupName         string
		updatedSecurityGroupRules []resources.Rule
		tmpRulesDir               string
		updatedRulesPath          string
		username                  string
		err                       error
	)

	BeforeEach(func() {
		username = helpers.LoginCF()
		securityGroupName = helpers.NewSecurityGroupName()
		securityGroup = resources.SecurityGroup{
			Name:  securityGroupName,
			Rules: []resources.Rule{},
		}
		helpers.CreateSecurityGroup(securityGroup)
	})

	BeforeEach(func() {
		// create the JSON rules file
		ports := "25,465,587"
		description := "Give me a place to stand on, and I shall spam the world!"
		updatedSecurityGroupRules = []resources.Rule{
			{
				Protocol:    "tcp",
				Destination: "0.0.0.0/0",
				Ports:       &ports,
				Description: &description,
			},
		}

		tmpRulesDir, err = ioutil.TempDir("", "spam-security-group")
		Expect(err).ToNot(HaveOccurred())

		updatedRulesPath = filepath.Join(tmpRulesDir, "spam-security-group.json")

		securityGroup, err := json.Marshal(updatedSecurityGroupRules)
		Expect(err).ToNot(HaveOccurred())

		err = ioutil.WriteFile(updatedRulesPath, securityGroup, 0666)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		helpers.DeleteSecurityGroup(securityGroup)
		err = os.RemoveAll(tmpRulesDir)
		Expect(err).ToNot(HaveOccurred())
	})

	It("updates the security group", func() {
		session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
		Eventually(session).Should(Say(`Updating security group %s as %s`, securityGroupName, username))
		Eventually(session).Should(Exit(0))

		session = helpers.CF("security-group", securityGroupName)
		Eventually(session).Should(Say(`Getting info for security group %s as %s\.\.\.`, securityGroupName, username))
		Eventually(session).Should(Say(`name:\s+%s`, securityGroupName))
		Eventually(session).Should(Say("rules:"))
		Eventually(session).Should(Say("\\["))
		Eventually(session).Should(Say("{"))
		Eventually(session).Should(Say(`"protocol": "tcp",`))
		Eventually(session).Should(Say(`"destination": "0.0.0.0/0",`))
		Eventually(session).Should(Say(`"ports": "25,465,587",`))
		Eventually(session).Should(Say(`"description": "Give me a place to stand on, and I shall spam the world!"`))
		Eventually(session).Should(Say("}"))
		Eventually(session).Should(Say("]"))
		Eventually(session).Should(Say("No spaces assigned"))

		Eventually(session).Should(Exit(0))
	})

	When("the security group does not exist or is not visible to the user", func() {
		BeforeEach(func() {
			helpers.DeleteSecurityGroup(securityGroup)
		})
		It("displays a 'not-found' security group error message and fails", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
			Eventually(session).Should(Say("FAILED"))
			Eventually(session.Err).Should(Say("Security group '%s' not found.", securityGroupName))
			Eventually(session).Should(Exit(1))
		})
	})

	When("the JSON file with the security group's rules does not exist", func() {
		It("displays a does-not-exist error message and displays help text", func() {
			session := helpers.CF("update-security-group", securityGroupName, "/non/existent/path")
			Eventually(session.Err).Should(Say("Incorrect Usage: The specified path '/non/existent/path' does not exist."))
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(1))
		})
	})

	When("the JSON file with the security group's rules is not JSON", func() {
		BeforeEach(func() {
			err = ioutil.WriteFile(updatedRulesPath, []byte("I'm definitely not JSON"), 0666)
			Expect(err).ToNot(HaveOccurred())
		})
		It("displays a an incorrect JSON format message and fails", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
			Eventually(session).Should(Say("FAILED"))
			Eventually(session.Err).Should(Say("Incorrect json format: %s", updatedRulesPath))
			Eventually(session).Should(Exit(1))
		})
	})

	When("there are unexpected arguments", func() {
		It("complains that there are unexpected arguments, fails, and prints the help text", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath, "unexpected-argument")
			Eventually(session.Err).Should(Say(`Incorrect Usage: unexpected argument "unexpected-argument"`))
			Eventually(session).Should(Say("FAILED"))
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(1))
		})
	})
})

func displaysUpdateSecurityGroupHelpText(session *Session) {
	Eventually(session).Should(Say("NAME:"))
	Eventually(session).Should(Say("update-security-group - Update a security group"))
	Eventually(session).Should(Say("USAGE:"))
	Eventually(session).Should(Say("cf update-security-group SECURITY_GROUP PATH_TO_JSON_RULES_FILE"))
	Eventually(session).Should(Say("The provided path can be an absolute or relative path to a file. The file should have"))
	Eventually(session).Should(Say("a single array with JSON objects inside describing the rules. The JSON Base Object is"))
	Eventually(session).Should(Say("omitted and only the square brackets and associated child object are required in the file."))
	Eventually(session).Should(Say("Valid json file example:"))
	Eventually(session).Should(Say("\\["))
	Eventually(session).Should(Say("{"))
	Eventually(session).Should(Say(`"protocol": "tcp",`))
	Eventually(session).Should(Say(`"destination": "10.0.11.0/24",`))
	Eventually(session).Should(Say(`"ports": "80,443",`))
	Eventually(session).Should(Say(`"description": "Allow http and https traffic from ZoneA"`))
	Eventually(session).Should(Say("}"))
	Eventually(session).Should(Say("]"))
	Eventually(session).Should(Say(`TIP: Changes require an app restart \(for running\) or restage \(for staging\) to apply to existing applications.`))
	Eventually(session).Should(Say("SEE ALSO:"))
	Eventually(session).Should(Say("restage, security-groups"))
}
<commit_msg>v7: Use SayPath matcher to assert on platform-independent paths<commit_after>package isolated

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"

	"code.cloudfoundry.org/cli/integration/helpers"
	"code.cloudfoundry.org/cli/resources"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("update-security-group command", func() {
	When("--help flag is set", func() {
		It("Displays command usage to output", func() {
			session := helpers.CF("update-security-group", "--help")
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(0))
		})
	})

	var (
		securityGroup             resources.SecurityGroup
		securityGroupName         string
		updatedSecurityGroupRules []resources.Rule
		tmpRulesDir               string
		updatedRulesPath          string
		username                  string
		err                       error
	)

	BeforeEach(func() {
		username = helpers.LoginCF()
		securityGroupName = helpers.NewSecurityGroupName()
		securityGroup = resources.SecurityGroup{
			Name:  securityGroupName,
			Rules: []resources.Rule{},
		}
		helpers.CreateSecurityGroup(securityGroup)
	})

	BeforeEach(func() {
		// create the JSON rules file
		ports := "25,465,587"
		description := "Give me a place to stand on, and I shall spam the world!"
		updatedSecurityGroupRules = []resources.Rule{
			{
				Protocol:    "tcp",
				Destination: "0.0.0.0/0",
				Ports:       &ports,
				Description: &description,
			},
		}

		tmpRulesDir, err = ioutil.TempDir("", "spam-security-group")
		Expect(err).ToNot(HaveOccurred())

		updatedRulesPath = filepath.Join(tmpRulesDir, "spam-security-group.json")

		securityGroup, err := json.Marshal(updatedSecurityGroupRules)
		Expect(err).ToNot(HaveOccurred())

		err = ioutil.WriteFile(updatedRulesPath, securityGroup, 0666)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		helpers.DeleteSecurityGroup(securityGroup)
		err = os.RemoveAll(tmpRulesDir)
		Expect(err).ToNot(HaveOccurred())
	})

	It("updates the security group", func() {
		session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
		Eventually(session).Should(Say(`Updating security group %s as %s`, securityGroupName, username))
		Eventually(session).Should(Exit(0))

		session = helpers.CF("security-group", securityGroupName)
		Eventually(session).Should(Say(`Getting info for security group %s as %s\.\.\.`, securityGroupName, username))
		Eventually(session).Should(Say(`name:\s+%s`, securityGroupName))
		Eventually(session).Should(Say("rules:"))
		Eventually(session).Should(Say("\\["))
		Eventually(session).Should(Say("{"))
		Eventually(session).Should(Say(`"protocol": "tcp",`))
		Eventually(session).Should(Say(`"destination": "0.0.0.0/0",`))
		Eventually(session).Should(Say(`"ports": "25,465,587",`))
		Eventually(session).Should(Say(`"description": "Give me a place to stand on, and I shall spam the world!"`))
		Eventually(session).Should(Say("}"))
		Eventually(session).Should(Say("]"))
		Eventually(session).Should(Say("No spaces assigned"))

		Eventually(session).Should(Exit(0))
	})

	When("the security group does not exist or is not visible to the user", func() {
		BeforeEach(func() {
			helpers.DeleteSecurityGroup(securityGroup)
		})
		It("displays a 'not-found' security group error message and fails", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
			Eventually(session).Should(Say("FAILED"))
			Eventually(session.Err).Should(Say("Security group '%s' not found.", securityGroupName))
			Eventually(session).Should(Exit(1))
		})
	})

	When("the JSON file with the security group's rules does not exist", func() {
		It("displays a does-not-exist error message and displays help text", func() {
			session := helpers.CF("update-security-group", securityGroupName, "/non/existent/path")
			Eventually(session.Err).Should(Say("Incorrect Usage: The specified path '/non/existent/path' does not exist."))
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(1))
		})
	})

	When("the JSON file with the security group's rules is not JSON", func() {
		BeforeEach(func() {
			err = ioutil.WriteFile(updatedRulesPath, []byte("I'm definitely not JSON"), 0666)
			Expect(err).ToNot(HaveOccurred())
		})
		It("displays a an incorrect JSON format message and fails", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath)
			Eventually(session).Should(Say("FAILED"))
			Eventually(session.Err).Should(helpers.SayPath("Incorrect json format: %s", updatedRulesPath))
			Eventually(session).Should(Exit(1))
		})
	})

	When("there are unexpected arguments", func() {
		It("complains that there are unexpected arguments, fails, and prints the help text", func() {
			session := helpers.CF("update-security-group", securityGroupName, updatedRulesPath, "unexpected-argument")
			Eventually(session.Err).Should(Say(`Incorrect Usage: unexpected argument "unexpected-argument"`))
			Eventually(session).Should(Say("FAILED"))
			displaysUpdateSecurityGroupHelpText(session)
			Eventually(session).Should(Exit(1))
		})
	})
})

func displaysUpdateSecurityGroupHelpText(session *Session) {
	Eventually(session).Should(Say("NAME:"))
	Eventually(session).Should(Say("update-security-group - Update a security group"))
	Eventually(session).Should(Say("USAGE:"))
	Eventually(session).Should(Say("cf update-security-group SECURITY_GROUP PATH_TO_JSON_RULES_FILE"))
	Eventually(session).Should(Say("The provided path can be an absolute or relative path to a file. The file should have"))
	Eventually(session).Should(Say("a single array with JSON objects inside describing the rules. The JSON Base Object is"))
	Eventually(session).Should(Say("omitted and only the square brackets and associated child object are required in the file."))
	Eventually(session).Should(Say("Valid json file example:"))
	Eventually(session).Should(Say("\\["))
	Eventually(session).Should(Say("{"))
	Eventually(session).Should(Say(`"protocol": "tcp",`))
	Eventually(session).Should(Say(`"destination": "10.0.11.0/24",`))
	Eventually(session).Should(Say(`"ports": "80,443",`))
	Eventually(session).Should(Say(`"description": "Allow http and https traffic from ZoneA"`))
	Eventually(session).Should(Say("}"))
	Eventually(session).Should(Say("]"))
	Eventually(session).Should(Say(`TIP: Changes require an app restart \(for running\) or restage \(for staging\) to apply to existing applications.`))
	Eventually(session).Should(Say("SEE ALSO:"))
	Eventually(session).Should(Say("restage, security-groups"))
}
<|endoftext|>


<commit_before>// Wrapper over jose utilities that implements JSON Web Signature (JWS) and JSON Web Encryption (JWE).

package jwt

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
)

const (
	prvKeyBits = 2048

	defaultSignatureAlgorithm = jose.RS256
	defaultKeyAlgorithm       = jose.RSA_OAEP
	defaultContentAlgorithm   = jose.A128GCM

	// See https://tools.ietf.org/html/draft-ietf-jose-json-web-key-41#section-4.2
	signatureJWKUse  = "sig"
	encryptionJWKUse = "enc"
)

func genKey() (*rsa.PrivateKey, error) {
	return rsa.GenerateKey(rand.Reader, prvKeyBits)
}

// Claims specify registered claim names specified in https://tools.ietf.org/html/rfc7519#section-4.1.
type Claims struct {
	Issuer    string      `json:"iss,omitempty"`
	Subject   string      `json:"sub,omitempty"`
	Audience  []string    `json:"aud,omitempty"`
	Expiry    NumericDate `json:"exp,omitempty"`
	NotBefore NumericDate `json:"nbf,omitempty"`
	IssuedAt  NumericDate `json:"iat,omitempty"`
	ID        string      `json:"jti,omitempty"`

	timeNow func() time.Time
}

// Validate checks claims in a token against expected values.
func (c Claims) Validate(e Claims) error {
	if e.Issuer != "" && e.Issuer != c.Issuer {
		return fmt.Errorf("Invalid issuer. Expected: %q, got: %q", e.Issuer, c.Issuer)
	}

	if e.Subject != "" && e.Subject != c.Subject {
		return fmt.Errorf("Invalid subject. Expected: %q, got: %q", e.Subject, c.Subject)
	}

	if e.ID != "" && e.ID != c.ID {
		return fmt.Errorf("Invalid ID. Expected: %q, got: %q", e.ID, c.ID)
	}

	if len(e.Audience) != 0 {
		if len(e.Audience) != len(c.Audience) {
			return fmt.Errorf("Invalid Audience. Expected len: %d, got: %d", len(e.Audience), len(c.Audience))
		}

		audienceLookup := map[string]struct{}{}
		for _, v := range c.Audience {
			audienceLookup[v] = struct{}{}
		}

		for _, v := range e.Audience {
			if _, ok := audienceLookup[v]; !ok {
				return fmt.Errorf("Invalid Audience. Expected: %q and  not found.", v)
			}
		}
	}
	return nil
}

// ValidateExpiryWithLeeway checks time based claims. A custom leeway may be specified for comparing
// time values. You may pass a zero value to check time values with no leeway, but you should note that numeric date
// values are rounded to the nearest second and sub-second precision is not supported.
func (c Claims) ValidateExpiryWithLeeway(leeway time.Duration) error {
	if c.timeNow == nil {
		c.timeNow = time.Now
	}

	now := c.timeNow()
	if now.IsZero() {
		return nil
	}

	if now.Add(leeway).Before(c.NotBefore.Time()) {
		return fmt.Errorf("Violated NotBefore. It is before specified time %q", c.NotBefore.Time().String())
	}

	if now.Add(-leeway).After(c.Expiry.Time()) {
		return fmt.Errorf("Exipred claim. Expiration time %q", c.Expiry.Time())

	}

	return nil
}

// jwtPayload defines special claim custom payload.
// Needed to avoid collisions. https://tools.ietf.org/html/rfc7519#section-4.2
type jwtPayload struct {
	Payload interface{} `json:"payload"`
}

// Builder is an builder that is able to construct nested JWT or JWS with custom payload claim or other claims.
// (Nested JSON Web Token is token that is signed and encrypted respectively).
// For each instance separate private RSA key is assigned, which is used to constructs and obtain all tokens.
type Builder struct {
	*SignedObtainer

	// Private key encryption and signatures used with every JWT handled by this builder.
	prvKey *rsa.PrivateKey

	signer    jose.Signer
	encrypter jose.Encrypter
}

// NewDefaultBuilder constructs new Builder that is able to construct and read all types of JSON Web tokens.
// Uses default signature, key and content algorithms. Private key is auto-generated.
func NewDefaultBuilder() (*Builder, error) {
	prvKey, err := genKey()
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not generate RSA key. Err: %v", err)
	}

	return NewBuilder(prvKey, defaultSignatureAlgorithm, defaultKeyAlgorithm, defaultContentAlgorithm)
}

// NewBuilder constructs new Builder that is able to construct and read all types of JSON Web tokens.
func NewBuilder(
	prvKey *rsa.PrivateKey,
	signatureAlgorithm jose.SignatureAlgorithm,
	keyAlgorithm jose.KeyAlgorithm,
	contentAlgorithm jose.ContentEncryption,
) (*Builder, error) {
	signer, err := jose.NewSigner(
		jose.SigningKey{
			Algorithm: signatureAlgorithm,
			Key:       prvKey,
		},
		nil,
	)
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not build signer. Err: %v", err)
	}

	encrypter, err := jose.NewEncrypter(
		contentAlgorithm,
		jose.Recipient{
			Algorithm: keyAlgorithm,
			Key:       &prvKey.PublicKey,
		},
		(&jose.EncrypterOptions{}).WithType("JWT").WithContentType("JWT"),
	)
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not build encrypter. Err: %v", err)
	}

	return &Builder{
		SignedObtainer: NewSignedObtainer(&prvKey.PublicKey, signatureAlgorithm),
		prvKey:         prvKey,
		signer:         signer,
		encrypter:      encrypter,
	}, nil
}

// SignedAndEncryptedJWT constructs nested JSON Web Token that is both signed and encrypted respectively
func (j *Builder) SignedAndEncryptedJWT() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.SignedAndEncrypted(j.signer, j.encrypter)}
}

// JWS constructs JSON Web Token that is only signed.
func (j *Builder) JWS() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.Signed(j.signer)}
}

// JWE constructs JSON Web Token that is only encrypted.
func (j *Builder) JWE() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.Encrypted(j.encrypter)}
}

// FromSignedAndEncryptedJWT parsed given token as nested JSON Web Token's and decrypts it.
// Returned ObtainerWrapper can be used to fetch claims. Signature verification is done by ObtainerWrapper.
func (j *Builder) FromSignedAndEncryptedJWT(token string) *ObtainerWrapper {
	nestedJWT, err := jwt.ParseSignedAndEncrypted(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as nested JSON Web Token. Err: %v", err)}
	}

	signedJWT, err := nestedJWT.Decrypt(j.prvKey)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not decrypt given nested JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: signedJWT, key: &j.prvKey.PublicKey}
}

// FromJWE decrypts JSON Web Token's.
// Returned ObtainerWrapper can be used to fetch claims. Decryption is done by ObtainerWrapper.
func (j *Builder) FromJWE(token string) *ObtainerWrapper {
	encrypted, err := jwt.ParseEncrypted(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as encrypted JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: encrypted, key: j.prvKey}
}

// SignedObtainer is struct that is able only to parse not encrypted JSON Web tokens.
type SignedObtainer struct {
	pubKey             *rsa.PublicKey
	signatureAlgorithm jose.SignatureAlgorithm

	timeNow func() time.Time
}

// NewSignedObtainer constructs SignedObtainer.
func NewSignedObtainer(publicKey *rsa.PublicKey, signatureAlgorithm jose.SignatureAlgorithm) *SignedObtainer {
	return &SignedObtainer{
		pubKey:             publicKey,
		signatureAlgorithm: signatureAlgorithm,
		timeNow:            time.Now,
	}
}

// PublicRSAKey gets Public RSA key used by this Obtainer.
func (j *SignedObtainer) PublicRSAKey() rsa.PublicKey {
	return *j.pubKey
}

// PublicJWK gets Public RSA key wrapped in JSON Web Key used by this Obtainer.
func (j *SignedObtainer) PublicJWK() jose.JSONWebKey {
	return jose.JSONWebKey{
		Key:       j.pubKey,
		Use:       signatureJWKUse,
		Algorithm: string(j.signatureAlgorithm),
	}
}

// FromJWS parses given JWS.
// Returned ObtainerWrapper can be used to fetch claims. Signature verification is done by ObtainerWrapper.
func (j *SignedObtainer) FromJWS(token string) *ObtainerWrapper {
	signed, err := jwt.ParseSigned(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as signed JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: signed, key: j.pubKey}
}

// VerifyStdClaims verifies standard "iss", "sub", "aud", "exp" claims from JWT RFC (https://tools.ietf.org/html/rfc7519).
func (j *SignedObtainer) VerifyStdClaims(claims Claims, expected Claims) error {
	err := claims.Validate(expected)
	if err != nil {
		return fmt.Errorf("JWT Builder: claims validation failed. Err: %v", err)
	}

	err = claims.ValidateExpiryWithLeeway(1 * time.Second)
	if err != nil {
		return fmt.Errorf("JWT Builder: claims validation failed. Err: %v", err)
	}
	return nil
}

// BuilderWrapper wraps specified engine and enabling packing and serializing claims into single token.
type BuilderWrapper struct {
	engine interface{}
}

// Claims encodes claims into JWE/JWS form. Multiple calls will merge claims
// into single JSON object.
func (b *BuilderWrapper) Claims(claims interface{}) *BuilderWrapper {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		b.engine = builder.Claims(claims)
	case jwt.NestedBuilder:
		b.engine = builder.Claims(claims)
	}
	return b
}

// Payload encodes payload into JWE/JWS form in a `payload` field. Multiple calls will override payload.
func (b *BuilderWrapper) Payload(payload interface{}) *BuilderWrapper {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		b.engine = builder.Claims(
			jwtPayload{
				Payload: payload,
			},
		)
	case jwt.NestedBuilder:
		b.engine = builder.Claims(
			jwtPayload{
				Payload: payload,
			},
		)
	}
	return b
}

// CompactSerialize serializes constructed token into compact form.
func (b *BuilderWrapper) CompactSerialize() (string, error) {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		return builder.CompactSerialize()
	case jwt.NestedBuilder:
		return builder.CompactSerialize()
	}
	return "", errors.New("JWT BuildWrapper internal error: Wrong type of wrapped engine")
}

// ObtainerWrapper wraps token and enables deserialization from token.
type ObtainerWrapper struct {
	err         error
	parsedToken *jwt.JSONWebToken
	key         interface{} // *rsa.PublicKey or rsa.PrivateKey
}

// Claims decodes claims from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) Claims(out interface{}) error {
	if o.err != nil {
		return o.err
	}
	err := o.parsedToken.Claims(o.key, &out)
	if err != nil {
		return fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into given type. Err: %v", err)
	}
	return nil
}

// Payload decodes payload from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) Payload(out interface{}) error {
	if o.err != nil {
		return o.err
	}
	payload := jwtPayload{
		Payload: out,
	}
	err := o.parsedToken.Claims(o.key, &payload)
	if err != nil {
		return fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into given payload type. Err: %v", err)
	}
	return nil
}

// StdClaims decodes standard registered claims from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) StdClaims() (Claims, error) {
	if o.err != nil {
		return Claims{}, o.err
	}
	var claims Claims
	err := o.parsedToken.Claims(o.key, &claims)
	if err != nil {
		return Claims{}, fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into standard claims. Err: %v", err)
	}
	return claims, nil
}

// NumericDate represents date and time as the number of seconds since the
// epoch, including leap seconds. Non-integer values can be represented
// in the serialized format, but we round to the nearest second.
type NumericDate int64

// NewNumericDate constructs NumericDate from time.Time value.
func NewNumericDate(t time.Time) NumericDate {
	if t.IsZero() {
		return NumericDate(0)
	}

	// While RFC 7519 technically states that NumericDate values may be
	// non-integer values, we don't bother serializing timestamps in
	// claims with sub-second accuracy and just round to the nearest
	// second instead.
	return NumericDate(t.Unix())
}

// MarshalJSON serializes the given NumericDate into its JSON representation.
func (n NumericDate) MarshalJSON() ([]byte, error) {
	return []byte(strconv.FormatInt(int64(n), 10)), nil
}

// UnmarshalJSON reads a date from its JSON representation.
func (n *NumericDate) UnmarshalJSON(b []byte) error {
	s := string(b)

	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return fmt.Errorf("Failed to unmarshall NumericDate. Err: %v", err)
	}

	*n = NumericDate(f)
	return nil
}

// Time returns time.Time representation of NumericDate.
func (n NumericDate) Time() time.Time {
	return time.Unix(int64(n), 0)
}
<commit_msg>Added KeyID for compatibility.<commit_after>// Wrapper over jose utilities that implements JSON Web Signature (JWS) and JSON Web Encryption (JWE).

package jwt

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
)

const (
	prvKeyBits = 2048

	defaultSignatureAlgorithm = jose.RS256
	defaultKeyAlgorithm       = jose.RSA_OAEP
	defaultContentAlgorithm   = jose.A128GCM

	// See https://tools.ietf.org/html/draft-ietf-jose-json-web-key-41#section-4.2
	signatureJWKUse  = "sig"
	encryptionJWKUse = "enc"
)

func genKey() (*rsa.PrivateKey, error) {
	return rsa.GenerateKey(rand.Reader, prvKeyBits)
}

// Claims specify registered claim names specified in https://tools.ietf.org/html/rfc7519#section-4.1.
type Claims struct {
	Issuer    string      `json:"iss,omitempty"`
	Subject   string      `json:"sub,omitempty"`
	Audience  []string    `json:"aud,omitempty"`
	Expiry    NumericDate `json:"exp,omitempty"`
	NotBefore NumericDate `json:"nbf,omitempty"`
	IssuedAt  NumericDate `json:"iat,omitempty"`
	ID        string      `json:"jti,omitempty"`

	timeNow func() time.Time
}

// Validate checks claims in a token against expected values.
func (c Claims) Validate(e Claims) error {
	if e.Issuer != "" && e.Issuer != c.Issuer {
		return fmt.Errorf("Invalid issuer. Expected: %q, got: %q", e.Issuer, c.Issuer)
	}

	if e.Subject != "" && e.Subject != c.Subject {
		return fmt.Errorf("Invalid subject. Expected: %q, got: %q", e.Subject, c.Subject)
	}

	if e.ID != "" && e.ID != c.ID {
		return fmt.Errorf("Invalid ID. Expected: %q, got: %q", e.ID, c.ID)
	}

	if len(e.Audience) != 0 {
		if len(e.Audience) != len(c.Audience) {
			return fmt.Errorf("Invalid Audience. Expected len: %d, got: %d", len(e.Audience), len(c.Audience))
		}

		audienceLookup := map[string]struct{}{}
		for _, v := range c.Audience {
			audienceLookup[v] = struct{}{}
		}

		for _, v := range e.Audience {
			if _, ok := audienceLookup[v]; !ok {
				return fmt.Errorf("Invalid Audience. Expected: %q and  not found.", v)
			}
		}
	}
	return nil
}

// ValidateExpiryWithLeeway checks time based claims. A custom leeway may be specified for comparing
// time values. You may pass a zero value to check time values with no leeway, but you should note that numeric date
// values are rounded to the nearest second and sub-second precision is not supported.
func (c Claims) ValidateExpiryWithLeeway(leeway time.Duration) error {
	if c.timeNow == nil {
		c.timeNow = time.Now
	}

	now := c.timeNow()
	if now.IsZero() {
		return nil
	}

	if now.Add(leeway).Before(c.NotBefore.Time()) {
		return fmt.Errorf("Violated NotBefore. It is before specified time %q", c.NotBefore.Time().String())
	}

	if now.Add(-leeway).After(c.Expiry.Time()) {
		return fmt.Errorf("Exipred claim. Expiration time %q", c.Expiry.Time())

	}

	return nil
}

// jwtPayload defines special claim custom payload.
// Needed to avoid collisions. https://tools.ietf.org/html/rfc7519#section-4.2
type jwtPayload struct {
	Payload interface{} `json:"payload"`
}

// Builder is an builder that is able to construct nested JWT or JWS with custom payload claim or other claims.
// (Nested JSON Web Token is token that is signed and encrypted respectively).
// For each instance separate private RSA key is assigned, which is used to constructs and obtain all tokens.
type Builder struct {
	*SignedObtainer

	// Private key encryption and signatures used with every JWT handled by this builder.
	prvKey *rsa.PrivateKey

	signer    jose.Signer
	encrypter jose.Encrypter
}

// NewDefaultBuilder constructs new Builder that is able to construct and read all types of JSON Web tokens.
// Uses default signature, key and content algorithms. Private key is auto-generated.
func NewDefaultBuilder() (*Builder, error) {
	prvKey, err := genKey()
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not generate RSA key. Err: %v", err)
	}

	return NewBuilder(prvKey, defaultSignatureAlgorithm, defaultKeyAlgorithm, defaultContentAlgorithm)
}

// NewBuilder constructs new Builder that is able to construct and read all types of JSON Web tokens.
func NewBuilder(
	prvKey *rsa.PrivateKey,
	signatureAlgorithm jose.SignatureAlgorithm,
	keyAlgorithm jose.KeyAlgorithm,
	contentAlgorithm jose.ContentEncryption,
) (*Builder, error) {
	signer, err := jose.NewSigner(
		jose.SigningKey{
			Algorithm: signatureAlgorithm,
			Key:       prvKey,
		},
		nil,
	)
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not build signer. Err: %v", err)
	}

	encrypter, err := jose.NewEncrypter(
		contentAlgorithm,
		jose.Recipient{
			Algorithm: keyAlgorithm,
			Key:       &prvKey.PublicKey,
		},
		(&jose.EncrypterOptions{}).WithType("JWT").WithContentType("JWT"),
	)
	if err != nil {
		return nil, fmt.Errorf("JWT Builder: Could not build encrypter. Err: %v", err)
	}

	return &Builder{
		SignedObtainer: NewSignedObtainer(&prvKey.PublicKey, signatureAlgorithm),
		prvKey:         prvKey,
		signer:         signer,
		encrypter:      encrypter,
	}, nil
}

// SignedAndEncryptedJWT constructs nested JSON Web Token that is both signed and encrypted respectively
func (j *Builder) SignedAndEncryptedJWT() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.SignedAndEncrypted(j.signer, j.encrypter)}
}

// JWS constructs JSON Web Token that is only signed.
func (j *Builder) JWS() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.Signed(j.signer)}
}

// JWE constructs JSON Web Token that is only encrypted.
func (j *Builder) JWE() *BuilderWrapper {
	return &BuilderWrapper{engine: jwt.Encrypted(j.encrypter)}
}

// FromSignedAndEncryptedJWT parsed given token as nested JSON Web Token's and decrypts it.
// Returned ObtainerWrapper can be used to fetch claims. Signature verification is done by ObtainerWrapper.
func (j *Builder) FromSignedAndEncryptedJWT(token string) *ObtainerWrapper {
	nestedJWT, err := jwt.ParseSignedAndEncrypted(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as nested JSON Web Token. Err: %v", err)}
	}

	signedJWT, err := nestedJWT.Decrypt(j.prvKey)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not decrypt given nested JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: signedJWT, key: &j.prvKey.PublicKey}
}

// FromJWE decrypts JSON Web Token's.
// Returned ObtainerWrapper can be used to fetch claims. Decryption is done by ObtainerWrapper.
func (j *Builder) FromJWE(token string) *ObtainerWrapper {
	encrypted, err := jwt.ParseEncrypted(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as encrypted JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: encrypted, key: j.prvKey}
}

// SignedObtainer is struct that is able only to parse not encrypted JSON Web tokens.
type SignedObtainer struct {
	pubKey             *rsa.PublicKey
	signatureAlgorithm jose.SignatureAlgorithm

	timeNow func() time.Time
}

// NewSignedObtainer constructs SignedObtainer.
func NewSignedObtainer(publicKey *rsa.PublicKey, signatureAlgorithm jose.SignatureAlgorithm) *SignedObtainer {
	return &SignedObtainer{
		pubKey:             publicKey,
		signatureAlgorithm: signatureAlgorithm,
		timeNow:            time.Now,
	}
}

// PublicRSAKey gets Public RSA key used by this Obtainer.
func (j *SignedObtainer) PublicRSAKey() rsa.PublicKey {
	return *j.pubKey
}

// PublicJWK gets Public RSA key wrapped in JSON Web Key used by this Obtainer.
func (j *SignedObtainer) PublicJWK() jose.JSONWebKey {
	return jose.JSONWebKey{
		Key:   j.pubKey,
		KeyID: "1",
		// Use parameter value is REQUIRED for all keys in the referenced JWK Set to indicate each key's intended usage.
		Use:       signatureJWKUse,
		Algorithm: string(j.signatureAlgorithm),
	}
}

// FromJWS parses given JWS.
// Returned ObtainerWrapper can be used to fetch claims. Signature verification is done by ObtainerWrapper.
func (j *SignedObtainer) FromJWS(token string) *ObtainerWrapper {
	signed, err := jwt.ParseSigned(token)
	if err != nil {
		return &ObtainerWrapper{err: fmt.Errorf("JWT Builder: Could not parse given token as signed JSON Web Token. Err: %v", err)}
	}

	return &ObtainerWrapper{parsedToken: signed, key: j.pubKey}
}

// VerifyStdClaims verifies standard "iss", "sub", "aud", "exp" claims from JWT RFC (https://tools.ietf.org/html/rfc7519).
func (j *SignedObtainer) VerifyStdClaims(claims Claims, expected Claims) error {
	err := claims.Validate(expected)
	if err != nil {
		return fmt.Errorf("JWT Builder: claims validation failed. Err: %v", err)
	}

	err = claims.ValidateExpiryWithLeeway(1 * time.Second)
	if err != nil {
		return fmt.Errorf("JWT Builder: claims validation failed. Err: %v", err)
	}
	return nil
}

// BuilderWrapper wraps specified engine and enabling packing and serializing claims into single token.
type BuilderWrapper struct {
	engine interface{}
}

// Claims encodes claims into JWE/JWS form. Multiple calls will merge claims
// into single JSON object.
func (b *BuilderWrapper) Claims(claims interface{}) *BuilderWrapper {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		b.engine = builder.Claims(claims)
	case jwt.NestedBuilder:
		b.engine = builder.Claims(claims)
	}
	return b
}

// Payload encodes payload into JWE/JWS form in a `payload` field. Multiple calls will override payload.
func (b *BuilderWrapper) Payload(payload interface{}) *BuilderWrapper {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		b.engine = builder.Claims(
			jwtPayload{
				Payload: payload,
			},
		)
	case jwt.NestedBuilder:
		b.engine = builder.Claims(
			jwtPayload{
				Payload: payload,
			},
		)
	}
	return b
}

// CompactSerialize serializes constructed token into compact form.
func (b *BuilderWrapper) CompactSerialize() (string, error) {
	switch builder := b.engine.(type) {
	case jwt.Builder:
		return builder.CompactSerialize()
	case jwt.NestedBuilder:
		return builder.CompactSerialize()
	}
	return "", errors.New("JWT BuildWrapper internal error: Wrong type of wrapped engine")
}

// ObtainerWrapper wraps token and enables deserialization from token.
type ObtainerWrapper struct {
	err         error
	parsedToken *jwt.JSONWebToken
	key         interface{} // *rsa.PublicKey or rsa.PrivateKey
}

// Claims decodes claims from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) Claims(out interface{}) error {
	if o.err != nil {
		return o.err
	}
	err := o.parsedToken.Claims(o.key, &out)
	if err != nil {
		return fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into given type. Err: %v", err)
	}
	return nil
}

// Payload decodes payload from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) Payload(out interface{}) error {
	if o.err != nil {
		return o.err
	}
	payload := jwtPayload{
		Payload: out,
	}
	err := o.parsedToken.Claims(o.key, &payload)
	if err != nil {
		return fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into given payload type. Err: %v", err)
	}
	return nil
}

// StdClaims decodes standard registered claims from JWE/JWS form. Multiple calls are allowed.
func (o *ObtainerWrapper) StdClaims() (Claims, error) {
	if o.err != nil {
		return Claims{}, o.err
	}
	var claims Claims
	err := o.parsedToken.Claims(o.key, &claims)
	if err != nil {
		return Claims{}, fmt.Errorf("JWT ObtainerWrapper: Could not deserialize JSON Web Token into standard claims. Err: %v", err)
	}
	return claims, nil
}

// NumericDate represents date and time as the number of seconds since the
// epoch, including leap seconds. Non-integer values can be represented
// in the serialized format, but we round to the nearest second.
type NumericDate int64

// NewNumericDate constructs NumericDate from time.Time value.
func NewNumericDate(t time.Time) NumericDate {
	if t.IsZero() {
		return NumericDate(0)
	}

	// While RFC 7519 technically states that NumericDate values may be
	// non-integer values, we don't bother serializing timestamps in
	// claims with sub-second accuracy and just round to the nearest
	// second instead.
	return NumericDate(t.Unix())
}

// MarshalJSON serializes the given NumericDate into its JSON representation.
func (n NumericDate) MarshalJSON() ([]byte, error) {
	return []byte(strconv.FormatInt(int64(n), 10)), nil
}

// UnmarshalJSON reads a date from its JSON representation.
func (n *NumericDate) UnmarshalJSON(b []byte) error {
	s := string(b)

	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return fmt.Errorf("Failed to unmarshall NumericDate. Err: %v", err)
	}

	*n = NumericDate(f)
	return nil
}

// Time returns time.Time representation of NumericDate.
func (n NumericDate) Time() time.Time {
	return time.Unix(int64(n), 0)
}
<|endoftext|>


<commit_before>package mtp

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/hanwen/go-mtpfs/usb"
)

// An MTP device.
type Device struct {
	h   *usb.DeviceHandle
	dev *usb.Device

	claimed bool

	// split off descriptor?
	devDescr    usb.DeviceDescriptor
	ifaceDescr  usb.InterfaceDescriptor
	sendEp      byte
	fetchEp     byte
	eventEp     byte
	configValue byte

	// In milliseconds. Defaults to 2 seconds.
	Timeout int

	// Print request/response codes.
	DebugPrint bool

	// Print data as it passes over the USB connection.
	DataPrint bool

	// If set, send header in separate write.
	SeparateHeader bool

	session *sessionData
}

type sessionData struct {
	tid uint32
	sid uint32
}

// RCError are return codes from the Container.Code field.
type RCError uint16

func (e RCError) Error() string {
	n, ok := RC_names[int(e)]
	if ok {
		return n
	}
	return fmt.Sprintf("RetCode %x", uint16(e))
}

// Close releases the interface, and closes the device.
func (d *Device) Close() error {
	if d.h == nil {
		return nil // or error?
	}

	if d.session != nil {
		err := d.CloseSession()
		if err != nil {
 			d.h.Reset()
		}
	}

	if d.claimed {
		d.h.ReleaseInterface(d.ifaceDescr.InterfaceNumber)
	}
	err := d.h.Close()
	d.h = nil
	return err
}

// Done releases the libusb device reference.
func (d *Device) Done() {
	d.dev.Unref()
	d.dev = nil
}

// Claims the USB interface of the device.
func (d *Device) claim() error {
	if d.h == nil {
		return fmt.Errorf("device not open")
	}

	err := d.h.ClaimInterface(d.ifaceDescr.InterfaceNumber)
	if err == nil {
		d.claimed = true
	}

	return err
}

// Open opens an MTP device.
func (d *Device) Open() error {
	if d.Timeout == 0 {
		d.Timeout = 2000
	}

	if d.h != nil {
		return fmt.Errorf("already open")
	}

	var err error
	d.h, err = d.dev.Open()
	if err != nil {
		return err
	}

	iface, err := d.h.GetStringDescriptorASCII(d.ifaceDescr.InterfaceStringIndex)
	if err != nil {
		d.Close()
		return err
	}

	if !strings.Contains(iface, "MTP") {
		d.Close()
		return fmt.Errorf("has no MTP in interface string")
	}

	return d.claim()
}

// Id is the manufacturer + product + serial
func (d *Device) ID() (string, error) {
	if d.h == nil {
		return "", fmt.Errorf("device not open")
	}

	var ids []string
	for _, b := range []byte{
		d.devDescr.Manufacturer,
		d.devDescr.Product,
		d.devDescr.SerialNumber} {
		s, err := d.h.GetStringDescriptorASCII(b)
		if err != nil {
			return "", err
		}
		ids = append(ids, s)
	}

	return strings.Join(ids, " "), nil
}

func (d *Device) sendReq(req *Container) error {
	c := usbBulkContainer{
		usbBulkHeader: usbBulkHeader{
			Length:        uint32(usbHdrLen + 4*len(req.Param)),
			Type:          USB_CONTAINER_COMMAND,
			Code:          req.Code,
			TransactionID: req.TransactionID,
		},
	}
	for i := range req.Param {
		c.Param[i] = req.Param[i]
	}

	var wData [usbBulkLen]byte
	buf := bytes.NewBuffer(wData[:0])

	err := binary.Write(buf, binary.LittleEndian, c)
	if err != nil {
		panic(err)
	}

	d.dataPrint(d.sendEp, buf.Bytes())
	_, err = d.h.BulkTransfer(d.sendEp, buf.Bytes(), d.Timeout)
	if err != nil {
		return err
	}
	return nil
}

const packetSize = 512

// Fetches one USB packet. The header is split off, and the remainder is returned.
// dest should be at least 512bytes.
func (d *Device) fetchPacket(dest []byte, header *usbBulkHeader) (rest []byte, err error) {
	n, err := d.h.BulkTransfer(d.fetchEp, dest[:packetSize], d.Timeout)
	if n > 0 {
		d.dataPrint(d.fetchEp, dest[:n])
	}

	if err != nil {
		return nil, err
	}

	buf := bytes.NewBuffer(dest[:n])
	err = binary.Read(buf, binary.LittleEndian, header)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func (d *Device) decodeRep(h *usbBulkHeader, rest []byte, rep *Container) error {
	if h.Type != USB_CONTAINER_RESPONSE {
		return fmt.Errorf("got type %d in response", h.Type)
	}

	rep.Code = h.Code
	if rep.Code != RC_OK {
		return RCError(rep.Code)
	}
	rep.TransactionID = h.TransactionID

	restLen := int(h.Length) - usbHdrLen
	if restLen != len(rest) {
		return fmt.Errorf("header specified 0x%x bytes, but have 0x%x",
			restLen, len(rest))
	}
	nParam := restLen / 4
	for i := 0; i < nParam; i++ {
		rep.Param = append(rep.Param, byteOrder.Uint32(rest[4*i:]))
	}

	return nil
}

// Runs a single MTP transaction. dest and src cannot be specified at
// the same time.  The request should fill out Code and Param as
// necessary. The response is provided here, but usually only the
// return code is of interest.  If the return code is an error, this
// function will return an RCError instance.
func (d *Device) RunTransaction(req *Container, rep *Container,
	dest io.Writer, src io.Reader, writeSize int64) error {
	if d.h == nil {
		return fmt.Errorf("mtp: cannot run operation %v, device is not open",
			OC_names[int(req.Code)])
	}
	if d.session != nil {
		req.SessionID = d.session.sid
		req.TransactionID = d.session.tid
		d.session.tid++
	}

	if d.DebugPrint {
		log.Printf("MTP request %s %v\n", OC_names[int(req.Code)], req.Param)
	}

	err := d.sendReq(req)

	if err != nil {
		if d.DebugPrint {
			log.Printf("MTP sendreq failed: %v\n", err)
		}
		return err
	}

	if src != nil {
		hdr := usbBulkHeader{
			Type:          USB_CONTAINER_DATA,
			Code:          req.Code,
			Length:        uint32(writeSize),
			TransactionID: req.TransactionID,
		}

		_, err := d.bulkWrite(&hdr, src, writeSize)
		if err != nil {
			d.Close()
			return err
		}
	}
	var data [packetSize]byte
	h := &usbBulkHeader{}
	rest, err := d.fetchPacket(data[:], h)
	if err != nil {
		return err
	}
	if h.Type == USB_CONTAINER_DATA {
		if dest == nil {
			d.Close()
			return fmt.Errorf("no sink for data")
		}
		if d.DebugPrint {
			log.Printf("MTP data 0x%x bytes", h.Length)
		}

		size := int(h.Length)
		dest.Write(rest)
		size -= len(rest) + usbHdrLen
		if size > 0 {
			_, err = d.bulkRead(dest)
			if err != nil {
				return err
			}
		}

		h = &usbBulkHeader{}
		rest, err = d.fetchPacket(data[:], h)
	}

	err = d.decodeRep(h, rest, rep)
	if d.DebugPrint {
		log.Printf("MTP response %s %v", RC_names[int(rep.Code)], rep.Param)
	}
	if err != nil {
		return err
	}
	if d.session != nil && rep.TransactionID != req.TransactionID {
		return fmt.Errorf("transaction ID mismatch got %x want %x",
			rep.TransactionID, req.TransactionID)
	}
	rep.SessionID = req.SessionID
	return nil
}

// Prints data going over the USB connection.
func (d *Device) dataPrint(ep byte, data []byte) {
	if !d.DataPrint {
		return
	}
	dir := "send"
	if 0 != ep&usb.ENDPOINT_IN {
		dir = "recv"
	}
	fmt.Fprintf(os.Stderr, "%s: 0x%x bytes with ep 0x%x:\n", dir, len(data), ep)
	hexDump(data)
}

// The linux usb stack can send 16kb per call, according to libusb.
const rwBufSize = 0x4000

// bulkWrite returns the number of non-header bytes written.
func (d *Device) bulkWrite(hdr *usbBulkHeader, r io.Reader, size int64) (n int64, err error) {
	if hdr != nil {
		if size+usbHdrLen > 0xFFFFFFFF {
			hdr.Length = 0xFFFFFFFF
		} else {
			hdr.Length = uint32(size + usbHdrLen)
		}

		var packetArr [packetSize]byte
		var packet []byte
		if d.SeparateHeader {
			packet = packetArr[:usbHdrLen]
		} else {
			packet = packetArr[:]
		}

		buf := bytes.NewBuffer(packet[:0])
		binary.Write(buf, byteOrder, hdr)
		cpSize := int64(len(packet) - usbHdrLen)
		if cpSize > size {
			cpSize = size
		}

		_, err = io.CopyN(buf, r, cpSize)
		d.dataPrint(d.sendEp, buf.Bytes())
		_, err = d.h.BulkTransfer(d.sendEp, buf.Bytes(), d.Timeout)
		if err != nil {
			return cpSize, err
		}
		size -= cpSize
		n += cpSize
	}

	var buf [rwBufSize]byte
	var lastTransfer int
	for size > 0 {
		var m int
		toread := buf[:]
		if int64(len(toread)) > size {
			toread = buf[:int(size)]
		}

		m, err = r.Read(toread)
		if err != nil {
			break
		}
		size -= int64(m)

		d.dataPrint(d.sendEp, buf[:m])
		lastTransfer, err = d.h.BulkTransfer(d.sendEp, buf[:m], d.Timeout)
		n += int64(lastTransfer)

		if err != nil || lastTransfer == 0 {
			break
		}
	}
	if lastTransfer%packetSize == 0 {
		// write a short packet just to be sure.
		d.h.BulkTransfer(d.sendEp, buf[:0], d.Timeout)
	}

	return n, err
}

func (d *Device) bulkRead(w io.Writer) (n int64, err error) {
	var buf [rwBufSize]byte
	var lastRead int
	for {
		toread := buf[:]
		lastRead, err = d.h.BulkTransfer(d.fetchEp, toread, d.Timeout)
		if err != nil {
			break
		}
		if lastRead > 0 {
			d.dataPrint(d.fetchEp, buf[:lastRead])

			w, err := w.Write(buf[:lastRead])
			n += int64(w)
			if err != nil {
				break
			}
		}
		if lastRead < len(toread) {
			// short read.
			break
		}
	}
	if lastRead%packetSize == 0 {
		// Null read.
		d.h.BulkTransfer(d.fetchEp, buf[:0], d.Timeout)
	}
	return n, err
}

// Configure is a robust version of OpenSession. On failure, it resets
// the device and reopens the device and the session.
func (d *Device) Configure() error {
	err := d.OpenSession()
	if err == RCError(RC_SessionAlreadyOpened) {
		// It's open, so close the session. Fortunately, this
		// even works without a transaction ID.
		d.CloseSession()
		err = d.OpenSession()
	}

	if err == usb.ERROR_IO || err == RCError(RC_SessionAlreadyOpened) ||
		err == usb.ERROR_TIMEOUT {
		d.h.Reset()
		d.Close()

		// Give the device some rest.
		time.Sleep(100 * time.Millisecond)
		if err := d.Open(); err != nil {
			return fmt.Errorf("opening after reset: %v", err)
		}
		if err := d.OpenSession(); err != nil {
			return fmt.Errorf("OpenSession after reset: %v", err)
		}
	}
	return nil
}
<commit_msg>Reset on any error.<commit_after>package mtp

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/hanwen/go-mtpfs/usb"
)

// An MTP device.
type Device struct {
	h   *usb.DeviceHandle
	dev *usb.Device

	claimed bool

	// split off descriptor?
	devDescr    usb.DeviceDescriptor
	ifaceDescr  usb.InterfaceDescriptor
	sendEp      byte
	fetchEp     byte
	eventEp     byte
	configValue byte

	// In milliseconds. Defaults to 2 seconds.
	Timeout int

	// Print request/response codes.
	DebugPrint bool

	// Print data as it passes over the USB connection.
	DataPrint bool

	// If set, send header in separate write.
	SeparateHeader bool

	session *sessionData
}

type sessionData struct {
	tid uint32
	sid uint32
}

// RCError are return codes from the Container.Code field.
type RCError uint16

func (e RCError) Error() string {
	n, ok := RC_names[int(e)]
	if ok {
		return n
	}
	return fmt.Sprintf("RetCode %x", uint16(e))
}

// Close releases the interface, and closes the device.
func (d *Device) Close() error {
	if d.h == nil {
		return nil // or error?
	}

	if d.session != nil {
		err := d.CloseSession()
		if err != nil {
 			d.h.Reset()
		}
	}

	if d.claimed {
		d.h.ReleaseInterface(d.ifaceDescr.InterfaceNumber)
	}
	err := d.h.Close()
	d.h = nil
	return err
}

// Done releases the libusb device reference.
func (d *Device) Done() {
	d.dev.Unref()
	d.dev = nil
}

// Claims the USB interface of the device.
func (d *Device) claim() error {
	if d.h == nil {
		return fmt.Errorf("device not open")
	}

	err := d.h.ClaimInterface(d.ifaceDescr.InterfaceNumber)
	if err == nil {
		d.claimed = true
	}

	return err
}

// Open opens an MTP device.
func (d *Device) Open() error {
	if d.Timeout == 0 {
		d.Timeout = 2000
	}

	if d.h != nil {
		return fmt.Errorf("already open")
	}

	var err error
	d.h, err = d.dev.Open()
	if err != nil {
		return err
	}

	iface, err := d.h.GetStringDescriptorASCII(d.ifaceDescr.InterfaceStringIndex)
	if err != nil {
		d.Close()
		return err
	}

	if !strings.Contains(iface, "MTP") {
		d.Close()
		return fmt.Errorf("has no MTP in interface string")
	}

	return d.claim()
}

// Id is the manufacturer + product + serial
func (d *Device) ID() (string, error) {
	if d.h == nil {
		return "", fmt.Errorf("device not open")
	}

	var ids []string
	for _, b := range []byte{
		d.devDescr.Manufacturer,
		d.devDescr.Product,
		d.devDescr.SerialNumber} {
		s, err := d.h.GetStringDescriptorASCII(b)
		if err != nil {
			return "", err
		}
		ids = append(ids, s)
	}

	return strings.Join(ids, " "), nil
}

func (d *Device) sendReq(req *Container) error {
	c := usbBulkContainer{
		usbBulkHeader: usbBulkHeader{
			Length:        uint32(usbHdrLen + 4*len(req.Param)),
			Type:          USB_CONTAINER_COMMAND,
			Code:          req.Code,
			TransactionID: req.TransactionID,
		},
	}
	for i := range req.Param {
		c.Param[i] = req.Param[i]
	}

	var wData [usbBulkLen]byte
	buf := bytes.NewBuffer(wData[:0])

	err := binary.Write(buf, binary.LittleEndian, c)
	if err != nil {
		panic(err)
	}

	d.dataPrint(d.sendEp, buf.Bytes())
	_, err = d.h.BulkTransfer(d.sendEp, buf.Bytes(), d.Timeout)
	if err != nil {
		return err
	}
	return nil
}

const packetSize = 512

// Fetches one USB packet. The header is split off, and the remainder is returned.
// dest should be at least 512bytes.
func (d *Device) fetchPacket(dest []byte, header *usbBulkHeader) (rest []byte, err error) {
	n, err := d.h.BulkTransfer(d.fetchEp, dest[:packetSize], d.Timeout)
	if n > 0 {
		d.dataPrint(d.fetchEp, dest[:n])
	}

	if err != nil {
		return nil, err
	}

	buf := bytes.NewBuffer(dest[:n])
	err = binary.Read(buf, binary.LittleEndian, header)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func (d *Device) decodeRep(h *usbBulkHeader, rest []byte, rep *Container) error {
	if h.Type != USB_CONTAINER_RESPONSE {
		return fmt.Errorf("got type %d in response", h.Type)
	}

	rep.Code = h.Code
	if rep.Code != RC_OK {
		return RCError(rep.Code)
	}
	rep.TransactionID = h.TransactionID

	restLen := int(h.Length) - usbHdrLen
	if restLen != len(rest) {
		return fmt.Errorf("header specified 0x%x bytes, but have 0x%x",
			restLen, len(rest))
	}
	nParam := restLen / 4
	for i := 0; i < nParam; i++ {
		rep.Param = append(rep.Param, byteOrder.Uint32(rest[4*i:]))
	}

	return nil
}

// Runs a single MTP transaction. dest and src cannot be specified at
// the same time.  The request should fill out Code and Param as
// necessary. The response is provided here, but usually only the
// return code is of interest.  If the return code is an error, this
// function will return an RCError instance.
func (d *Device) RunTransaction(req *Container, rep *Container,
	dest io.Writer, src io.Reader, writeSize int64) error {
	if d.h == nil {
		return fmt.Errorf("mtp: cannot run operation %v, device is not open",
			OC_names[int(req.Code)])
	}
	if d.session != nil {
		req.SessionID = d.session.sid
		req.TransactionID = d.session.tid
		d.session.tid++
	}

	if d.DebugPrint {
		log.Printf("MTP request %s %v\n", OC_names[int(req.Code)], req.Param)
	}

	err := d.sendReq(req)

	if err != nil {
		if d.DebugPrint {
			log.Printf("MTP sendreq failed: %v\n", err)
		}
		return err
	}

	if src != nil {
		hdr := usbBulkHeader{
			Type:          USB_CONTAINER_DATA,
			Code:          req.Code,
			Length:        uint32(writeSize),
			TransactionID: req.TransactionID,
		}

		_, err := d.bulkWrite(&hdr, src, writeSize)
		if err != nil {
			d.Close()
			return err
		}
	}
	var data [packetSize]byte
	h := &usbBulkHeader{}
	rest, err := d.fetchPacket(data[:], h)
	if err != nil {
		return err
	}
	if h.Type == USB_CONTAINER_DATA {
		if dest == nil {
			d.Close()
			return fmt.Errorf("no sink for data")
		}
		if d.DebugPrint {
			log.Printf("MTP data 0x%x bytes", h.Length)
		}

		size := int(h.Length)
		dest.Write(rest)
		size -= len(rest) + usbHdrLen
		if size > 0 {
			_, err = d.bulkRead(dest)
			if err != nil {
				return err
			}
		}

		h = &usbBulkHeader{}
		rest, err = d.fetchPacket(data[:], h)
	}

	err = d.decodeRep(h, rest, rep)
	if d.DebugPrint {
		log.Printf("MTP response %s %v", RC_names[int(rep.Code)], rep.Param)
	}
	if err != nil {
		return err
	}
	if d.session != nil && rep.TransactionID != req.TransactionID {
		return fmt.Errorf("transaction ID mismatch got %x want %x",
			rep.TransactionID, req.TransactionID)
	}
	rep.SessionID = req.SessionID
	return nil
}

// Prints data going over the USB connection.
func (d *Device) dataPrint(ep byte, data []byte) {
	if !d.DataPrint {
		return
	}
	dir := "send"
	if 0 != ep&usb.ENDPOINT_IN {
		dir = "recv"
	}
	fmt.Fprintf(os.Stderr, "%s: 0x%x bytes with ep 0x%x:\n", dir, len(data), ep)
	hexDump(data)
}

// The linux usb stack can send 16kb per call, according to libusb.
const rwBufSize = 0x4000

// bulkWrite returns the number of non-header bytes written.
func (d *Device) bulkWrite(hdr *usbBulkHeader, r io.Reader, size int64) (n int64, err error) {
	if hdr != nil {
		if size+usbHdrLen > 0xFFFFFFFF {
			hdr.Length = 0xFFFFFFFF
		} else {
			hdr.Length = uint32(size + usbHdrLen)
		}

		var packetArr [packetSize]byte
		var packet []byte
		if d.SeparateHeader {
			packet = packetArr[:usbHdrLen]
		} else {
			packet = packetArr[:]
		}

		buf := bytes.NewBuffer(packet[:0])
		binary.Write(buf, byteOrder, hdr)
		cpSize := int64(len(packet) - usbHdrLen)
		if cpSize > size {
			cpSize = size
		}

		_, err = io.CopyN(buf, r, cpSize)
		d.dataPrint(d.sendEp, buf.Bytes())
		_, err = d.h.BulkTransfer(d.sendEp, buf.Bytes(), d.Timeout)
		if err != nil {
			return cpSize, err
		}
		size -= cpSize
		n += cpSize
	}

	var buf [rwBufSize]byte
	var lastTransfer int
	for size > 0 {
		var m int
		toread := buf[:]
		if int64(len(toread)) > size {
			toread = buf[:int(size)]
		}

		m, err = r.Read(toread)
		if err != nil {
			break
		}
		size -= int64(m)

		d.dataPrint(d.sendEp, buf[:m])
		lastTransfer, err = d.h.BulkTransfer(d.sendEp, buf[:m], d.Timeout)
		n += int64(lastTransfer)

		if err != nil || lastTransfer == 0 {
			break
		}
	}
	if lastTransfer%packetSize == 0 {
		// write a short packet just to be sure.
		d.h.BulkTransfer(d.sendEp, buf[:0], d.Timeout)
	}

	return n, err
}

func (d *Device) bulkRead(w io.Writer) (n int64, err error) {
	var buf [rwBufSize]byte
	var lastRead int
	for {
		toread := buf[:]
		lastRead, err = d.h.BulkTransfer(d.fetchEp, toread, d.Timeout)
		if err != nil {
			break
		}
		if lastRead > 0 {
			d.dataPrint(d.fetchEp, buf[:lastRead])

			w, err := w.Write(buf[:lastRead])
			n += int64(w)
			if err != nil {
				break
			}
		}
		if lastRead < len(toread) {
			// short read.
			break
		}
	}
	if lastRead%packetSize == 0 {
		// Null read.
		d.h.BulkTransfer(d.fetchEp, buf[:0], d.Timeout)
	}
	return n, err
}

// Configure is a robust version of OpenSession. On failure, it resets
// the device and reopens the device and the session.
func (d *Device) Configure() error {
	if d.h == nil {
		err := d.Open()
		if err != nil {
			return err
		}
	}
	err := d.OpenSession()
	if err == RCError(RC_SessionAlreadyOpened) {
		// It's open, so close the session. Fortunately, this
		// even works without a transaction ID.
		d.CloseSession()
		err = d.OpenSession()
	}

 	if err != nil {
		log.Printf("OpenSession failed: %v; attempting reset", err)
		d.h.Reset()
		d.Close()

		// Give the device some rest.
		time.Sleep(100 * time.Millisecond)
		if err := d.Open(); err != nil {
			return fmt.Errorf("opening after reset: %v", err)
		}
		if err := d.OpenSession(); err != nil {
			return fmt.Errorf("OpenSession after reset: %v", err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before><commit_msg>better output<commit_after><|endoftext|>


<commit_before>package kite

import (
	"errors"
	"fmt"
	"runtime/debug"
	"strings"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/cache"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/protocol"
	"github.com/koding/kite/sockjsclient"
)

// Request contains information about the incoming request.
type Request struct {
	// Method defines the method name which is invoked by the incoming request
	Method string

	// Args defines the incoming arguments for the given method
	Args *dnode.Partial

	// LocalKite defines a context for the local kite
	LocalKite *Kite

	// Client defines a context for the remote kite
	Client *Client

	// Username defines the username which the incoming request is bound to.
	// This is authenticated and validated if authentication is enabled.
	Username string

	// Auth stores the authentication information for the incoming request and
	// the type of authentication. This is not used when authentication is disabled
	Auth *Auth

	// Context holds a context that used by the current ServeKite handler. Any
	// items added to the Context can be fetched from other handlers in the
	// chain. This is useful with PreHandle and PostHandle handlers to pass
	// data between handlers.
	Context cache.Cache
}

// Response is the type of the object that is returned from request handlers
// and the type of only argument that is passed to callback functions.
type Response struct {
	Error  *Error      `json:"error" dnode:"-"`
	Result interface{} `json:"result"`
}

// runMethod is called when a method is received from remote Kite.
func (c *Client) runMethod(method *Method, args *dnode.Partial) {
	var (
		callFunc func(interface{}, *Error)
		request  *Request
	)

	// Recover dnode argument errors and send them back. The caller can use
	// functions like MustString(), MustSlice()... without the fear of panic.
	defer func() {
		if r := recover(); r != nil {
			debug.PrintStack()
			kiteErr := createError(r)
			c.LocalKite.Log.Error(kiteErr.Error()) // let's log it too :)
			callFunc(nil, kiteErr)
		}
	}()

	// The request that will be constructed from incoming dnode message.
	request, callFunc = c.newRequest(method.name, args)
	if method.authenticate {
		if err := request.authenticate(); err != nil {
			callFunc(nil, err)
			return
		}
	} else {
		// if not valided accept any username it sends, also useful for test
		// cases.
		request.Username = request.Client.Kite.Username
	}

	method.mu.Lock()
	if !method.initialized {
		method.preHandlers = append(method.preHandlers, c.LocalKite.preHandlers...)
		method.postHandlers = append(method.postHandlers, c.LocalKite.postHandlers...)
		method.initialized = true
	}
	method.mu.Unlock()

	// Call the handler functions.
	result, err := method.ServeKite(request)

	callFunc(result, createError(err))
}

// runCallback is called when a callback method call is received from remote Kite.
func (c *Client) runCallback(callback func(*dnode.Partial), args *dnode.Partial) {
	// Do not panic no matter what.
	defer func() {
		if err := recover(); err != nil {
			c.LocalKite.Log.Warning("Error in calling the callback function : %v", err)
		}
	}()

	// Call the callback function.
	callback(args)
}

// newRequest returns a new *Request from the method and arguments passed.
func (c *Client) newRequest(method string, args *dnode.Partial) (*Request, func(interface{}, *Error)) {
	// Parse dnode method arguments: [options]
	var options callOptions
	args.One().MustUnmarshal(&options)

	// Notify the handlers registered with Kite.OnFirstRequest().
	if _, ok := c.session.(*sockjsclient.WebsocketSession); !ok {
		c.firstRequestHandlersNotified.Do(func() {
			c.Kite = options.Kite
			c.LocalKite.callOnFirstRequestHandlers(c)
		})
	}

	request := &Request{
		Method:    method,
		Args:      options.WithArgs,
		LocalKite: c.LocalKite,
		Client:    c,
		Auth:      options.Auth,
		Context:   cache.NewMemory(),
	}

	// Call response callback function, send back our response
	callFunc := func(result interface{}, err *Error) {
		if options.ResponseCallback.Caller == nil {
			return
		}

		// Only argument to the callback.
		response := Response{
			Result: result,
			Error:  err,
		}

		if err := options.ResponseCallback.Call(response); err != nil {
			c.LocalKite.Log.Error(err.Error())
		}
	}

	return request, callFunc
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.
	// Following cast means, session is opened by the client.
	if _, ok := r.Client.session.(*sockjsclient.WebsocketSession); ok {
		return nil
	}

	if r.Auth == nil {
		return &Error{
			Type:    "authenticationError",
			Message: "No authentication information is provided",
		}
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Auth.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Auth.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Replace username of the remote Kite with the username that client send
	// us. This prevents a Kite to impersonate someone else's Kite.
	r.Client.SetUsername(r.Username)
	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Auth.Key, r.LocalKite.RSAKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	// check if we have an audience and it matches our own signature
	audience, ok := token.Claims["aud"].(string)
	if ok && audience != "/" {
		if checkAudience(k.Kite().String(), audience); err != nil {
			return err
		}
	}

	// We don't check for exp and nbf claims here because jwt-go package
	// already checks them.
	username, ok := token.Claims["sub"].(string)
	if !ok {
		return errors.New("Username is not present in token")
	}

	// replace the requester username so we reflect the validated
	r.Username = username

	return nil
}

func checkAudience(kiteRepr, audience string) error {
	a, err := protocol.KiteFromString(audience)
	if err != nil {
		return err
	}

	// it doesn't make sense to return an error if the audience is fully empty
	if a.Username == "" {
		return nil
	}

	// this is good so our kites can also work behind load balancers
	threePart := fmt.Sprintf("/%s/%s/%s", a.Username, a.Environment, a.Name)

	// now check if the first three fields are matching our own fields
	if !strings.HasPrefix(kiteRepr, threePart) {
		return fmt.Errorf("Invalid audience in token. Have: '%s' Must be a part of: '%s'",
			audience, kiteRepr)
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from kite key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateSimpleKiteKey authenticates user from the given kite key and
// returns the authenticated username. It's the same as AuthenticateFromKiteKey
// but can be used without the need for a *kite.Request.
func (k *Kite) AuthenticateSimpleKiteKey(key string) (string, error) {
	token, err := jwt.Parse(key, kitekey.GetKontrolKey)
	if err != nil {
		return "", err
	}

	if !token.Valid {
		return "", errors.New("Invalid signature in token")
	}

	username, ok := token.Claims["sub"].(string)
	if !ok {
		return "", errors.New("Username is not present in token")
	}

	// return authenticated username
	return username, nil
}
<commit_msg>request: add XHR based session option too<commit_after>package kite

import (
	"errors"
	"fmt"
	"runtime/debug"
	"strings"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/cache"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/protocol"
	"github.com/koding/kite/sockjsclient"
)

// Request contains information about the incoming request.
type Request struct {
	// Method defines the method name which is invoked by the incoming request
	Method string

	// Args defines the incoming arguments for the given method
	Args *dnode.Partial

	// LocalKite defines a context for the local kite
	LocalKite *Kite

	// Client defines a context for the remote kite
	Client *Client

	// Username defines the username which the incoming request is bound to.
	// This is authenticated and validated if authentication is enabled.
	Username string

	// Auth stores the authentication information for the incoming request and
	// the type of authentication. This is not used when authentication is disabled
	Auth *Auth

	// Context holds a context that used by the current ServeKite handler. Any
	// items added to the Context can be fetched from other handlers in the
	// chain. This is useful with PreHandle and PostHandle handlers to pass
	// data between handlers.
	Context cache.Cache
}

// Response is the type of the object that is returned from request handlers
// and the type of only argument that is passed to callback functions.
type Response struct {
	Error  *Error      `json:"error" dnode:"-"`
	Result interface{} `json:"result"`
}

// runMethod is called when a method is received from remote Kite.
func (c *Client) runMethod(method *Method, args *dnode.Partial) {
	var (
		callFunc func(interface{}, *Error)
		request  *Request
	)

	// Recover dnode argument errors and send them back. The caller can use
	// functions like MustString(), MustSlice()... without the fear of panic.
	defer func() {
		if r := recover(); r != nil {
			debug.PrintStack()
			kiteErr := createError(r)
			c.LocalKite.Log.Error(kiteErr.Error()) // let's log it too :)
			callFunc(nil, kiteErr)
		}
	}()

	// The request that will be constructed from incoming dnode message.
	request, callFunc = c.newRequest(method.name, args)
	if method.authenticate {
		if err := request.authenticate(); err != nil {
			callFunc(nil, err)
			return
		}
	} else {
		// if not valided accept any username it sends, also useful for test
		// cases.
		request.Username = request.Client.Kite.Username
	}

	method.mu.Lock()
	if !method.initialized {
		method.preHandlers = append(method.preHandlers, c.LocalKite.preHandlers...)
		method.postHandlers = append(method.postHandlers, c.LocalKite.postHandlers...)
		method.initialized = true
	}
	method.mu.Unlock()

	// Call the handler functions.
	result, err := method.ServeKite(request)

	callFunc(result, createError(err))
}

// runCallback is called when a callback method call is received from remote Kite.
func (c *Client) runCallback(callback func(*dnode.Partial), args *dnode.Partial) {
	// Do not panic no matter what.
	defer func() {
		if err := recover(); err != nil {
			c.LocalKite.Log.Warning("Error in calling the callback function : %v", err)
		}
	}()

	// Call the callback function.
	callback(args)
}

// newRequest returns a new *Request from the method and arguments passed.
func (c *Client) newRequest(method string, args *dnode.Partial) (*Request, func(interface{}, *Error)) {
	// Parse dnode method arguments: [options]
	var options callOptions
	args.One().MustUnmarshal(&options)

	// Notify the handlers registered with Kite.OnFirstRequest().
	if _, ok := c.session.(*sockjsclient.WebsocketSession); !ok {
		c.firstRequestHandlersNotified.Do(func() {
			c.Kite = options.Kite
			c.LocalKite.callOnFirstRequestHandlers(c)
		})
	}

	request := &Request{
		Method:    method,
		Args:      options.WithArgs,
		LocalKite: c.LocalKite,
		Client:    c,
		Auth:      options.Auth,
		Context:   cache.NewMemory(),
	}

	// Call response callback function, send back our response
	callFunc := func(result interface{}, err *Error) {
		if options.ResponseCallback.Caller == nil {
			return
		}

		// Only argument to the callback.
		response := Response{
			Result: result,
			Error:  err,
		}

		if err := options.ResponseCallback.Call(response); err != nil {
			c.LocalKite.Log.Error(err.Error())
		}
	}

	return request, callFunc
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.  Following casts
	// means, session is opened by the client.
	if _, ok := r.Client.session.(*sockjsclient.WebsocketSession); ok {
		return nil
	}

	if _, ok := r.Client.session.(*sockjsclient.XHRSession); ok {
		return nil
	}

	if r.Auth == nil {
		return &Error{
			Type:    "authenticationError",
			Message: "No authentication information is provided",
		}
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Auth.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Auth.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Replace username of the remote Kite with the username that client send
	// us. This prevents a Kite to impersonate someone else's Kite.
	r.Client.SetUsername(r.Username)
	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Auth.Key, r.LocalKite.RSAKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	// check if we have an audience and it matches our own signature
	audience, ok := token.Claims["aud"].(string)
	if ok && audience != "/" {
		if checkAudience(k.Kite().String(), audience); err != nil {
			return err
		}
	}

	// We don't check for exp and nbf claims here because jwt-go package
	// already checks them.
	username, ok := token.Claims["sub"].(string)
	if !ok {
		return errors.New("Username is not present in token")
	}

	// replace the requester username so we reflect the validated
	r.Username = username

	return nil
}

func checkAudience(kiteRepr, audience string) error {
	a, err := protocol.KiteFromString(audience)
	if err != nil {
		return err
	}

	// it doesn't make sense to return an error if the audience is fully empty
	if a.Username == "" {
		return nil
	}

	// this is good so our kites can also work behind load balancers
	threePart := fmt.Sprintf("/%s/%s/%s", a.Username, a.Environment, a.Name)

	// now check if the first three fields are matching our own fields
	if !strings.HasPrefix(kiteRepr, threePart) {
		return fmt.Errorf("Invalid audience in token. Have: '%s' Must be a part of: '%s'",
			audience, kiteRepr)
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from kite key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateSimpleKiteKey authenticates user from the given kite key and
// returns the authenticated username. It's the same as AuthenticateFromKiteKey
// but can be used without the need for a *kite.Request.
func (k *Kite) AuthenticateSimpleKiteKey(key string) (string, error) {
	token, err := jwt.Parse(key, kitekey.GetKontrolKey)
	if err != nil {
		return "", err
	}

	if !token.Valid {
		return "", errors.New("Invalid signature in token")
	}

	username, ok := token.Claims["sub"].(string)
	if !ok {
		return "", errors.New("Username is not present in token")
	}

	// return authenticated username
	return username, nil
}
<|endoftext|>


<commit_before>// +build linux

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

Copyright 2016 Intel Corporation

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 cpu provides implementation of snap-plugin-collector-cpu plugin
package cpu

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/intelsdi-x/snap-plugin-utilities/ns"
	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"
	"github.com/intelsdi-x/snap/core/ctypes"
)

const (
	//vendor namespace part
	vendor = "intel"

	//fs namespace part
	fs = "procfs"

	//pluginName namespace part
	pluginName = "cpu"

	// version of cpu plugin
	version = 2

	//pluginType type of plugin
	pluginType = plugin.CollectorPluginType

	//userProcStat "user" metric from /proc/stat
	userProcStat = "user"

	//userProcStat "nice" metric from /proc/stat
	niceProcStat = "nice"

	//systemProcStat "system" metric from /proc/stat
	systemProcStat = "system"

	//idleProcStat "idle" metric from /proc/stat
	idleProcStat = "idle"

	//iowaitProcStat "iowait" metric from /proc/stat
	iowaitProcStat = "iowait"

	//irqProcStat "irq" metric from /proc/stat
	irqProcStat = "irq"

	//softirqProcStat "softirq" metric from /proc/stat
	softirqProcStat = "softirq"

	//stealProcStat "steal" metric from /proc/stat
	stealProcStat = "steal"

	//guestProcStat "guest" metric from /proc/stat
	guestProcStat = "guest"

	//guestNiceProcStat "guest_nice" metric from /proc/stat
	guestNiceProcStat = "guest_nice"

	//activeProcStat "active" snap metric
	activeProcStat = "active"

	//utilizationProcStat "utilization" snap metric
	utilizationProcStat = "utilization"

	//jiffiesRepresentationType jiffies representation type
	jiffiesRepresentationType = "jiffies"

	//percentageRepresentationType percentage representation type
	percentageRepresentationType = "percentage"

	//maxNamespaceSize max size of namespace for metrics
	maxNamespaceSize = 5

	//allCPU string indentifier for aggregation metrics (for all CPUs)
	allCPU = "all"

	//firstCPU string indentifier for the first CPU
	firstCPU = "0"

	//secondCPU string indentifier for the second CPU
	secondCPU = "1"

	//cpuStr string indentifier for /proc/stat line which have desired CPU metrics
	cpuStr = "cpu"
)

//cpuInfo source of data for metrics
var cpuInfo = "/proc/stat"

// GetMetricTypes returns list of available metric types
// It returns error in case retrieval was not successful
func (p *Plugin) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	if !p.initialized {
		if err := p.init(cfg.Table()); err != nil {
			return nil, err
		}
	}
	metricTypes := []plugin.MetricType{}

	namespaces := []string{}

	err := ns.FromMap(p.stats, filepath.Join(vendor, fs, pluginName), &namespaces)

	if err != nil {
		return nil, err
	}

	for _, namespace := range namespaces {
		namespace = strings.TrimRight(namespace, string(os.PathSeparator))
		metricType := plugin.MetricType{
			Namespace_: core.NewNamespace(strings.Split(namespace, string(os.PathSeparator))...)}
		metricTypes = append(metricTypes, metricType)
	}
	return metricTypes, nil
}

// CollectMetrics returns list of requested metric values
// It returns error in case retrieval was not successful
func (p *Plugin) CollectMetrics(metricTypes []plugin.MetricType) ([]plugin.MetricType, error) {
	metrics := []plugin.MetricType{}
	for _, metricType := range metricTypes {
		if !p.initialized {
			if err := p.init(metricType.Config().Table()); err != nil {
				return nil, err
			}
		}
		ns := metricType.Namespace()
		if len(ns) != maxNamespaceSize {
			return nil, fmt.Errorf("Incorrect namespace length (len = %d)", len(ns))
		}

		val, err := getMapValueByNamespace(p.stats, ns[3:].Strings())

		if err != nil {
			return metrics, err
		}

		metric := plugin.MetricType{
			Namespace_: ns,
			Data_:      val,
			Timestamp_: time.Now(),
		}
		metrics = append(metrics, metric)
	}
	return metrics, nil
}

// GetConfigPolicy returns config policy
// It returns error in case retrieval was not successful
func (p *Plugin) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	return cpolicy.New(), nil
}

//Meta returns meta data for plugin
func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		pluginName,
		version,
		pluginType,
		[]string{},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1))
}

func (p *Plugin) init(cfg map[string]ctypes.ConfigValue) error {
	path := cpuInfo
	if procPath, ok := cfg["proc_path"]; ok {
		path = procPath.(ctypes.ConfigValueStr).Value + "/stat"
	}
	fh, err := os.Open(path)
	if err != nil {
		return err
	}
	defer fh.Close()

	var procStatMetricsNumber int
	p.cpuMetricsNumber, procStatMetricsNumber, err = getInitialProcStatData()
	if err != nil {
		return err
	}

	//initialize metric names arrays
	p.procStatMetricsNames = []string{userProcStat, niceProcStat, systemProcStat, idleProcStat,
		iowaitProcStat, irqProcStat, softirqProcStat, stealProcStat, guestProcStat, guestNiceProcStat}
	snapSpecificMetricsNames := []string{activeProcStat, utilizationProcStat}

	//build snapMetricsNames to support different kernels
	//var snapMetricsNames []string
	p.snapMetricsNames = append(p.snapMetricsNames, p.procStatMetricsNames[0:procStatMetricsNumber]...)
	p.snapMetricsNames = append(p.snapMetricsNames, snapSpecificMetricsNames...)
	p.stats = make(map[string]interface{})
	p.prevMetricsSum = make(map[string]float64)
	if err := getStats(p.stats, p.prevMetricsSum, p.cpuMetricsNumber,
		p.snapMetricsNames, p.procStatMetricsNames); err != nil {
		return err
	}
	p.initialized = true
	return nil
}

// New creates instance of interface info plugin
func New() *Plugin {
	host, err := os.Hostname()
	if err != nil {
		host = "localhost"
	}
	p := &Plugin{host: host}
	return p
}

//Plugin cpu plugin struct which gathers plugin specific data
type Plugin struct {
	initialized          bool
	host                 string
	cpuMetricsNumber     int // number of cpu + "all" metric
	stats                map[string]interface{}
	prevMetricsSum       map[string]float64
	procStatMetricsNames []string
	snapMetricsNames     []string
}

//getStats gets metrics from /proc/stat output and calculates snap specific metrics
func getStats(stats map[string]interface{}, prevMetricsSum map[string]float64, cpuMetricsNumber int,
	snapMetricsNames []string, procStatMetricsNames []string) (err error) {
	fh, err := os.Open(cpuInfo)
	if err != nil {
		return err
	}
	defer fh.Close()

	scanner := bufio.NewScanner(fh)
	for i := 0; i < cpuMetricsNumber; i++ {
		scanErr := scanner.Scan()
		if !scanErr {
			return fmt.Errorf("Wrong %s format", cpuInfo)
		}
		fields := strings.Fields(scanner.Text())

		if len(fields) < 2 {
			return fmt.Errorf("Wrong %s format", cpuInfo)
		}

		cpuID := strings.TrimSpace(fields[0])
		if cpuID == cpuStr {
			if err != nil {
				return err
			}
			cpuID = allCPU //change CPU identifier for aggregation metrics
		} else {
			cpuID = cpuID[len(cpuID)-1:] //get the last character
		}
		metrics := fields[1:]

		if len(metrics) != len(procStatMetricsNames) {
			return fmt.Errorf("Wrong data length. Expected {%d} is {%d}",
				len(procStatMetricsNames), len(metrics))
		}

		//sum of new data in line
		currDataSum, err := strTabSum(metrics)
		if err != nil {
			return err
		}

		metricStats := make(map[string]interface{})
		for j := range snapMetricsNames {

			metricName := snapMetricsNames[j]
			var currVal float64
			//data collecting, there is an assumption that firstly metrics from /proc/stat/
			//are gathered then snap specific metrics (e.g. active and utilization are calculated)
			if metricName == activeProcStat {
				idleVal, err := getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(idleProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}
				currVal = currDataSum - idleVal
			} else if metricName == utilizationProcStat {
				nonActiveVal, err := getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(idleProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}

				currVal = currDataSum - nonActiveVal

				nonActiveVal, err = getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(iowaitProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}

				currVal = currVal - nonActiveVal
			} else {
				currVal, err = strconv.ParseFloat(metrics[j], 64)
				if err != nil {
					return err
				}
			}

			metricStats[getNamespaceMetricPart(metricName, percentageRepresentationType)] = float64(0)

			if mapKeyExists(cpuID, prevMetricsSum) {
				diffSum := currDataSum - prevMetricsSum[cpuID]
				if diffSum != 0 {
					prevVal, err := getMapFloatValueByNamespace(stats,
						[]string{cpuID, getNamespaceMetricPart(metricName, jiffiesRepresentationType)})
					if err != nil {
						return err
					}
					metricStats[getNamespaceMetricPart(metricName, percentageRepresentationType)] =
						float64(100 * (currVal - prevVal) / diffSum)
				}
			}
			metricStats[getNamespaceMetricPart(metricName, jiffiesRepresentationType)] = currVal
		}
		stats[cpuID] = metricStats
		prevMetricsSum[cpuID] = currDataSum
	}
	return nil
}

//getNamespaceMetricPart builds part of namespace specific for metric and representation type
func getNamespaceMetricPart(metricName string, representationType string) (s string) {
	s = metricName + "_" + representationType
	return s
}

//mapKeyExists checks if element with given key exists in map
func mapKeyExists(key string, m map[string]float64) bool {
	var ret = false
	if _, ok := m[key]; ok {
		ret = true
	}
	return ret
}

//strTabSum adds string data as float
func strTabSum(metrics []string) (sum float64, err error) {
	sum = 0
	for i := range metrics {
		val, err := strconv.ParseFloat(metrics[i], 64)
		if err != nil {
			return sum, err
		}
		sum += val
	}
	return sum, err
}

//getMapFloatValueByNamespace gets value as float from map by namespace given in array of strings
func getMapFloatValueByNamespace(m map[string]interface{}, ns []string) (val float64, err error) {
	var interfaceVal interface{}
	interfaceVal, err = getMapValueByNamespace(m, ns)
	if err != nil {
		return val, err
	}
	var errBool bool
	val, errBool = interfaceVal.(float64)
	if !errBool {
		return val, fmt.Errorf("Parsing error")
	}
	return val, err
}

//getMapValueByNamespace gets value from map by namespace given in array of strings
func getMapValueByNamespace(m map[string]interface{}, ns []string) (val interface{}, err error) {
	if len(ns) == 0 {
		return val, fmt.Errorf("Namespace length equal to zero")
	}

	current := ns[0]

	if len(ns) == 1 {
		if val, ok := m[current]; ok {
			return val, err
		}
		return val, fmt.Errorf("Key does not exist in map {key %s}", current)
	}

	if v, ok := m[current].(map[string]interface{}); ok {
		val, err = getMapValueByNamespace(v, ns[1:])
		return val, err
	}
	return val, err
}

//getInitialProcStatData gets number of CPUs and number of metrics available in /proc/stat output
func getInitialProcStatData() (cpuMetricsNumber int, procStatMetricNumber int, err error) {
	fh, err := os.Open(cpuInfo)
	if err != nil {
		return cpuMetricsNumber, procStatMetricNumber, err
	}
	defer fh.Close()

	scanner := bufio.NewScanner(fh)
	//read the first line to start the loop
	scanErr := scanner.Scan()
	if scanErr != true {
		return cpuMetricsNumber, procStatMetricNumber, fmt.Errorf("Cannot read from /proc/stat")
	}

	procStatLine := strings.Fields(scanner.Text())
	procStatMetricNumber = len(procStatLine) - 1 //without cpu identifier (e.g. cpu or cpu0)

	cpuMetricsNumber = 0
	for strings.Contains(procStatLine[0], cpuStr) {
		//check line length, compare current length with length of the first line
		//length of the first line = len(procStatMetricNumber) + CPU identifier
		if len(procStatLine) != (procStatMetricNumber + 1) {
			return cpuMetricsNumber, procStatMetricNumber, fmt.Errorf("Incorrect /proc/stat output")
		}
		cpuMetricsNumber++
		//read the next line to be able to check loop condition
		scanErr = scanner.Scan()
		if scanErr != true {
			break //no more lines to read
		}
		procStatLine = strings.Fields(scanner.Text())
	}
	return cpuMetricsNumber, procStatMetricNumber, err
}
<commit_msg>Uniformize configuration across plugins<commit_after>// +build linux

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

Copyright 2016 Intel Corporation

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 cpu provides implementation of snap-plugin-collector-cpu plugin
package cpu

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"
	"github.com/intelsdi-x/snap/core/ctypes"

	"github.com/intelsdi-x/snap-plugin-utilities/ns"
)

const (
	//vendor namespace part
	vendor = "intel"

	//fs namespace part
	fs = "procfs"

	//pluginName namespace part
	pluginName = "cpu"

	// version of cpu plugin
	version = 2

	//pluginType type of plugin
	pluginType = plugin.CollectorPluginType

	//userProcStat "user" metric from /proc/stat
	userProcStat = "user"

	//userProcStat "nice" metric from /proc/stat
	niceProcStat = "nice"

	//systemProcStat "system" metric from /proc/stat
	systemProcStat = "system"

	//idleProcStat "idle" metric from /proc/stat
	idleProcStat = "idle"

	//iowaitProcStat "iowait" metric from /proc/stat
	iowaitProcStat = "iowait"

	//irqProcStat "irq" metric from /proc/stat
	irqProcStat = "irq"

	//softirqProcStat "softirq" metric from /proc/stat
	softirqProcStat = "softirq"

	//stealProcStat "steal" metric from /proc/stat
	stealProcStat = "steal"

	//guestProcStat "guest" metric from /proc/stat
	guestProcStat = "guest"

	//guestNiceProcStat "guest_nice" metric from /proc/stat
	guestNiceProcStat = "guest_nice"

	//activeProcStat "active" snap metric
	activeProcStat = "active"

	//utilizationProcStat "utilization" snap metric
	utilizationProcStat = "utilization"

	//jiffiesRepresentationType jiffies representation type
	jiffiesRepresentationType = "jiffies"

	//percentageRepresentationType percentage representation type
	percentageRepresentationType = "percentage"

	//maxNamespaceSize max size of namespace for metrics
	maxNamespaceSize = 5

	//allCPU string indentifier for aggregation metrics (for all CPUs)
	allCPU = "all"

	//firstCPU string indentifier for the first CPU
	firstCPU = "0"

	//secondCPU string indentifier for the second CPU
	secondCPU = "1"

	//cpuStr string indentifier for /proc/stat line which have desired CPU metrics
	cpuStr = "cpu"
)

//cpuInfo source of data for metrics
var cpuInfo = "/proc/stat"

// GetMetricTypes returns list of available metric types
// It returns error in case retrieval was not successful
func (p *Plugin) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	if !p.initialized {
		if err := p.init(cfg.Table()); err != nil {
			return nil, err
		}
	}
	metricTypes := []plugin.MetricType{}

	namespaces := []string{}

	err := ns.FromMap(p.stats, filepath.Join(vendor, fs, pluginName), &namespaces)

	if err != nil {
		return nil, err
	}

	for _, namespace := range namespaces {
		namespace = strings.TrimRight(namespace, string(os.PathSeparator))
		metricType := plugin.MetricType{
			Namespace_: core.NewNamespace(strings.Split(namespace, string(os.PathSeparator))...)}
		metricTypes = append(metricTypes, metricType)
	}
	return metricTypes, nil
}

// CollectMetrics returns list of requested metric values
// It returns error in case retrieval was not successful
func (p *Plugin) CollectMetrics(metricTypes []plugin.MetricType) ([]plugin.MetricType, error) {
	metrics := []plugin.MetricType{}
	for _, metricType := range metricTypes {
		if !p.initialized {
			if err := p.init(metricType.Config().Table()); err != nil {
				return nil, err
			}
		}
		ns := metricType.Namespace()
		if len(ns) != maxNamespaceSize {
			return nil, fmt.Errorf("Incorrect namespace length (len = %d)", len(ns))
		}

		val, err := getMapValueByNamespace(p.stats, ns[3:].Strings())

		if err != nil {
			return metrics, err
		}

		metric := plugin.MetricType{
			Namespace_: ns,
			Data_:      val,
			Timestamp_: time.Now(),
		}
		metrics = append(metrics, metric)
	}
	return metrics, nil
}

// GetConfigPolicy returns config policy
// It returns error in case retrieval was not successful
func (p *Plugin) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("proc_path", false, "/proc")
	node := cpolicy.NewPolicyNode()
	node.Add(rule)
	cp.Add([]string{vendor, fs, pluginName}, node)
	return cp, nil
}

//Meta returns meta data for plugin
func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		pluginName,
		version,
		pluginType,
		[]string{},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1))
}

func (p *Plugin) init(cfg map[string]ctypes.ConfigValue) error {
	path := cpuInfo
	if procPath, ok := cfg["proc_path"]; ok {
		path = procPath.(ctypes.ConfigValueStr).Value + "/stat"
	}
	fh, err := os.Open(path)
	if err != nil {
		return err
	}
	defer fh.Close()

	var procStatMetricsNumber int
	p.cpuMetricsNumber, procStatMetricsNumber, err = getInitialProcStatData()
	if err != nil {
		return err
	}

	//initialize metric names arrays
	p.procStatMetricsNames = []string{userProcStat, niceProcStat, systemProcStat, idleProcStat,
		iowaitProcStat, irqProcStat, softirqProcStat, stealProcStat, guestProcStat, guestNiceProcStat}
	snapSpecificMetricsNames := []string{activeProcStat, utilizationProcStat}

	//build snapMetricsNames to support different kernels
	//var snapMetricsNames []string
	p.snapMetricsNames = append(p.snapMetricsNames, p.procStatMetricsNames[0:procStatMetricsNumber]...)
	p.snapMetricsNames = append(p.snapMetricsNames, snapSpecificMetricsNames...)
	p.stats = make(map[string]interface{})
	p.prevMetricsSum = make(map[string]float64)
	if err := getStats(p.stats, p.prevMetricsSum, p.cpuMetricsNumber,
		p.snapMetricsNames, p.procStatMetricsNames); err != nil {
		return err
	}
	p.initialized = true
	return nil
}

// New creates instance of interface info plugin
func New() *Plugin {
	host, err := os.Hostname()
	if err != nil {
		host = "localhost"
	}
	p := &Plugin{host: host}
	return p
}

//Plugin cpu plugin struct which gathers plugin specific data
type Plugin struct {
	initialized          bool
	host                 string
	cpuMetricsNumber     int // number of cpu + "all" metric
	stats                map[string]interface{}
	prevMetricsSum       map[string]float64
	procStatMetricsNames []string
	snapMetricsNames     []string
}

//getStats gets metrics from /proc/stat output and calculates snap specific metrics
func getStats(stats map[string]interface{}, prevMetricsSum map[string]float64, cpuMetricsNumber int,
	snapMetricsNames []string, procStatMetricsNames []string) (err error) {
	fh, err := os.Open(cpuInfo)
	if err != nil {
		return err
	}
	defer fh.Close()

	scanner := bufio.NewScanner(fh)
	for i := 0; i < cpuMetricsNumber; i++ {
		scanErr := scanner.Scan()
		if !scanErr {
			return fmt.Errorf("Wrong %s format", cpuInfo)
		}
		fields := strings.Fields(scanner.Text())

		if len(fields) < 2 {
			return fmt.Errorf("Wrong %s format", cpuInfo)
		}

		cpuID := strings.TrimSpace(fields[0])
		if cpuID == cpuStr {
			if err != nil {
				return err
			}
			cpuID = allCPU //change CPU identifier for aggregation metrics
		} else {
			cpuID = cpuID[len(cpuID)-1:] //get the last character
		}
		metrics := fields[1:]

		if len(metrics) != len(procStatMetricsNames) {
			return fmt.Errorf("Wrong data length. Expected {%d} is {%d}",
				len(procStatMetricsNames), len(metrics))
		}

		//sum of new data in line
		currDataSum, err := strTabSum(metrics)
		if err != nil {
			return err
		}

		metricStats := make(map[string]interface{})
		for j := range snapMetricsNames {

			metricName := snapMetricsNames[j]
			var currVal float64
			//data collecting, there is an assumption that firstly metrics from /proc/stat/
			//are gathered then snap specific metrics (e.g. active and utilization are calculated)
			if metricName == activeProcStat {
				idleVal, err := getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(idleProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}
				currVal = currDataSum - idleVal
			} else if metricName == utilizationProcStat {
				nonActiveVal, err := getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(idleProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}

				currVal = currDataSum - nonActiveVal

				nonActiveVal, err = getMapFloatValueByNamespace(metricStats,
					[]string{getNamespaceMetricPart(iowaitProcStat, jiffiesRepresentationType)})
				if err != nil {
					return err
				}

				currVal = currVal - nonActiveVal
			} else {
				currVal, err = strconv.ParseFloat(metrics[j], 64)
				if err != nil {
					return err
				}
			}

			metricStats[getNamespaceMetricPart(metricName, percentageRepresentationType)] = float64(0)

			if mapKeyExists(cpuID, prevMetricsSum) {
				diffSum := currDataSum - prevMetricsSum[cpuID]
				if diffSum != 0 {
					prevVal, err := getMapFloatValueByNamespace(stats,
						[]string{cpuID, getNamespaceMetricPart(metricName, jiffiesRepresentationType)})
					if err != nil {
						return err
					}
					metricStats[getNamespaceMetricPart(metricName, percentageRepresentationType)] =
						float64(100 * (currVal - prevVal) / diffSum)
				}
			}
			metricStats[getNamespaceMetricPart(metricName, jiffiesRepresentationType)] = currVal
		}
		stats[cpuID] = metricStats
		prevMetricsSum[cpuID] = currDataSum
	}
	return nil
}

//getNamespaceMetricPart builds part of namespace specific for metric and representation type
func getNamespaceMetricPart(metricName string, representationType string) (s string) {
	s = metricName + "_" + representationType
	return s
}

//mapKeyExists checks if element with given key exists in map
func mapKeyExists(key string, m map[string]float64) bool {
	var ret = false
	if _, ok := m[key]; ok {
		ret = true
	}
	return ret
}

//strTabSum adds string data as float
func strTabSum(metrics []string) (sum float64, err error) {
	sum = 0
	for i := range metrics {
		val, err := strconv.ParseFloat(metrics[i], 64)
		if err != nil {
			return sum, err
		}
		sum += val
	}
	return sum, err
}

//getMapFloatValueByNamespace gets value as float from map by namespace given in array of strings
func getMapFloatValueByNamespace(m map[string]interface{}, ns []string) (val float64, err error) {
	var interfaceVal interface{}
	interfaceVal, err = getMapValueByNamespace(m, ns)
	if err != nil {
		return val, err
	}
	var errBool bool
	val, errBool = interfaceVal.(float64)
	if !errBool {
		return val, fmt.Errorf("Parsing error")
	}
	return val, err
}

//getMapValueByNamespace gets value from map by namespace given in array of strings
func getMapValueByNamespace(m map[string]interface{}, ns []string) (val interface{}, err error) {
	if len(ns) == 0 {
		return val, fmt.Errorf("Namespace length equal to zero")
	}

	current := ns[0]

	if len(ns) == 1 {
		if val, ok := m[current]; ok {
			return val, err
		}
		return val, fmt.Errorf("Key does not exist in map {key %s}", current)
	}

	if v, ok := m[current].(map[string]interface{}); ok {
		val, err = getMapValueByNamespace(v, ns[1:])
		return val, err
	}
	return val, err
}

//getInitialProcStatData gets number of CPUs and number of metrics available in /proc/stat output
func getInitialProcStatData() (cpuMetricsNumber int, procStatMetricNumber int, err error) {
	fh, err := os.Open(cpuInfo)
	if err != nil {
		return cpuMetricsNumber, procStatMetricNumber, err
	}
	defer fh.Close()

	scanner := bufio.NewScanner(fh)
	//read the first line to start the loop
	scanErr := scanner.Scan()
	if scanErr != true {
		return cpuMetricsNumber, procStatMetricNumber, fmt.Errorf("Cannot read from /proc/stat")
	}

	procStatLine := strings.Fields(scanner.Text())
	procStatMetricNumber = len(procStatLine) - 1 //without cpu identifier (e.g. cpu or cpu0)

	cpuMetricsNumber = 0
	for strings.Contains(procStatLine[0], cpuStr) {
		//check line length, compare current length with length of the first line
		//length of the first line = len(procStatMetricNumber) + CPU identifier
		if len(procStatLine) != (procStatMetricNumber + 1) {
			return cpuMetricsNumber, procStatMetricNumber, fmt.Errorf("Incorrect /proc/stat output")
		}
		cpuMetricsNumber++
		//read the next line to be able to check loop condition
		scanErr = scanner.Scan()
		if scanErr != true {
			break //no more lines to read
		}
		procStatLine = strings.Fields(scanner.Text())
	}
	return cpuMetricsNumber, procStatMetricNumber, err
}
<|endoftext|>