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

package rsa

import (
	"big"
	"crypto/subtle"
	"io"
	"os"
)

// This file implements encryption and decryption using PKCS#1 v1.5 padding.

// EncryptPKCS1v15 encrypts the given message with RSA and the padding scheme from PKCS#1 v1.5.
// The message must be no longer than the length of the public modulus minus 11 bytes.
// WARNING: use of this function to encrypt plaintexts other than session keys
// is dangerous. Use RSA OAEP in new protocols.
func EncryptPKCS1v15(rand io.Reader, pub *PublicKey, msg []byte) (out []byte, err os.Error) {
	k := (pub.N.BitLen() + 7) / 8
	if len(msg) > k-11 {
		err = MessageTooLongError{}
		return
	}

	// EM = 0x02 || PS || 0x00 || M
	em := make([]byte, k-1)
	em[0] = 2
	ps, mm := em[1:len(em)-len(msg)-1], em[len(em)-len(msg):]
	err = nonZeroRandomBytes(ps, rand)
	if err != nil {
		return
	}
	em[len(em)-len(msg)-1] = 0
	copy(mm, msg)

	m := new(big.Int).SetBytes(em)
	c := encrypt(new(big.Int), pub, m)
	out = c.Bytes()
	return
}

// DecryptPKCS1v15 decrypts a plaintext using RSA and the padding scheme from PKCS#1 v1.5.
// If rand != nil, it uses RSA blinding to avoid timing side-channel attacks.
func DecryptPKCS1v15(rand io.Reader, priv *PrivateKey, ciphertext []byte) (out []byte, err os.Error) {
	valid, out, err := decryptPKCS1v15(rand, priv, ciphertext)
	if err == nil && valid == 0 {
		err = DecryptionError{}
	}

	return
}

// DecryptPKCS1v15SessionKey decrypts a session key using RSA and the padding scheme from PKCS#1 v1.5.
// If rand != nil, it uses RSA blinding to avoid timing side-channel attacks.
// It returns an error if the ciphertext is the wrong length or if the
// ciphertext is greater than the public modulus. Otherwise, no error is
// returned. If the padding is valid, the resulting plaintext message is copied
// into key. Otherwise, key is unchanged. These alternatives occur in constant
// time. It is intended that the user of this function generate a random
// session key beforehand and continue the protocol with the resulting value.
// This will remove any possibility that an attacker can learn any information
// about the plaintext.
// See ``Chosen Ciphertext Attacks Against Protocols Based on the RSA
// Encryption Standard PKCS #1'', Daniel Bleichenbacher, Advances in Cryptology
// (Crypto '98),
func DecryptPKCS1v15SessionKey(rand io.Reader, priv *PrivateKey, ciphertext []byte, key []byte) (err os.Error) {
	k := (priv.N.BitLen() + 7) / 8
	if k-(len(key)+3+8) < 0 {
		err = DecryptionError{}
		return
	}

	valid, msg, err := decryptPKCS1v15(rand, priv, ciphertext)
	if err != nil {
		return
	}

	valid &= subtle.ConstantTimeEq(int32(len(msg)), int32(len(key)))
	subtle.ConstantTimeCopy(valid, key, msg)
	return
}

func decryptPKCS1v15(rand io.Reader, priv *PrivateKey, ciphertext []byte) (valid int, msg []byte, err os.Error) {
	k := (priv.N.BitLen() + 7) / 8
	if k < 11 {
		err = DecryptionError{}
		return
	}

	c := new(big.Int).SetBytes(ciphertext)
	m, err := decrypt(rand, priv, c)
	if err != nil {
		return
	}

	em := leftPad(m.Bytes(), k)
	firstByteIsZero := subtle.ConstantTimeByteEq(em[0], 0)
	secondByteIsTwo := subtle.ConstantTimeByteEq(em[1], 2)

	// The remainder of the plaintext must be a string of non-zero random
	// octets, followed by a 0, followed by the message.
	//   lookingForIndex: 1 iff we are still looking for the zero.
	//   index: the offset of the first zero byte.
	var lookingForIndex, index int
	lookingForIndex = 1

	for i := 2; i < len(em); i++ {
		equals0 := subtle.ConstantTimeByteEq(em[i], 0)
		index = subtle.ConstantTimeSelect(lookingForIndex&equals0, i, index)
		lookingForIndex = subtle.ConstantTimeSelect(equals0, 0, lookingForIndex)
	}

	valid = firstByteIsZero & secondByteIsTwo & (^lookingForIndex & 1)
	msg = em[index+1:]
	return
}

// nonZeroRandomBytes fills the given slice with non-zero random octets.
func nonZeroRandomBytes(s []byte, rand io.Reader) (err os.Error) {
	_, err = io.ReadFull(rand, s)
	if err != nil {
		return
	}

	for i := 0; i < len(s); i++ {
		for s[i] == 0 {
			_, err = rand.Read(s[i : i+1])
			if err != nil {
				return
			}
			// In tests, the PRNG may return all zeros so we do
			// this to break the loop.
			s[i] ^= 0x42
		}
	}

	return
}

// Due to the design of PKCS#1 v1.5, we need to know the exact hash function in
// use. A generic hash.Hash will not do.
type PKCS1v15Hash int

const (
	HashMD5 PKCS1v15Hash = iota
	HashSHA1
	HashSHA256
	HashSHA384
	HashSHA512
	HashMD5SHA1 // combined MD5 and SHA1 hash used for RSA signing in TLS.
)

// These are ASN1 DER structures:
//   DigestInfo ::= SEQUENCE {
//     digestAlgorithm AlgorithmIdentifier,
//     digest OCTET STRING
//   }
// For performance, we don't use the generic ASN1 encoder. Rather, we
// precompute a prefix of the digest value that makes a valid ASN1 DER string
// with the correct contents.
var hashPrefixes = [][]byte{
	// HashMD5
	{0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	// HashSHA1
	{0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	// HashSHA256
	{0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	// HashSHA384
	{0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	// HashSHA512
	{0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	// HashMD5SHA1
	{}, // A special TLS case which doesn't use an ASN1 prefix.
}

// SignPKCS1v15 calcuates the signature of hashed using RSASSA-PSS-SIGN from RSA PKCS#1 v1.5.
// Note that hashed must be the result of hashing the input message using the
// given hash function.
func SignPKCS1v15(rand io.Reader, priv *PrivateKey, hash PKCS1v15Hash, hashed []byte) (s []byte, err os.Error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (priv.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, MessageTooLongError{}
	}

	// EM = 0x00 || 0x01 || PS || 0x00 || T
	em := make([]byte, k)
	em[1] = 1
	for i := 2; i < k-tLen-1; i++ {
		em[i] = 0xff
	}
	copy(em[k-tLen:k-hashLen], prefix)
	copy(em[k-hashLen:k], hashed)

	m := new(big.Int).SetBytes(em)
	c, err := decrypt(rand, priv, m)
	if err == nil {
		s = c.Bytes()
	}
	return
}

// VerifyPKCS1v15 verifies an RSA PKCS#1 v1.5 signature.
// hashed is the result of hashing the input message using the given hash
// function and sig is the signature. A valid signature is indicated by
// returning a nil error.
func VerifyPKCS1v15(pub *PublicKey, hash PKCS1v15Hash, hashed []byte, sig []byte) (err os.Error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (pub.N.BitLen() + 7) / 8
	if k < tLen+11 {
		err = VerificationError{}
		return
	}

	c := new(big.Int).SetBytes(sig)
	m := encrypt(new(big.Int), pub, c)
	em := leftPad(m.Bytes(), k)
	// EM = 0x00 || 0x01 || PS || 0x00 || T

	ok := subtle.ConstantTimeByteEq(em[0], 0)
	ok &= subtle.ConstantTimeByteEq(em[1], 1)
	ok &= subtle.ConstantTimeCompare(em[k-hashLen:k], hashed)
	ok &= subtle.ConstantTimeCompare(em[k-tLen:k-hashLen], prefix)
	ok &= subtle.ConstantTimeByteEq(em[k-tLen-1], 0)

	for i := 2; i < k-tLen-1; i++ {
		ok &= subtle.ConstantTimeByteEq(em[i], 0xff)
	}

	if ok != 1 {
		return VerificationError{}
	}

	return nil
}

func pkcs1v15HashInfo(hash PKCS1v15Hash, inLen int) (hashLen int, prefix []byte, err os.Error) {
	switch hash {
	case HashMD5:
		hashLen = 16
	case HashSHA1:
		hashLen = 20
	case HashSHA256:
		hashLen = 32
	case HashSHA384:
		hashLen = 48
	case HashSHA512:
		hashLen = 64
	case HashMD5SHA1:
		hashLen = 36
	default:
		return 0, nil, os.ErrorString("unknown hash function")
	}

	if inLen != hashLen {
		return 0, nil, os.ErrorString("input must be hashed message")
	}

	prefix = hashPrefixes[int(hash)]
	return
}
<commit_msg>crypto/rsa: correct docstring for SignPKCS1v15.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rsa

import (
	"big"
	"crypto/subtle"
	"io"
	"os"
)

// This file implements encryption and decryption using PKCS#1 v1.5 padding.

// EncryptPKCS1v15 encrypts the given message with RSA and the padding scheme from PKCS#1 v1.5.
// The message must be no longer than the length of the public modulus minus 11 bytes.
// WARNING: use of this function to encrypt plaintexts other than session keys
// is dangerous. Use RSA OAEP in new protocols.
func EncryptPKCS1v15(rand io.Reader, pub *PublicKey, msg []byte) (out []byte, err os.Error) {
	k := (pub.N.BitLen() + 7) / 8
	if len(msg) > k-11 {
		err = MessageTooLongError{}
		return
	}

	// EM = 0x02 || PS || 0x00 || M
	em := make([]byte, k-1)
	em[0] = 2
	ps, mm := em[1:len(em)-len(msg)-1], em[len(em)-len(msg):]
	err = nonZeroRandomBytes(ps, rand)
	if err != nil {
		return
	}
	em[len(em)-len(msg)-1] = 0
	copy(mm, msg)

	m := new(big.Int).SetBytes(em)
	c := encrypt(new(big.Int), pub, m)
	out = c.Bytes()
	return
}

// DecryptPKCS1v15 decrypts a plaintext using RSA and the padding scheme from PKCS#1 v1.5.
// If rand != nil, it uses RSA blinding to avoid timing side-channel attacks.
func DecryptPKCS1v15(rand io.Reader, priv *PrivateKey, ciphertext []byte) (out []byte, err os.Error) {
	valid, out, err := decryptPKCS1v15(rand, priv, ciphertext)
	if err == nil && valid == 0 {
		err = DecryptionError{}
	}

	return
}

// DecryptPKCS1v15SessionKey decrypts a session key using RSA and the padding scheme from PKCS#1 v1.5.
// If rand != nil, it uses RSA blinding to avoid timing side-channel attacks.
// It returns an error if the ciphertext is the wrong length or if the
// ciphertext is greater than the public modulus. Otherwise, no error is
// returned. If the padding is valid, the resulting plaintext message is copied
// into key. Otherwise, key is unchanged. These alternatives occur in constant
// time. It is intended that the user of this function generate a random
// session key beforehand and continue the protocol with the resulting value.
// This will remove any possibility that an attacker can learn any information
// about the plaintext.
// See ``Chosen Ciphertext Attacks Against Protocols Based on the RSA
// Encryption Standard PKCS #1'', Daniel Bleichenbacher, Advances in Cryptology
// (Crypto '98),
func DecryptPKCS1v15SessionKey(rand io.Reader, priv *PrivateKey, ciphertext []byte, key []byte) (err os.Error) {
	k := (priv.N.BitLen() + 7) / 8
	if k-(len(key)+3+8) < 0 {
		err = DecryptionError{}
		return
	}

	valid, msg, err := decryptPKCS1v15(rand, priv, ciphertext)
	if err != nil {
		return
	}

	valid &= subtle.ConstantTimeEq(int32(len(msg)), int32(len(key)))
	subtle.ConstantTimeCopy(valid, key, msg)
	return
}

func decryptPKCS1v15(rand io.Reader, priv *PrivateKey, ciphertext []byte) (valid int, msg []byte, err os.Error) {
	k := (priv.N.BitLen() + 7) / 8
	if k < 11 {
		err = DecryptionError{}
		return
	}

	c := new(big.Int).SetBytes(ciphertext)
	m, err := decrypt(rand, priv, c)
	if err != nil {
		return
	}

	em := leftPad(m.Bytes(), k)
	firstByteIsZero := subtle.ConstantTimeByteEq(em[0], 0)
	secondByteIsTwo := subtle.ConstantTimeByteEq(em[1], 2)

	// The remainder of the plaintext must be a string of non-zero random
	// octets, followed by a 0, followed by the message.
	//   lookingForIndex: 1 iff we are still looking for the zero.
	//   index: the offset of the first zero byte.
	var lookingForIndex, index int
	lookingForIndex = 1

	for i := 2; i < len(em); i++ {
		equals0 := subtle.ConstantTimeByteEq(em[i], 0)
		index = subtle.ConstantTimeSelect(lookingForIndex&equals0, i, index)
		lookingForIndex = subtle.ConstantTimeSelect(equals0, 0, lookingForIndex)
	}

	valid = firstByteIsZero & secondByteIsTwo & (^lookingForIndex & 1)
	msg = em[index+1:]
	return
}

// nonZeroRandomBytes fills the given slice with non-zero random octets.
func nonZeroRandomBytes(s []byte, rand io.Reader) (err os.Error) {
	_, err = io.ReadFull(rand, s)
	if err != nil {
		return
	}

	for i := 0; i < len(s); i++ {
		for s[i] == 0 {
			_, err = rand.Read(s[i : i+1])
			if err != nil {
				return
			}
			// In tests, the PRNG may return all zeros so we do
			// this to break the loop.
			s[i] ^= 0x42
		}
	}

	return
}

// Due to the design of PKCS#1 v1.5, we need to know the exact hash function in
// use. A generic hash.Hash will not do.
type PKCS1v15Hash int

const (
	HashMD5 PKCS1v15Hash = iota
	HashSHA1
	HashSHA256
	HashSHA384
	HashSHA512
	HashMD5SHA1 // combined MD5 and SHA1 hash used for RSA signing in TLS.
)

// These are ASN1 DER structures:
//   DigestInfo ::= SEQUENCE {
//     digestAlgorithm AlgorithmIdentifier,
//     digest OCTET STRING
//   }
// For performance, we don't use the generic ASN1 encoder. Rather, we
// precompute a prefix of the digest value that makes a valid ASN1 DER string
// with the correct contents.
var hashPrefixes = [][]byte{
	// HashMD5
	{0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	// HashSHA1
	{0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	// HashSHA256
	{0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	// HashSHA384
	{0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	// HashSHA512
	{0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	// HashMD5SHA1
	{}, // A special TLS case which doesn't use an ASN1 prefix.
}

// SignPKCS1v15 calculates the signature of hashed using RSASSA-PKCS1-V1_5-SIGN from RSA PKCS#1 v1.5.
// Note that hashed must be the result of hashing the input message using the
// given hash function.
func SignPKCS1v15(rand io.Reader, priv *PrivateKey, hash PKCS1v15Hash, hashed []byte) (s []byte, err os.Error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (priv.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, MessageTooLongError{}
	}

	// EM = 0x00 || 0x01 || PS || 0x00 || T
	em := make([]byte, k)
	em[1] = 1
	for i := 2; i < k-tLen-1; i++ {
		em[i] = 0xff
	}
	copy(em[k-tLen:k-hashLen], prefix)
	copy(em[k-hashLen:k], hashed)

	m := new(big.Int).SetBytes(em)
	c, err := decrypt(rand, priv, m)
	if err == nil {
		s = c.Bytes()
	}
	return
}

// VerifyPKCS1v15 verifies an RSA PKCS#1 v1.5 signature.
// hashed is the result of hashing the input message using the given hash
// function and sig is the signature. A valid signature is indicated by
// returning a nil error.
func VerifyPKCS1v15(pub *PublicKey, hash PKCS1v15Hash, hashed []byte, sig []byte) (err os.Error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (pub.N.BitLen() + 7) / 8
	if k < tLen+11 {
		err = VerificationError{}
		return
	}

	c := new(big.Int).SetBytes(sig)
	m := encrypt(new(big.Int), pub, c)
	em := leftPad(m.Bytes(), k)
	// EM = 0x00 || 0x01 || PS || 0x00 || T

	ok := subtle.ConstantTimeByteEq(em[0], 0)
	ok &= subtle.ConstantTimeByteEq(em[1], 1)
	ok &= subtle.ConstantTimeCompare(em[k-hashLen:k], hashed)
	ok &= subtle.ConstantTimeCompare(em[k-tLen:k-hashLen], prefix)
	ok &= subtle.ConstantTimeByteEq(em[k-tLen-1], 0)

	for i := 2; i < k-tLen-1; i++ {
		ok &= subtle.ConstantTimeByteEq(em[i], 0xff)
	}

	if ok != 1 {
		return VerificationError{}
	}

	return nil
}

func pkcs1v15HashInfo(hash PKCS1v15Hash, inLen int) (hashLen int, prefix []byte, err os.Error) {
	switch hash {
	case HashMD5:
		hashLen = 16
	case HashSHA1:
		hashLen = 20
	case HashSHA256:
		hashLen = 32
	case HashSHA384:
		hashLen = 48
	case HashSHA512:
		hashLen = 64
	case HashMD5SHA1:
		hashLen = 36
	default:
		return 0, nil, os.ErrorString("unknown hash function")
	}

	if inLen != hashLen {
		return 0, nil, os.ErrorString("input must be hashed message")
	}

	prefix = hashPrefixes[int(hash)]
	return
}
<|endoftext|>


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

// Package debug contains facilities for programs to debug themselves while
// they are running.
package debug

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
)

// PrintStack prints to standard error the stack trace returned by Stack.
func PrintStack() {
	os.Stderr.Write(stack())
}

// Stack returns a formatted stack trace of the goroutine that calls it.
// For each routine, it includes the source line information and PC value,
// then attempts to discover, for Go functions, the calling function or
// method and the text of the line containing the invocation.
//
// This function is deprecated. Use package runtime's Stack instead.
func Stack() []byte {
	return stack()
}

// stack implements Stack, skipping 2 frames
func stack() []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := 2; ; i++ { // Caller we care about is the user, 2 frames up
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		line-- // in stack trace, lines are 1-indexed but our array is 0-indexed
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.Trim(lines[n], " \t")
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}
<commit_msg>runtime/debug: fix incorrect Stack output if package path contains a dot<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package debug contains facilities for programs to debug themselves while
// they are running.
package debug

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
	slash     = []byte("/")
)

// PrintStack prints to standard error the stack trace returned by Stack.
func PrintStack() {
	os.Stderr.Write(stack())
}

// Stack returns a formatted stack trace of the goroutine that calls it.
// For each routine, it includes the source line information and PC value,
// then attempts to discover, for Go functions, the calling function or
// method and the text of the line containing the invocation.
//
// This function is deprecated. Use package runtime's Stack instead.
func Stack() []byte {
	return stack()
}

// stack implements Stack, skipping 2 frames
func stack() []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := 2; ; i++ { // Caller we care about is the user, 2 frames up
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		line-- // in stack trace, lines are 1-indexed but our array is 0-indexed
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.Trim(lines[n], " \t")
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	// Since the package path might contains dots (e.g. code.google.com/...),
	// we first remove the path prefix if there is one.
	if lastslash := bytes.LastIndex(name, slash); lastslash >= 0 {
		name = name[lastslash+1:]
	}
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}
<|endoftext|>


<commit_before>package worker_test

import (
	"errors"
	"testing"
	"time"

	"restic/worker"
)

const concurrency = 10

var errTooLarge = errors.New("too large")

func square(job worker.Job, done <-chan struct{}) (interface{}, error) {
	n := job.Data.(int)
	if n > 2000 {
		return nil, errTooLarge
	}
	return n * n, nil
}

func newBufferedPool(bufsize int, n int, f worker.Func) (chan worker.Job, chan worker.Job, *worker.Pool) {
	inCh := make(chan worker.Job, bufsize)
	outCh := make(chan worker.Job, bufsize)

	return inCh, outCh, worker.New(n, f, inCh, outCh)
}

func TestPool(t *testing.T) {
	inCh, outCh, p := newBufferedPool(200, concurrency, square)

	for i := 0; i < 150; i++ {
		inCh <- worker.Job{Data: i}
	}

	close(inCh)
	p.Wait()

	for res := range outCh {
		if res.Error != nil {
			t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
			continue
		}

		n := res.Data.(int)
		m := res.Result.(int)

		if m != n*n {
			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
		}
	}
}

func TestPoolErrors(t *testing.T) {
	inCh, outCh, p := newBufferedPool(200, concurrency, square)

	for i := 0; i < 150; i++ {
		inCh <- worker.Job{Data: i + 1900}
	}

	close(inCh)
	p.Wait()

	for res := range outCh {
		n := res.Data.(int)

		if n > 2000 {
			if res.Error == nil {
				t.Errorf("expected error not found, result is %v", res)
				continue
			}

			if res.Error != errTooLarge {
				t.Errorf("unexpected error found, result is %v", res)
			}

			continue
		} else {
			if res.Error != nil {
				t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
				continue
			}
		}

		m := res.Result.(int)
		if m != n*n {
			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
		}
	}
}

var errCancelled = errors.New("cancelled")

type Job struct {
	suc chan struct{}
	d   time.Duration
}

func wait(job worker.Job, done <-chan struct{}) (interface{}, error) {
	j := job.Data.(Job)
	select {
	case j.suc <- struct{}{}:
		return time.Now(), nil
	case <-time.After(j.d):
		return time.Now(), nil
	case <-done:
		return nil, errCancelled
	}
}

func TestPoolCancel(t *testing.T) {
	jobCh, resCh, p := newBufferedPool(20, concurrency, wait)

	suc := make(chan struct{}, 1)
	for i := 0; i < 20; i++ {
		jobCh <- worker.Job{Data: Job{suc: suc, d: time.Second}}
	}

	<-suc
	p.Cancel()
	p.Wait()

	foundResult := false
	foundCancelError := false
	for res := range resCh {
		if res.Error == nil {
			foundResult = true
		}

		if res.Error == errCancelled {
			foundCancelError = true
		}
	}

	if !foundResult {
		t.Error("did not find one expected result")
	}

	if !foundCancelError {
		t.Error("did not find one expected cancel error")
	}
}
<commit_msg>Fix flaky worker cancel test<commit_after>package worker_test

import (
	"errors"
	"testing"
	"time"

	"restic/worker"
)

const concurrency = 10

var errTooLarge = errors.New("too large")

func square(job worker.Job, done <-chan struct{}) (interface{}, error) {
	n := job.Data.(int)
	if n > 2000 {
		return nil, errTooLarge
	}
	return n * n, nil
}

func newBufferedPool(bufsize int, n int, f worker.Func) (chan worker.Job, chan worker.Job, *worker.Pool) {
	inCh := make(chan worker.Job, bufsize)
	outCh := make(chan worker.Job, bufsize)

	return inCh, outCh, worker.New(n, f, inCh, outCh)
}

func TestPool(t *testing.T) {
	inCh, outCh, p := newBufferedPool(200, concurrency, square)

	for i := 0; i < 150; i++ {
		inCh <- worker.Job{Data: i}
	}

	close(inCh)
	p.Wait()

	for res := range outCh {
		if res.Error != nil {
			t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
			continue
		}

		n := res.Data.(int)
		m := res.Result.(int)

		if m != n*n {
			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
		}
	}
}

func TestPoolErrors(t *testing.T) {
	inCh, outCh, p := newBufferedPool(200, concurrency, square)

	for i := 0; i < 150; i++ {
		inCh <- worker.Job{Data: i + 1900}
	}

	close(inCh)
	p.Wait()

	for res := range outCh {
		n := res.Data.(int)

		if n > 2000 {
			if res.Error == nil {
				t.Errorf("expected error not found, result is %v", res)
				continue
			}

			if res.Error != errTooLarge {
				t.Errorf("unexpected error found, result is %v", res)
			}

			continue
		} else {
			if res.Error != nil {
				t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
				continue
			}
		}

		m := res.Result.(int)
		if m != n*n {
			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
		}
	}
}

var errCancelled = errors.New("cancelled")

type Job struct {
	suc chan struct{}
}

func wait(job worker.Job, done <-chan struct{}) (interface{}, error) {
	j := job.Data.(Job)
	select {
	case j.suc <- struct{}{}:
		return time.Now(), nil
	case <-done:
		return nil, errCancelled
	}
}

func TestPoolCancel(t *testing.T) {
	jobCh, resCh, p := newBufferedPool(20, concurrency, wait)

	suc := make(chan struct{}, 1)
	for i := 0; i < 20; i++ {
		jobCh <- worker.Job{Data: Job{suc: suc}}
	}

	<-suc
	p.Cancel()
	p.Wait()

	foundResult := false
	foundCancelError := false
	for res := range resCh {
		if res.Error == nil {
			foundResult = true
		}

		if res.Error == errCancelled {
			foundCancelError = true
		}
	}

	if !foundResult {
		t.Error("did not find one expected result")
	}

	if !foundCancelError {
		t.Error("did not find one expected cancel error")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/wowqhb/ringbuffer/ringbuffer"
	"fmt"
	"time"
	"bytes"
	"strconv"
)

func main() {
	rbuffer := ringbuffer.RingBuffer{}
	rbuffer.RingBufferInit(3)
	fmt.Println(rbuffer.GetCurrentReadIndex())
	fmt.Println(rbuffer.GetCurrentWriteIndex())
	/*bytes := make([]byte, 20)
	bytes[1] = byte(20)
	fmt.Println(&bytes)
	ok := rbuffer.WriteBuffer(&bytes)
	fmt.Println(ok)
	retP, ok := rbuffer.ReadBuffer()
	fmt.Println(ok)
	fmt.Println(retP)
	bytes[0] = byte(19)
	fmt.Println(bytes)
	fmt.Println(retP)*/
	go writegoroutine(&rbuffer)
	go readgoroutine(&rbuffer)

	time.Sleep(60 * time.Second)

}

func readgoroutine(rbuffer  *ringbuffer.RingBuffer) {
	for {
		fmt.Println(strconv.FormatInt(rbuffer.GetCurrentWriteIndex(),10) + "::::" +strconv.FormatInt(rbuffer.GetCurrentReadIndex(),10))
		retP, ok := rbuffer.ReadBuffer()
		if ok {
			fmt.Print("read::")
			fmt.Print(retP)
			fmt.Print(" =>> ")
			fmt.Println(ok)
		}else {
			fmt.Print("read::nil =>> ")
			fmt.Println(ok)
		}
		//time.Sleep(10 * time.Millisecond)
	}
}

func writegoroutine(rbuffer *ringbuffer.RingBuffer) {
	for {
		fmt.Println(strconv.FormatInt(rbuffer.GetCurrentWriteIndex(),10) + "::::" +strconv.FormatInt(rbuffer.GetCurrentReadIndex(),10))
		time_ := time.Now().String()
		bytes := bytes.NewBufferString(time_).Bytes()
		ok := rbuffer.WriteBuffer(&bytes)
		fmt.Print("write::" + time_ + " =>> ")
		fmt.Println(ok)
		//time.Sleep(10 * time.Millisecond)
	}
}<commit_msg>提交测试代码<commit_after>package main

import (
	"github.com/wowqhb/ringbuffer/ringbuffer"
	"fmt"
	"time"
	"bytes"
//"strconv"
	"strconv"
)

func main() {
	rbuffer := ringbuffer.RingBuffer{}
	rbuffer.RingBufferInit(uint64(8))
	fmt.Println(rbuffer.GetCurrentReadIndex())
	fmt.Println(rbuffer.GetCurrentWriteIndex())
	/*bytes := make([]byte, 20)
	bytes[1] = byte(20)
	fmt.Println(&bytes)
	ok := rbuffer.WriteBuffer(&bytes)
	fmt.Println(ok)
	retP, ok := rbuffer.ReadBuffer()
	fmt.Println(ok)
	fmt.Println(retP)
	bytes[0] = byte(19)
	fmt.Println(bytes)
	fmt.Println(retP)*/
	go writegoroutine(&rbuffer)
	go readgoroutine(&rbuffer)

	time.Sleep(60 * time.Second)

}

func readgoroutine(rbuffer  *ringbuffer.RingBuffer) {
	for {
		retP, ok := rbuffer.ReadBuffer()
		if ok {
			if retP == nil {
				//fmt.Println(strconv.FormatUint(rbuffer.GetCurrentReadIndex() - 1, 10) + "::READ::nil =>> " + strconv.FormatBool(ok))
			}else {
				fmt.Println(strconv.FormatUint(rbuffer.GetCurrentReadIndex() - 1, 10) + "::READ::" + string(*retP) + " =>> " + strconv.FormatBool(ok))
			}

		}else {
			//fmt.Println(strconv.FormatUint(rbuffer.GetCurrentReadIndex(), 10) + "::READ::nil =>> " + strconv.FormatBool(ok))
		}
	}
}

func writegoroutine(rbuffer *ringbuffer.RingBuffer) {
	for {
		time_ := strconv.FormatUint(rbuffer.GetCurrentWriteIndex(), 10);
		bytes := bytes.NewBufferString(time_).Bytes()
		ok := rbuffer.WriteBuffer(&bytes)
		windex := rbuffer.GetCurrentWriteIndex()
		if ok {
			windex = rbuffer.GetCurrentWriteIndex() - 1
			fmt.Println(strconv.FormatUint(windex, 10) + "::WRITE::" + time_ + " =>> " + strconv.FormatBool(ok))
		}
		//fmt.Println(strconv.FormatUint(windex, 10) + "::WRITE::" + time_ + " =>> " + strconv.FormatBool(ok))
	}
}<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/lovoo/goka"
	"github.com/lovoo/goka/codec"
)

var (
	brokers             = []string{"127.0.0.1:9092"}
	topic   goka.Stream = "user-click"
	group   goka.Group  = "mini-group"
)

// A user is the object that is stored in the processor's group table
type user struct {
	// number of clicks the user has performed.
	Clicks int
}

// This codec allows marshalling (encode) and unmarshalling (decode) the user to and from the
// group table.
type userCodec struct{}

// Encodes a user into []byte
func (jc *userCodec) Encode(value interface{}) ([]byte, error) {
	if _, isUser := value.(*user); !isUser {
		return nil, fmt.Errorf("Codec requires value *user, got %T", value)
	}
	return json.Marshal(value)
}

// Decodes a user from []byte to it's go representation.
func (jc *userCodec) Decode(data []byte) (interface{}, error) {
	var (
		c   user
		err error
	)
	err = json.Unmarshal(data, &c)
	if err != nil {
		return nil, fmt.Errorf("Error unmarshaling user: %v", err)
	}
	return &c, nil
}

func runEmitter() {
	emitter, err := goka.NewEmitter(brokers, topic,
		new(codec.String))
	if err != nil {
		panic(err)
	}
	defer emitter.Finish()

	t := time.NewTicker(100 * time.Millisecond)
	defer t.Stop()

	var i int
	for range t.C {
		key := fmt.Sprintf("user-%d", i%10)
		value := fmt.Sprintf("%s", time.Now())
		emitter.EmitSync(key, value)
		i++
	}
}

func process(ctx goka.Context, msg interface{}) {
	var u *user
	if val := ctx.Value(); val != nil {
		u = val.(*user)
	} else {
		u = new(user)
	}

	u.Clicks++
	ctx.SetValue(u)
	fmt.Printf("[proc] key: %s clicks: %d, msg: %v\n", ctx.Key(), u.Clicks, msg)
}

func runProcessor() {
	g := goka.DefineGroup(group,
		goka.Input(topic, new(codec.String), process),
		goka.Persist(new(userCodec)),
	)
	tmc := goka.NewTopicManagerConfig()
	tmc.Table.Replication = 1
	tmc.Stream.Replication = 1
	p, err := goka.NewProcessor(brokers,
		g,
		goka.WithTopicManagerBuilder(goka.TopicManagerBuilderWithTopicManagerConfig(tmc)),
		goka.WithConsumerGroupBuilder(goka.DefaultConsumerGroupBuilder),
	)
	if err != nil {
		panic(err)
	}

	p.Run(context.Background())
}

func runView() {
	view, err := goka.NewView(brokers,
		goka.GroupTable(group),
		new(userCodec),
	)
	if err != nil {
		panic(err)
	}

	root := mux.NewRouter()
	root.HandleFunc("/{key}", func(w http.ResponseWriter, r *http.Request) {
		value, _ := view.Get(mux.Vars(r)["key"])
		data, _ := json.Marshal(value)
		w.Write(data)
	})
	fmt.Println("View opened at http://localhost:9095/")
	go http.ListenAndServe(":9095", root)

	view.Run(context.Background())
}

func main() {
	go runEmitter()
	go runProcessor()
	runView()
}
<commit_msg>Fix initialization of topics in example 2<commit_after>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/lovoo/goka"
	"github.com/lovoo/goka/codec"
)

var (
	brokers             = []string{"127.0.0.1:9092"}
	topic   goka.Stream = "user-click"
	group   goka.Group  = "mini-group"

	tmc *goka.TopicManagerConfig
)

// A user is the object that is stored in the processor's group table
type user struct {
	// number of clicks the user has performed.
	Clicks int
}

// This codec allows marshalling (encode) and unmarshalling (decode) the user to and from the
// group table.
type userCodec struct{}

func init() {
	tmc = goka.NewTopicManagerConfig()
	tmc.Table.Replication = 1
	tmc.Stream.Replication = 1
}

// Encodes a user into []byte
func (jc *userCodec) Encode(value interface{}) ([]byte, error) {
	if _, isUser := value.(*user); !isUser {
		return nil, fmt.Errorf("Codec requires value *user, got %T", value)
	}
	return json.Marshal(value)
}

// Decodes a user from []byte to it's go representation.
func (jc *userCodec) Decode(data []byte) (interface{}, error) {
	var (
		c   user
		err error
	)
	err = json.Unmarshal(data, &c)
	if err != nil {
		return nil, fmt.Errorf("Error unmarshaling user: %v", err)
	}
	return &c, nil
}

func runEmitter() {
	emitter, err := goka.NewEmitter(brokers, topic,
		new(codec.String))
	if err != nil {
		panic(err)
	}
	defer emitter.Finish()

	t := time.NewTicker(100 * time.Millisecond)
	defer t.Stop()

	var i int
	for range t.C {
		key := fmt.Sprintf("user-%d", i%10)
		value := fmt.Sprintf("%s", time.Now())
		emitter.EmitSync(key, value)
		i++
	}
}

func process(ctx goka.Context, msg interface{}) {
	var u *user
	if val := ctx.Value(); val != nil {
		u = val.(*user)
	} else {
		u = new(user)
	}

	u.Clicks++
	ctx.SetValue(u)
	fmt.Printf("[proc] key: %s clicks: %d, msg: %v\n", ctx.Key(), u.Clicks, msg)
}

func runProcessor(initialized chan bool) {
	g := goka.DefineGroup(group,
		goka.Input(topic, new(codec.String), process),
		goka.Persist(new(userCodec)),
	)
	p, err := goka.NewProcessor(brokers,
		g,
		goka.WithTopicManagerBuilder(goka.TopicManagerBuilderWithTopicManagerConfig(tmc)),
		goka.WithConsumerGroupBuilder(goka.DefaultConsumerGroupBuilder),
	)
	if err != nil {
		panic(err)
	}

	initialized <- true

	p.Run(context.Background())
}

func runView(initialized chan bool) {
	<-initialized

	view, err := goka.NewView(brokers,
		goka.GroupTable(group),
		new(userCodec),
	)
	if err != nil {
		panic(err)
	}

	root := mux.NewRouter()
	root.HandleFunc("/{key}", func(w http.ResponseWriter, r *http.Request) {
		value, _ := view.Get(mux.Vars(r)["key"])
		data, _ := json.Marshal(value)
		w.Write(data)
	})
	fmt.Println("View opened at http://localhost:9095/")
	go http.ListenAndServe(":9095", root)

	view.Run(context.Background())
}

func main() {
	tm, err := goka.NewTopicManager(brokers, goka.DefaultConfig(), tmc)
	if err != nil {
		log.Fatalf("Error creating topic manager: %v", err)
	}
	err = tm.EnsureStreamExists(string(topic), 8)
	if err != nil {
		log.Printf("Error creating kafka topic %s: %v", topic, err)
	}

	// When this example is run the first time, wait for creation of all internal topics (this will be
	// done by function runProcessor)
	initialized := make(chan bool)

	go runEmitter()
	go runProcessor(initialized)
	runView(initialized)
}
<|endoftext|>


<commit_before>package rpc

import (
	"bufio"
	"cgl.tideland.biz/asserts"
	"github.com/mitchellh/packer/packer"
	"io"
	"net/rpc"
	"testing"
	"time"
)

type testCommunicator struct {
	startCalled bool
	startCmd    *packer.RemoteCmd

	uploadCalled bool
	uploadPath   string
	uploadData   string

	downloadCalled bool
	downloadPath   string
}

func (t *testCommunicator) Start(cmd *packer.RemoteCmd) error {
	t.startCalled = true
	t.startCmd = cmd
	return nil
}

func (t *testCommunicator) Upload(path string, reader io.Reader) (err error) {
	t.uploadCalled = true
	t.uploadPath = path
	t.uploadData, err = bufio.NewReader(reader).ReadString('\n')
	return
}

func (t *testCommunicator) Download(path string, writer io.Writer) error {
	t.downloadCalled = true
	t.downloadPath = path
	writer.Write([]byte("download\n"))

	return nil
}

func TestCommunicatorRPC(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	// Create the interface to test
	c := new(testCommunicator)

	// Start the server
	server := rpc.NewServer()
	RegisterCommunicator(server, c)
	address := serveSingleConn(server)

	// Create the client over RPC and run some methods to verify it works
	client, err := rpc.Dial("tcp", address)
	assert.Nil(err, "should be able to connect")
	remote := Communicator(client)

	// The remote command we'll use
	stdin_r, stdin_w := io.Pipe()
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()

	var cmd packer.RemoteCmd
	cmd.Command = "foo"
	cmd.Stdin = stdin_r
	cmd.Stdout = stdout_w
	cmd.Stderr = stderr_w

	// Test Start
	err = remote.Start(&cmd)
	assert.Nil(err, "should not have an error")

	// Test that we can read from stdout
	c.startCmd.Stdout.Write([]byte("outfoo\n"))
	bufOut := bufio.NewReader(stdout_r)
	data, err := bufOut.ReadString('\n')
	assert.Nil(err, "should have no problem reading stdout")
	assert.Equal(data, "outfoo\n", "should be correct stdout")

	// Test that we can read from stderr
	c.startCmd.Stderr.Write([]byte("errfoo\n"))
	bufErr := bufio.NewReader(stderr_r)
	data, err = bufErr.ReadString('\n')
	assert.Nil(err, "should have no problem reading stderr")
	assert.Equal(data, "errfoo\n", "should be correct stderr")

	// Test that we can write to stdin
	stdin_w.Write([]byte("infoo\n"))
	bufIn := bufio.NewReader(c.startCmd.Stdin)
	data, err = bufIn.ReadString('\n')
	assert.Nil(err, "should have no problem reading stdin")
	assert.Equal(data, "infoo\n", "should be correct stdin")

	// Test that we can get the exit status properly
	c.startCmd.SetExited(42)

	for i := 0; i < 5; i++ {
		if cmd.Exited {
			assert.Equal(cmd.ExitStatus, 42, "should have proper exit status")
			break
		}

		time.Sleep(50 * time.Millisecond)
	}

	assert.True(cmd.Exited, "should have exited")

	// Test that we can upload things
	uploadR, uploadW := io.Pipe()
	go uploadW.Write([]byte("uploadfoo\n"))
	err = remote.Upload("foo", uploadR)
	assert.Nil(err, "should not error")
	assert.True(c.uploadCalled, "should be called")
	assert.Equal(c.uploadPath, "foo", "should be correct path")
	assert.Equal(c.uploadData, "uploadfoo\n", "should have the proper data")

	// Test that we can download things
	downloadR, downloadW := io.Pipe()
	downloadDone := make(chan bool)
	var downloadData string
	var downloadErr error

	go func() {
		bufDownR := bufio.NewReader(downloadR)
		downloadData, downloadErr = bufDownR.ReadString('\n')
		downloadDone <- true
	}()

	err = remote.Download("bar", downloadW)
	assert.Nil(err, "should not error")
	assert.True(c.downloadCalled, "should have called download")
	assert.Equal(c.downloadPath, "bar", "should have correct download path")

	<-downloadDone
	assert.Nil(downloadErr, "should not error reading download data")
	assert.Equal(downloadData, "download\n", "should have the proper data")
}

func TestCommunicator_ImplementsCommunicator(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	var r packer.Communicator
	c := Communicator(nil)

	assert.Implementor(c, &r, "should be a Communicator")
}
<commit_msg>packer/rpc: get rid of data races in tests<commit_after>package rpc

import (
	"bufio"
	"cgl.tideland.biz/asserts"
	"github.com/mitchellh/packer/packer"
	"io"
	"net/rpc"
	"testing"
	"time"
)

type testCommunicator struct {
	startCalled bool
	startCmd    *packer.RemoteCmd

	uploadCalled bool
	uploadPath   string
	uploadData   string

	downloadCalled bool
	downloadPath   string
}

func (t *testCommunicator) Start(cmd *packer.RemoteCmd) error {
	t.startCalled = true
	t.startCmd = cmd
	return nil
}

func (t *testCommunicator) Upload(path string, reader io.Reader) (err error) {
	t.uploadCalled = true
	t.uploadPath = path
	t.uploadData, err = bufio.NewReader(reader).ReadString('\n')
	return
}

func (t *testCommunicator) Download(path string, writer io.Writer) error {
	t.downloadCalled = true
	t.downloadPath = path
	writer.Write([]byte("download\n"))

	return nil
}

func TestCommunicatorRPC(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	// Create the interface to test
	c := new(testCommunicator)

	// Start the server
	server := rpc.NewServer()
	RegisterCommunicator(server, c)
	address := serveSingleConn(server)

	// Create the client over RPC and run some methods to verify it works
	client, err := rpc.Dial("tcp", address)
	assert.Nil(err, "should be able to connect")
	remote := Communicator(client)

	// The remote command we'll use
	stdin_r, stdin_w := io.Pipe()
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()

	var cmd packer.RemoteCmd
	cmd.Command = "foo"
	cmd.Stdin = stdin_r
	cmd.Stdout = stdout_w
	cmd.Stderr = stderr_w

	// Test Start
	err = remote.Start(&cmd)
	assert.Nil(err, "should not have an error")

	// Test that we can read from stdout
	c.startCmd.Stdout.Write([]byte("outfoo\n"))
	bufOut := bufio.NewReader(stdout_r)
	data, err := bufOut.ReadString('\n')
	assert.Nil(err, "should have no problem reading stdout")
	assert.Equal(data, "outfoo\n", "should be correct stdout")

	// Test that we can read from stderr
	c.startCmd.Stderr.Write([]byte("errfoo\n"))
	bufErr := bufio.NewReader(stderr_r)
	data, err = bufErr.ReadString('\n')
	assert.Nil(err, "should have no problem reading stderr")
	assert.Equal(data, "errfoo\n", "should be correct stderr")

	// Test that we can write to stdin
	stdin_w.Write([]byte("infoo\n"))
	bufIn := bufio.NewReader(c.startCmd.Stdin)
	data, err = bufIn.ReadString('\n')
	assert.Nil(err, "should have no problem reading stdin")
	assert.Equal(data, "infoo\n", "should be correct stdin")

	// Test that we can get the exit status properly
	c.startCmd.SetExited(42)

	for i := 0; i < 5; i++ {
		cmd.Lock()
		exited := cmd.Exited
		cmd.Unlock()
		if exited {
			assert.Equal(cmd.ExitStatus, 42, "should have proper exit status")
			break
		}

		time.Sleep(50 * time.Millisecond)
	}

	assert.True(cmd.Exited, "should have exited")

	// Test that we can upload things
	uploadR, uploadW := io.Pipe()
	go uploadW.Write([]byte("uploadfoo\n"))
	err = remote.Upload("foo", uploadR)
	assert.Nil(err, "should not error")
	assert.True(c.uploadCalled, "should be called")
	assert.Equal(c.uploadPath, "foo", "should be correct path")
	assert.Equal(c.uploadData, "uploadfoo\n", "should have the proper data")

	// Test that we can download things
	downloadR, downloadW := io.Pipe()
	downloadDone := make(chan bool)
	var downloadData string
	var downloadErr error

	go func() {
		bufDownR := bufio.NewReader(downloadR)
		downloadData, downloadErr = bufDownR.ReadString('\n')
		downloadDone <- true
	}()

	err = remote.Download("bar", downloadW)
	assert.Nil(err, "should not error")
	assert.True(c.downloadCalled, "should have called download")
	assert.Equal(c.downloadPath, "bar", "should have correct download path")

	<-downloadDone
	assert.Nil(downloadErr, "should not error reading download data")
	assert.Equal(downloadData, "download\n", "should have the proper data")
}

func TestCommunicator_ImplementsCommunicator(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	var r packer.Communicator
	c := Communicator(nil)

	assert.Implementor(c, &r, "should be a Communicator")
}
<|endoftext|>


<commit_before>package dockerbuild

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

	docker "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/credentialprovider"
	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"

	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	"github.com/openshift/origin/pkg/util/docker/dockerfile/builder"
)

const (
	dockerbuildLong = `
Build a Dockerfile into a single layer

Builds the provided directory with a Dockerfile into a single layered image.
Requires that you have a working connection to a Docker engine.`

	dockerbuildExample = `  # Build the current directory into a single layer and tag
  $ %[1]s dockerbuild . myimage:latest`
)

type DockerbuildOptions struct {
	Out io.Writer
	Err io.Writer

	Client *docker.Client

	Directory      string
	Tag            string
	DockerfilePath string
	AllowPull      bool
	Keyring        credentialprovider.DockerKeyring
	Arguments      cmdutil.Environment
}

func NewCmdDockerbuild(fullName string, f *clientcmd.Factory, out, errOut io.Writer) *cobra.Command {
	options := &DockerbuildOptions{
		Out: out,
		Err: errOut,
	}
	cmd := &cobra.Command{
		Use:     "dockerbuild DIRECTORY TAG [--dockerfile=PATH]",
		Short:   "Perform a direct Docker build",
		Long:    dockerbuildLong,
		Example: fmt.Sprintf(dockerbuildExample, fullName),
		Run: func(cmd *cobra.Command, args []string) {
			kcmdutil.CheckErr(options.Complete(f, cmd, args))
			kcmdutil.CheckErr(options.Validate())
			if err := options.Run(); err != nil {
				// TODO: move met to kcmdutil
				if err == cmdutil.ErrExit {
					os.Exit(1)
				}
				kcmdutil.CheckErr(err)
			}
		},
	}

	cmd.Flags().StringVar(&options.DockerfilePath, "dockerfile", options.DockerfilePath, "An optional path to a Dockerfile to use.")
	cmd.Flags().BoolVar(&options.AllowPull, "allow-pull", true, "Pull the images that are not present.")
	cmd.MarkFlagFilename("dockerfile")

	return cmd
}

func (o *DockerbuildOptions) Complete(f *clientcmd.Factory, cmd *cobra.Command, args []string) error {
	paths, envArgs, ok := cmdutil.SplitEnvironmentFromResources(args)
	if !ok {
		return kcmdutil.UsageError(cmd, "context directory must be specified before environment changes: %s", strings.Join(args, " "))
	}
	if len(paths) != 2 {
		return kcmdutil.UsageError(cmd, "the directory to build and tag must be specified")
	}
	o.Arguments, _, _ = cmdutil.ParseEnvironmentArguments(envArgs)
	o.Directory = paths[0]
	o.Tag = paths[1]
	if len(o.DockerfilePath) == 0 {
		o.DockerfilePath = filepath.Join(o.Directory, "Dockerfile")
	}
	client, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}
	o.Client = client

	o.Keyring = credentialprovider.NewDockerKeyring()

	return nil
}

func (o *DockerbuildOptions) Validate() error {
	return nil
}

func (o *DockerbuildOptions) Run() error {
	f, err := os.Open(o.DockerfilePath)
	if err != nil {
		return err
	}
	defer f.Close()
	e := builder.NewClientExecutor(o.Client)
	e.Out, e.ErrOut = o.Out, o.Err
	e.AllowPull = o.AllowPull
	e.Directory = o.Directory
	e.Tag = o.Tag
	e.AuthFn = o.Keyring.Lookup
	e.LogFn = func(format string, args ...interface{}) {
		if glog.V(2) {
			glog.Infof("Builder: "+format, args...)
		} else {
			fmt.Fprintf(e.ErrOut, "# %s\n", fmt.Sprintf(format, args...))
		}
	}
	return stripLeadingError(e.Build(f, o.Arguments))
}

func stripLeadingError(err error) error {
	if err == nil {
		return err
	}
	if strings.HasPrefix(err.Error(), "Error: ") {
		return fmt.Errorf(strings.TrimPrefix(err.Error(), "Error: "))
	}
	return err
}
<commit_msg>Interrupting a dockerbuild command should clean up leftovers<commit_after>package dockerbuild

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

	docker "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/credentialprovider"
	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/util/interrupt"

	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	"github.com/openshift/origin/pkg/util/docker/dockerfile/builder"
)

const (
	dockerbuildLong = `
Build a Dockerfile into a single layer

Builds the provided directory with a Dockerfile into a single layered image.
Requires that you have a working connection to a Docker engine.`

	dockerbuildExample = `  # Build the current directory into a single layer and tag
  $ %[1]s dockerbuild . myimage:latest`
)

type DockerbuildOptions struct {
	Out io.Writer
	Err io.Writer

	Client *docker.Client

	Directory      string
	Tag            string
	DockerfilePath string
	AllowPull      bool
	Keyring        credentialprovider.DockerKeyring
	Arguments      cmdutil.Environment
}

func NewCmdDockerbuild(fullName string, f *clientcmd.Factory, out, errOut io.Writer) *cobra.Command {
	options := &DockerbuildOptions{
		Out: out,
		Err: errOut,
	}
	cmd := &cobra.Command{
		Use:     "dockerbuild DIRECTORY TAG [--dockerfile=PATH]",
		Short:   "Perform a direct Docker build",
		Long:    dockerbuildLong,
		Example: fmt.Sprintf(dockerbuildExample, fullName),
		Run: func(cmd *cobra.Command, args []string) {
			kcmdutil.CheckErr(options.Complete(f, cmd, args))
			kcmdutil.CheckErr(options.Validate())
			if err := options.Run(); err != nil {
				// TODO: move met to kcmdutil
				if err == cmdutil.ErrExit {
					os.Exit(1)
				}
				kcmdutil.CheckErr(err)
			}
		},
	}

	cmd.Flags().StringVar(&options.DockerfilePath, "dockerfile", options.DockerfilePath, "An optional path to a Dockerfile to use.")
	cmd.Flags().BoolVar(&options.AllowPull, "allow-pull", true, "Pull the images that are not present.")
	cmd.MarkFlagFilename("dockerfile")

	return cmd
}

func (o *DockerbuildOptions) Complete(f *clientcmd.Factory, cmd *cobra.Command, args []string) error {
	paths, envArgs, ok := cmdutil.SplitEnvironmentFromResources(args)
	if !ok {
		return kcmdutil.UsageError(cmd, "context directory must be specified before environment changes: %s", strings.Join(args, " "))
	}
	if len(paths) != 2 {
		return kcmdutil.UsageError(cmd, "the directory to build and tag must be specified")
	}
	o.Arguments, _, _ = cmdutil.ParseEnvironmentArguments(envArgs)
	o.Directory = paths[0]
	o.Tag = paths[1]
	if len(o.DockerfilePath) == 0 {
		o.DockerfilePath = filepath.Join(o.Directory, "Dockerfile")
	}
	client, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}
	o.Client = client

	o.Keyring = credentialprovider.NewDockerKeyring()

	return nil
}

func (o *DockerbuildOptions) Validate() error {
	return nil
}

func (o *DockerbuildOptions) Run() error {
	f, err := os.Open(o.DockerfilePath)
	if err != nil {
		return err
	}
	defer f.Close()
	e := builder.NewClientExecutor(o.Client)
	e.Out, e.ErrOut = o.Out, o.Err
	e.AllowPull = o.AllowPull
	e.Directory = o.Directory
	e.Tag = o.Tag
	e.AuthFn = o.Keyring.Lookup
	e.LogFn = func(format string, args ...interface{}) {
		if glog.V(2) {
			glog.Infof("Builder: "+format, args...)
		} else {
			fmt.Fprintf(e.ErrOut, "--> %s\n", fmt.Sprintf(format, args...))
		}
	}
	safe := interrupt.New(func(os.Signal) { os.Exit(1) }, func() {
		glog.V(5).Infof("invoking cleanup")
		if err := e.Cleanup(); err != nil {
			fmt.Fprintf(o.Err, "error: Unable to clean up build: %v\n", err)
		}
	})
	return safe.Run(func() error { return stripLeadingError(e.Build(f, o.Arguments)) })
}

func stripLeadingError(err error) error {
	if err == nil {
		return err
	}
	if strings.HasPrefix(err.Error(), "Error: ") {
		return fmt.Errorf(strings.TrimPrefix(err.Error(), "Error: "))
	}
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package metricsutil

import (
	"fmt"
	"io"

	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/printers"
	metricsapi "k8s.io/metrics/pkg/apis/metrics/v1alpha1"
)

var (
	MeasuredResources = []api.ResourceName{
		api.ResourceCPU,
		api.ResourceMemory,
	}
	NodeColumns     = []string{"NAME", "CPU(cores)", "CPU%", "MEMORY(bytes)", "MEMORY%"}
	PodColumns      = []string{"NAME", "CPU(cores)", "MEMORY(bytes)"}
	NamespaceColumn = "NAMESPACE"
	PodColumn       = "POD"
)

type ResourceMetricsInfo struct {
	Name      string
	Metrics   api.ResourceList
	Available api.ResourceList
}

type TopCmdPrinter struct {
	out io.Writer
}

func NewTopCmdPrinter(out io.Writer) *TopCmdPrinter {
	return &TopCmdPrinter{out: out}
}

func (printer *TopCmdPrinter) PrintNodeMetrics(metrics []metricsapi.NodeMetrics, availableResources map[string]api.ResourceList) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	printColumnNames(w, NodeColumns)
	var usage api.ResourceList
	for _, m := range metrics {
		err := api.Scheme.Convert(&m.Usage, &usage, nil)
		if err != nil {
			return err
		}
		printMetricsLine(w, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   usage,
			Available: availableResources[m.Name],
		})
	}
	return nil
}

func (printer *TopCmdPrinter) PrintPodMetrics(metrics []metricsapi.PodMetrics, printContainers bool, withNamespace bool) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	if withNamespace {
		printValue(w, NamespaceColumn)
	}
	if printContainers {
		printValue(w, PodColumn)
	}
	printColumnNames(w, PodColumns)
	for _, m := range metrics {
		err := printSinglePodMetrics(w, &m, printContainers, withNamespace)
		if err != nil {
			return err
		}
	}
	return nil
}

func printColumnNames(out io.Writer, names []string) {
	for _, name := range names {
		printValue(out, name)
	}
	fmt.Fprint(out, "\n")
}

func printSinglePodMetrics(out io.Writer, m *metricsapi.PodMetrics, printContainersOnly bool, withNamespace bool) error {
	containers := make(map[string]api.ResourceList)
	podMetrics := make(api.ResourceList)
	for _, res := range MeasuredResources {
		podMetrics[res], _ = resource.ParseQuantity("0")
	}

	for _, c := range m.Containers {
		var usage api.ResourceList
		err := api.Scheme.Convert(&c.Usage, &usage, nil)
		if err != nil {
			return err
		}
		containers[c.Name] = usage
		if !printContainersOnly {
			for _, res := range MeasuredResources {
				quantity := podMetrics[res]
				quantity.Add(usage[res])
				podMetrics[res] = quantity
			}
		}
	}
	if printContainersOnly {
		for contName := range containers {
			if withNamespace {
				printValue(out, m.Namespace)
			}
			printValue(out, m.Name)
			printMetricsLine(out, &ResourceMetricsInfo{
				Name:      contName,
				Metrics:   containers[contName],
				Available: api.ResourceList{},
			})
		}
	} else {
		if withNamespace {
			printValue(out, m.Namespace)
		}
		printMetricsLine(out, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   podMetrics,
			Available: api.ResourceList{},
		})
	}
	return nil
}

func printMetricsLine(out io.Writer, metrics *ResourceMetricsInfo) {
	printValue(out, metrics.Name)
	printAllResourceUsages(out, metrics)
	fmt.Fprint(out, "\n")
}

func printValue(out io.Writer, value interface{}) {
	fmt.Fprintf(out, "%v\t", value)
}

func printAllResourceUsages(out io.Writer, metrics *ResourceMetricsInfo) {
	for _, res := range MeasuredResources {
		quantity := metrics.Metrics[res]
		printSingleResourceUsage(out, res, quantity)
		fmt.Fprint(out, "\t")
		if available, found := metrics.Available[res]; found {
			fraction := float64(quantity.MilliValue()) / float64(available.MilliValue()) * 100
			fmt.Fprintf(out, "%d%%\t", int64(fraction))
		}
	}
}

func printSingleResourceUsage(out io.Writer, resourceType api.ResourceName, quantity resource.Quantity) {
	switch resourceType {
	case api.ResourceCPU:
		fmt.Fprintf(out, "%vm", quantity.MilliValue())
	case api.ResourceMemory:
		fmt.Fprintf(out, "%vMi", quantity.Value()/(1024*1024))
	default:
		fmt.Fprintf(out, "%v", quantity.Value())
	}
}
<commit_msg>sort top pod and top node output by namespace/name<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package metricsutil

import (
	"fmt"
	"io"
	"sort"

	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/printers"
	metricsapi "k8s.io/metrics/pkg/apis/metrics/v1alpha1"
)

var (
	MeasuredResources = []api.ResourceName{
		api.ResourceCPU,
		api.ResourceMemory,
	}
	NodeColumns     = []string{"NAME", "CPU(cores)", "CPU%", "MEMORY(bytes)", "MEMORY%"}
	PodColumns      = []string{"NAME", "CPU(cores)", "MEMORY(bytes)"}
	NamespaceColumn = "NAMESPACE"
	PodColumn       = "POD"
)

type ResourceMetricsInfo struct {
	Name      string
	Metrics   api.ResourceList
	Available api.ResourceList
}

type TopCmdPrinter struct {
	out io.Writer
}

func NewTopCmdPrinter(out io.Writer) *TopCmdPrinter {
	return &TopCmdPrinter{out: out}
}

func (printer *TopCmdPrinter) PrintNodeMetrics(metrics []metricsapi.NodeMetrics, availableResources map[string]api.ResourceList) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	sort.Slice(metrics, func(i, j int) bool {
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, NodeColumns)
	var usage api.ResourceList
	for _, m := range metrics {
		err := api.Scheme.Convert(&m.Usage, &usage, nil)
		if err != nil {
			return err
		}
		printMetricsLine(w, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   usage,
			Available: availableResources[m.Name],
		})
	}
	return nil
}

func (printer *TopCmdPrinter) PrintPodMetrics(metrics []metricsapi.PodMetrics, printContainers bool, withNamespace bool) error {
	if len(metrics) == 0 {
		return nil
	}
	w := printers.GetNewTabWriter(printer.out)
	defer w.Flush()

	if withNamespace {
		printValue(w, NamespaceColumn)
	}
	if printContainers {
		printValue(w, PodColumn)
	}

	sort.Slice(metrics, func(i, j int) bool {
		if withNamespace && metrics[i].Namespace != metrics[j].Namespace {
			return metrics[i].Namespace < metrics[j].Namespace
		}
		return metrics[i].Name < metrics[j].Name
	})

	printColumnNames(w, PodColumns)
	for _, m := range metrics {
		err := printSinglePodMetrics(w, &m, printContainers, withNamespace)
		if err != nil {
			return err
		}
	}
	return nil
}

func printColumnNames(out io.Writer, names []string) {
	for _, name := range names {
		printValue(out, name)
	}
	fmt.Fprint(out, "\n")
}

func printSinglePodMetrics(out io.Writer, m *metricsapi.PodMetrics, printContainersOnly bool, withNamespace bool) error {
	containers := make(map[string]api.ResourceList)
	podMetrics := make(api.ResourceList)
	for _, res := range MeasuredResources {
		podMetrics[res], _ = resource.ParseQuantity("0")
	}

	for _, c := range m.Containers {
		var usage api.ResourceList
		err := api.Scheme.Convert(&c.Usage, &usage, nil)
		if err != nil {
			return err
		}
		containers[c.Name] = usage
		if !printContainersOnly {
			for _, res := range MeasuredResources {
				quantity := podMetrics[res]
				quantity.Add(usage[res])
				podMetrics[res] = quantity
			}
		}
	}
	if printContainersOnly {
		for contName := range containers {
			if withNamespace {
				printValue(out, m.Namespace)
			}
			printValue(out, m.Name)
			printMetricsLine(out, &ResourceMetricsInfo{
				Name:      contName,
				Metrics:   containers[contName],
				Available: api.ResourceList{},
			})
		}
	} else {
		if withNamespace {
			printValue(out, m.Namespace)
		}
		printMetricsLine(out, &ResourceMetricsInfo{
			Name:      m.Name,
			Metrics:   podMetrics,
			Available: api.ResourceList{},
		})
	}
	return nil
}

func printMetricsLine(out io.Writer, metrics *ResourceMetricsInfo) {
	printValue(out, metrics.Name)
	printAllResourceUsages(out, metrics)
	fmt.Fprint(out, "\n")
}

func printValue(out io.Writer, value interface{}) {
	fmt.Fprintf(out, "%v\t", value)
}

func printAllResourceUsages(out io.Writer, metrics *ResourceMetricsInfo) {
	for _, res := range MeasuredResources {
		quantity := metrics.Metrics[res]
		printSingleResourceUsage(out, res, quantity)
		fmt.Fprint(out, "\t")
		if available, found := metrics.Available[res]; found {
			fraction := float64(quantity.MilliValue()) / float64(available.MilliValue()) * 100
			fmt.Fprintf(out, "%d%%\t", int64(fraction))
		}
	}
}

func printSingleResourceUsage(out io.Writer, resourceType api.ResourceName, quantity resource.Quantity) {
	switch resourceType {
	case api.ResourceCPU:
		fmt.Fprintf(out, "%vm", quantity.MilliValue())
	case api.ResourceMemory:
		fmt.Fprintf(out, "%vMi", quantity.Value()/(1024*1024))
	default:
		fmt.Fprintf(out, "%v", quantity.Value())
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

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

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

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

package resources

import (
	"fmt"
	"math"
	"strconv"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	network "knative.dev/networking/pkg"
	pkgnet "knative.dev/networking/pkg/apis/networking"
	"knative.dev/pkg/metrics"
	"knative.dev/pkg/profiling"
	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/deployment"
	"knative.dev/serving/pkg/networking"
	"knative.dev/serving/pkg/queue"
	"knative.dev/serving/pkg/queue/readiness"
	"knative.dev/serving/pkg/reconciler/revision/config"
)

const (
	localAddress             = "127.0.0.1"
	requestQueueHTTPPortName = "queue-port"
	profilingPortName        = "profiling-port"
)

var (
	queueHTTPPort = corev1.ContainerPort{
		Name:          requestQueueHTTPPortName,
		ContainerPort: networking.BackendHTTPPort,
	}
	queueHTTP2Port = corev1.ContainerPort{
		Name:          requestQueueHTTPPortName,
		ContainerPort: networking.BackendHTTP2Port,
	}
	queueNonServingPorts = []corev1.ContainerPort{{
		// Provides health checks and lifecycle hooks.
		Name:          v1.QueueAdminPortName,
		ContainerPort: networking.QueueAdminPort,
	}, {
		Name:          v1.AutoscalingQueueMetricsPortName,
		ContainerPort: networking.AutoscalingQueueMetricsPort,
	}, {
		Name:          v1.UserQueueMetricsPortName,
		ContainerPort: networking.UserQueueMetricsPort,
	}}

	profilingPort = corev1.ContainerPort{
		Name:          profilingPortName,
		ContainerPort: profiling.ProfilingPort,
	}

	queueSecurityContext = &corev1.SecurityContext{
		AllowPrivilegeEscalation: ptr.Bool(false),
	}
)

func createQueueResources(cfg *deployment.Config, annotations map[string]string, userContainer *corev1.Container) corev1.ResourceRequirements {
	resourceRequests := corev1.ResourceList{}
	resourceLimits := corev1.ResourceList{}

	for _, r := range []struct {
		Name    corev1.ResourceName
		Request *resource.Quantity
		Limit   *resource.Quantity
	}{{
		Name:    corev1.ResourceCPU,
		Request: cfg.QueueSidecarCPURequest,
		Limit:   cfg.QueueSidecarCPULimit,
	}, {
		Name:    corev1.ResourceMemory,
		Request: cfg.QueueSidecarMemoryRequest,
		Limit:   cfg.QueueSidecarMemoryLimit,
	}, {
		Name:    corev1.ResourceEphemeralStorage,
		Request: cfg.QueueSidecarEphemeralStorageRequest,
		Limit:   cfg.QueueSidecarEphemeralStorageLimit,
	}} {
		if r.Request != nil {
			resourceRequests[r.Name] = *r.Request
		}
		if r.Limit != nil {
			resourceLimits[r.Name] = *r.Limit
		}
	}

	var requestCPU, limitCPU, requestMemory, limitMemory resource.Quantity

	if resourceFraction, ok := fractionFromPercentage(annotations, serving.QueueSideCarResourcePercentageAnnotation); ok {
		if ok, requestCPU = computeResourceRequirements(userContainer.Resources.Requests.Cpu(), resourceFraction, queueContainerRequestCPU); ok {
			resourceRequests[corev1.ResourceCPU] = requestCPU
		}

		if ok, limitCPU = computeResourceRequirements(userContainer.Resources.Limits.Cpu(), resourceFraction, queueContainerLimitCPU); ok {
			resourceLimits[corev1.ResourceCPU] = limitCPU
		}

		if ok, requestMemory = computeResourceRequirements(userContainer.Resources.Requests.Memory(), resourceFraction, queueContainerRequestMemory); ok {
			resourceRequests[corev1.ResourceMemory] = requestMemory
		}

		if ok, limitMemory = computeResourceRequirements(userContainer.Resources.Limits.Memory(), resourceFraction, queueContainerLimitMemory); ok {
			resourceLimits[corev1.ResourceMemory] = limitMemory
		}
	}

	resources := corev1.ResourceRequirements{
		Requests: resourceRequests,
	}
	if len(resourceLimits) != 0 {
		resources.Limits = resourceLimits
	}

	return resources
}

func computeResourceRequirements(resourceQuantity *resource.Quantity, fraction float64, boundary resourceBoundary) (bool, resource.Quantity) {
	if resourceQuantity.IsZero() {
		return false, resource.Quantity{}
	}

	// In case the resourceQuantity MilliValue overflows int64 we use MaxInt64
	// https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go
	scaledValue := resourceQuantity.Value()
	scaledMilliValue := int64(math.MaxInt64 - 1)
	if scaledValue < (math.MaxInt64 / 1000) {
		scaledMilliValue = resourceQuantity.MilliValue()
	}

	// float64(math.MaxInt64) > math.MaxInt64, to avoid overflow
	percentageValue := float64(scaledMilliValue) * fraction
	newValue := int64(math.MaxInt64)
	if percentageValue < math.MaxInt64 {
		newValue = int64(percentageValue)
	}

	newquantity := boundary.applyBoundary(*resource.NewMilliQuantity(newValue, resource.BinarySI))
	return true, newquantity
}

func fractionFromPercentage(m map[string]string, k string) (float64, bool) {
	value, err := strconv.ParseFloat(m[k], 64)
	return value / 100, err == nil
}

func makeQueueProbe(in *corev1.Probe) *corev1.Probe {
	if in == nil || in.PeriodSeconds == 0 {
		out := &corev1.Probe{
			Handler: corev1.Handler{
				Exec: &corev1.ExecAction{
					Command: []string{"/ko-app/queue", "-probe-period", "0"},
				},
			},
			// The exec probe enables us to retry failed probes quickly to get sub-second
			// resolution and achieve faster cold-starts.  However, for draining pods as
			// part of the K8s lifecycle this period will bound the tail of how quickly
			// we can remove a Pod's endpoint from the K8s service.
			//
			// The trade-off here is that exec probes cost CPU to run, and for idle pods
			// (e.g. due to minScale) we see ~50m/{period} of idle CPU usage in the
			// queue-proxy.  So while setting this to 1s results in slightly faster drains
			// it also means that in the steady state the queue-proxy is consuming 10x
			// more CPU due to probes than with a period of 10s.
			//
			// See also: https://github.com/knative/serving/issues/8147
			PeriodSeconds: 10,
			// We keep the connection open for a while because we're
			// actively probing the user-container on that endpoint and
			// thus don't want to be limited by K8s granularity here.
			TimeoutSeconds: 10,
		}

		if in != nil {
			out.InitialDelaySeconds = in.InitialDelaySeconds
		}
		return out
	}

	timeout := time.Second
	if in.TimeoutSeconds > 1 {
		timeout = time.Duration(in.TimeoutSeconds) * time.Second
	}

	return &corev1.Probe{
		Handler: corev1.Handler{
			Exec: &corev1.ExecAction{
				Command: []string{"/ko-app/queue", "-probe-period", timeout.String()},
			},
		},
		PeriodSeconds:       in.PeriodSeconds,
		TimeoutSeconds:      int32(timeout.Seconds()),
		SuccessThreshold:    in.SuccessThreshold,
		FailureThreshold:    in.FailureThreshold,
		InitialDelaySeconds: in.InitialDelaySeconds,
	}
}

// makeQueueContainer creates the container spec for the queue sidecar.
func makeQueueContainer(rev *v1.Revision, cfg *config.Config) (*corev1.Container, error) {
	configName := ""
	if owner := metav1.GetControllerOf(rev); owner != nil && owner.Kind == "Configuration" {
		configName = owner.Name
	}
	serviceName := rev.Labels[serving.ServiceLabelKey]

	userPort := getUserPort(rev)

	var loggingLevel string
	if ll, ok := cfg.Logging.LoggingLevel["queueproxy"]; ok {
		loggingLevel = ll.String()
	}

	ts := int64(0)
	if rev.Spec.TimeoutSeconds != nil {
		ts = *rev.Spec.TimeoutSeconds
	}

	ports := queueNonServingPorts
	if cfg.Observability.EnableProfiling {
		ports = append(ports, profilingPort)
	}
	// We need to configure only one serving port for the Queue proxy, since
	// we know the protocol that is being used by this application.
	servingPort := queueHTTPPort
	if rev.GetProtocol() == pkgnet.ProtocolH2C {
		servingPort = queueHTTP2Port
	}
	ports = append(ports, servingPort)

	container := rev.Spec.GetContainer()
	rp := container.ReadinessProbe.DeepCopy()

	applyReadinessProbeDefaults(rp, userPort)

	probeJSON, err := readiness.EncodeProbe(rp)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize readiness probe: %w", err)
	}

	return &corev1.Container{
		Name:            QueueContainerName,
		Image:           cfg.Deployment.QueueSidecarImage,
		Resources:       createQueueResources(cfg.Deployment, rev.GetAnnotations(), container),
		Ports:           ports,
		ReadinessProbe:  makeQueueProbe(rp),
		SecurityContext: queueSecurityContext,
		Env: []corev1.EnvVar{{
			Name:  "SERVING_NAMESPACE",
			Value: rev.Namespace,
		}, {
			Name:  "SERVING_SERVICE",
			Value: serviceName,
		}, {
			Name:  "SERVING_CONFIGURATION",
			Value: configName,
		}, {
			Name:  "SERVING_REVISION",
			Value: rev.Name,
		}, {
			Name:  "QUEUE_SERVING_PORT",
			Value: strconv.Itoa(int(servingPort.ContainerPort)),
		}, {
			Name:  "CONTAINER_CONCURRENCY",
			Value: strconv.Itoa(int(rev.Spec.GetContainerConcurrency())),
		}, {
			Name:  "REVISION_TIMEOUT_SECONDS",
			Value: strconv.Itoa(int(ts)),
		}, {
			Name: "SERVING_POD",
			ValueFrom: &corev1.EnvVarSource{
				FieldRef: &corev1.ObjectFieldSelector{
					FieldPath: "metadata.name",
				},
			},
		}, {
			Name: "SERVING_POD_IP",
			ValueFrom: &corev1.EnvVarSource{
				FieldRef: &corev1.ObjectFieldSelector{
					FieldPath: "status.podIP",
				},
			},
		}, {
			Name:  "SERVING_LOGGING_CONFIG",
			Value: cfg.Logging.LoggingConfig,
		}, {
			Name:  "SERVING_LOGGING_LEVEL",
			Value: loggingLevel,
		}, {
			Name:  "SERVING_REQUEST_LOG_TEMPLATE",
			Value: cfg.Observability.RequestLogTemplate,
		}, {
			Name:  "SERVING_ENABLE_REQUEST_LOG",
			Value: strconv.FormatBool(cfg.Observability.EnableRequestLog),
		}, {
			Name:  "SERVING_REQUEST_METRICS_BACKEND",
			Value: cfg.Observability.RequestMetricsBackend,
		}, {
			Name:  "TRACING_CONFIG_BACKEND",
			Value: string(cfg.Tracing.Backend),
		}, {
			Name:  "TRACING_CONFIG_ZIPKIN_ENDPOINT",
			Value: cfg.Tracing.ZipkinEndpoint,
		}, {
			Name:  "TRACING_CONFIG_STACKDRIVER_PROJECT_ID",
			Value: cfg.Tracing.StackdriverProjectID,
		}, {
			Name:  "TRACING_CONFIG_DEBUG",
			Value: strconv.FormatBool(cfg.Tracing.Debug),
		}, {
			Name:  "TRACING_CONFIG_SAMPLE_RATE",
			Value: fmt.Sprint(cfg.Tracing.SampleRate),
		}, {
			Name:  "USER_PORT",
			Value: strconv.Itoa(int(userPort)),
		}, {
			Name:  system.NamespaceEnvKey,
			Value: system.Namespace(),
		}, {
			Name:  metrics.DomainEnv,
			Value: metrics.Domain(),
		}, {
			Name:  "SERVING_READINESS_PROBE",
			Value: probeJSON,
		}, {
			Name:  "ENABLE_PROFILING",
			Value: strconv.FormatBool(cfg.Observability.EnableProfiling),
		}, {
			Name:  "SERVING_ENABLE_PROBE_REQUEST_LOG",
			Value: strconv.FormatBool(cfg.Observability.EnableProbeRequestLog),
		}, {
			Name:  "METRICS_COLLECTOR_ADDRESS",
			Value: cfg.Observability.MetricsCollectorAddress,
		}},
	}, nil
}

func applyReadinessProbeDefaults(p *corev1.Probe, port int32) {
	switch {
	case p == nil:
		return
	case p.HTTPGet != nil:
		p.HTTPGet.Host = localAddress
		p.HTTPGet.Port = intstr.FromInt(int(port))

		if p.HTTPGet.Scheme == "" {
			p.HTTPGet.Scheme = corev1.URISchemeHTTP
		}

		p.HTTPGet.HTTPHeaders = append(p.HTTPGet.HTTPHeaders, corev1.HTTPHeader{
			Name:  network.KubeletProbeHeaderName,
			Value: queue.Name,
		})
	case p.TCPSocket != nil:
		p.TCPSocket.Host = localAddress
		p.TCPSocket.Port = intstr.FromInt(int(port))
	case p.Exec != nil:
		// User-defined ExecProbe will still be run on user-container.
		// Use TCP probe in queue-proxy.
		p.TCPSocket = &corev1.TCPSocketAction{
			Host: localAddress,
			Port: intstr.FromInt(int(port)),
		}
		p.Exec = nil
	}

	if p.PeriodSeconds > 0 && p.TimeoutSeconds < 1 {
		p.TimeoutSeconds = 1
	}
}
<commit_msg>Narrow down the security profile of queue-proxy containers (#9974)<commit_after>/*
Copyright 2018 The Knative Authors

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

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

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

package resources

import (
	"fmt"
	"math"
	"strconv"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	network "knative.dev/networking/pkg"
	pkgnet "knative.dev/networking/pkg/apis/networking"
	"knative.dev/pkg/metrics"
	"knative.dev/pkg/profiling"
	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/deployment"
	"knative.dev/serving/pkg/networking"
	"knative.dev/serving/pkg/queue"
	"knative.dev/serving/pkg/queue/readiness"
	"knative.dev/serving/pkg/reconciler/revision/config"
)

const (
	localAddress             = "127.0.0.1"
	requestQueueHTTPPortName = "queue-port"
	profilingPortName        = "profiling-port"
)

var (
	queueHTTPPort = corev1.ContainerPort{
		Name:          requestQueueHTTPPortName,
		ContainerPort: networking.BackendHTTPPort,
	}
	queueHTTP2Port = corev1.ContainerPort{
		Name:          requestQueueHTTPPortName,
		ContainerPort: networking.BackendHTTP2Port,
	}
	queueNonServingPorts = []corev1.ContainerPort{{
		// Provides health checks and lifecycle hooks.
		Name:          v1.QueueAdminPortName,
		ContainerPort: networking.QueueAdminPort,
	}, {
		Name:          v1.AutoscalingQueueMetricsPortName,
		ContainerPort: networking.AutoscalingQueueMetricsPort,
	}, {
		Name:          v1.UserQueueMetricsPortName,
		ContainerPort: networking.UserQueueMetricsPort,
	}}

	profilingPort = corev1.ContainerPort{
		Name:          profilingPortName,
		ContainerPort: profiling.ProfilingPort,
	}

	queueSecurityContext = &corev1.SecurityContext{
		AllowPrivilegeEscalation: ptr.Bool(false),
		ReadOnlyRootFilesystem:   ptr.Bool(false), // Needed to be able to generate the unix socket.
		RunAsNonRoot:             ptr.Bool(true),
		Capabilities: &corev1.Capabilities{
			Drop: []corev1.Capability{"all"},
		},
	}
)

func createQueueResources(cfg *deployment.Config, annotations map[string]string, userContainer *corev1.Container) corev1.ResourceRequirements {
	resourceRequests := corev1.ResourceList{}
	resourceLimits := corev1.ResourceList{}

	for _, r := range []struct {
		Name    corev1.ResourceName
		Request *resource.Quantity
		Limit   *resource.Quantity
	}{{
		Name:    corev1.ResourceCPU,
		Request: cfg.QueueSidecarCPURequest,
		Limit:   cfg.QueueSidecarCPULimit,
	}, {
		Name:    corev1.ResourceMemory,
		Request: cfg.QueueSidecarMemoryRequest,
		Limit:   cfg.QueueSidecarMemoryLimit,
	}, {
		Name:    corev1.ResourceEphemeralStorage,
		Request: cfg.QueueSidecarEphemeralStorageRequest,
		Limit:   cfg.QueueSidecarEphemeralStorageLimit,
	}} {
		if r.Request != nil {
			resourceRequests[r.Name] = *r.Request
		}
		if r.Limit != nil {
			resourceLimits[r.Name] = *r.Limit
		}
	}

	var requestCPU, limitCPU, requestMemory, limitMemory resource.Quantity

	if resourceFraction, ok := fractionFromPercentage(annotations, serving.QueueSideCarResourcePercentageAnnotation); ok {
		if ok, requestCPU = computeResourceRequirements(userContainer.Resources.Requests.Cpu(), resourceFraction, queueContainerRequestCPU); ok {
			resourceRequests[corev1.ResourceCPU] = requestCPU
		}

		if ok, limitCPU = computeResourceRequirements(userContainer.Resources.Limits.Cpu(), resourceFraction, queueContainerLimitCPU); ok {
			resourceLimits[corev1.ResourceCPU] = limitCPU
		}

		if ok, requestMemory = computeResourceRequirements(userContainer.Resources.Requests.Memory(), resourceFraction, queueContainerRequestMemory); ok {
			resourceRequests[corev1.ResourceMemory] = requestMemory
		}

		if ok, limitMemory = computeResourceRequirements(userContainer.Resources.Limits.Memory(), resourceFraction, queueContainerLimitMemory); ok {
			resourceLimits[corev1.ResourceMemory] = limitMemory
		}
	}

	resources := corev1.ResourceRequirements{
		Requests: resourceRequests,
	}
	if len(resourceLimits) != 0 {
		resources.Limits = resourceLimits
	}

	return resources
}

func computeResourceRequirements(resourceQuantity *resource.Quantity, fraction float64, boundary resourceBoundary) (bool, resource.Quantity) {
	if resourceQuantity.IsZero() {
		return false, resource.Quantity{}
	}

	// In case the resourceQuantity MilliValue overflows int64 we use MaxInt64
	// https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go
	scaledValue := resourceQuantity.Value()
	scaledMilliValue := int64(math.MaxInt64 - 1)
	if scaledValue < (math.MaxInt64 / 1000) {
		scaledMilliValue = resourceQuantity.MilliValue()
	}

	// float64(math.MaxInt64) > math.MaxInt64, to avoid overflow
	percentageValue := float64(scaledMilliValue) * fraction
	newValue := int64(math.MaxInt64)
	if percentageValue < math.MaxInt64 {
		newValue = int64(percentageValue)
	}

	newquantity := boundary.applyBoundary(*resource.NewMilliQuantity(newValue, resource.BinarySI))
	return true, newquantity
}

func fractionFromPercentage(m map[string]string, k string) (float64, bool) {
	value, err := strconv.ParseFloat(m[k], 64)
	return value / 100, err == nil
}

func makeQueueProbe(in *corev1.Probe) *corev1.Probe {
	if in == nil || in.PeriodSeconds == 0 {
		out := &corev1.Probe{
			Handler: corev1.Handler{
				Exec: &corev1.ExecAction{
					Command: []string{"/ko-app/queue", "-probe-period", "0"},
				},
			},
			// The exec probe enables us to retry failed probes quickly to get sub-second
			// resolution and achieve faster cold-starts.  However, for draining pods as
			// part of the K8s lifecycle this period will bound the tail of how quickly
			// we can remove a Pod's endpoint from the K8s service.
			//
			// The trade-off here is that exec probes cost CPU to run, and for idle pods
			// (e.g. due to minScale) we see ~50m/{period} of idle CPU usage in the
			// queue-proxy.  So while setting this to 1s results in slightly faster drains
			// it also means that in the steady state the queue-proxy is consuming 10x
			// more CPU due to probes than with a period of 10s.
			//
			// See also: https://github.com/knative/serving/issues/8147
			PeriodSeconds: 10,
			// We keep the connection open for a while because we're
			// actively probing the user-container on that endpoint and
			// thus don't want to be limited by K8s granularity here.
			TimeoutSeconds: 10,
		}

		if in != nil {
			out.InitialDelaySeconds = in.InitialDelaySeconds
		}
		return out
	}

	timeout := time.Second
	if in.TimeoutSeconds > 1 {
		timeout = time.Duration(in.TimeoutSeconds) * time.Second
	}

	return &corev1.Probe{
		Handler: corev1.Handler{
			Exec: &corev1.ExecAction{
				Command: []string{"/ko-app/queue", "-probe-period", timeout.String()},
			},
		},
		PeriodSeconds:       in.PeriodSeconds,
		TimeoutSeconds:      int32(timeout.Seconds()),
		SuccessThreshold:    in.SuccessThreshold,
		FailureThreshold:    in.FailureThreshold,
		InitialDelaySeconds: in.InitialDelaySeconds,
	}
}

// makeQueueContainer creates the container spec for the queue sidecar.
func makeQueueContainer(rev *v1.Revision, cfg *config.Config) (*corev1.Container, error) {
	configName := ""
	if owner := metav1.GetControllerOf(rev); owner != nil && owner.Kind == "Configuration" {
		configName = owner.Name
	}
	serviceName := rev.Labels[serving.ServiceLabelKey]

	userPort := getUserPort(rev)

	var loggingLevel string
	if ll, ok := cfg.Logging.LoggingLevel["queueproxy"]; ok {
		loggingLevel = ll.String()
	}

	ts := int64(0)
	if rev.Spec.TimeoutSeconds != nil {
		ts = *rev.Spec.TimeoutSeconds
	}

	ports := queueNonServingPorts
	if cfg.Observability.EnableProfiling {
		ports = append(ports, profilingPort)
	}
	// We need to configure only one serving port for the Queue proxy, since
	// we know the protocol that is being used by this application.
	servingPort := queueHTTPPort
	if rev.GetProtocol() == pkgnet.ProtocolH2C {
		servingPort = queueHTTP2Port
	}
	ports = append(ports, servingPort)

	container := rev.Spec.GetContainer()
	rp := container.ReadinessProbe.DeepCopy()

	applyReadinessProbeDefaults(rp, userPort)

	probeJSON, err := readiness.EncodeProbe(rp)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize readiness probe: %w", err)
	}

	return &corev1.Container{
		Name:            QueueContainerName,
		Image:           cfg.Deployment.QueueSidecarImage,
		Resources:       createQueueResources(cfg.Deployment, rev.GetAnnotations(), container),
		Ports:           ports,
		ReadinessProbe:  makeQueueProbe(rp),
		SecurityContext: queueSecurityContext,
		Env: []corev1.EnvVar{{
			Name:  "SERVING_NAMESPACE",
			Value: rev.Namespace,
		}, {
			Name:  "SERVING_SERVICE",
			Value: serviceName,
		}, {
			Name:  "SERVING_CONFIGURATION",
			Value: configName,
		}, {
			Name:  "SERVING_REVISION",
			Value: rev.Name,
		}, {
			Name:  "QUEUE_SERVING_PORT",
			Value: strconv.Itoa(int(servingPort.ContainerPort)),
		}, {
			Name:  "CONTAINER_CONCURRENCY",
			Value: strconv.Itoa(int(rev.Spec.GetContainerConcurrency())),
		}, {
			Name:  "REVISION_TIMEOUT_SECONDS",
			Value: strconv.Itoa(int(ts)),
		}, {
			Name: "SERVING_POD",
			ValueFrom: &corev1.EnvVarSource{
				FieldRef: &corev1.ObjectFieldSelector{
					FieldPath: "metadata.name",
				},
			},
		}, {
			Name: "SERVING_POD_IP",
			ValueFrom: &corev1.EnvVarSource{
				FieldRef: &corev1.ObjectFieldSelector{
					FieldPath: "status.podIP",
				},
			},
		}, {
			Name:  "SERVING_LOGGING_CONFIG",
			Value: cfg.Logging.LoggingConfig,
		}, {
			Name:  "SERVING_LOGGING_LEVEL",
			Value: loggingLevel,
		}, {
			Name:  "SERVING_REQUEST_LOG_TEMPLATE",
			Value: cfg.Observability.RequestLogTemplate,
		}, {
			Name:  "SERVING_ENABLE_REQUEST_LOG",
			Value: strconv.FormatBool(cfg.Observability.EnableRequestLog),
		}, {
			Name:  "SERVING_REQUEST_METRICS_BACKEND",
			Value: cfg.Observability.RequestMetricsBackend,
		}, {
			Name:  "TRACING_CONFIG_BACKEND",
			Value: string(cfg.Tracing.Backend),
		}, {
			Name:  "TRACING_CONFIG_ZIPKIN_ENDPOINT",
			Value: cfg.Tracing.ZipkinEndpoint,
		}, {
			Name:  "TRACING_CONFIG_STACKDRIVER_PROJECT_ID",
			Value: cfg.Tracing.StackdriverProjectID,
		}, {
			Name:  "TRACING_CONFIG_DEBUG",
			Value: strconv.FormatBool(cfg.Tracing.Debug),
		}, {
			Name:  "TRACING_CONFIG_SAMPLE_RATE",
			Value: fmt.Sprint(cfg.Tracing.SampleRate),
		}, {
			Name:  "USER_PORT",
			Value: strconv.Itoa(int(userPort)),
		}, {
			Name:  system.NamespaceEnvKey,
			Value: system.Namespace(),
		}, {
			Name:  metrics.DomainEnv,
			Value: metrics.Domain(),
		}, {
			Name:  "SERVING_READINESS_PROBE",
			Value: probeJSON,
		}, {
			Name:  "ENABLE_PROFILING",
			Value: strconv.FormatBool(cfg.Observability.EnableProfiling),
		}, {
			Name:  "SERVING_ENABLE_PROBE_REQUEST_LOG",
			Value: strconv.FormatBool(cfg.Observability.EnableProbeRequestLog),
		}, {
			Name:  "METRICS_COLLECTOR_ADDRESS",
			Value: cfg.Observability.MetricsCollectorAddress,
		}},
	}, nil
}

func applyReadinessProbeDefaults(p *corev1.Probe, port int32) {
	switch {
	case p == nil:
		return
	case p.HTTPGet != nil:
		p.HTTPGet.Host = localAddress
		p.HTTPGet.Port = intstr.FromInt(int(port))

		if p.HTTPGet.Scheme == "" {
			p.HTTPGet.Scheme = corev1.URISchemeHTTP
		}

		p.HTTPGet.HTTPHeaders = append(p.HTTPGet.HTTPHeaders, corev1.HTTPHeader{
			Name:  network.KubeletProbeHeaderName,
			Value: queue.Name,
		})
	case p.TCPSocket != nil:
		p.TCPSocket.Host = localAddress
		p.TCPSocket.Port = intstr.FromInt(int(port))
	case p.Exec != nil:
		// User-defined ExecProbe will still be run on user-container.
		// Use TCP probe in queue-proxy.
		p.TCPSocket = &corev1.TCPSocketAction{
			Host: localAddress,
			Port: intstr.FromInt(int(port)),
		}
		p.Exec = nil
	}

	if p.PeriodSeconds > 0 && p.TimeoutSeconds < 1 {
		p.TimeoutSeconds = 1
	}
}
<|endoftext|>


<commit_before>package alerting

import (
	"context"
	"fmt"
	"math/rand"
	"net/http"

	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
)

// NotificationTestCommand initiates an test
// execution of an alert notification.
type NotificationTestCommand struct {
	OrgID          int64
	ID             int64
	State          models.AlertStateType
	Name           string
	Type           string
	Settings       *simplejson.Json
	SecureSettings map[string]string
}

var (
	logger = log.New("alerting.testnotification")
)

func (s *AlertNotificationService) HandleNotificationTestCommand(ctx context.Context, cmd *NotificationTestCommand) error {
	notificationSvc := newNotificationService(nil, nil)

	model := models.AlertNotification{
		Id:       cmd.ID,
		Name:     cmd.Name,
		Type:     cmd.Type,
		Settings: cmd.Settings,
	}

	notifier, err := s.createNotifier(ctx, &model, cmd.SecureSettings)
	if err != nil {
		return err
	}

	return notificationSvc.sendNotifications(createTestEvalContext(cmd), notifierStateSlice{{notifier: notifier}})
}

func createTestEvalContext(cmd *NotificationTestCommand) *EvalContext {
	testRule := &Rule{
		DashboardID: 1,
		PanelID:     1,
		Name:        "Test notification",
		Message:     "Someone is testing the alert notification within Grafana.",
		State:       models.AlertStateAlerting,
		ID:          rand.Int63(),
	}

	ctx := NewEvalContext(context.Background(), testRule, fakeRequestValidator{})
	if cmd.Settings.Get("uploadImage").MustBool(true) {
		ctx.ImagePublicURL = "https://grafana.com/assets/img/blog/mixed_styles.png"
	}
	ctx.IsTestRun = true
	ctx.Firing = true
	ctx.Error = fmt.Errorf("this is only a test")
	ctx.EvalMatches = evalMatchesBasedOnState()

	return ctx
}

func evalMatchesBasedOnState() []*EvalMatch {
	matches := make([]*EvalMatch, 0)
	matches = append(matches, &EvalMatch{
		Metric: "High value",
		Value:  null.FloatFrom(100),
	})

	matches = append(matches, &EvalMatch{
		Metric: "Higher Value",
		Value:  null.FloatFrom(200),
	})

	return matches
}

type fakeRequestValidator struct{}

func (fakeRequestValidator) Validate(_ string, _ *http.Request) error {
	return nil
}
<commit_msg>fix: pass OrgId to alertNotification model (#43678)<commit_after>package alerting

import (
	"context"
	"fmt"
	"math/rand"
	"net/http"

	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
)

// NotificationTestCommand initiates an test
// execution of an alert notification.
type NotificationTestCommand struct {
	OrgID          int64
	ID             int64
	State          models.AlertStateType
	Name           string
	Type           string
	Settings       *simplejson.Json
	SecureSettings map[string]string
}

var (
	logger = log.New("alerting.testnotification")
)

func (s *AlertNotificationService) HandleNotificationTestCommand(ctx context.Context, cmd *NotificationTestCommand) error {
	notificationSvc := newNotificationService(nil, nil)

	model := models.AlertNotification{
		Id:       cmd.ID,
		OrgId:    cmd.OrgID,
		Name:     cmd.Name,
		Type:     cmd.Type,
		Settings: cmd.Settings,
	}

	notifier, err := s.createNotifier(ctx, &model, cmd.SecureSettings)
	if err != nil {
		return err
	}

	return notificationSvc.sendNotifications(createTestEvalContext(cmd), notifierStateSlice{{notifier: notifier}})
}

func createTestEvalContext(cmd *NotificationTestCommand) *EvalContext {
	testRule := &Rule{
		DashboardID: 1,
		PanelID:     1,
		Name:        "Test notification",
		Message:     "Someone is testing the alert notification within Grafana.",
		State:       models.AlertStateAlerting,
		ID:          rand.Int63(),
	}

	ctx := NewEvalContext(context.Background(), testRule, fakeRequestValidator{})
	if cmd.Settings.Get("uploadImage").MustBool(true) {
		ctx.ImagePublicURL = "https://grafana.com/assets/img/blog/mixed_styles.png"
	}
	ctx.IsTestRun = true
	ctx.Firing = true
	ctx.Error = fmt.Errorf("this is only a test")
	ctx.EvalMatches = evalMatchesBasedOnState()

	return ctx
}

func evalMatchesBasedOnState() []*EvalMatch {
	matches := make([]*EvalMatch, 0)
	matches = append(matches, &EvalMatch{
		Metric: "High value",
		Value:  null.FloatFrom(100),
	})

	matches = append(matches, &EvalMatch{
		Metric: "Higher Value",
		Value:  null.FloatFrom(200),
	})

	return matches
}

type fakeRequestValidator struct{}

func (fakeRequestValidator) Validate(_ string, _ *http.Request) error {
	return nil
}
<|endoftext|>