<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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 webpush provides helper functions for sending encrpyted payloads
// using the Web Push protocol.
//
// Sending a message:
//   import (
//     "github.com/googlechrome/push-encryption-go/webpush"
//   )
//
//   func main() {
//     // The values that make up the Subscription struct come from the browser
//     sub := &webpush.Subscription{endpoint, key, auth}
//     webpush.Send(nil, sub, "Yay! Web Push!", "")
//   }
//
// You can turn a JSON string representation of a PushSubscription object you
// collected from the browser into a Subscription struct with a helper function.
//
//   var exampleJSON = []byte(`{"endpoint": "...", "keys": {"p256dh": "...", "auth": "..."}}`)
//   sub, err := SubscriptionFromJSON(exampleJSON)
//
// If the push service requires an authentication header (notably Google Cloud
// Messaging, used by Chrome) then you can add that as a fourth parameter:
//
//   if strings.Contains(sub.Endpoint, "https://android.googleapis.com/gcm/send/") {
//     webpush.Send(nil, sub, "A message for Chrome", myGCMKey)
//   }
package webpush

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"strings"
)

const (
	maxPayloadLength = 4078
)

var (
	authInfo = []byte("Content-Encoding: auth\x00")
	curve    = elliptic.P256()

	// Generate a random key pair to be used for the encryption. Overridable for
	// testing.
	randomKey = func() (priv []byte, pub []byte, err error) {
		priv, x, y, err := elliptic.GenerateKey(curve, rand.Reader)
		if err != nil {
			return nil, nil, err
		}

		return priv, elliptic.Marshal(curve, x, y), nil
	}

	// Generate a random salt for the encryption. Overridable for testing.
	randomSalt = func() ([]byte, error) {
		salt := make([]byte, 16)
		_, err := rand.Read(salt)
		if err != nil {
			return nil, err
		}
		return salt, nil
	}
)

// Subscription holds the useful values from a PushSubscription object acquired
// from the browser
type Subscription struct {
	// Endpoint is the URL to send the Web Push message to. Comes from the
	// endpoint field of the PushSubscription.
	Endpoint string
	// Key is the client's public key. From the keys.p256dh field.
	Key []byte
	// Auth is a value used by the client to validate the encryption. From the
	// keys.auth field.
	Auth []byte
}

// SubscriptionFromJSON is a convenience function that takes a JSON encoded
// PushSubscription object acquired from the browser and returns a pointer to a
// Subscription
func SubscriptionFromJSON(b []byte) (*Subscription, error) {
	var sub struct {
		Endpoint string
		Keys     struct {
			P256dh string
			Auth   string
		}
	}
	if err := json.Unmarshal(b, &sub); err != nil {
		return nil, err
	}

	b64 := base64.URLEncoding.WithPadding(base64.NoPadding)

	// Chrome < 52 incorrectly adds padding when Base64 encoding the values, so
	// we need to strip that out
	key, err := b64.DecodeString(strings.TrimRight(sub.Keys.P256dh, "="))
	if err != nil {
		return nil, err
	}

	auth, err := b64.DecodeString(strings.TrimRight(sub.Keys.Auth, "="))
	if err != nil {
		return nil, err
	}

	return &Subscription{sub.Endpoint, key, auth}, nil
}

// EncryptionResult stores the result of encrypting a message. The ciphertext is
// the actual encrypted message, while the salt and server public key are
// required to be sent to the client so that the message can be decrypted.
type EncryptionResult struct {
	Ciphertext      []byte
	Salt            []byte
	ServerPublicKey []byte
}

// Encrypt a message such that it can be sent using the Web Push protocol.
// You can find out more about the various pieces:
//    - https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding
//    - https://en.wikipedia.org/wiki/Elliptic_curve_Diffie%E2%80%93Hellman
//    - https://tools.ietf.org/html/draft-ietf-webpush-encryption
func Encrypt(sub *Subscription, message string) (*EncryptionResult, error) {
	plaintext := []byte(message)
	if len(plaintext) > maxPayloadLength {
		return nil, fmt.Errorf("Payload is too large. The max number of bytes is %d, input is %d bytes.", maxPayloadLength, len(plaintext))
	}

	if len(sub.Key) == 0 {
		return nil, fmt.Errorf("Subscription must include the client's public key")
	}

	if len(sub.Auth) == 0 {
		return nil, fmt.Errorf("Subscription must include the client's auth value")
	}

	salt, err := randomSalt()
	if err != nil {
		return nil, err
	}

	// Use ECDH to derive a shared secret between us and the client. We generate
	// a fresh private/public key pair at random every time we encrypt.
	serverPrivateKey, serverPublicKey, err := randomKey()
	if err != nil {
		return nil, err
	}
	secret, err := sharedSecret(curve, sub.Key, serverPrivateKey)
	if err != nil {
		return nil, err
	}

	// Derive a Pseudo-Random Key (prk) that can be used to further derive our
	// other encryption parameters. These derivations are described in
	// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00
	prk := hkdf(sub.Auth, secret, authInfo, 32)

	// Derive the Content Encryption Key and nonce
	ctx := newContext(sub.Key, serverPublicKey)
	cek := newCEK(ctx, salt, prk)
	nonce := newNonce(ctx, salt, prk)

	// Do the actual encryption
	ciphertext, err := encrypt(plaintext, cek, nonce)
	if err != nil {
		return nil, err
	}

	// Return all of the values needed to construct a Web Push HTTP request.
	return &EncryptionResult{ciphertext, salt, serverPublicKey}, nil
}

func newCEK(ctx, salt, prk []byte) []byte {
	info := newInfo("aesgcm", ctx)
	return hkdf(salt, prk, info, 16)
}

func newNonce(ctx, salt, prk []byte) []byte {
	info := newInfo("nonce", ctx)
	return hkdf(salt, prk, info, 12)
}

// Creates a context for deriving encyption parameters, as described in
// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00.
// The 'context' in this case is just the public keys of both client and server.
// The keys should always be 65 bytes each. The format of the keys is
// described in section 4.3.6 of the (sadly not freely linkable) ANSI X9.62
// specification.
func newContext(clientPublicKey, serverPublicKey []byte) []byte {
	// The context format is:
	// 0x00 || length(clientPublicKey) || clientPublicKey ||
	//         length(serverPublicKey) || serverPublicKey
	// The lengths are 16-bit, Big Endian, unsigned integers so take 2 bytes each.
	cplen := uint16(len(clientPublicKey))
	cplenbuf := make([]byte, 2)
	binary.BigEndian.PutUint16(cplenbuf, cplen)

	splen := uint16(len(serverPublicKey))
	splenbuf := make([]byte, 2)
	binary.BigEndian.PutUint16(splenbuf, splen)

	var ctx []byte
	ctx = append(ctx, 0)
	ctx = append(ctx, cplenbuf...)
	ctx = append(ctx, []byte(clientPublicKey)...)
	ctx = append(ctx, splenbuf...)
	ctx = append(ctx, []byte(serverPublicKey)...)

	return ctx
}

// Returns an info record. See sections 3.2 and 3.3 of
// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00.
// The context argument should match what newContext creates
func newInfo(infoType string, context []byte) []byte {
	var info []byte
	info = append(info, []byte("Content-Encoding: ")...)
	info = append(info, []byte(infoType)...)
	info = append(info, 0)
	info = append(info, []byte("P-256")...)
	info = append(info, context...)
	return info
}

// HMAC-based Extract-and-Expand Key Derivation Function (HKDF)
//
// This is used to derive a secure encryption key from a mostly-secure shared
// secret.
//
// This is a partial implementation of HKDF tailored to our specific purposes.
// In particular, for us the value of N will always be 1, and thus T always
// equals HMAC-Hash(PRK, info | 0x01). This is true because the maximum output
// length we need/allow is 32.
//
// See https://www.rfc-editor.org/rfc/rfc5869.txt
func hkdf(salt, ikm, info []byte, length int) []byte {
	// HMAC length for SHA256 is 32 bytes, so that is the maximum result length.
	if length > 32 {
		panic("Can only produce HKDF outputs up to 32 bytes long")
	}

	// Extract
	mac := hmac.New(sha256.New, salt)
	mac.Write(ikm)
	prk := mac.Sum(nil)

	// Expand
	mac = hmac.New(sha256.New, prk)
	mac.Write(info)
	mac.Write([]byte{1})
	return mac.Sum(nil)[0:length]
}

// Encrypt the plaintext message using AES128/GCM
func encrypt(plaintext, key, nonce []byte) ([]byte, error) {
	// Add padding. There is a uint16 size followed by that number of bytes of
	// padding.
	// TODO: Right now we leave the size at zero. We should add a padding option
	// that allows the payload size to be obscured.
	padding := make([]byte, 2)
	data := append(padding, plaintext...)

	c, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(c)
	if err != nil {
		return nil, err
	}

	return gcm.Seal([]byte{}, nonce, data, nil), nil
}

// Given the coordinates of a party A's public key and the bytes of party B's
// private key, compute a shared secret.
func sharedSecret(curve elliptic.Curve, pub, priv []byte) ([]byte, error) {
	publicX, publicY := elliptic.Unmarshal(curve, pub)
	if publicX == nil {
		return nil, fmt.Errorf("Couldn't unmarshal public key. Not a valid point on the curve.")
	}
	x, _ := curve.ScalarMult(publicX, publicY, priv)
	return x.Bytes(), nil
}
<commit_msg>Fix golint issues<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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 webpush provides helper functions for sending encrpyted payloads
// using the Web Push protocol.
//
// Sending a message:
//   import (
//     "github.com/googlechrome/push-encryption-go/webpush"
//   )
//
//   func main() {
//     // The values that make up the Subscription struct come from the browser
//     sub := &webpush.Subscription{endpoint, key, auth}
//     webpush.Send(nil, sub, "Yay! Web Push!", "")
//   }
//
// You can turn a JSON string representation of a PushSubscription object you
// collected from the browser into a Subscription struct with a helper function.
//
//   var exampleJSON = []byte(`{"endpoint": "...", "keys": {"p256dh": "...", "auth": "..."}}`)
//   sub, err := SubscriptionFromJSON(exampleJSON)
//
// If the push service requires an authentication header (notably Google Cloud
// Messaging, used by Chrome) then you can add that as a fourth parameter:
//
//   if strings.Contains(sub.Endpoint, "https://android.googleapis.com/gcm/send/") {
//     webpush.Send(nil, sub, "A message for Chrome", myGCMKey)
//   }
package webpush

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

const (
	maxPayloadLength = 4078
)

var (
	authInfo = []byte("Content-Encoding: auth\x00")
	curve    = elliptic.P256()

	// Generate a random key pair to be used for the encryption. Overridable for
	// testing.
	randomKey = func() (priv []byte, pub []byte, err error) {
		priv, x, y, err := elliptic.GenerateKey(curve, rand.Reader)
		if err != nil {
			return nil, nil, err
		}

		return priv, elliptic.Marshal(curve, x, y), nil
	}

	// Generate a random salt for the encryption. Overridable for testing.
	randomSalt = func() ([]byte, error) {
		salt := make([]byte, 16)
		_, err := rand.Read(salt)
		if err != nil {
			return nil, err
		}
		return salt, nil
	}
)

// Subscription holds the useful values from a PushSubscription object acquired
// from the browser
type Subscription struct {
	// Endpoint is the URL to send the Web Push message to. Comes from the
	// endpoint field of the PushSubscription.
	Endpoint string
	// Key is the client's public key. From the keys.p256dh field.
	Key []byte
	// Auth is a value used by the client to validate the encryption. From the
	// keys.auth field.
	Auth []byte
}

// SubscriptionFromJSON is a convenience function that takes a JSON encoded
// PushSubscription object acquired from the browser and returns a pointer to a
// Subscription
func SubscriptionFromJSON(b []byte) (*Subscription, error) {
	var sub struct {
		Endpoint string
		Keys     struct {
			P256dh string
			Auth   string
		}
	}
	if err := json.Unmarshal(b, &sub); err != nil {
		return nil, err
	}

	b64 := base64.URLEncoding.WithPadding(base64.NoPadding)

	// Chrome < 52 incorrectly adds padding when Base64 encoding the values, so
	// we need to strip that out
	key, err := b64.DecodeString(strings.TrimRight(sub.Keys.P256dh, "="))
	if err != nil {
		return nil, err
	}

	auth, err := b64.DecodeString(strings.TrimRight(sub.Keys.Auth, "="))
	if err != nil {
		return nil, err
	}

	return &Subscription{sub.Endpoint, key, auth}, nil
}

// EncryptionResult stores the result of encrypting a message. The ciphertext is
// the actual encrypted message, while the salt and server public key are
// required to be sent to the client so that the message can be decrypted.
type EncryptionResult struct {
	Ciphertext      []byte
	Salt            []byte
	ServerPublicKey []byte
}

// Encrypt a message such that it can be sent using the Web Push protocol.
// You can find out more about the various pieces:
//    - https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding
//    - https://en.wikipedia.org/wiki/Elliptic_curve_Diffie%E2%80%93Hellman
//    - https://tools.ietf.org/html/draft-ietf-webpush-encryption
func Encrypt(sub *Subscription, message string) (*EncryptionResult, error) {
	plaintext := []byte(message)
	if len(plaintext) > maxPayloadLength {
		return nil, fmt.Errorf("Payload is too large. The max number of bytes is %d, input is %d bytes", maxPayloadLength, len(plaintext))
	}

	if len(sub.Key) == 0 {
		return nil, errors.New("Subscription must include the client's public key")
	}

	if len(sub.Auth) == 0 {
		return nil, errors.New("Subscription must include the client's auth value")
	}

	salt, err := randomSalt()
	if err != nil {
		return nil, err
	}

	// Use ECDH to derive a shared secret between us and the client. We generate
	// a fresh private/public key pair at random every time we encrypt.
	serverPrivateKey, serverPublicKey, err := randomKey()
	if err != nil {
		return nil, err
	}
	secret, err := sharedSecret(curve, sub.Key, serverPrivateKey)
	if err != nil {
		return nil, err
	}

	// Derive a Pseudo-Random Key (prk) that can be used to further derive our
	// other encryption parameters. These derivations are described in
	// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00
	prk := hkdf(sub.Auth, secret, authInfo, 32)

	// Derive the Content Encryption Key and nonce
	ctx := newContext(sub.Key, serverPublicKey)
	cek := newCEK(ctx, salt, prk)
	nonce := newNonce(ctx, salt, prk)

	// Do the actual encryption
	ciphertext, err := encrypt(plaintext, cek, nonce)
	if err != nil {
		return nil, err
	}

	// Return all of the values needed to construct a Web Push HTTP request.
	return &EncryptionResult{ciphertext, salt, serverPublicKey}, nil
}

func newCEK(ctx, salt, prk []byte) []byte {
	info := newInfo("aesgcm", ctx)
	return hkdf(salt, prk, info, 16)
}

func newNonce(ctx, salt, prk []byte) []byte {
	info := newInfo("nonce", ctx)
	return hkdf(salt, prk, info, 12)
}

// Creates a context for deriving encyption parameters, as described in
// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00.
// The 'context' in this case is just the public keys of both client and server.
// The keys should always be 65 bytes each. The format of the keys is
// described in section 4.3.6 of the (sadly not freely linkable) ANSI X9.62
// specification.
func newContext(clientPublicKey, serverPublicKey []byte) []byte {
	// The context format is:
	// 0x00 || length(clientPublicKey) || clientPublicKey ||
	//         length(serverPublicKey) || serverPublicKey
	// The lengths are 16-bit, Big Endian, unsigned integers so take 2 bytes each.
	cplen := uint16(len(clientPublicKey))
	cplenbuf := make([]byte, 2)
	binary.BigEndian.PutUint16(cplenbuf, cplen)

	splen := uint16(len(serverPublicKey))
	splenbuf := make([]byte, 2)
	binary.BigEndian.PutUint16(splenbuf, splen)

	var ctx []byte
	ctx = append(ctx, 0)
	ctx = append(ctx, cplenbuf...)
	ctx = append(ctx, []byte(clientPublicKey)...)
	ctx = append(ctx, splenbuf...)
	ctx = append(ctx, []byte(serverPublicKey)...)

	return ctx
}

// Returns an info record. See sections 3.2 and 3.3 of
// https://tools.ietf.org/html/draft-ietf-httpbis-encryption-encoding-00.
// The context argument should match what newContext creates
func newInfo(infoType string, context []byte) []byte {
	var info []byte
	info = append(info, []byte("Content-Encoding: ")...)
	info = append(info, []byte(infoType)...)
	info = append(info, 0)
	info = append(info, []byte("P-256")...)
	info = append(info, context...)
	return info
}

// HMAC-based Extract-and-Expand Key Derivation Function (HKDF)
//
// This is used to derive a secure encryption key from a mostly-secure shared
// secret.
//
// This is a partial implementation of HKDF tailored to our specific purposes.
// In particular, for us the value of N will always be 1, and thus T always
// equals HMAC-Hash(PRK, info | 0x01). This is true because the maximum output
// length we need/allow is 32.
//
// See https://www.rfc-editor.org/rfc/rfc5869.txt
func hkdf(salt, ikm, info []byte, length int) []byte {
	// HMAC length for SHA256 is 32 bytes, so that is the maximum result length.
	if length > 32 {
		panic("Can only produce HKDF outputs up to 32 bytes long")
	}

	// Extract
	mac := hmac.New(sha256.New, salt)
	mac.Write(ikm)
	prk := mac.Sum(nil)

	// Expand
	mac = hmac.New(sha256.New, prk)
	mac.Write(info)
	mac.Write([]byte{1})
	return mac.Sum(nil)[0:length]
}

// Encrypt the plaintext message using AES128/GCM
func encrypt(plaintext, key, nonce []byte) ([]byte, error) {
	// Add padding. There is a uint16 size followed by that number of bytes of
	// padding.
	// TODO: Right now we leave the size at zero. We should add a padding option
	// that allows the payload size to be obscured.
	padding := make([]byte, 2)
	data := append(padding, plaintext...)

	c, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(c)
	if err != nil {
		return nil, err
	}

	return gcm.Seal([]byte{}, nonce, data, nil), nil
}

// Given the coordinates of a party A's public key and the bytes of party B's
// private key, compute a shared secret.
func sharedSecret(curve elliptic.Curve, pub, priv []byte) ([]byte, error) {
	publicX, publicY := elliptic.Unmarshal(curve, pub)
	if publicX == nil {
		return nil, errors.New("Couldn't unmarshal public key. Not a valid point on the curve")
	}
	x, _ := curve.ScalarMult(publicX, publicY, priv)
	return x.Bytes(), nil
}
<|endoftext|>


<commit_before>package webservice

import (
	"crypto/tls"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/flexiant/concerto/config"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
)

const contentDispositionRegex = "filename=\\\"([^\\\"]*){1}\\\""

type Webservice struct {
	config *config.Config
	client *http.Client
}

func NewWebService() (*Webservice, error) {
	config, err := config.ConcertoServerConfiguration()
	if err != nil {
		return nil, err
	}

	client, err := httpClient(config)
	if err != nil {
		return nil, err
	}

	return &Webservice{config, client}, nil
}

func httpClient(config *config.Config) (*http.Client, error) {

	// Loads Clients Certificates and creates and 509KeyPair
	cert, err := tls.LoadX509KeyPair(config.Certificate.Cert, config.Certificate.Key)
	if err != nil {
		return nil, err
	}

	// Creates a client with specific transport configurations
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{Certificates: []tls.Certificate{cert}, InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: transport}

	return client, nil
}

func (w *Webservice) Post(endpoint string, json []byte) error {
	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	output := strings.NewReader(string(json))
	response, err := w.client.Post(w.config.ApiEndpoint+endpoint, "application/json", output)

	log.Debugf("Posting: %s", output)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	log.Debugf("Status code: %s", response.Status)
	return nil
}

func (w *Webservice) Get(endpoint string) ([]byte, error) {

	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	response, err := w.client.Get(w.config.ApiEndpoint + endpoint)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	log.Debugf("Status code: %s", response.Status)
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func (w *Webservice) GetFile(endpoint string, directoryPath string) (string, error) {

	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	response, err := w.client.Get(w.config.ApiEndpoint + endpoint)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	log.Debugf("Status code: %s", response.Status)

	r, err := regexp.Compile(contentDispositionRegex)
	if err != nil {
		return "", err
	}

	fileName := r.FindStringSubmatch(response.Header.Get("Content-Disposition"))[1]
	if err != nil {
		return "", err
	}
	realFileName := fmt.Sprintf("%s/%s", directoryPath, fileName)

	output, err := os.Create(realFileName)
	if err != nil {
		return "", err
	}
	defer output.Close()

	n, err := io.Copy(output, response.Body)
	if err != nil {
		return "", err
	}

	log.Debugf("%#v bytes downloaded", n)
	return realFileName, nil
}
<commit_msg>Minor updates to fix windows implementation<commit_after>package webservice

import (
	"crypto/tls"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/flexiant/concerto/config"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
)

const contentDispositionRegex = "filename=\\\"([^\\\"]*){1}\\\""

type Webservice struct {
	config *config.Config
	client *http.Client
}

func NewWebService() (*Webservice, error) {
	config, err := config.ConcertoServerConfiguration()
	if err != nil {
		return nil, err
	}

	client, err := httpClient(config)
	if err != nil {
		return nil, err
	}

	return &Webservice{config, client}, nil
}

func httpClient(config *config.Config) (*http.Client, error) {

	// Loads Clients Certificates and creates and 509KeyPair
	cert, err := tls.LoadX509KeyPair(config.Certificate.Cert, config.Certificate.Key)
	if err != nil {
		return nil, err
	}

	// Creates a client with specific transport configurations
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{Certificates: []tls.Certificate{cert}, InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: transport}

	return client, nil
}

func (w *Webservice) Post(endpoint string, json []byte) error {
	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	output := strings.NewReader(string(json))
	response, err := w.client.Post(w.config.ApiEndpoint+endpoint, "application/json", output)

	log.Debugf("Posting: %s", output)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	log.Debugf("Status code: %s", response.Status)
	return nil
}

func (w *Webservice) Get(endpoint string) ([]byte, error) {

	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	response, err := w.client.Get(w.config.ApiEndpoint + endpoint)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	log.Debugf("Status code: %s", response.Status)
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	log.Debugf("Response: %s", string(body))
	return body, nil
}

func (w *Webservice) GetFile(endpoint string, directoryPath string) (string, error) {

	log.Debugf("Connecting: %s%s", w.config.ApiEndpoint, endpoint)
	response, err := w.client.Get(w.config.ApiEndpoint + endpoint)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	log.Debugf("Status code: %s", response.Status)

	r, err := regexp.Compile(contentDispositionRegex)
	if err != nil {
		return "", err
	}

	fileName := r.FindStringSubmatch(response.Header.Get("Content-Disposition"))[1]
	if err != nil {
		return "", err
	}
	realFileName := fmt.Sprintf("%s/%s", directoryPath, fileName)

	output, err := os.Create(realFileName)
	if err != nil {
		return "", err
	}
	defer output.Close()

	n, err := io.Copy(output, response.Body)
	if err != nil {
		return "", err
	}

	log.Debugf("%#v bytes downloaded", n)
	return realFileName, nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/apex/log"
	MQTT "github.com/eclipse/paho.mqtt.golang"
)

const QoS = 0x00

// Client connects to the MQTT server and can publish/subscribe on uplink, downlink and activations from devices
type Client interface {
	Connect() error
	Disconnect()

	IsConnected() bool

	// Uplink pub/sub
	PublishUplink(payload UplinkMessage) Token
	SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token
	SubscribeAppUplink(appID string, handler UplinkHandler) Token
	SubscribeUplink(handler UplinkHandler) Token
	UnsubscribeDeviceUplink(appID string, devID string) Token
	UnsubscribeAppUplink(appID string) Token
	UnsubscribeUplink() Token

	// Downlink pub/sub
	PublishDownlink(payload DownlinkMessage) Token
	SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token
	SubscribeAppDownlink(appID string, handler DownlinkHandler) Token
	SubscribeDownlink(handler DownlinkHandler) Token
	UnsubscribeDeviceDownlink(appID string, devID string) Token
	UnsubscribeAppDownlink(appID string) Token
	UnsubscribeDownlink() Token

	// Activation pub/sub
	PublishActivation(payload Activation) Token
	SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token
	SubscribeAppActivations(appID string, handler ActivationHandler) Token
	SubscribeActivations(handler ActivationHandler) Token
	UnsubscribeDeviceActivations(appID string, devID string) Token
	UnsubscribeAppActivations(appID string) Token
	UnsubscribeActivations() Token
}

type Token interface {
	Wait() bool
	WaitTimeout(time.Duration) bool
	Error() error
}

type simpleToken struct {
	err error
}

// Wait always returns true
func (t *simpleToken) Wait() bool {
	return true
}

// WaitTimeout always returns true
func (t *simpleToken) WaitTimeout(_ time.Duration) bool {
	return true
}

// Error contains the error if present
func (t *simpleToken) Error() error {
	return t.err
}

type UplinkHandler func(client Client, appID string, devID string, req UplinkMessage)
type DownlinkHandler func(client Client, appID string, devID string, req DownlinkMessage)
type ActivationHandler func(client Client, appID string, devID string, req Activation)

type defaultClient struct {
	mqtt          MQTT.Client
	ctx           log.Interface
	subscriptions map[string]MQTT.MessageHandler
}

func NewClient(ctx log.Interface, id, username, password string, brokers ...string) Client {
	mqttOpts := MQTT.NewClientOptions()

	for _, broker := range brokers {
		mqttOpts.AddBroker(broker)
	}

	mqttOpts.SetClientID(fmt.Sprintf("%s-%s", id, random.String(16)))
	mqttOpts.SetUsername(username)
	mqttOpts.SetPassword(password)

	// TODO: Some tuning of these values probably won't hurt:
	mqttOpts.SetKeepAlive(30 * time.Second)
	mqttOpts.SetPingTimeout(10 * time.Second)

	mqttOpts.SetCleanSession(true)

	mqttOpts.SetDefaultPublishHandler(func(client MQTT.Client, msg MQTT.Message) {
		ctx.WithField("message", msg).Warn("Received unhandled message")
	})

	mqttOpts.SetConnectionLostHandler(func(client MQTT.Client, err error) {
		ctx.WithError(err).Warn("Disconnected, reconnecting...")
	})

	ttnClient := &defaultClient{
		mqtt:          MQTT.NewClient(mqttOpts),
		ctx:           ctx,
		subscriptions: make(map[string]MQTT.MessageHandler),
	}

	mqttOpts.SetOnConnectHandler(func(client MQTT.Client) {
		ctx.Debug("Connected to MQTT")
		for topic, handler := range ttnClient.subscriptions {
			ttnClient.subscribe(topic, handler) // re-subscribe
		}
	})

	return ttnClient
}

var (
	// ConnectRetries says how many times the client should retry a failed connection
	ConnectRetries = 10
	// ConnectRetryDelay says how long the client should wait between retries
	ConnectRetryDelay = time.Second
)

func (c *defaultClient) Connect() error {
	if c.mqtt.IsConnected() {
		return nil
	}
	var err error
	for retries := 0; retries < ConnectRetries; retries++ {
		c.ctx.Debug("Connecting to MQTT...")
		token := c.mqtt.Connect()
		token.Wait()
		err = token.Error()
		if err == nil {
			break
		}
		<-time.After(ConnectRetryDelay)
	}
	if err != nil {
		return fmt.Errorf("Could not connect: %s", err)
	}
	return nil
}

func (c *defaultClient) subscribe(topic string, handler MQTT.MessageHandler) Token {
	c.subscriptions[topic] = handler
	return c.mqtt.Subscribe(topic, QoS, handler)
}

func (c *defaultClient) unsubscribe(topic string) Token {
	delete(c.subscriptions, topic)
	return c.mqtt.Unsubscribe(topic)
}

func (c *defaultClient) Disconnect() {
	if !c.mqtt.IsConnected() {
		return
	}
	c.ctx.Debug("Disconnecting from MQTT")
	c.mqtt.Disconnect(25)
}

func (c *defaultClient) IsConnected() bool {
	return c.mqtt.IsConnected()
}

func (c *defaultClient) PublishUplink(dataUp UplinkMessage) Token {
	topic := DeviceTopic{dataUp.AppID, dataUp.DevID, Uplink}
	msg, err := json.Marshal(dataUp)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid uplink topic")
			return
		}

		// Unmarshal the payload
		dataUp := &UplinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataUp)

		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal uplink")
			return
		}

		// Call the uplink handler
		handler(c, topic.AppID, topic.DevID, *dataUp)
	})
}

func (c *defaultClient) SubscribeAppUplink(appID string, handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink(appID, "", handler)
}

func (c *defaultClient) SubscribeUplink(handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceUplink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.unsubscribe(topic.String())
}
func (c *defaultClient) UnsubscribeAppUplink(appID string) Token {
	return c.UnsubscribeDeviceUplink(appID, "")
}
func (c *defaultClient) UnsubscribeUplink() Token {
	return c.UnsubscribeDeviceUplink("", "")
}

func (c *defaultClient) PublishDownlink(dataDown DownlinkMessage) Token {
	topic := DeviceTopic{dataDown.AppID, dataDown.DevID, Downlink}
	msg, err := json.Marshal(dataDown)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Downlink topic")
			return
		}

		// Unmarshal the payload
		dataDown := &DownlinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataDown)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Downlink")
			return
		}

		// Call the Downlink handler
		handler(c, topic.AppID, topic.DevID, *dataDown)
	})
}

func (c *defaultClient) SubscribeAppDownlink(appID string, handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink(appID, "", handler)
}

func (c *defaultClient) SubscribeDownlink(handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceDownlink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.unsubscribe(topic.String())
}
func (c *defaultClient) UnsubscribeAppDownlink(appID string) Token {
	return c.UnsubscribeDeviceDownlink(appID, "")
}
func (c *defaultClient) UnsubscribeDownlink() Token {
	return c.UnsubscribeDeviceDownlink("", "")
}

func (c *defaultClient) PublishActivation(activation Activation) Token {
	topic := DeviceTopic{activation.AppID, activation.DevID, Activations}
	msg, err := json.Marshal(activation)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Activations topic")
			return
		}

		// Unmarshal the payload
		activation := &Activation{}
		err = json.Unmarshal(msg.Payload(), activation)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Activation")
			return
		}

		// Call the Activation handler
		handler(c, topic.AppID, topic.DevID, *activation)
	})
}

func (c *defaultClient) SubscribeAppActivations(appID string, handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations(appID, "", handler)
}

func (c *defaultClient) SubscribeActivations(handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceActivations(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.unsubscribe(topic.String())
}

func (c *defaultClient) UnsubscribeAppActivations(appID string) Token {
	return c.UnsubscribeDeviceActivations(appID, "")
}

func (c *defaultClient) UnsubscribeActivations() Token {
	return c.UnsubscribeDeviceActivations("", "")
}
<commit_msg>Add MQTT OnConnectHandler before making client<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/apex/log"
	MQTT "github.com/eclipse/paho.mqtt.golang"
)

const QoS = 0x00

// Client connects to the MQTT server and can publish/subscribe on uplink, downlink and activations from devices
type Client interface {
	Connect() error
	Disconnect()

	IsConnected() bool

	// Uplink pub/sub
	PublishUplink(payload UplinkMessage) Token
	SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token
	SubscribeAppUplink(appID string, handler UplinkHandler) Token
	SubscribeUplink(handler UplinkHandler) Token
	UnsubscribeDeviceUplink(appID string, devID string) Token
	UnsubscribeAppUplink(appID string) Token
	UnsubscribeUplink() Token

	// Downlink pub/sub
	PublishDownlink(payload DownlinkMessage) Token
	SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token
	SubscribeAppDownlink(appID string, handler DownlinkHandler) Token
	SubscribeDownlink(handler DownlinkHandler) Token
	UnsubscribeDeviceDownlink(appID string, devID string) Token
	UnsubscribeAppDownlink(appID string) Token
	UnsubscribeDownlink() Token

	// Activation pub/sub
	PublishActivation(payload Activation) Token
	SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token
	SubscribeAppActivations(appID string, handler ActivationHandler) Token
	SubscribeActivations(handler ActivationHandler) Token
	UnsubscribeDeviceActivations(appID string, devID string) Token
	UnsubscribeAppActivations(appID string) Token
	UnsubscribeActivations() Token
}

type Token interface {
	Wait() bool
	WaitTimeout(time.Duration) bool
	Error() error
}

type simpleToken struct {
	err error
}

// Wait always returns true
func (t *simpleToken) Wait() bool {
	return true
}

// WaitTimeout always returns true
func (t *simpleToken) WaitTimeout(_ time.Duration) bool {
	return true
}

// Error contains the error if present
func (t *simpleToken) Error() error {
	return t.err
}

type UplinkHandler func(client Client, appID string, devID string, req UplinkMessage)
type DownlinkHandler func(client Client, appID string, devID string, req DownlinkMessage)
type ActivationHandler func(client Client, appID string, devID string, req Activation)

type defaultClient struct {
	mqtt          MQTT.Client
	ctx           log.Interface
	subscriptions map[string]MQTT.MessageHandler
}

func NewClient(ctx log.Interface, id, username, password string, brokers ...string) Client {
	mqttOpts := MQTT.NewClientOptions()

	for _, broker := range brokers {
		mqttOpts.AddBroker(broker)
	}

	mqttOpts.SetClientID(fmt.Sprintf("%s-%s", id, random.String(16)))
	mqttOpts.SetUsername(username)
	mqttOpts.SetPassword(password)

	// TODO: Some tuning of these values probably won't hurt:
	mqttOpts.SetKeepAlive(30 * time.Second)
	mqttOpts.SetPingTimeout(10 * time.Second)

	mqttOpts.SetCleanSession(true)

	mqttOpts.SetDefaultPublishHandler(func(client MQTT.Client, msg MQTT.Message) {
		ctx.WithField("message", msg).Warn("Received unhandled message")
	})

	var reconnecting bool

	mqttOpts.SetConnectionLostHandler(func(client MQTT.Client, err error) {
		ctx.WithError(err).Warn("Disconnected, reconnecting...")
		reconnecting = true
	})

	ttnClient := &defaultClient{
		ctx:           ctx,
		subscriptions: make(map[string]MQTT.MessageHandler),
	}

	mqttOpts.SetOnConnectHandler(func(client MQTT.Client) {
		ctx.Info("Connected to MQTT")
		if reconnecting {
			for topic, handler := range ttnClient.subscriptions {
				ctx.Infof("Re-subscribing to %s", topic)
				ttnClient.subscribe(topic, handler)
			}
			reconnecting = false
		}
	})

	ttnClient.mqtt = MQTT.NewClient(mqttOpts)

	return ttnClient
}

var (
	// ConnectRetries says how many times the client should retry a failed connection
	ConnectRetries = 10
	// ConnectRetryDelay says how long the client should wait between retries
	ConnectRetryDelay = time.Second
)

func (c *defaultClient) Connect() error {
	if c.mqtt.IsConnected() {
		return nil
	}
	var err error
	for retries := 0; retries < ConnectRetries; retries++ {
		c.ctx.Debug("Connecting to MQTT...")
		token := c.mqtt.Connect()
		token.Wait()
		err = token.Error()
		if err == nil {
			break
		}
		<-time.After(ConnectRetryDelay)
	}
	if err != nil {
		return fmt.Errorf("Could not connect: %s", err)
	}
	return nil
}

func (c *defaultClient) subscribe(topic string, handler MQTT.MessageHandler) Token {
	c.subscriptions[topic] = handler
	return c.mqtt.Subscribe(topic, QoS, handler)
}

func (c *defaultClient) unsubscribe(topic string) Token {
	delete(c.subscriptions, topic)
	return c.mqtt.Unsubscribe(topic)
}

func (c *defaultClient) Disconnect() {
	if !c.mqtt.IsConnected() {
		return
	}
	c.ctx.Debug("Disconnecting from MQTT")
	c.mqtt.Disconnect(25)
}

func (c *defaultClient) IsConnected() bool {
	return c.mqtt.IsConnected()
}

func (c *defaultClient) PublishUplink(dataUp UplinkMessage) Token {
	topic := DeviceTopic{dataUp.AppID, dataUp.DevID, Uplink}
	msg, err := json.Marshal(dataUp)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid uplink topic")
			return
		}

		// Unmarshal the payload
		dataUp := &UplinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataUp)

		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal uplink")
			return
		}

		// Call the uplink handler
		handler(c, topic.AppID, topic.DevID, *dataUp)
	})
}

func (c *defaultClient) SubscribeAppUplink(appID string, handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink(appID, "", handler)
}

func (c *defaultClient) SubscribeUplink(handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceUplink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.unsubscribe(topic.String())
}
func (c *defaultClient) UnsubscribeAppUplink(appID string) Token {
	return c.UnsubscribeDeviceUplink(appID, "")
}
func (c *defaultClient) UnsubscribeUplink() Token {
	return c.UnsubscribeDeviceUplink("", "")
}

func (c *defaultClient) PublishDownlink(dataDown DownlinkMessage) Token {
	topic := DeviceTopic{dataDown.AppID, dataDown.DevID, Downlink}
	msg, err := json.Marshal(dataDown)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Downlink topic")
			return
		}

		// Unmarshal the payload
		dataDown := &DownlinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataDown)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Downlink")
			return
		}

		// Call the Downlink handler
		handler(c, topic.AppID, topic.DevID, *dataDown)
	})
}

func (c *defaultClient) SubscribeAppDownlink(appID string, handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink(appID, "", handler)
}

func (c *defaultClient) SubscribeDownlink(handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceDownlink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.unsubscribe(topic.String())
}
func (c *defaultClient) UnsubscribeAppDownlink(appID string) Token {
	return c.UnsubscribeDeviceDownlink(appID, "")
}
func (c *defaultClient) UnsubscribeDownlink() Token {
	return c.UnsubscribeDeviceDownlink("", "")
}

func (c *defaultClient) PublishActivation(activation Activation) Token {
	topic := DeviceTopic{activation.AppID, activation.DevID, Activations}
	msg, err := json.Marshal(activation)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

func (c *defaultClient) SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Activations topic")
			return
		}

		// Unmarshal the payload
		activation := &Activation{}
		err = json.Unmarshal(msg.Payload(), activation)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Activation")
			return
		}

		// Call the Activation handler
		handler(c, topic.AppID, topic.DevID, *activation)
	})
}

func (c *defaultClient) SubscribeAppActivations(appID string, handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations(appID, "", handler)
}

func (c *defaultClient) SubscribeActivations(handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations("", "", handler)
}

func (c *defaultClient) UnsubscribeDeviceActivations(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.unsubscribe(topic.String())
}

func (c *defaultClient) UnsubscribeAppActivations(appID string) Token {
	return c.UnsubscribeDeviceActivations(appID, "")
}

func (c *defaultClient) UnsubscribeActivations() Token {
	return c.UnsubscribeDeviceActivations("", "")
}
<|endoftext|>


<commit_before>package dropbox

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	ma "gx/ipfs/QmWWQ2Txc2c6tqjsBpzg5Ar652cHPGNsQQp2SejkNmkUMb/go-multiaddr"
	peer "gx/ipfs/QmZoWKhxUmZ2seW4BzX6fJkNR8hh9PsGModr7q171yq2SS/go-libp2p-peer"
	mh "gx/ipfs/QmZyZDi491cCNTLfAhwcaDii2Kg4pwKRkhqQzURGDvY6ua/go-multihash"

	"github.com/dropbox/dropbox-sdk-go-unofficial"
	"github.com/dropbox/dropbox-sdk-go-unofficial/files"
	"github.com/dropbox/dropbox-sdk-go-unofficial/sharing"
)

type DropBoxStorage struct {
	apiToken string
}

func NewDropBoxStorage(apiToken string) (*DropBoxStorage, error) {
	api := dropbox.Client(apiToken, dropbox.Options{Verbose: true})
	if _, err := api.GetCurrentAccount(); err != nil {
		return nil, err
	}
	return &DropBoxStorage{
		apiToken: apiToken,
	}, nil
}

func (s *DropBoxStorage) Store(peerID peer.ID, ciphertext []byte) (ma.Multiaddr, error) {
	api := dropbox.Client(s.apiToken, dropbox.Options{Verbose: true})
	hash := sha256.Sum256(ciphertext)
	hex := hex.EncodeToString(hash[:])

	// Upload ciphertext
	uploadArg := files.NewCommitInfo("/" + hex)
	r := bytes.NewReader(ciphertext)
	_, err := api.Upload(uploadArg, r)
	if err != nil {
		return nil, err
	}

	// Set public sharing
	sharingArg := sharing.NewCreateSharedLinkArg("/" + hex)
	res, err := api.CreateSharedLink(sharingArg)
	if err != nil {
		return nil, err
	}

	// Create encoded multiaddr
	url := res.Url[:len(res.Url)-1] + "1"
	b, err := mh.Encode([]byte(url), mh.SHA1)
	if err != nil {
		return nil, err
	}
	m, err := mh.Cast(b)
	if err != nil {
		return nil, err
	}

	addr, err := ma.NewMultiaddr("/ipfs/" + m.B58String() + "/https/")
	if err != nil {
		return nil, err
	}
	return addr, nil
}
<commit_msg>[cleanup] avoid variable/package name collission (hex -> hexStr)<commit_after>package dropbox

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	ma "gx/ipfs/QmWWQ2Txc2c6tqjsBpzg5Ar652cHPGNsQQp2SejkNmkUMb/go-multiaddr"
	peer "gx/ipfs/QmZoWKhxUmZ2seW4BzX6fJkNR8hh9PsGModr7q171yq2SS/go-libp2p-peer"
	mh "gx/ipfs/QmZyZDi491cCNTLfAhwcaDii2Kg4pwKRkhqQzURGDvY6ua/go-multihash"

	"github.com/dropbox/dropbox-sdk-go-unofficial"
	"github.com/dropbox/dropbox-sdk-go-unofficial/files"
	"github.com/dropbox/dropbox-sdk-go-unofficial/sharing"
)

type DropBoxStorage struct {
	apiToken string
}

func NewDropBoxStorage(apiToken string) (*DropBoxStorage, error) {
	api := dropbox.Client(apiToken, dropbox.Options{Verbose: true})
	if _, err := api.GetCurrentAccount(); err != nil {
		return nil, err
	}
	return &DropBoxStorage{
		apiToken: apiToken,
	}, nil
}

func (s *DropBoxStorage) Store(peerID peer.ID, ciphertext []byte) (ma.Multiaddr, error) {
	api := dropbox.Client(s.apiToken, dropbox.Options{Verbose: true})
	hash := sha256.Sum256(ciphertext)
	hexStr := hex.EncodeToString(hash[:])

	// Upload ciphertext
	uploadArg := files.NewCommitInfo("/" + hexStr)
	r := bytes.NewReader(ciphertext)
	_, err := api.Upload(uploadArg, r)
	if err != nil {
		return nil, err
	}

	// Set public sharing
	sharingArg := sharing.NewCreateSharedLinkArg("/" + hexStr)
	res, err := api.CreateSharedLink(sharingArg)
	if err != nil {
		return nil, err
	}

	// Create encoded multiaddr
	url := res.Url[:len(res.Url)-1] + "1"
	b, err := mh.Encode([]byte(url), mh.SHA1)
	if err != nil {
		return nil, err
	}
	m, err := mh.Cast(b)
	if err != nil {
		return nil, err
	}

	addr, err := ma.NewMultiaddr("/ipfs/" + m.B58String() + "/https/")
	if err != nil {
		return nil, err
	}
	return addr, nil
}
<|endoftext|>


<commit_before>package solr

const (
	VERSION = "0.0"
)
<commit_msg>#11 step up version<commit_after>package solr

const (
	VERSION = "0.1"
)
<|endoftext|>


<commit_before>package analyzer

import (
	"fmt"
	"math"
	"time"

	"github.com/m-lab/signal-searcher/sequencer"
)

// An Incident represents a piece of a timeseries that contains a user-visible
// problem.
type Incident struct {
	Start, End         time.Time
	AffectedCount      int
	Severity           float64
	GoodPeriodDownload float64
	BadPeriodDownload  float64
}

// URL converts an incident (along with provided Metadata) into a viz URL.
func (i *Incident) URL(m sequencer.Meta) string {
	twelveBeforeStart := i.Start.AddDate(-1, 0, 0)
	return fmt.Sprintf(
		"https://viz.measurementlab.net/location/%s?aggr=month&isps=%s&start=%s&end=%s",
		m.Loc, m.ASN, twelveBeforeStart.Format("2006-01-02"), i.End.Format("2006-01-02"))
}

type arrayIncident struct {
	start, end         int
	severity           float64
	goodPeriodDownload float64
	badPeriodDownload  float64
}

// severity of merged incident should be the max of the incidents merged
func mergeArrayIncidents(a []arrayIncident) (merged []arrayIncident) {
	if len(a) <= 1 {
		return a
	}
	current := a[0]
	for i := 1; i < len(a); i++ {
		if current.end+1 == a[i].end {
			current.end = a[i].end
			current.severity = math.Max(a[i].severity, current.severity)
		} else {
			merged = append(merged, current)
			current = a[i]
		}
	}
	merged = append(merged, current)
	return
}

// FindPerformanceDrops discovers time periods of a year or greater where
// performance showed more than a 30% average drop.
func FindPerformanceDrops(s *sequencer.Sequence) []Incident {
	dates, data := s.SortedSlices()
	var previous, current sequencer.Datum
	for i := 0; i < 12; i++ {
		previous.Download += data[i].Download
	}
	for i := 12; i < 24; i++ {
		current.Download += data[i].Download
	}
	var arrayIncidents []arrayIncident
	for i := 24; i < len(data); i++ {
		// Update the running sums
		previous.Download = previous.Download - data[i-24].Download + data[i-12].Download
		current.Download = current.Download - data[i-12].Download + data[i].Download
		if previous.Download*.7 > current.Download {
			arrayIncidents = append(arrayIncidents, arrayIncident{start: i - 12, end: i, severity: 1.0 - current.Download/previous.Download, goodPeriodDownload: previous.Download / 12, badPeriodDownload: current.Download / 12})
		}
	}
	arrayIncidents = mergeArrayIncidents(arrayIncidents)
	incidents := []Incident{}
	for _, ai := range arrayIncidents {
		newIncident := Incident{Start: dates[ai.start], End: dates[ai.end], Severity: ai.severity, GoodPeriodDownload: ai.goodPeriodDownload, BadPeriodDownload: ai.badPeriodDownload}
		for i := ai.start; i < ai.end; i++ {
			newIncident.AffectedCount += data[i].Count
		}
		incidents = append(incidents, newIncident)
	}
	return incidents
}
<commit_msg>add conditional for good/bad periods when merging array incidents<commit_after>package analyzer

import (
	"fmt"
	"time"

	"github.com/m-lab/signal-searcher/sequencer"
)

// An Incident represents a piece of a timeseries that contains a user-visible
// problem.
type Incident struct {
	Start, End         time.Time
	AffectedCount      int
	Severity           float64
	GoodPeriodDownload float64
	BadPeriodDownload  float64
}

// URL converts an incident (along with provided Metadata) into a viz URL.
func (i *Incident) URL(m sequencer.Meta) string {
	twelveBeforeStart := i.Start.AddDate(-1, 0, 0)
	return fmt.Sprintf(
		"https://viz.measurementlab.net/location/%s?aggr=month&isps=%s&start=%s&end=%s",
		m.Loc, m.ASN, twelveBeforeStart.Format("2006-01-02"), i.End.Format("2006-01-02"))
}

type arrayIncident struct {
	start, end         int
	severity           float64
	goodPeriodDownload float64
	badPeriodDownload  float64
}

// severity of merged incident should be the max of the incidents merged
func mergeArrayIncidents(a []arrayIncident) (merged []arrayIncident) {
	if len(a) <= 1 {
		return a
	}
	current := a[0]
	for i := 1; i < len(a); i++ {
		if current.end+1 == a[i].end {
			current.end = a[i].end
			if a[i].severity > current.severity {
				current.goodPeriodDownload = a[i].goodPeriodDownload
				current.badPeriodDownload = a[i].badPeriodDownload
				current.severity = a[i].severity
			}
		} else {
			merged = append(merged, current)
			current = a[i]
		}
	}
	merged = append(merged, current)
	return
}

// FindPerformanceDrops discovers time periods of a year or greater where
// performance showed more than a 30% average drop.
func FindPerformanceDrops(s *sequencer.Sequence) []Incident {
	dates, data := s.SortedSlices()
	var previous, current sequencer.Datum
	for i := 0; i < 12; i++ {
		previous.Download += data[i].Download
	}
	for i := 12; i < 24; i++ {
		current.Download += data[i].Download
	}
	var arrayIncidents []arrayIncident
	for i := 24; i < len(data); i++ {
		// Update the running sums
		previous.Download = previous.Download - data[i-24].Download + data[i-12].Download
		current.Download = current.Download - data[i-12].Download + data[i].Download
		if previous.Download*.7 > current.Download {
			arrayIncidents = append(arrayIncidents, arrayIncident{start: i - 12, end: i, severity: 1.0 - current.Download/previous.Download, goodPeriodDownload: previous.Download / 12, badPeriodDownload: current.Download / 12})
		}
	}
	arrayIncidents = mergeArrayIncidents(arrayIncidents)
	incidents := []Incident{}
	for _, ai := range arrayIncidents {
		newIncident := Incident{Start: dates[ai.start], End: dates[ai.end], Severity: ai.severity, GoodPeriodDownload: ai.goodPeriodDownload, BadPeriodDownload: ai.badPeriodDownload}
		for i := ai.start; i < ai.end; i++ {
			newIncident.AffectedCount += data[i].Count
		}
		incidents = append(incidents, newIncident)
	}
	return incidents
}
<|endoftext|>


<commit_before>package mstate

import (
	"errors"
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/trivial"
)

// ResolvedMode describes the way state transition errors
// are resolved.
type ResolvedMode int

const (
	ResolvedNone ResolvedMode = iota
	ResolvedRetryHooks
	ResolvedNoHooks
	nResolvedModes
)

// AssignmentPolicy controls what machine a unit will be assigned to.
type AssignmentPolicy string

const (
	// AssignLocal indicates that all service units should be assigned
	// to machine 0.
	AssignLocal AssignmentPolicy = "local"
	// AssignUnused indicates that every service unit should be assigned
	// to a dedicated machine, and that new machines should be launched
	// if required.
	AssignUnused AssignmentPolicy = "unused"
)

// NeedsUpgrade describes if a unit needs an
// upgrade and if this is forced.
type NeedsUpgrade struct {
	Upgrade bool
	Force   bool
}

// Port identifies a network port number for a particular protocol.
type Port struct {
	Protocol string `yaml:"proto"`
	Number   int    `yaml:"port"`
}

// UnitSettings holds information about a service unit's settings within a
// relation.
type UnitSettings struct {
	Version  int
	Settings map[string]interface{}
}

// unitDoc represents the internal state of a unit in MongoDB.
type unitDoc struct {
	Name           string `bson:"_id"`
	Service        string
	Principal      string
	PublicAddress  string
	PrivateAddress string
	MachineId      *int
	Resolved       ResolvedMode
	NeedsUpgrade   *NeedsUpgrade
	Life           Life
}

// Unit represents the state of a service unit.
type Unit struct {
	st  *State
	doc unitDoc
}

func newUnit(st *State, udoc *unitDoc) *Unit {
	return &Unit{
		st:  st,
		doc: *udoc,
	}
}

// ServiceName returns the service name.
func (u *Unit) ServiceName() string {
	return u.doc.Service
}

// String returns the unit as string.
func (u *Unit) String() string {
	return u.doc.Name
}

// Name returns the unit name.
func (u *Unit) Name() string {
	return u.doc.Name
}

// Resolved returns the resolved mode for the unit.
func (u *Unit) Resolved() (mode ResolvedMode, err error) {
	return u.doc.Resolved, nil
}

// IsPrincipal returns whether the unit is deployed in its own container,
// and can therefore have subordinate services deployed alongside it.
func (u *Unit) IsPrincipal() bool {
	return u.doc.Principal == ""
}

// PublicAddress returns the public address of the unit.
func (u *Unit) PublicAddress() (string, error) {
	if u.doc.PublicAddress == "" {
		return "", fmt.Errorf("public address of unit %q not found", u)
	}
	return u.doc.PublicAddress, nil
}

// PrivateAddress returns the public address of the unit.
func (u *Unit) PrivateAddress() (string, error) {
	if u.doc.PrivateAddress == "" {
		return "", fmt.Errorf("private address of unit %q not found", u)
	}
	return u.doc.PrivateAddress, nil
}

func (u *Unit) Refresh() error {
	err := u.st.units.FindId(u.doc.Name).One(&u.doc)
	if err != nil {
		return fmt.Errorf("cannot refresh unit %q: %v", u, err)
	}
	return nil
}

// AssignedMachineId returns the id of the assigned machine.
func (u *Unit) AssignedMachineId() (id int, err error) {
	defer trivial.ErrorContextf(&err, "cannot get machine id of unit %q", u)
	if u.IsPrincipal() {
		if u.doc.MachineId == nil {
			return 0, errors.New("unit not assigned to machine")
		}
		return *u.doc.MachineId, nil
	}
	pudoc := unitDoc{}
	sel := bson.D{{"_id", u.doc.Principal}, {"life", Alive}}
	err = u.st.units.Find(sel).One(&pudoc)
	if err != nil {
		return 0, err
	}
	if pudoc.MachineId == nil {
		return 0, errors.New("unit not assigned to machine")
	}
	return *pudoc.MachineId, nil
}

// AssignToMachine assigns this unit to a given machine.
func (u *Unit) AssignToMachine(m *Machine) (err error) {
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"machineid", nil}},
			bson.D{{"machineid", m.Id()}},
		}},
		{"life", Alive},
	}
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     sel,
		Change:     bson.D{{"$set", bson.D{{"machineid", m.Id()}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot assign unit %q to machine %s: %v", u, m, err)
	}
	u.doc.MachineId = &m.doc.Id
	return nil
}

// UnassignFromMachine removes the assignment between this unit and the
// machine it's assigned to.
func (u *Unit) UnassignFromMachine() (err error) {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"machineid", nil}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot unassign unit %q from machine: %v", u, err)
	}
	u.doc.MachineId = nil
	return nil
}

// SetPublicAddress sets the public address of the unit.
func (u *Unit) SetPublicAddress(address string) error {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"publicaddress", address}}}},
	}}
	err := u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot set public address of unit %q: %v", u, err)
	}
	u.doc.PublicAddress = address
	return nil
}

// SetPrivateAddress sets the public address of the unit.
func (u *Unit) SetPrivateAddress(address string) error {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"privateaddress", address}}}},
	}}
	err := u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot set private address of unit %q: %v", u, err)
	}
	u.doc.PrivateAddress = address
	return nil
}

// SetResolved marks the unit as having had any previous state transition
// problems resolved, and informs the unit that it may attempt to
// reestablish normal workflow. The resolved mode parameter informs
// whether to attempt to reexecute previous failed hooks or to continue
// as if they had succeeded before.
func (u *Unit) SetResolved(mode ResolvedMode) (err error) {
	defer trivial.ErrorContextf(&err, "cannot set resolved mode for unit %q", u)
	if !(0 <= mode && mode < nResolvedModes) {
		return fmt.Errorf("invalid error resolution mode: %v", mode)
	}
	change := bson.D{{"$set", bson.D{{"resolved", mode}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"resolved", ResolvedNone},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("flag already set")
	}
	if err != nil {
		return err
	}
	u.doc.Resolved = mode
	return nil
}

// ClearResolved removes any resolved setting on the unit.
func (u *Unit) ClearResolved() error {
	change := bson.D{{"$set", bson.D{{"resolved", ResolvedNone}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot clear resolved mode for unit %q: %v", u, err)
	}
	u.doc.Resolved = ResolvedNone
	return nil
}

// NeedsUpgrade returns whether the unit needs an upgrade
// and if it does, if this is forced.
func (u *Unit) NeedsUpgrade() (*NeedsUpgrade, error) {
	if u.doc.NeedsUpgrade == nil {
		return &NeedsUpgrade{Upgrade: false, Force: false}, nil
	}
	return u.doc.NeedsUpgrade, nil
}

// SetNeedsUpgrade informs the unit that it should perform
// a regular or forced upgrade.
func (u *Unit) SetNeedsUpgrade(force bool) (err error) {
	defer trivial.ErrorContextf(&err, "cannot inform unit %q about upgrade", u)
	nu := &NeedsUpgrade{Upgrade: true, Force: force}
	change := bson.D{{"$set", bson.D{{"needsupgrade", nu}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"needsupgrade", nil}},
			bson.D{{"needsupgrade", nu}},
		}},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("upgrade already enabled")
	}
	if err != nil {
		return err
	}
	u.doc.NeedsUpgrade = nu
	return nil
}

// ClearNeedsUpgrade resets the upgrade notification. It is typically
// done by the unit agent before beginning the upgrade.
func (u *Unit) ClearNeedsUpgrade() error {
	change := bson.D{{"$set", bson.D{{"needsupgrade", nil}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("upgrade notification for unit %q cannot be reset: %v", u, err)
	}
	u.doc.NeedsUpgrade = nil
	return nil
}
<commit_msg>mstate: use txn for unit SetResolved<commit_after>package mstate

import (
	"errors"
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/trivial"
)

// ResolvedMode describes the way state transition errors
// are resolved.
type ResolvedMode int

const (
	ResolvedNone ResolvedMode = iota
	ResolvedRetryHooks
	ResolvedNoHooks
	nResolvedModes
)

// AssignmentPolicy controls what machine a unit will be assigned to.
type AssignmentPolicy string

const (
	// AssignLocal indicates that all service units should be assigned
	// to machine 0.
	AssignLocal AssignmentPolicy = "local"
	// AssignUnused indicates that every service unit should be assigned
	// to a dedicated machine, and that new machines should be launched
	// if required.
	AssignUnused AssignmentPolicy = "unused"
)

// NeedsUpgrade describes if a unit needs an
// upgrade and if this is forced.
type NeedsUpgrade struct {
	Upgrade bool
	Force   bool
}

// Port identifies a network port number for a particular protocol.
type Port struct {
	Protocol string `yaml:"proto"`
	Number   int    `yaml:"port"`
}

// UnitSettings holds information about a service unit's settings within a
// relation.
type UnitSettings struct {
	Version  int
	Settings map[string]interface{}
}

// unitDoc represents the internal state of a unit in MongoDB.
type unitDoc struct {
	Name           string `bson:"_id"`
	Service        string
	Principal      string
	PublicAddress  string
	PrivateAddress string
	MachineId      *int
	Resolved       ResolvedMode
	NeedsUpgrade   *NeedsUpgrade
	Life           Life
}

// Unit represents the state of a service unit.
type Unit struct {
	st  *State
	doc unitDoc
}

func newUnit(st *State, udoc *unitDoc) *Unit {
	return &Unit{
		st:  st,
		doc: *udoc,
	}
}

// ServiceName returns the service name.
func (u *Unit) ServiceName() string {
	return u.doc.Service
}

// String returns the unit as string.
func (u *Unit) String() string {
	return u.doc.Name
}

// Name returns the unit name.
func (u *Unit) Name() string {
	return u.doc.Name
}

// Resolved returns the resolved mode for the unit.
func (u *Unit) Resolved() (mode ResolvedMode, err error) {
	return u.doc.Resolved, nil
}

// IsPrincipal returns whether the unit is deployed in its own container,
// and can therefore have subordinate services deployed alongside it.
func (u *Unit) IsPrincipal() bool {
	return u.doc.Principal == ""
}

// PublicAddress returns the public address of the unit.
func (u *Unit) PublicAddress() (string, error) {
	if u.doc.PublicAddress == "" {
		return "", fmt.Errorf("public address of unit %q not found", u)
	}
	return u.doc.PublicAddress, nil
}

// PrivateAddress returns the public address of the unit.
func (u *Unit) PrivateAddress() (string, error) {
	if u.doc.PrivateAddress == "" {
		return "", fmt.Errorf("private address of unit %q not found", u)
	}
	return u.doc.PrivateAddress, nil
}

func (u *Unit) Refresh() error {
	err := u.st.units.FindId(u.doc.Name).One(&u.doc)
	if err != nil {
		return fmt.Errorf("cannot refresh unit %q: %v", u, err)
	}
	return nil
}

// AssignedMachineId returns the id of the assigned machine.
func (u *Unit) AssignedMachineId() (id int, err error) {
	defer trivial.ErrorContextf(&err, "cannot get machine id of unit %q", u)
	if u.IsPrincipal() {
		if u.doc.MachineId == nil {
			return 0, errors.New("unit not assigned to machine")
		}
		return *u.doc.MachineId, nil
	}
	pudoc := unitDoc{}
	sel := bson.D{{"_id", u.doc.Principal}, {"life", Alive}}
	err = u.st.units.Find(sel).One(&pudoc)
	if err != nil {
		return 0, err
	}
	if pudoc.MachineId == nil {
		return 0, errors.New("unit not assigned to machine")
	}
	return *pudoc.MachineId, nil
}

// AssignToMachine assigns this unit to a given machine.
func (u *Unit) AssignToMachine(m *Machine) (err error) {
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"machineid", nil}},
			bson.D{{"machineid", m.Id()}},
		}},
		{"life", Alive},
	}
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     sel,
		Change:     bson.D{{"$set", bson.D{{"machineid", m.Id()}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot assign unit %q to machine %s: %v", u, m, err)
	}
	u.doc.MachineId = &m.doc.Id
	return nil
}

// UnassignFromMachine removes the assignment between this unit and the
// machine it's assigned to.
func (u *Unit) UnassignFromMachine() (err error) {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"machineid", nil}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot unassign unit %q from machine: %v", u, err)
	}
	u.doc.MachineId = nil
	return nil
}

// SetPublicAddress sets the public address of the unit.
func (u *Unit) SetPublicAddress(address string) error {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"publicaddress", address}}}},
	}}
	err := u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot set public address of unit %q: %v", u, err)
	}
	u.doc.PublicAddress = address
	return nil
}

// SetPrivateAddress sets the public address of the unit.
func (u *Unit) SetPrivateAddress(address string) error {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"privateaddress", address}}}},
	}}
	err := u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot set private address of unit %q: %v", u, err)
	}
	u.doc.PrivateAddress = address
	return nil
}

// SetResolved marks the unit as having had any previous state transition
// problems resolved, and informs the unit that it may attempt to
// reestablish normal workflow. The resolved mode parameter informs
// whether to attempt to reexecute previous failed hooks or to continue
// as if they had succeeded before.
func (u *Unit) SetResolved(mode ResolvedMode) (err error) {
	defer trivial.ErrorContextf(&err, "cannot set resolved mode for unit %q", u)
	if !(0 <= mode && mode < nResolvedModes) {
		return fmt.Errorf("invalid error resolution mode: %v", mode)
	}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"resolved", ResolvedNone},
	}
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     sel,
		Change:    	bson.D{{"$set", bson.D{{"resolved", mode}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err == txn.ErrAborted {
		return errors.New("flag already set")
	}
	if err != nil {
		return err
	}
	u.doc.Resolved = mode
	return nil
}

// ClearResolved removes any resolved setting on the unit.
func (u *Unit) ClearResolved() error {
	change := bson.D{{"$set", bson.D{{"resolved", ResolvedNone}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot clear resolved mode for unit %q: %v", u, err)
	}
	u.doc.Resolved = ResolvedNone
	return nil
}

// NeedsUpgrade returns whether the unit needs an upgrade
// and if it does, if this is forced.
func (u *Unit) NeedsUpgrade() (*NeedsUpgrade, error) {
	if u.doc.NeedsUpgrade == nil {
		return &NeedsUpgrade{Upgrade: false, Force: false}, nil
	}
	return u.doc.NeedsUpgrade, nil
}

// SetNeedsUpgrade informs the unit that it should perform
// a regular or forced upgrade.
func (u *Unit) SetNeedsUpgrade(force bool) (err error) {
	defer trivial.ErrorContextf(&err, "cannot inform unit %q about upgrade", u)
	nu := &NeedsUpgrade{Upgrade: true, Force: force}
	change := bson.D{{"$set", bson.D{{"needsupgrade", nu}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"needsupgrade", nil}},
			bson.D{{"needsupgrade", nu}},
		}},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("upgrade already enabled")
	}
	if err != nil {
		return err
	}
	u.doc.NeedsUpgrade = nu
	return nil
}

// ClearNeedsUpgrade resets the upgrade notification. It is typically
// done by the unit agent before beginning the upgrade.
func (u *Unit) ClearNeedsUpgrade() error {
	change := bson.D{{"$set", bson.D{{"needsupgrade", nil}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("upgrade notification for unit %q cannot be reset: %v", u, err)
	}
	u.doc.NeedsUpgrade = nil
	return nil
}
<|endoftext|>


<commit_before>package projects

import (
	"encoding/json"
	"net/http"
	"strconv"

	h "github.com/ernestio/api-gateway/helpers"
	"github.com/ernestio/api-gateway/models"
)

// Update : responds to PUT /projects/:id: by updating
// an existing project
func Update(au models.User, project string, body []byte) (int, []byte) {
	var err error
	var d models.Project
	var existing models.Project
	var r models.Role
	var roles []models.Role

	if d.Map(body) != nil {
		return 400, models.NewJSONError("Invalid input")
	}

	if err = existing.FindByName(project); err != nil {
		id, err := strconv.Atoi(project)
		if err = existing.FindByID(id); err != nil {
			return 404, models.NewJSONError("Project not found")
		}
	}

	if st, res := h.IsAuthorizedToResource(&au, h.UpdateProject, d.GetType(), d.Name); st != 200 {
		return st, res
	}

	if err := r.FindAllByResource(existing.GetID(), existing.GetType(), &roles); err == nil {
		existing.Members = roles
	}

	existing.Credentials = d.Credentials

	if err = existing.Save(); err != nil {
		h.L.Error(err.Error())
		return 500, models.NewJSONError("Internal server error")
	}

	if d.Members == nil {
		if body, err = json.Marshal(d); err != nil {
			return 500, models.NewJSONError("Internal server error")
		}

		return http.StatusOK, body
	}

	for _, r := range d.Members {
		if r.ID == 0 {
			if !au.IsAdmin() {
				if ok := au.IsOwner(r.ResourceType, r.ResourceID); !ok {
					return 403, models.NewJSONError("You're not authorized to perform this action")
				}
			}

			err = r.Save()
			if err != nil {
				h.L.Error(err.Error())
				return http.StatusBadRequest, models.NewJSONError(err.Error())
			}

			continue
		}

		for _, er := range existing.Members {
			// update role
			if r.ID == er.ID && r.Role != er.Role {
				if !au.IsAdmin() {
					if ok := au.IsOwner(r.ResourceType, r.ResourceID); !ok {
						return 403, models.NewJSONError("You're not authorized to perform this action")
					}
				}

				err = r.Save()
				if err != nil {
					h.L.Error(err.Error())
					return http.StatusBadRequest, models.NewJSONError(err.Error())
				}
			}
		}
	}

	for _, er := range existing.Members {
		var exists bool

		for _, r := range d.Members {
			if r.ID == er.ID {
				exists = true
			}
		}

		// delete roles
		if !exists {
			if !au.IsAdmin() {
				if ok := au.IsOwner(er.ResourceType, er.ResourceID); !ok {
					return 403, models.NewJSONError("You're not authorized to perform this action")
				}
			}

			err = er.Delete()
			if err != nil {
				h.L.Error(err.Error())
				return http.StatusBadRequest, models.NewJSONError(err.Error())
			}
		}
	}

	if body, err = json.Marshal(d); err != nil {
		return 500, models.NewJSONError("Internal server error")
	}

	return http.StatusOK, body
}
<commit_msg>fixing project update response on project not found<commit_after>package projects

import (
	"encoding/json"
	"net/http"

	h "github.com/ernestio/api-gateway/helpers"
	"github.com/ernestio/api-gateway/models"
)

// Update : responds to PUT /projects/:id: by updating
// an existing project
func Update(au models.User, project string, body []byte) (int, []byte) {
	var err error
	var d models.Project
	var existing models.Project
	var r models.Role
	var roles []models.Role

	if d.Map(body) != nil {
		return 400, models.NewJSONError("Invalid input")
	}

	if err = existing.FindByName(project); err != nil {
		return 404, models.NewJSONError("Project not found")
	}

	if st, res := h.IsAuthorizedToResource(&au, h.UpdateProject, d.GetType(), d.Name); st != 200 {
		return st, res
	}

	if err = r.FindAllByResource(existing.GetID(), existing.GetType(), &roles); err == nil {
		existing.Members = roles
	}

	existing.Credentials = d.Credentials

	if err = existing.Save(); err != nil {
		h.L.Error(err.Error())
		return 500, models.NewJSONError("Internal server error")
	}

	if d.Members == nil {
		if body, err = json.Marshal(d); err != nil {
			return 500, models.NewJSONError("Internal server error")
		}

		return http.StatusOK, body
	}

	for _, r := range d.Members {
		if r.ID == 0 {
			if !au.IsAdmin() {
				if ok := au.IsOwner(r.ResourceType, r.ResourceID); !ok {
					return 403, models.NewJSONError("You're not authorized to perform this action")
				}
			}

			err = r.Save()
			if err != nil {
				h.L.Error(err.Error())
				return http.StatusBadRequest, models.NewJSONError(err.Error())
			}

			continue
		}

		for _, er := range existing.Members {
			// update role
			if r.ID == er.ID && r.Role != er.Role {
				if !au.IsAdmin() {
					if ok := au.IsOwner(r.ResourceType, r.ResourceID); !ok {
						return 403, models.NewJSONError("You're not authorized to perform this action")
					}
				}

				err = r.Save()
				if err != nil {
					h.L.Error(err.Error())
					return http.StatusBadRequest, models.NewJSONError(err.Error())
				}
			}
		}
	}

	for _, er := range existing.Members {
		var exists bool

		for _, r := range d.Members {
			if r.ID == er.ID {
				exists = true
			}
		}

		// delete roles
		if !exists {
			if !au.IsAdmin() {
				if ok := au.IsOwner(er.ResourceType, er.ResourceID); !ok {
					return 403, models.NewJSONError("You're not authorized to perform this action")
				}
			}

			err = er.Delete()
			if err != nil {
				h.L.Error(err.Error())
				return http.StatusBadRequest, models.NewJSONError(err.Error())
			}
		}
	}

	if body, err = json.Marshal(d); err != nil {
		return 500, models.NewJSONError("Internal server error")
	}

	return http.StatusOK, body
}
<|endoftext|>


<commit_before>package cell_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"code.cloudfoundry.org/archiver/compressor"
	archive_helper "code.cloudfoundry.org/archiver/extractor/test_helper"
	"code.cloudfoundry.org/bbs/models"
	ssh_helpers "code.cloudfoundry.org/diego-ssh/helpers"
	"code.cloudfoundry.org/diego-ssh/routes"
	"code.cloudfoundry.org/inigo/fixtures"
	"code.cloudfoundry.org/inigo/helpers"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/ifrit/grouper"
	"golang.org/x/crypto/ssh"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("SSH", func() {
	verifySSH := func(address, processGuid string, index int) {
		clientConfig := &ssh.ClientConfig{
			User: fmt.Sprintf("diego:%s/%d", processGuid, index),
			Auth: []ssh.AuthMethod{ssh.Password("")},
		}

		client, err := ssh.Dial("tcp", address, clientConfig)
		Expect(err).NotTo(HaveOccurred())

		session, err := client.NewSession()
		Expect(err).NotTo(HaveOccurred())

		output, err := session.Output("env")
		Expect(err).NotTo(HaveOccurred())

		Expect(string(output)).To(ContainSubstring("USER=root"))
		Expect(string(output)).To(ContainSubstring("TEST=foobar"))
		Expect(string(output)).To(ContainSubstring(fmt.Sprintf("INSTANCE_INDEX=%d", index)))
	}

	var (
		processGuid         string
		fileServerStaticDir string

		runtime ifrit.Process
		address string

		lrp models.DesiredLRP
	)

	BeforeEach(func() {
		processGuid = helpers.GenerateGuid()
		address = componentMaker.Addresses.SSHProxy

		var fileServer ifrit.Runner
		fileServer, fileServerStaticDir = componentMaker.FileServer()
		runtime = ginkgomon.Invoke(grouper.NewParallel(os.Kill, grouper.Members{
			{"router", componentMaker.Router()},
			{"file-server", fileServer},
			{"rep", componentMaker.Rep()},
			{"auctioneer", componentMaker.Auctioneer()},
			{"route-emitter", componentMaker.RouteEmitter()},
			{"ssh-proxy", componentMaker.SSHProxy()},
		}))

		tgCompressor := compressor.NewTgz()
		err := tgCompressor.Compress(componentMaker.Artifacts.Executables["sshd"], filepath.Join(fileServerStaticDir, "sshd.tgz"))
		Expect(err).NotTo(HaveOccurred())

		archive_helper.CreateZipArchive(
			filepath.Join(fileServerStaticDir, "lrp.zip"),
			fixtures.GoServerApp(),
		)

		sshRoute := routes.SSHRoute{
			ContainerPort:   3456,
			PrivateKey:      componentMaker.SSHConfig.PrivateKeyPem,
			HostFingerprint: ssh_helpers.MD5Fingerprint(componentMaker.SSHConfig.HostKey.PublicKey()),
		}

		sshRoutePayload, err := json.Marshal(sshRoute)
		Expect(err).NotTo(HaveOccurred())

		sshRouteMessage := json.RawMessage(sshRoutePayload)

		envVars := []*models.EnvironmentVariable{
			{Name: "TEST", Value: "foobar"},
		}

		lrp = models.DesiredLRP{
			ProcessGuid:        processGuid,
			Domain:             "inigo",
			Instances:          2,
			Privileged:         true,
			LegacyDownloadUser: "vcap",
			Setup: models.WrapAction(models.Serial(
				&models.DownloadAction{
					Artifact: "sshd",
					From:     fmt.Sprintf("http://%s/v1/static/%s", componentMaker.Addresses.FileServer, "sshd.tgz"),
					To:       "/tmp/diego",
					CacheKey: "sshd",
					User:     "root",
				},
				&models.DownloadAction{
					Artifact: "go-server",
					From:     fmt.Sprintf("http://%s/v1/static/%s", componentMaker.Addresses.FileServer, "lrp.zip"),
					To:       "/tmp/diego",
					CacheKey: "lrp-cache-key",
					User:     "root",
				},
			)),
			Action: models.WrapAction(models.Codependent(
				&models.RunAction{
					User: "root",
					Path: "/tmp/diego/sshd",
					Args: []string{
						"-address=0.0.0.0:3456",
						"-hostKey=" + componentMaker.SSHConfig.HostKeyPem,
						"-authorizedKey=" + componentMaker.SSHConfig.AuthorizedKey,
						"-inheritDaemonEnv",
					},
				},
				&models.RunAction{
					User: "root",
					Path: "/tmp/diego/go-server",
					Env:  []*models.EnvironmentVariable{{"PORT", "9999"}},
				},
			)),
			Monitor: models.WrapAction(&models.RunAction{
				User: "root",
				Path: "nc",
				Args: []string{"-z", "127.0.0.1", "3456"},
			}),
			StartTimeoutMs: 60000,
			RootFs:         "preloaded:" + helpers.PreloadedStacks[0],
			MemoryMb:       128,
			DiskMb:         128,
			Ports:          []uint32{3456},
			Routes: &models.Routes{
				routes.DIEGO_SSH: &sshRouteMessage,
			},
			EnvironmentVariables: envVars,
		}
	})

	JustBeforeEach(func() {
		logger := lagertest.NewTestLogger("test")
		logger.Info("desired-ssh-lrp", lager.Data{"lrp": lrp})

		err := bbsClient.DesireLRP(logger, &lrp)
		Expect(err).NotTo(HaveOccurred())

		Eventually(func() []*models.ActualLRPGroup {
			lrps, err := bbsClient.ActualLRPGroupsByProcessGuid(logger, processGuid)
			Expect(err).NotTo(HaveOccurred())
			return lrps
		}).Should(HaveLen(2))

		Eventually(
			helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 0, nil),
		).Should(Equal(models.ActualLRPStateRunning))

		Eventually(
			helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 1, nil),
		).Should(Equal(models.ActualLRPStateRunning))
	})

	AfterEach(func() {
		helpers.StopProcesses(runtime)
	})

	Context("when valid process guid and index are used in the username", func() {
		It("can ssh to appropriate app instance container", func() {
			verifySSH(address, processGuid, 0)
			verifySSH(address, processGuid, 1)
		})

		It("supports local port fowarding", func() {
			clientConfig := &ssh.ClientConfig{
				User: fmt.Sprintf("diego:%s/%d", processGuid, 0),
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}

			client, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).NotTo(HaveOccurred())

			httpClient := &http.Client{
				Transport: &http.Transport{
					Dial: client.Dial,
				},
				Timeout: 20 * time.Second,
			}

			resp, err := httpClient.Get("http://localhost:9999/yo")
			Expect(err).NotTo(HaveOccurred())
			defer resp.Body.Close()
			Expect(resp.StatusCode).To(Equal(http.StatusOK))

			contents, err := ioutil.ReadAll(resp.Body)
			Expect(err).NotTo(HaveOccurred())
			Expect(contents).To(ContainSubstring("sup dawg"))
		})

		Context("when invalid password is used", func() {
			var clientConfig *ssh.ClientConfig

			BeforeEach(func() {
				clientConfig = &ssh.ClientConfig{
					User: "diego:" + processGuid + "/0",
					Auth: []ssh.AuthMethod{ssh.Password("invalid:password")},
				}
			})

			It("returns an error", func() {
				Eventually(
					helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 0, nil),
				).Should(Equal(models.ActualLRPStateRunning))

				_, err := ssh.Dial("tcp", address, clientConfig)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when a bare-bones docker image is used as the root filesystem", func() {
			BeforeEach(func() {
				lrp.StartTimeoutMs = 120000
				lrp.RootFs = "docker:///cloudfoundry/diego-docker-app"

				// busybox nc doesn't support -z
				lrp.Monitor = models.WrapAction(&models.RunAction{
					User: "root",
					Path: "sh",
					Args: []string{
						"-c",
						"echo -n '' | telnet localhost 3456 >/dev/null 2>&1 && echo -n '' | telnet localhost 9999 >/dev/null 2>&1 && true",
					},
				})
			})

			It("can ssh to appropriate app instance container", func() {
				verifySSH(address, processGuid, 0)
				verifySSH(address, processGuid, 1)
			})

			It("supports local port fowarding", func() {
				clientConfig := &ssh.ClientConfig{
					User: fmt.Sprintf("diego:%s/%d", processGuid, 0),
					Auth: []ssh.AuthMethod{ssh.Password("")},
				}

				client, err := ssh.Dial("tcp", address, clientConfig)
				Expect(err).NotTo(HaveOccurred())

				httpClient := &http.Client{
					Transport: &http.Transport{
						Dial: client.Dial,
					},
					Timeout: 20 * time.Second,
				}

				resp, err := httpClient.Get("http://localhost:9999/yo")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(http.StatusOK))

				contents, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(contents).To(ContainSubstring("sup dawg"))
			})
		})
	})

	Context("when non-existent index is used as part of username", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "diego:" + processGuid + "/3",
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})

	Context("when non-existent process guid is used as part of username", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "diego:not-existing-process-guid/0",
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})

	Context("when invalid username format is used", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "root",
				Auth: []ssh.AuthMethod{ssh.Password("some-password")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})
})
<commit_msg>Improve inigo ssh tests<commit_after>package cell_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"code.cloudfoundry.org/archiver/compressor"
	archive_helper "code.cloudfoundry.org/archiver/extractor/test_helper"
	"code.cloudfoundry.org/bbs/models"
	ssh_helpers "code.cloudfoundry.org/diego-ssh/helpers"
	"code.cloudfoundry.org/diego-ssh/routes"
	"code.cloudfoundry.org/inigo/fixtures"
	"code.cloudfoundry.org/inigo/helpers"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"code.cloudfoundry.org/localip"
	"code.cloudfoundry.org/rep/cmd/rep/config"
	"github.com/cloudfoundry/sonde-go/events"
	"github.com/gogo/protobuf/proto"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/ifrit/grouper"
	"golang.org/x/crypto/ssh"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("SSH", func() {
	verifySSH := func(address, processGuid string, index int) {
		clientConfig := &ssh.ClientConfig{
			User: fmt.Sprintf("diego:%s/%d", processGuid, index),
			Auth: []ssh.AuthMethod{ssh.Password("")},
		}

		client, err := ssh.Dial("tcp", address, clientConfig)
		Expect(err).NotTo(HaveOccurred())

		session, err := client.NewSession()
		Expect(err).NotTo(HaveOccurred())

		output, err := session.Output("env")
		Expect(err).NotTo(HaveOccurred())

		Expect(string(output)).To(ContainSubstring("USER=root"))
		Expect(string(output)).To(ContainSubstring("TEST=foobar"))
		Expect(string(output)).To(ContainSubstring(fmt.Sprintf("INSTANCE_INDEX=%d", index)))
	}

	var (
		processGuid         string
		fileServerStaticDir string

		runtime ifrit.Process
		address string

		lrp models.DesiredLRP
	)

	BeforeEach(func() {
		port, err := localip.LocalPort()
		Expect(err).NotTo(HaveOccurred())
		addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", port))
		Expect(err).NotTo(HaveOccurred())
		udpConn, err := net.ListenUDP("udp", addr)
		Expect(err).NotTo(HaveOccurred())

		metronAgent := ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
			logger := logger.Session("metron-agent")
			close(ready)
			logger.Info("starting", lager.Data{"port": addr.Port})
			msgs := make(chan []byte)
			errCh := make(chan error)
			go func() {
				for {
					bs := make([]byte, 102400)
					n, _, err := udpConn.ReadFromUDP(bs)
					if err != nil {
						errCh <- err
						return
					}
					msgs <- bs[:n]
				}
			}()
		loop:
			for {
				select {
				case <-signals:
					logger.Info("signaled")
					break loop
				case err := <-errCh:
					return err
				case msg := <-msgs:
					var envelope events.Envelope
					err := proto.Unmarshal(msg, &envelope)
					if err != nil {
						logger.Error("error-parsing-message", err)
						continue
					}
					if envelope.GetEventType() != events.Envelope_LogMessage {
						continue
					}
					logger.Info("received-data", lager.Data{"message": string(envelope.GetLogMessage().GetMessage())})
				}
			}
			udpConn.Close()
			return nil
		})

		processGuid = helpers.GenerateGuid()
		address = componentMaker.Addresses.SSHProxy

		var fileServer ifrit.Runner
		fileServer, fileServerStaticDir = componentMaker.FileServer()
		repConfig := func(cfg *config.RepConfig) {
			cfg.DropsondePort = addr.Port
		}
		runtime = ginkgomon.Invoke(grouper.NewParallel(os.Kill, grouper.Members{
			{"router", componentMaker.Router()},
			{"file-server", fileServer},
			{"metron-agent", metronAgent},
			{"rep", componentMaker.Rep(repConfig)},
			{"auctioneer", componentMaker.Auctioneer()},
			{"route-emitter", componentMaker.RouteEmitter()},
			{"ssh-proxy", componentMaker.SSHProxy()},
		}))

		tgCompressor := compressor.NewTgz()
		err = tgCompressor.Compress(componentMaker.Artifacts.Executables["sshd"], filepath.Join(fileServerStaticDir, "sshd.tgz"))
		Expect(err).NotTo(HaveOccurred())

		archive_helper.CreateZipArchive(
			filepath.Join(fileServerStaticDir, "lrp.zip"),
			fixtures.GoServerApp(),
		)

		sshRoute := routes.SSHRoute{
			ContainerPort:   3456,
			PrivateKey:      componentMaker.SSHConfig.PrivateKeyPem,
			HostFingerprint: ssh_helpers.MD5Fingerprint(componentMaker.SSHConfig.HostKey.PublicKey()),
		}

		sshRoutePayload, err := json.Marshal(sshRoute)
		Expect(err).NotTo(HaveOccurred())

		sshRouteMessage := json.RawMessage(sshRoutePayload)

		envVars := []*models.EnvironmentVariable{
			{Name: "TEST", Value: "foobar"},
		}

		lrp = models.DesiredLRP{
			LogGuid:            processGuid,
			ProcessGuid:        processGuid,
			Domain:             "inigo",
			Instances:          2,
			Privileged:         true,
			LegacyDownloadUser: "vcap",
			Setup: models.WrapAction(models.Serial(
				&models.DownloadAction{
					Artifact: "sshd",
					From:     fmt.Sprintf("http://%s/v1/static/%s", componentMaker.Addresses.FileServer, "sshd.tgz"),
					To:       "/tmp/diego",
					CacheKey: "sshd",
					User:     "root",
				},
				&models.DownloadAction{
					Artifact: "go-server",
					From:     fmt.Sprintf("http://%s/v1/static/%s", componentMaker.Addresses.FileServer, "lrp.zip"),
					To:       "/tmp/diego",
					CacheKey: "lrp-cache-key",
					User:     "root",
				},
			)),
			Action: models.WrapAction(models.Codependent(
				&models.RunAction{
					User: "root",
					Path: "/tmp/diego/sshd",
					Args: []string{
						"-address=0.0.0.0:3456",
						"-hostKey=" + componentMaker.SSHConfig.HostKeyPem,
						"-authorizedKey=" + componentMaker.SSHConfig.AuthorizedKey,
						"-inheritDaemonEnv",
						"-logLevel=debug",
					},
				},
				&models.RunAction{
					User: "root",
					Path: "/tmp/diego/go-server",
					Env:  []*models.EnvironmentVariable{{"PORT", "9999"}},
				},
			)),
			Monitor: models.WrapAction(&models.RunAction{
				User: "root",
				Path: "nc",
				Args: []string{"-z", "127.0.0.1", "3456"},
			}),
			StartTimeoutMs: 60000,
			RootFs:         "preloaded:" + helpers.PreloadedStacks[0],
			MemoryMb:       128,
			DiskMb:         128,
			Ports:          []uint32{3456},
			Routes: &models.Routes{
				routes.DIEGO_SSH: &sshRouteMessage,
			},
			EnvironmentVariables: envVars,
		}
	})

	JustBeforeEach(func() {
		logger := lagertest.NewTestLogger("test")
		logger.Info("desired-ssh-lrp", lager.Data{"lrp": lrp})

		err := bbsClient.DesireLRP(logger, &lrp)
		Expect(err).NotTo(HaveOccurred())

		Eventually(func() []*models.ActualLRPGroup {
			lrps, err := bbsClient.ActualLRPGroupsByProcessGuid(logger, processGuid)
			Expect(err).NotTo(HaveOccurred())
			return lrps
		}).Should(HaveLen(2))

		Eventually(
			helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 0, nil),
		).Should(Equal(models.ActualLRPStateRunning))

		Eventually(
			helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 1, nil),
		).Should(Equal(models.ActualLRPStateRunning))
	})

	AfterEach(func() {
		helpers.StopProcesses(runtime)
	})

	Context("when valid process guid and index are used in the username", func() {
		It("can ssh to appropriate app instance container", func() {
			verifySSH(address, processGuid, 0)
			verifySSH(address, processGuid, 1)
		})

		It("supports local port fowarding", func() {
			clientConfig := &ssh.ClientConfig{
				User: fmt.Sprintf("diego:%s/%d", processGuid, 0),
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}

			client, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).NotTo(HaveOccurred())

			httpClient := &http.Client{
				Transport: &http.Transport{
					Dial: client.Dial,
				},
				Timeout: 5 * time.Second,
			}

			resp, err := httpClient.Get("http://127.0.0.1:9999/yo")
			Expect(err).NotTo(HaveOccurred())
			defer resp.Body.Close()
			Expect(resp.StatusCode).To(Equal(http.StatusOK))

			contents, err := ioutil.ReadAll(resp.Body)
			Expect(err).NotTo(HaveOccurred())
			Expect(contents).To(ContainSubstring("sup dawg"))
		})

		Context("when invalid password is used", func() {
			var clientConfig *ssh.ClientConfig

			BeforeEach(func() {
				clientConfig = &ssh.ClientConfig{
					User: "diego:" + processGuid + "/0",
					Auth: []ssh.AuthMethod{ssh.Password("invalid:password")},
				}
			})

			It("returns an error", func() {
				Eventually(
					helpers.LRPInstanceStatePoller(logger, bbsClient, processGuid, 0, nil),
				).Should(Equal(models.ActualLRPStateRunning))

				_, err := ssh.Dial("tcp", address, clientConfig)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when a bare-bones docker image is used as the root filesystem", func() {
			BeforeEach(func() {
				lrp.StartTimeoutMs = 120000
				lrp.RootFs = "docker:///cloudfoundry/diego-docker-app"

				// busybox nc doesn't support -z
				lrp.Monitor = models.WrapAction(&models.RunAction{
					User: "root",
					Path: "sh",
					Args: []string{
						"-c",
						"echo -n '' | telnet localhost 3456 >/dev/null 2>&1 && echo -n '' | telnet localhost 9999 >/dev/null 2>&1 && true",
					},
				})
			})

			It("can ssh to appropriate app instance container", func() {
				verifySSH(address, processGuid, 0)
				verifySSH(address, processGuid, 1)
			})

			It("supports local port fowarding", func() {
				clientConfig := &ssh.ClientConfig{
					User: fmt.Sprintf("diego:%s/%d", processGuid, 0),
					Auth: []ssh.AuthMethod{ssh.Password("")},
				}

				client, err := ssh.Dial("tcp", address, clientConfig)
				Expect(err).NotTo(HaveOccurred())

				httpClient := &http.Client{
					Transport: &http.Transport{
						Dial: client.Dial,
					},
					Timeout: 5 * time.Second,
				}

				resp, err := httpClient.Get("http://127.0.0.1:9999/yo")
				Expect(err).NotTo(HaveOccurred())
				defer resp.Body.Close()
				Expect(resp.StatusCode).To(Equal(http.StatusOK))

				contents, err := ioutil.ReadAll(resp.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(contents).To(ContainSubstring("sup dawg"))
			})
		})
	})

	Context("when non-existent index is used as part of username", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "diego:" + processGuid + "/3",
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})

	Context("when non-existent process guid is used as part of username", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "diego:not-existing-process-guid/0",
				Auth: []ssh.AuthMethod{ssh.Password("")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})

	Context("when invalid username format is used", func() {
		var clientConfig *ssh.ClientConfig

		BeforeEach(func() {
			clientConfig = &ssh.ClientConfig{
				User: "root",
				Auth: []ssh.AuthMethod{ssh.Password("some-password")},
			}
		})

		It("returns an error", func() {
			_, err := ssh.Dial("tcp", address, clientConfig)
			Expect(err).To(HaveOccurred())
		})
	})
})
<|endoftext|>