<commit_before>package encoding

import (
	"crypto/md5"
	"encoding/binary"
	"hash/fnv"
	"io"
	"io/ioutil"
	"sync"
	"time"
	"v2ray.com/core/common/dice"

	"golang.org/x/crypto/chacha20poly1305"

	"v2ray.com/core/common"
	"v2ray.com/core/common/bitmask"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/crypto"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/task"
	"v2ray.com/core/proxy/vmess"
)

type sessionId struct {
	user  [16]byte
	key   [16]byte
	nonce [16]byte
}

// SessionHistory keeps track of historical session ids, to prevent replay attacks.
type SessionHistory struct {
	sync.RWMutex
	cache map[sessionId]time.Time
	task  *task.Periodic
}

// NewSessionHistory creates a new SessionHistory object.
func NewSessionHistory() *SessionHistory {
	h := &SessionHistory{
		cache: make(map[sessionId]time.Time, 128),
	}
	h.task = &task.Periodic{
		Interval: time.Second * 30,
		Execute:  h.removeExpiredEntries,
	}
	return h
}

// Close implements common.Closable.
func (h *SessionHistory) Close() error {
	return h.task.Close()
}

func (h *SessionHistory) addIfNotExits(session sessionId) bool {
	h.Lock()

	if expire, found := h.cache[session]; found && expire.After(time.Now()) {
		h.Unlock()
		return false
	}

	h.cache[session] = time.Now().Add(time.Minute * 3)
	h.Unlock()
	common.Must(h.task.Start())
	return true
}

func (h *SessionHistory) removeExpiredEntries() error {
	now := time.Now()

	h.Lock()
	defer h.Unlock()

	if len(h.cache) == 0 {
		return newError("nothing to do")
	}

	for session, expire := range h.cache {
		if expire.Before(now) {
			delete(h.cache, session)
		}
	}

	if len(h.cache) == 0 {
		h.cache = make(map[sessionId]time.Time, 128)
	}

	return nil
}

// ServerSession keeps information for a session in VMess server.
type ServerSession struct {
	userValidator   *vmess.TimedUserValidator
	sessionHistory  *SessionHistory
	requestBodyKey  [16]byte
	requestBodyIV   [16]byte
	responseBodyKey [16]byte
	responseBodyIV  [16]byte
	responseWriter  io.Writer
	responseHeader  byte
}

// NewServerSession creates a new ServerSession, using the given UserValidator.
// The ServerSession instance doesn't take ownership of the validator.
func NewServerSession(validator *vmess.TimedUserValidator, sessionHistory *SessionHistory) *ServerSession {
	return &ServerSession{
		userValidator:  validator,
		sessionHistory: sessionHistory,
	}
}

func parseSecurityType(b byte) protocol.SecurityType {
	if _, f := protocol.SecurityType_name[int32(b)]; f {
		st := protocol.SecurityType(b)
		// For backward compatibility.
		if st == protocol.SecurityType_UNKNOWN {
			st = protocol.SecurityType_LEGACY
		}
		return st
	}
	return protocol.SecurityType_UNKNOWN
}

// DecodeRequestHeader decodes and returns (if successful) a RequestHeader from an input stream.
func (s *ServerSession) DecodeRequestHeader(reader io.Reader) (*protocol.RequestHeader, error) {
	buffer := buf.New()
	behaviorRand := dice.NewDeterministicDice(int64(s.userValidator.GetBehaviorSeed()))
	DrainSize := behaviorRand.Roll(387) + 16 + 38
	readSizeRemain := DrainSize

	drainConnection := func(e error) error {
		//We read a deterministic generated length of data before closing the connection to offset padding read pattern
		readSizeRemain -= int(buffer.Len())
		if readSizeRemain > 0 {
			err := s.DrainConnN(reader, readSizeRemain)
			if err != nil {
				return newError("failed to drain connection").Base(err).Base(e)
			}
			return newError("connection drained DrainSize = ", DrainSize).Base(e)
		}
		return e
	}

	defer func() {
		buffer.Release()
	}()

	if _, err := buffer.ReadFullFrom(reader, protocol.IDBytesLen); err != nil {
		return nil, newError("failed to read request header").Base(err)
	}

	user, timestamp, valid := s.userValidator.Get(buffer.Bytes())
	if !valid {
		return nil, newError("invalid user")
	}

	iv := hashTimestamp(md5.New(), timestamp)
	vmessAccount := user.Account.(*vmess.MemoryAccount)

	aesStream := crypto.NewAesDecryptionStream(vmessAccount.ID.CmdKey(), iv[:])
	decryptor := crypto.NewCryptionReader(aesStream, reader)

	readSizeRemain -= int(buffer.Len())
	buffer.Clear()
	if _, err := buffer.ReadFullFrom(decryptor, 38); err != nil {
		return nil, newError("failed to read request header").Base(err)
	}

	request := &protocol.RequestHeader{
		User:    user,
		Version: buffer.Byte(0),
	}

	copy(s.requestBodyIV[:], buffer.BytesRange(1, 17))   // 16 bytes
	copy(s.requestBodyKey[:], buffer.BytesRange(17, 33)) // 16 bytes
	var sid sessionId
	copy(sid.user[:], vmessAccount.ID.Bytes())
	sid.key = s.requestBodyKey
	sid.nonce = s.requestBodyIV
	if !s.sessionHistory.addIfNotExits(sid) {
		return nil, newError("duplicated session id, possibly under replay attack")
	}

	s.responseHeader = buffer.Byte(33)             // 1 byte
	request.Option = bitmask.Byte(buffer.Byte(34)) // 1 byte
	padingLen := int(buffer.Byte(35) >> 4)
	request.Security = parseSecurityType(buffer.Byte(35) & 0x0F)
	// 1 bytes reserved
	request.Command = protocol.RequestCommand(buffer.Byte(37))

	switch request.Command {
	case protocol.RequestCommandMux:
		request.Address = net.DomainAddress("v1.mux.cool")
		request.Port = 0
	case protocol.RequestCommandTCP, protocol.RequestCommandUDP:
		if addr, port, err := addrParser.ReadAddressPort(buffer, decryptor); err == nil {
			request.Address = addr
			request.Port = port
		}
	}

	if padingLen > 0 {
		if _, err := buffer.ReadFullFrom(decryptor, int32(padingLen)); err != nil {
			return nil, newError("failed to read padding").Base(err)
		}
	}

	if _, err := buffer.ReadFullFrom(decryptor, 4); err != nil {
		return nil, newError("failed to read checksum").Base(err)
	}

	fnv1a := fnv.New32a()
	common.Must2(fnv1a.Write(buffer.BytesTo(-4)))
	actualHash := fnv1a.Sum32()
	expectedHash := binary.BigEndian.Uint32(buffer.BytesFrom(-4))

	if actualHash != expectedHash {
		//It is possible that we are under attack described in https://github.com/v2ray/v2ray-core/issues/2523
		return nil, drainConnection(newError("invalid auth"))
	}

	if request.Address == nil {
		return nil, newError("invalid remote address")
	}

	if request.Security == protocol.SecurityType_UNKNOWN || request.Security == protocol.SecurityType_AUTO {
		return nil, newError("unknown security type: ", request.Security)
	}

	return request, nil
}

// DecodeRequestBody returns Reader from which caller can fetch decrypted body.
func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reader io.Reader) buf.Reader {
	var sizeParser crypto.ChunkSizeDecoder = crypto.PlainChunkSizeParser{}
	if request.Option.Has(protocol.RequestOptionChunkMasking) {
		sizeParser = NewShakeSizeParser(s.requestBodyIV[:])
	}
	var padding crypto.PaddingLengthGenerator
	if request.Option.Has(protocol.RequestOptionGlobalPadding) {
		padding = sizeParser.(crypto.PaddingLengthGenerator)
	}

	switch request.Security {
	case protocol.SecurityType_NONE:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			if request.Command.TransferType() == protocol.TransferTypeStream {
				return crypto.NewChunkStreamReader(sizeParser, reader)
			}

			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(NoOpAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket, padding)
		}

		return buf.NewReader(reader)
	case protocol.SecurityType_LEGACY:
		aesStream := crypto.NewAesDecryptionStream(s.requestBodyKey[:], s.requestBodyIV[:])
		cryptionReader := crypto.NewCryptionReader(aesStream, reader)
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(FnvAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationReader(auth, sizeParser, cryptionReader, request.Command.TransferType(), padding)
		}

		return buf.NewReader(cryptionReader)
	case protocol.SecurityType_AES128_GCM:
		aead := crypto.NewAesGcm(s.requestBodyKey[:])

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType(), padding)
	case protocol.SecurityType_CHACHA20_POLY1305:
		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.requestBodyKey[:]))

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType(), padding)
	default:
		panic("Unknown security type.")
	}
}

// EncodeResponseHeader writes encoded response header into the given writer.
func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, writer io.Writer) {
	s.responseBodyKey = md5.Sum(s.requestBodyKey[:])
	s.responseBodyIV = md5.Sum(s.requestBodyIV[:])

	aesStream := crypto.NewAesEncryptionStream(s.responseBodyKey[:], s.responseBodyIV[:])
	encryptionWriter := crypto.NewCryptionWriter(aesStream, writer)
	s.responseWriter = encryptionWriter

	common.Must2(encryptionWriter.Write([]byte{s.responseHeader, byte(header.Option)}))
	err := MarshalCommand(header.Command, encryptionWriter)
	if err != nil {
		common.Must2(encryptionWriter.Write([]byte{0x00, 0x00}))
	}
}

// EncodeResponseBody returns a Writer that auto-encrypt content written by caller.
func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writer io.Writer) buf.Writer {
	var sizeParser crypto.ChunkSizeEncoder = crypto.PlainChunkSizeParser{}
	if request.Option.Has(protocol.RequestOptionChunkMasking) {
		sizeParser = NewShakeSizeParser(s.responseBodyIV[:])
	}
	var padding crypto.PaddingLengthGenerator
	if request.Option.Has(protocol.RequestOptionGlobalPadding) {
		padding = sizeParser.(crypto.PaddingLengthGenerator)
	}

	switch request.Security {
	case protocol.SecurityType_NONE:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			if request.Command.TransferType() == protocol.TransferTypeStream {
				return crypto.NewChunkStreamWriter(sizeParser, writer)
			}

			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(NoOpAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket, padding)
		}

		return buf.NewWriter(writer)
	case protocol.SecurityType_LEGACY:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(FnvAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationWriter(auth, sizeParser, s.responseWriter, request.Command.TransferType(), padding)
		}

		return &buf.SequentialWriter{Writer: s.responseWriter}
	case protocol.SecurityType_AES128_GCM:
		aead := crypto.NewAesGcm(s.responseBodyKey[:])

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType(), padding)
	case protocol.SecurityType_CHACHA20_POLY1305:
		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.responseBodyKey[:]))

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType(), padding)
	default:
		panic("Unknown security type.")
	}
}

func (s *ServerSession) DrainConnN(reader io.Reader, n int) error {
	_, err := io.CopyN(ioutil.Discard, reader, int64(n))
	return err
}
<commit_msg>Increase Drain Pool<commit_after>package encoding

import (
	"crypto/md5"
	"encoding/binary"
	"hash/fnv"
	"io"
	"io/ioutil"
	"sync"
	"time"
	"v2ray.com/core/common/dice"

	"golang.org/x/crypto/chacha20poly1305"

	"v2ray.com/core/common"
	"v2ray.com/core/common/bitmask"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/crypto"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/task"
	"v2ray.com/core/proxy/vmess"
)

type sessionId struct {
	user  [16]byte
	key   [16]byte
	nonce [16]byte
}

// SessionHistory keeps track of historical session ids, to prevent replay attacks.
type SessionHistory struct {
	sync.RWMutex
	cache map[sessionId]time.Time
	task  *task.Periodic
}

// NewSessionHistory creates a new SessionHistory object.
func NewSessionHistory() *SessionHistory {
	h := &SessionHistory{
		cache: make(map[sessionId]time.Time, 128),
	}
	h.task = &task.Periodic{
		Interval: time.Second * 30,
		Execute:  h.removeExpiredEntries,
	}
	return h
}

// Close implements common.Closable.
func (h *SessionHistory) Close() error {
	return h.task.Close()
}

func (h *SessionHistory) addIfNotExits(session sessionId) bool {
	h.Lock()

	if expire, found := h.cache[session]; found && expire.After(time.Now()) {
		h.Unlock()
		return false
	}

	h.cache[session] = time.Now().Add(time.Minute * 3)
	h.Unlock()
	common.Must(h.task.Start())
	return true
}

func (h *SessionHistory) removeExpiredEntries() error {
	now := time.Now()

	h.Lock()
	defer h.Unlock()

	if len(h.cache) == 0 {
		return newError("nothing to do")
	}

	for session, expire := range h.cache {
		if expire.Before(now) {
			delete(h.cache, session)
		}
	}

	if len(h.cache) == 0 {
		h.cache = make(map[sessionId]time.Time, 128)
	}

	return nil
}

// ServerSession keeps information for a session in VMess server.
type ServerSession struct {
	userValidator   *vmess.TimedUserValidator
	sessionHistory  *SessionHistory
	requestBodyKey  [16]byte
	requestBodyIV   [16]byte
	responseBodyKey [16]byte
	responseBodyIV  [16]byte
	responseWriter  io.Writer
	responseHeader  byte
}

// NewServerSession creates a new ServerSession, using the given UserValidator.
// The ServerSession instance doesn't take ownership of the validator.
func NewServerSession(validator *vmess.TimedUserValidator, sessionHistory *SessionHistory) *ServerSession {
	return &ServerSession{
		userValidator:  validator,
		sessionHistory: sessionHistory,
	}
}

func parseSecurityType(b byte) protocol.SecurityType {
	if _, f := protocol.SecurityType_name[int32(b)]; f {
		st := protocol.SecurityType(b)
		// For backward compatibility.
		if st == protocol.SecurityType_UNKNOWN {
			st = protocol.SecurityType_LEGACY
		}
		return st
	}
	return protocol.SecurityType_UNKNOWN
}

// DecodeRequestHeader decodes and returns (if successful) a RequestHeader from an input stream.
func (s *ServerSession) DecodeRequestHeader(reader io.Reader) (*protocol.RequestHeader, error) {
	buffer := buf.New()
	behaviorRand := dice.NewDeterministicDice(int64(s.userValidator.GetBehaviorSeed()))
	DrainSize := behaviorRand.Roll(3266) + 16 + 38 + dice.Roll(behaviorRand.Roll(64))
	readSizeRemain := DrainSize

	drainConnection := func(e error) error {
		//We read a deterministic generated length of data before closing the connection to offset padding read pattern
		readSizeRemain -= int(buffer.Len())
		if readSizeRemain > 0 {
			err := s.DrainConnN(reader, readSizeRemain)
			if err != nil {
				return newError("failed to drain connection").Base(err).Base(e)
			}
			return newError("connection drained DrainSize = ", DrainSize).Base(e)
		}
		return e
	}

	defer func() {
		buffer.Release()
	}()

	if _, err := buffer.ReadFullFrom(reader, protocol.IDBytesLen); err != nil {
		return nil, newError("failed to read request header").Base(err)
	}

	user, timestamp, valid := s.userValidator.Get(buffer.Bytes())
	if !valid {
		return nil, newError("invalid user")
	}

	iv := hashTimestamp(md5.New(), timestamp)
	vmessAccount := user.Account.(*vmess.MemoryAccount)

	aesStream := crypto.NewAesDecryptionStream(vmessAccount.ID.CmdKey(), iv[:])
	decryptor := crypto.NewCryptionReader(aesStream, reader)

	readSizeRemain -= int(buffer.Len())
	buffer.Clear()
	if _, err := buffer.ReadFullFrom(decryptor, 38); err != nil {
		return nil, newError("failed to read request header").Base(err)
	}

	request := &protocol.RequestHeader{
		User:    user,
		Version: buffer.Byte(0),
	}

	copy(s.requestBodyIV[:], buffer.BytesRange(1, 17))   // 16 bytes
	copy(s.requestBodyKey[:], buffer.BytesRange(17, 33)) // 16 bytes
	var sid sessionId
	copy(sid.user[:], vmessAccount.ID.Bytes())
	sid.key = s.requestBodyKey
	sid.nonce = s.requestBodyIV
	if !s.sessionHistory.addIfNotExits(sid) {
		return nil, newError("duplicated session id, possibly under replay attack")
	}

	s.responseHeader = buffer.Byte(33)             // 1 byte
	request.Option = bitmask.Byte(buffer.Byte(34)) // 1 byte
	padingLen := int(buffer.Byte(35) >> 4)
	request.Security = parseSecurityType(buffer.Byte(35) & 0x0F)
	// 1 bytes reserved
	request.Command = protocol.RequestCommand(buffer.Byte(37))

	switch request.Command {
	case protocol.RequestCommandMux:
		request.Address = net.DomainAddress("v1.mux.cool")
		request.Port = 0
	case protocol.RequestCommandTCP, protocol.RequestCommandUDP:
		if addr, port, err := addrParser.ReadAddressPort(buffer, decryptor); err == nil {
			request.Address = addr
			request.Port = port
		}
	}

	if padingLen > 0 {
		if _, err := buffer.ReadFullFrom(decryptor, int32(padingLen)); err != nil {
			return nil, newError("failed to read padding").Base(err)
		}
	}

	if _, err := buffer.ReadFullFrom(decryptor, 4); err != nil {
		return nil, newError("failed to read checksum").Base(err)
	}

	fnv1a := fnv.New32a()
	common.Must2(fnv1a.Write(buffer.BytesTo(-4)))
	actualHash := fnv1a.Sum32()
	expectedHash := binary.BigEndian.Uint32(buffer.BytesFrom(-4))

	if actualHash != expectedHash {
		//It is possible that we are under attack described in https://github.com/v2ray/v2ray-core/issues/2523
		return nil, drainConnection(newError("invalid auth"))
	}

	if request.Address == nil {
		return nil, newError("invalid remote address")
	}

	if request.Security == protocol.SecurityType_UNKNOWN || request.Security == protocol.SecurityType_AUTO {
		return nil, newError("unknown security type: ", request.Security)
	}

	return request, nil
}

// DecodeRequestBody returns Reader from which caller can fetch decrypted body.
func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reader io.Reader) buf.Reader {
	var sizeParser crypto.ChunkSizeDecoder = crypto.PlainChunkSizeParser{}
	if request.Option.Has(protocol.RequestOptionChunkMasking) {
		sizeParser = NewShakeSizeParser(s.requestBodyIV[:])
	}
	var padding crypto.PaddingLengthGenerator
	if request.Option.Has(protocol.RequestOptionGlobalPadding) {
		padding = sizeParser.(crypto.PaddingLengthGenerator)
	}

	switch request.Security {
	case protocol.SecurityType_NONE:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			if request.Command.TransferType() == protocol.TransferTypeStream {
				return crypto.NewChunkStreamReader(sizeParser, reader)
			}

			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(NoOpAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket, padding)
		}

		return buf.NewReader(reader)
	case protocol.SecurityType_LEGACY:
		aesStream := crypto.NewAesDecryptionStream(s.requestBodyKey[:], s.requestBodyIV[:])
		cryptionReader := crypto.NewCryptionReader(aesStream, reader)
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(FnvAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationReader(auth, sizeParser, cryptionReader, request.Command.TransferType(), padding)
		}

		return buf.NewReader(cryptionReader)
	case protocol.SecurityType_AES128_GCM:
		aead := crypto.NewAesGcm(s.requestBodyKey[:])

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType(), padding)
	case protocol.SecurityType_CHACHA20_POLY1305:
		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.requestBodyKey[:]))

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType(), padding)
	default:
		panic("Unknown security type.")
	}
}

// EncodeResponseHeader writes encoded response header into the given writer.
func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, writer io.Writer) {
	s.responseBodyKey = md5.Sum(s.requestBodyKey[:])
	s.responseBodyIV = md5.Sum(s.requestBodyIV[:])

	aesStream := crypto.NewAesEncryptionStream(s.responseBodyKey[:], s.responseBodyIV[:])
	encryptionWriter := crypto.NewCryptionWriter(aesStream, writer)
	s.responseWriter = encryptionWriter

	common.Must2(encryptionWriter.Write([]byte{s.responseHeader, byte(header.Option)}))
	err := MarshalCommand(header.Command, encryptionWriter)
	if err != nil {
		common.Must2(encryptionWriter.Write([]byte{0x00, 0x00}))
	}
}

// EncodeResponseBody returns a Writer that auto-encrypt content written by caller.
func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writer io.Writer) buf.Writer {
	var sizeParser crypto.ChunkSizeEncoder = crypto.PlainChunkSizeParser{}
	if request.Option.Has(protocol.RequestOptionChunkMasking) {
		sizeParser = NewShakeSizeParser(s.responseBodyIV[:])
	}
	var padding crypto.PaddingLengthGenerator
	if request.Option.Has(protocol.RequestOptionGlobalPadding) {
		padding = sizeParser.(crypto.PaddingLengthGenerator)
	}

	switch request.Security {
	case protocol.SecurityType_NONE:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			if request.Command.TransferType() == protocol.TransferTypeStream {
				return crypto.NewChunkStreamWriter(sizeParser, writer)
			}

			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(NoOpAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket, padding)
		}

		return buf.NewWriter(writer)
	case protocol.SecurityType_LEGACY:
		if request.Option.Has(protocol.RequestOptionChunkStream) {
			auth := &crypto.AEADAuthenticator{
				AEAD:                    new(FnvAuthenticator),
				NonceGenerator:          crypto.GenerateEmptyBytes(),
				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
			}
			return crypto.NewAuthenticationWriter(auth, sizeParser, s.responseWriter, request.Command.TransferType(), padding)
		}

		return &buf.SequentialWriter{Writer: s.responseWriter}
	case protocol.SecurityType_AES128_GCM:
		aead := crypto.NewAesGcm(s.responseBodyKey[:])

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType(), padding)
	case protocol.SecurityType_CHACHA20_POLY1305:
		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.responseBodyKey[:]))

		auth := &crypto.AEADAuthenticator{
			AEAD:                    aead,
			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV[:], uint32(aead.NonceSize())),
			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
		}
		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType(), padding)
	default:
		panic("Unknown security type.")
	}
}

func (s *ServerSession) DrainConnN(reader io.Reader, n int) error {
	_, err := io.CopyN(ioutil.Discard, reader, int64(n))
	return err
}
<|endoftext|>


<commit_before>package GoMetricMiddleware

import (
	"net/http"
	"time"

	"strconv"

	"github.com/prometheus/client_golang/prometheus"
)

type loggingResponseWriter struct {
	http.ResponseWriter
	statusCode int
}

func newLoggingResponseWriter(w http.ResponseWriter) *loggingResponseWriter {
	return &loggingResponseWriter{w, http.StatusOK}
}

func (lrw *loggingResponseWriter) WriteHeader(code int) {
	lrw.statusCode = code
	lrw.ResponseWriter.WriteHeader(code)
}

func MetricsMiddleware(namespace string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestsDuration := prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Namespace: namespace,
				Name:      "request_duration_seconds",
				Help:      "The duration of the requests to the Statistics service.",
			},
		)

		requestsCurrent := prometheus.NewGauge(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "requests_current",
				Help:      "The current number of requests to " + namespace,
			},
		)

		requestsStatus := prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "requests_total",
				Help:      "The total number of requests to the " + namespace + " by status, method and path.",
			},
			[]string{"code", "method", "path"},
		)

		clientErrors := prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "errors",
				Help:      "The total number of " + namespace + " client errors",
			})

		start := time.Now()
		requestsCurrent.Inc()

		lrw := newLoggingResponseWriter(w)
		next.ServeHTTP(lrw, r)

		statusCode := lrw.statusCode

		requestsStatus.WithLabelValues(strconv.Itoa(statusCode), r.Method, r.URL.Path).Inc()

		if statusCode != http.StatusOK && statusCode != http.StatusNoContent {
			clientErrors.Inc()
		}

		requestsCurrent.Dec()
		requestsDuration.Observe(float64(time.Since(start).Seconds()))
	})
}
<commit_msg>Initing the prometheus.<commit_after>package GoMetricMiddleware

import (
	"net/http"
	"time"

	"strconv"

	"github.com/prometheus/client_golang/prometheus"
)

type loggingResponseWriter struct {
	http.ResponseWriter
	statusCode int
}

var requestsDuration prometheus.Histogram
var requestsCurrent prometheus.Gauge
var requestsStatus prometheus.CounterVec
var clientErrors prometheus.Counter

func newLoggingResponseWriter(w http.ResponseWriter) *loggingResponseWriter {
	return &loggingResponseWriter{w, http.StatusOK}
}

func (lrw *loggingResponseWriter) WriteHeader(code int) {
	lrw.statusCode = code
	lrw.ResponseWriter.WriteHeader(code)
}

func MetricsMiddleware(namespace string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestsDuration := prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Namespace: namespace,
				Name:      "request_duration_seconds",
				Help:      "The duration of the requests to the Statistics service.",
			},
		)

		requestsCurrent := prometheus.NewGauge(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "requests_current",
				Help:      "The current number of requests to " + namespace,
			},
		)

		requestsStatus := prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "requests_total",
				Help:      "The total number of requests to the " + namespace + " by status, method and path.",
			},
			[]string{"code", "method", "path"},
		)

		clientErrors := prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "errors",
				Help:      "The total number of " + namespace + " client errors",
			})

		start := time.Now()
		requestsCurrent.Inc()

		lrw := newLoggingResponseWriter(w)
		next.ServeHTTP(lrw, r)

		statusCode := lrw.statusCode

		requestsStatus.WithLabelValues(strconv.Itoa(statusCode), r.Method, r.URL.Path).Inc()

		if statusCode != http.StatusOK && statusCode != http.StatusNoContent {
			clientErrors.Inc()
		}

		requestsCurrent.Dec()
		requestsDuration.Observe(float64(time.Since(start).Seconds()))
	})
}

func init() {
	prometheus.MustRegister(requestsDuration)
	prometheus.MustRegister(requestsCurrent)
	prometheus.MustRegister(requestsStatus)
	prometheus.MustRegister(clientErrors)
}
<|endoftext|>


<commit_before>package negronistatsd

import (
	"net/http"
	"strconv"
	"time"

	"github.com/cactus/go-statsd-client/statsd"
	"github.com/codegangsta/negroni"
)

// Middleware is a middleware handler that logs the request as it goes in and the response as it goes out.
type Middleware struct {
	StatsdClient *statsd.Client
	// Name is the name of the application as recorded in latency metrics
}

// NewMiddleware returns a new *Middleware, yay!
func NewMiddleware(statsd *statsd.Client) *Middleware {
	return &Middleware{StatsdClient: statsd}
}

func (s *Middleware) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	start := time.Now()

	next(rw, r)

	// latency in milliseconds
	latency := int64(time.Since(start) / time.Millisecond)
	res := rw.(negroni.ResponseWriter)

	s.StatsdClient.Inc("requests."+strconv.Itoa(res.Status()), 1, 1.0)
	s.StatsdClient.Timing("request_time", latency, 1.0)

}
<commit_msg>use current statsd library types<commit_after>package negronistatsd

import (
	"net/http"
	"strconv"
	"time"

	"github.com/cactus/go-statsd-client/statsd"
	"github.com/codegangsta/negroni"
)

// Middleware is a middleware handler that logs the request as it goes in and the response as it goes out.
type Middleware struct {
	StatsdClient statsd.Statter
	// Name is the name of the application as recorded in latency metrics
}

// NewMiddleware returns a new *Middleware, yay!
func NewMiddleware(statsd statsd.Statter) *Middleware {
	return &Middleware{StatsdClient: statsd}
}

func (s *Middleware) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	start := time.Now()

	next(rw, r)

	// latency in milliseconds
	latency := int64(time.Since(start) / time.Millisecond)
	res := rw.(negroni.ResponseWriter)

	s.StatsdClient.Inc("requests."+strconv.Itoa(res.Status()), 1, 1.0)
	s.StatsdClient.Timing("request_time", latency, 1.0)

}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/led/job"
)

// EditRecipeBundleOpts are user-provided options for the recipe bundling
// process.
type EditRecipeBundleOpts struct {
	// Path on disk to the repo to extract the recipes from. May be a subdirectory
	// of the repo, as long as `git rev-parse --show-toplevel` can find the root
	// of the repository.
	//
	// If empty, uses the current working directory.
	RepoDir string

	// Overrides is a mapping of recipe project id (e.g. "recipe_engine") to
	// a local path to a checkout of that repo (e.g. "/path/to/recipes-py.git").
	//
	// When the bundle is created, this local repo will be used instead of the
	// pinned version of this recipe project id. This is helpful for preparing
	// bundles which have code changes in multiple recipe repos.
	Overrides map[string]string

	// DebugSleep is the amount of time to wait after the recipe completes
	// execution (either success or failure). This is injected into the generated
	// recipe bundle as a 'sleep X' command after the invocation of the recipe
	// itself.
	DebugSleep time.Duration
}

// RecipeDirectory is a very unfortunate constant which is here for
// a combination of reasons:
//   1) swarming doesn't allow you to 'checkout' an isolate relative to any path
//      in the task (other than the task root). This means that whatever value
//      we pick for EditRecipeBundle must be used EVERYWHERE the isolated hash
//      is used.
//   2) Currently the 'recipe_engine/led' module will blindly take the isolated
//      input and 'inject' it into further uses of led. This module currently
//      doesn't specify the checkout dir, relying on kitchen's default value of
//      (you guessed it) "kitchen-checkout".
//
// In order to fix this (and it will need to be fixed for bbagent support):
//   * The 'recipe_engine/led' module needs to accept 'checkout-dir' as
//     a parameter in its input properties.
//   * led needs to start passing the checkout dir to the led module's input
//     properties.
//   * `led edit` needs a way to manipulate the checkout directory in a job
//   * The 'recipe_engine/led' module needs to set this in the job
//     alongside the isolate hash when it's doing the injection.
//
// For now, we just hard-code it.
//
// TODO(crbug.com/1072117): Fix this, it's weird.
const RecipeDirectory = "kitchen-checkout"

// EditRecipeBundle overrides the recipe bundle in the given job with one
// located on disk.
//
// It isolates the recipes from the repository in the given working directory
// into the UserPayload under the directory "kitchen-checkout/". If there's an
// existing directory in the UserPayload at that location, it will be removed.
func EditRecipeBundle(ctx context.Context, authClient *http.Client, authOpts auth.Options, jd *job.Definition, opts *EditRecipeBundleOpts) error {
	if jd.GetSwarming() != nil {
		return errors.New("ledcmd.EditRecipeBundle is only available for Buildbucket tasks")
	}

	if opts == nil {
		opts = &EditRecipeBundleOpts{}
	}

	recipesPy, err := findRecipesPy(ctx, opts.RepoDir)
	if err != nil {
		return err
	}
	logging.Debugf(ctx, "using recipes.py: %q", recipesPy)

	err = EditIsolated(ctx, authClient, authOpts, jd, func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "bundling recipes")
		bundlePath := filepath.Join(dir, RecipeDirectory)
		// Remove existing bundled recipes, if any. Ignore the error.
		os.RemoveAll(bundlePath)
		if err := opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath); err != nil {
			return err
		}
		logging.Infof(ctx, "isolating recipes")
		return nil
	})
	if err != nil {
		return err
	}

	return jd.HighLevelEdit(func(je job.HighLevelEditor) {
		je.TaskPayloadSource("", "")
		je.TaskPayloadPath(RecipeDirectory)
	})
}

func logCmd(ctx context.Context, inDir string, arg0 string, args ...string) *exec.Cmd {
	ret := exec.CommandContext(ctx, arg0, args...)
	ret.Dir = inDir
	logging.Debugf(ctx, "Running (from %q) - %s %v", inDir, arg0, args)
	return ret
}

func cmdErr(cmd *exec.Cmd, err error, reason string) error {
	if err != nil {
		ee, _ := err.(*exec.ExitError)
		outErr := ""
		if ee != nil {
			outErr = strings.TrimSpace(string(ee.Stderr))
			if len(outErr) > 128 {
				outErr = outErr[:128] + "..."
			}
		}
		err = errors.Annotate(err, "running %q: %s: %s", strings.Join(cmd.Args, " "), reason, outErr).Err()
	}
	return err
}

func appendText(path, fmtStr string, items ...interface{}) error {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = fmt.Fprintf(file, fmtStr, items...)
	return err
}

func (opts *EditRecipeBundleOpts) prepBundle(ctx context.Context, inDir, recipesPy, toDirectory string) (err error) {
	args := []string{
		recipesPy,
	}
	if logging.GetLevel(ctx) < logging.Info {
		args = append(args, "-v")
	}
	for projID, path := range opts.Overrides {
		args = append(args, "-O", fmt.Sprintf("%s=%s", projID, path))
	}
	args = append(args, "bundle", "--destination", filepath.Join(toDirectory))
	cmd := logCmd(ctx, inDir, "python", args...)
	if logging.GetLevel(ctx) < logging.Info {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	if err = cmdErr(cmd, cmd.Run(), "creating bundle"); err != nil {
		return
	}
	if opts.DebugSleep != 0 {
		for _, basename := range []string{"recipes", "luciexe"} {
			fname := filepath.Join(toDirectory, basename)
			seconds := opts.DebugSleep / time.Second
			msg := "echo ENTERING DEBUG SLEEP. SSH to the bot to debug."

			if err = appendText(fname, "\n%s\nsleep %d\n", msg, seconds); err != nil {
				return
			}
			// Wait for a bogus event that won't occur... Windows sucks, amirite?
			if err = appendText(fname+".bat", "\r\n%s\r\nwaitfor /t %d DebugSessionEnd\r\n", msg, seconds); err != nil {
				return
			}
		}
	}

	return
}

// findRecipesPy locates the current repo's `recipes.py`. It does this by:
//   * invoking git to find the repo root
//   * loading the recipes.cfg at infra/config/recipes.cfg
//   * stat'ing the recipes.py implied by the recipes_path in that cfg file.
//
// Failure will return an error.
//
// On success, the absolute path to recipes.py is returned.
func findRecipesPy(ctx context.Context, inDir string) (string, error) {
	cmd := logCmd(ctx, inDir, "git", "rev-parse", "--show-toplevel")
	out, err := cmd.Output()
	if err = cmdErr(cmd, err, "finding git repo"); err != nil {
		return "", err
	}

	repoRoot := strings.TrimSpace(string(out))

	pth := filepath.Join(repoRoot, "infra", "config", "recipes.cfg")
	switch st, err := os.Stat(pth); {
	case err != nil:
		return "", errors.Annotate(err, "reading recipes.cfg").Err()

	case !st.Mode().IsRegular():
		return "", errors.Reason("%q is not a regular file", pth).Err()
	}

	type recipesJSON struct {
		RecipesPath string `json:"recipes_path"`
	}
	rj := &recipesJSON{}

	f, err := os.Open(pth)
	if err != nil {
		return "", errors.Reason("reading recipes.cfg: %q", pth).Err()
	}
	defer f.Close()

	if err := json.NewDecoder(f).Decode(rj); err != nil {
		return "", errors.Reason("parsing recipes.cfg: %q", pth).Err()
	}

	return filepath.Join(
		repoRoot, filepath.FromSlash(rj.RecipesPath), "recipes.py"), nil
}
<commit_msg>[led] Use new envvar to set debug sleep.<commit_after>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/led/job"
)

// EditRecipeBundleOpts are user-provided options for the recipe bundling
// process.
type EditRecipeBundleOpts struct {
	// Path on disk to the repo to extract the recipes from. May be a subdirectory
	// of the repo, as long as `git rev-parse --show-toplevel` can find the root
	// of the repository.
	//
	// If empty, uses the current working directory.
	RepoDir string

	// Overrides is a mapping of recipe project id (e.g. "recipe_engine") to
	// a local path to a checkout of that repo (e.g. "/path/to/recipes-py.git").
	//
	// When the bundle is created, this local repo will be used instead of the
	// pinned version of this recipe project id. This is helpful for preparing
	// bundles which have code changes in multiple recipe repos.
	Overrides map[string]string

	// DebugSleep is the amount of time to wait after the recipe completes
	// execution (either success or failure). This is injected into the generated
	// recipe bundle as a 'sleep X' command after the invocation of the recipe
	// itself.
	DebugSleep time.Duration
}

// RecipeDirectory is a very unfortunate constant which is here for
// a combination of reasons:
//   1) swarming doesn't allow you to 'checkout' an isolate relative to any path
//      in the task (other than the task root). This means that whatever value
//      we pick for EditRecipeBundle must be used EVERYWHERE the isolated hash
//      is used.
//   2) Currently the 'recipe_engine/led' module will blindly take the isolated
//      input and 'inject' it into further uses of led. This module currently
//      doesn't specify the checkout dir, relying on kitchen's default value of
//      (you guessed it) "kitchen-checkout".
//
// In order to fix this (and it will need to be fixed for bbagent support):
//   * The 'recipe_engine/led' module needs to accept 'checkout-dir' as
//     a parameter in its input properties.
//   * led needs to start passing the checkout dir to the led module's input
//     properties.
//   * `led edit` needs a way to manipulate the checkout directory in a job
//   * The 'recipe_engine/led' module needs to set this in the job
//     alongside the isolate hash when it's doing the injection.
//
// For now, we just hard-code it.
//
// TODO(crbug.com/1072117): Fix this, it's weird.
const RecipeDirectory = "kitchen-checkout"

// EditRecipeBundle overrides the recipe bundle in the given job with one
// located on disk.
//
// It isolates the recipes from the repository in the given working directory
// into the UserPayload under the directory "kitchen-checkout/". If there's an
// existing directory in the UserPayload at that location, it will be removed.
func EditRecipeBundle(ctx context.Context, authClient *http.Client, authOpts auth.Options, jd *job.Definition, opts *EditRecipeBundleOpts) error {
	if jd.GetSwarming() != nil {
		return errors.New("ledcmd.EditRecipeBundle is only available for Buildbucket tasks")
	}

	if opts == nil {
		opts = &EditRecipeBundleOpts{}
	}

	recipesPy, err := findRecipesPy(ctx, opts.RepoDir)
	if err != nil {
		return err
	}
	logging.Debugf(ctx, "using recipes.py: %q", recipesPy)

	err = EditIsolated(ctx, authClient, authOpts, jd, func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "bundling recipes")
		bundlePath := filepath.Join(dir, RecipeDirectory)
		// Remove existing bundled recipes, if any. Ignore the error.
		os.RemoveAll(bundlePath)
		if err := opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath); err != nil {
			return err
		}
		logging.Infof(ctx, "isolating recipes")
		return nil
	})
	if err != nil {
		return err
	}

	return jd.HighLevelEdit(func(je job.HighLevelEditor) {
		je.TaskPayloadSource("", "")
		je.TaskPayloadPath(RecipeDirectory)
		if opts.DebugSleep != 0 {
			je.Env(map[string]string{
				"RECIPES_DEBUG_SLEEP": fmt.Sprintf("%f", opts.DebugSleep.Seconds()),
			})
		}
	})
}

func logCmd(ctx context.Context, inDir string, arg0 string, args ...string) *exec.Cmd {
	ret := exec.CommandContext(ctx, arg0, args...)
	ret.Dir = inDir
	logging.Debugf(ctx, "Running (from %q) - %s %v", inDir, arg0, args)
	return ret
}

func cmdErr(cmd *exec.Cmd, err error, reason string) error {
	if err != nil {
		ee, _ := err.(*exec.ExitError)
		outErr := ""
		if ee != nil {
			outErr = strings.TrimSpace(string(ee.Stderr))
			if len(outErr) > 128 {
				outErr = outErr[:128] + "..."
			}
		}
		err = errors.Annotate(err, "running %q: %s: %s", strings.Join(cmd.Args, " "), reason, outErr).Err()
	}
	return err
}

func appendText(path, fmtStr string, items ...interface{}) error {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = fmt.Fprintf(file, fmtStr, items...)
	return err
}

func (opts *EditRecipeBundleOpts) prepBundle(ctx context.Context, inDir, recipesPy, toDirectory string) (err error) {
	args := []string{
		recipesPy,
	}
	if logging.GetLevel(ctx) < logging.Info {
		args = append(args, "-v")
	}
	for projID, path := range opts.Overrides {
		args = append(args, "-O", fmt.Sprintf("%s=%s", projID, path))
	}
	args = append(args, "bundle", "--destination", filepath.Join(toDirectory))
	cmd := logCmd(ctx, inDir, "python", args...)
	if logging.GetLevel(ctx) < logging.Info {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	if err = cmdErr(cmd, cmd.Run(), "creating bundle"); err != nil {
		return
	}
	return
}

// findRecipesPy locates the current repo's `recipes.py`. It does this by:
//   * invoking git to find the repo root
//   * loading the recipes.cfg at infra/config/recipes.cfg
//   * stat'ing the recipes.py implied by the recipes_path in that cfg file.
//
// Failure will return an error.
//
// On success, the absolute path to recipes.py is returned.
func findRecipesPy(ctx context.Context, inDir string) (string, error) {
	cmd := logCmd(ctx, inDir, "git", "rev-parse", "--show-toplevel")
	out, err := cmd.Output()
	if err = cmdErr(cmd, err, "finding git repo"); err != nil {
		return "", err
	}

	repoRoot := strings.TrimSpace(string(out))

	pth := filepath.Join(repoRoot, "infra", "config", "recipes.cfg")
	switch st, err := os.Stat(pth); {
	case err != nil:
		return "", errors.Annotate(err, "reading recipes.cfg").Err()

	case !st.Mode().IsRegular():
		return "", errors.Reason("%q is not a regular file", pth).Err()
	}

	type recipesJSON struct {
		RecipesPath string `json:"recipes_path"`
	}
	rj := &recipesJSON{}

	f, err := os.Open(pth)
	if err != nil {
		return "", errors.Reason("reading recipes.cfg: %q", pth).Err()
	}
	defer f.Close()

	if err := json.NewDecoder(f).Decode(rj); err != nil {
		return "", errors.Reason("parsing recipes.cfg: %q", pth).Err()
	}

	return filepath.Join(
		repoRoot, filepath.FromSlash(rj.RecipesPath), "recipes.py"), nil
}
<|endoftext|>


<commit_before>package ginprometheus

import (
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	log "github.com/sirupsen/logrus"
)

var defaultMetricPath = "/metrics"

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz prometheus.Summary

	MetricsPath string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string) *Prometheus {
	p := &Prometheus{
		MetricsPath: defaultMetricPath,
	}

	p.registerMetrics(subsystem)

	return p
}

func (p *Prometheus) registerMetrics(subsystem string) {

	p.reqCnt = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Subsystem: subsystem,
			Name:      "requests_total",
			Help:      "How many HTTP requests processed, partitioned by status code and HTTP method.",
		},
		[]string{"code", "method", "handler"},
	)

	if err := prometheus.Register(p.reqCnt); err != nil {
		log.Info("reqCnt could not be registered: ", err)
	} else {
		log.Info("reqCnt registered.")
	}

	p.reqDur = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "request_duration_seconds",
			Help:      "The HTTP request latencies in seconds.",
		},
	)

	if err := prometheus.Register(p.reqDur); err != nil {
		log.Info("reqDur could not be registered: ", err)
	} else {
		log.Info("reqDur registered.")
	}

	p.reqSz = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "request_size_bytes",
			Help:      "The HTTP request sizes in bytes.",
		},
	)

	if err := prometheus.Register(p.reqSz); err != nil {
		log.Info("reqSz could not be registered: ", err)
	} else {
		log.Info("reqSz registered.")
	}

	p.resSz = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "response_size_bytes",
			Help:      "The HTTP response sizes in bytes.",
		},
	)

	if err := prometheus.Register(p.resSz); err != nil {
		log.Info("resSz could not be registered: ", err)
	} else {
		log.Info("resSz registered.")
	}

}

// Use adds the middleware to a gin engine.
func (p *Prometheus) Use(e *gin.Engine) {
	e.Use(p.handlerFunc())
	e.GET(p.MetricsPath, prometheusHandler())
}

// UseWithAuth adds the middleware to a gin engine with BasicAuth.
func (p *Prometheus) UseWithAuth(e *gin.Engine, accounts gin.Accounts) {
	e.Use(p.handlerFunc())
	e.GET(p.MetricsPath, gin.BasicAuth(accounts), prometheusHandler())
}

func (p *Prometheus) handlerFunc() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Request.URL.String() == p.MetricsPath {
			c.Next()
			return
		}

		start := time.Now()

		reqSz := make(chan int)
		go computeApproximateRequestSize(c.Request, reqSz)

		c.Next()

		status := strconv.Itoa(c.Writer.Status())
		elapsed := float64(time.Since(start)) / float64(time.Second)
		resSz := float64(c.Writer.Size())

		p.reqDur.Observe(elapsed)
		p.reqCnt.WithLabelValues(status, c.Request.Method, c.HandlerName()).Inc()
		p.reqSz.Observe(float64(<-reqSz))
		p.resSz.Observe(resSz)
	}
}

func prometheusHandler() gin.HandlerFunc {
	h := promhttp.Handler()
	return func(c *gin.Context) {
		h.ServeHTTP(c.Writer, c.Request)
	}
}

// From https://github.com/DanielHeckrath/gin-prometheus/blob/master/gin_prometheus.go
func computeApproximateRequestSize(r *http.Request, out chan int) {
	s := 0
	if r.URL != nil {
		s = len(r.URL.String())
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	out <- s
}
<commit_msg>Add the possibility of listening in a different address for metrics exposing<commit_after>package ginprometheus

import (
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	log "github.com/sirupsen/logrus"
)

var defaultMetricPath = "/metrics"

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz prometheus.Summary
	router               *gin.Engine
	listenAddress        string

	MetricsPath string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string) *Prometheus {
	p := &Prometheus{
		MetricsPath: defaultMetricPath,
	}

	p.registerMetrics(subsystem)

	return p
}

func (p *Prometheus) SetListenAddress(address string) {
	p.listenAddress = address
	if p.listenAddress != "" {
		p.router = gin.Default()
	}
}

func (p *Prometheus) setMetricsPath(e *gin.Engine) {

	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, prometheusHandler())
	}
}

func (p *Prometheus) setMetricsPathWithAuth(e *gin.Engine, accounts gin.Accounts) {

	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, gin.BasicAuth(accounts), prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, gin.BasicAuth(accounts), prometheusHandler())
	}

}

func (p *Prometheus) runServer() {
	if p.listenAddress != "" {
		go p.router.Run(p.listenAddress)
	}
}

func (p *Prometheus) registerMetrics(subsystem string) {

	p.reqCnt = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Subsystem: subsystem,
			Name:      "requests_total",
			Help:      "How many HTTP requests processed, partitioned by status code and HTTP method.",
		},
		[]string{"code", "method", "handler"},
	)

	if err := prometheus.Register(p.reqCnt); err != nil {
		log.Info("reqCnt could not be registered: ", err)
	} else {
		log.Info("reqCnt registered.")
	}

	p.reqDur = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "request_duration_seconds",
			Help:      "The HTTP request latencies in seconds.",
		},
	)

	if err := prometheus.Register(p.reqDur); err != nil {
		log.Info("reqDur could not be registered: ", err)
	} else {
		log.Info("reqDur registered.")
	}

	p.reqSz = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "request_size_bytes",
			Help:      "The HTTP request sizes in bytes.",
		},
	)

	if err := prometheus.Register(p.reqSz); err != nil {
		log.Info("reqSz could not be registered: ", err)
	} else {
		log.Info("reqSz registered.")
	}

	p.resSz = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Subsystem: subsystem,
			Name:      "response_size_bytes",
			Help:      "The HTTP response sizes in bytes.",
		},
	)

	if err := prometheus.Register(p.resSz); err != nil {
		log.Info("resSz could not be registered: ", err)
	} else {
		log.Info("resSz registered.")
	}

}

// Use adds the middleware to a gin engine.
func (p *Prometheus) Use(e *gin.Engine) {
	e.Use(p.handlerFunc())
	p.setMetricsPath(e)
}

// UseWithAuth adds the middleware to a gin engine with BasicAuth.
func (p *Prometheus) UseWithAuth(e *gin.Engine, accounts gin.Accounts) {
	e.Use(p.handlerFunc())
	p.setMetricsPathWithAuth(e, accounts)
}

func (p *Prometheus) handlerFunc() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Request.URL.String() == p.MetricsPath {
			c.Next()
			return
		}

		start := time.Now()

		reqSz := make(chan int)
		go computeApproximateRequestSize(c.Request, reqSz)

		c.Next()

		status := strconv.Itoa(c.Writer.Status())
		elapsed := float64(time.Since(start)) / float64(time.Second)
		resSz := float64(c.Writer.Size())

		p.reqDur.Observe(elapsed)
		p.reqCnt.WithLabelValues(status, c.Request.Method, c.HandlerName()).Inc()
		p.reqSz.Observe(float64(<-reqSz))
		p.resSz.Observe(resSz)
	}
}

func prometheusHandler() gin.HandlerFunc {
	h := promhttp.Handler()
	return func(c *gin.Context) {
		h.ServeHTTP(c.Writer, c.Request)
	}
}

// From https://github.com/DanielHeckrath/gin-prometheus/blob/master/gin_prometheus.go
func computeApproximateRequestSize(r *http.Request, out chan int) {
	s := 0
	if r.URL != nil {
		s = len(r.URL.String())
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	out <- s
}
<|endoftext|>


<commit_before>package server

// Copyright (c) 2014-2016 Shawn Boyette <shawn@firepear.net>. All
// rights reserved.  Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

// Socket code for petrel

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"time"

	"firepear.net/qsplit"
	"firepear.net/petrel"
)

var pp = petrel.Perrs

// sockAccept monitors the listener socket and spawns connections for
// clients.
func (h *Handler) sockAccept() {
	defer h.w.Done()
	var cn uint
	for cn = 1; true; cn++ {
		c, err := h.l.Accept()
		if err != nil {
			select {
			case <-h.q:
				// h.Quit() was invoked; close up shop
				h.genMsg(0, 0, pp["quit"], "", nil)
				return
			default:
				// we've had a networking error
				h.genMsg(0, 0, pp["listenerfail"], "", err)
				return
			}
		}
		// we have a new client
		h.w.Add(1)
		go h.connHandler(c, cn)
	}
}

// connHandler dispatches commands from, and sends reponses to, a client. It
// is launched, per-connection, from sockAccept().
func (h *Handler) connHandler(c net.Conn, cn uint) {
	defer h.w.Done()
	defer c.Close()
	// request counter for this connection
	var reqnum uint

	if h.li {
		h.genMsg(cn, reqnum, pp["connect"], c.RemoteAddr().String(), nil)
	} else {
		h.genMsg(cn, reqnum, pp["connect"], "", nil)
	}

	for {
		reqnum++
		// read the request
		req, perr, xtra, err := h.connRead(c, cn, reqnum)
		if perr != "" {
			h.genMsg(cn, reqnum, pp[perr], xtra, err)
			if pp[perr].Xmit != nil {
				h.send(c, cn, reqnum, pp[perr].Xmit)
			}
			//TODO send "you've been disconnected" msg
			return
		}
		if len(req) == 0 {
			h.genMsg(cn, reqnum, pp["nilreq"], "", nil)
			h.send(c, cn, reqnum, pp["nilreq"].Xmit)
			continue
		}

		// dispatch the request and get the reply
		reply, perr, xtra, err := h.reqDispatch(c, cn, reqnum, req)
		if perr != "" {
			h.genMsg(cn, reqnum, pp[perr], xtra, err)
			if pp[perr].Xmit != nil {
				h.send(c, cn, reqnum, pp[perr].Xmit)
			}
			continue
		}

		// send reply
		err = h.send(c, cn, reqnum, reply)
		if err != nil {
			return
		}
		h.genMsg(cn, reqnum, pp["success"], "", nil)
	}
}

// connRead does all network reads and assembles the request. If it
// returns an error, then the connection terminates because the state
// of the connection cannot be known.
func (h *Handler) connRead(c net.Conn, cn, reqnum uint) ([]byte, string, string, error) {
	// buffer 0 holds the message length
	b0 := make([]byte, 4)
	// buffer 1: network reads go here, 128B at a time
	b1 := make([]byte, 128)
	// buffer 2: data accumulates here; requests pulled from here
	var b2 []byte
	// message length
	var mlen int32
	// bytes read so far
	var bread int32

	// get the response message length
	if h.t > 0 {
		c.SetReadDeadline(time.Now().Add(h.t))
	}
	n, err := c.Read(b0)
	if err != nil {
		if err == io.EOF {
			return nil, "disconnect", "", err
		} else {
			return nil, "netreaderr", "no message length", err
		}
	}
	if  n != 4 {
		return nil, "netreaderr", "short read on message length", err
	}
	buf := bytes.NewReader(b0)
	err = binary.Read(buf, binary.BigEndian, &mlen)
	if err != nil {
		return nil, "internalerr", "could not decode message length", err
	}

	for bread < mlen {
		// if there are less than 128 bytes remaining to read in this
		// message, resize b1 to fit. this avoids reading across a
		// message boundary.
		if x := mlen - bread; x < 128 {
			b1 = make([]byte, x)
		}
		if h.t > 0 {
			c.SetReadDeadline(time.Now().Add(h.t))
		}
		n, err = c.Read(b1)
		if err != nil {
			if err == io.EOF {
				return nil, "disconnect", "", err
			} else {
				return nil, "netreaderr", "failed to read req from socket", err
			}
		}
		if n == 0 {
			// short-circuit just in case this ever manages to happen
			return b2[:mlen], "", "", err
		}
		bread += int32(n)
		if h.rl > 0 && bread > h.rl {
			return nil, "reqlen", "", nil
		}
		b2 = append(b2, b1[:n]...)
	}
	return b2[:mlen], "", "", err
}

// reqDispatch turns the request into a command and arguments, and
// dispatches these components to a handler.
func (h *Handler) reqDispatch(c net.Conn, cn, reqnum uint, req []byte) ([]byte, string, string, error) {
	// get chunk locations
	cl := qsplit.LocationsOnce(req)
	dcmd := string(req[cl[0]:cl[1]])
	// now get the args
	var dargs []byte
	if cl[2] != -1 {
		dargs = req[cl[2]:]
	}
	// send error if we don't recognize the command
	dfunc, ok := h.d[dcmd]
	if !ok {
		return nil, "badreq", fmt.Sprintf("'%s'", dcmd), nil
	}
	// ok, we know the command and we have its dispatch
	// func. call it and send response
	h.genMsg(cn, reqnum, pp["dispatch"], dcmd, nil)
	var rs [][]byte // req, split by word
	switch dfunc.mode {
	case "args":
		rs = qsplit.ToBytes(dargs)
	case "blob":
		rs = rs[:0]
		rs = append(rs, dargs)
	}
	resp, err := dfunc.df(rs)
	if err != nil {
		return nil, "reqerr", "", err
	}
	return resp, "", "", nil
}

// send handles all network writes.
func (h *Handler) send(c net.Conn, cn, reqnum uint, resp []byte) (string, string, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, int32(len(resp)))
	if err != nil {
		return "internalerr", "could not encode message length", err
	}
	resp = append(buf.Bytes(), resp...)
	if h.t > 0 {
		c.SetReadDeadline(time.Now().Add(h.t))
	}
	_, err = c.Write(resp)
	if err != nil {
		return "netwriteerr", "", err
	}
	return "", "", err
}
<commit_msg>all sends now have internal message generation and cause handler abort on failure, closing #47<commit_after>package server

// Copyright (c) 2014-2016 Shawn Boyette <shawn@firepear.net>. All
// rights reserved.  Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

// Socket code for petrel

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"time"

	"firepear.net/qsplit"
	"firepear.net/petrel"
)

var pp = petrel.Perrs

// sockAccept monitors the listener socket and spawns connections for
// clients.
func (h *Handler) sockAccept() {
	defer h.w.Done()
	var cn uint
	for cn = 1; true; cn++ {
		c, err := h.l.Accept()
		if err != nil {
			select {
			case <-h.q:
				// h.Quit() was invoked; close up shop
				h.genMsg(0, 0, pp["quit"], "", nil)
				return
			default:
				// we've had a networking error
				h.genMsg(0, 0, pp["listenerfail"], "", err)
				return
			}
		}
		// we have a new client
		h.w.Add(1)
		go h.connHandler(c, cn)
	}
}

// connHandler dispatches commands from, and sends reponses to, a client. It
// is launched, per-connection, from sockAccept().
func (h *Handler) connHandler(c net.Conn, cn uint) {
	defer h.w.Done()
	defer c.Close()
	// request counter for this connection
	var reqnum uint

	if h.li {
		h.genMsg(cn, reqnum, pp["connect"], c.RemoteAddr().String(), nil)
	} else {
		h.genMsg(cn, reqnum, pp["connect"], "", nil)
	}

	for {
		reqnum++
		// read the request
		req, perr, xtra, err := h.connRead(c, cn, reqnum)
		if perr != "" {
			h.genMsg(cn, reqnum, pp[perr], xtra, err)
			if pp[perr].Xmit != nil {
				err = h.send(c, cn, reqnum, pp[perr].Xmit)
				if err != nil {
					return
				}
			}
			//TODO send "you've been disconnected" msg
			return
		}
		if len(req) == 0 {
			h.genMsg(cn, reqnum, pp["nilreq"], "", nil)
			err = h.send(c, cn, reqnum, pp["nilreq"].Xmit)
			if err != nil {
				return
			}
			continue
		}

		// dispatch the request and get the reply
		reply, perr, xtra, err := h.reqDispatch(c, cn, reqnum, req)
		if perr != "" {
			h.genMsg(cn, reqnum, pp[perr], xtra, err)
			if pp[perr].Xmit != nil {
				err = h.send(c, cn, reqnum, pp[perr].Xmit)
				if err != nil {
					return
				}
			}
			continue
		}

		// send reply
		err = h.send(c, cn, reqnum, reply)
		if err != nil {
			return
		}
		h.genMsg(cn, reqnum, pp["success"], "", nil)
	}
}

// connRead does all network reads and assembles the request. If it
// returns an error, then the connection terminates because the state
// of the connection cannot be known.
func (h *Handler) connRead(c net.Conn, cn, reqnum uint) ([]byte, string, string, error) {
	// buffer 0 holds the message length
	b0 := make([]byte, 4)
	// buffer 1: network reads go here, 128B at a time
	b1 := make([]byte, 128)
	// buffer 2: data accumulates here; requests pulled from here
	var b2 []byte
	// message length
	var mlen int32
	// bytes read so far
	var bread int32

	// get the response message length
	if h.t > 0 {
		c.SetReadDeadline(time.Now().Add(h.t))
	}
	n, err := c.Read(b0)
	if err != nil {
		if err == io.EOF {
			return nil, "disconnect", "", err
		} else {
			return nil, "netreaderr", "no message length", err
		}
	}
	if  n != 4 {
		return nil, "netreaderr", "short read on message length", err
	}
	buf := bytes.NewReader(b0)
	err = binary.Read(buf, binary.BigEndian, &mlen)
	if err != nil {
		return nil, "internalerr", "could not decode message length", err
	}

	for bread < mlen {
		// if there are less than 128 bytes remaining to read in this
		// message, resize b1 to fit. this avoids reading across a
		// message boundary.
		if x := mlen - bread; x < 128 {
			b1 = make([]byte, x)
		}
		if h.t > 0 {
			c.SetReadDeadline(time.Now().Add(h.t))
		}
		n, err = c.Read(b1)
		if err != nil {
			if err == io.EOF {
				return nil, "disconnect", "", err
			} else {
				return nil, "netreaderr", "failed to read req from socket", err
			}
		}
		if n == 0 {
			// short-circuit just in case this ever manages to happen
			return b2[:mlen], "", "", err
		}
		bread += int32(n)
		if h.rl > 0 && bread > h.rl {
			return nil, "reqlen", "", nil
		}
		b2 = append(b2, b1[:n]...)
	}
	return b2[:mlen], "", "", err
}

// reqDispatch turns the request into a command and arguments, and
// dispatches these components to a handler.
func (h *Handler) reqDispatch(c net.Conn, cn, reqnum uint, req []byte) ([]byte, string, string, error) {
	// get chunk locations
	cl := qsplit.LocationsOnce(req)
	dcmd := string(req[cl[0]:cl[1]])
	// now get the args
	var dargs []byte
	if cl[2] != -1 {
		dargs = req[cl[2]:]
	}
	// send error if we don't recognize the command
	dfunc, ok := h.d[dcmd]
	if !ok {
		return nil, "badreq", fmt.Sprintf("'%s'", dcmd), nil
	}
	// ok, we know the command and we have its dispatch
	// func. call it and send response
	h.genMsg(cn, reqnum, pp["dispatch"], dcmd, nil)
	var rs [][]byte // req, split by word
	switch dfunc.mode {
	case "args":
		rs = qsplit.ToBytes(dargs)
	case "blob":
		rs = rs[:0]
		rs = append(rs, dargs)
	}
	resp, err := dfunc.df(rs)
	if err != nil {
		return nil, "reqerr", "", err
	}
	return resp, "", "", nil
}

// send handles all network writes.
func (h *Handler) send(c net.Conn, cn, reqnum uint, resp []byte) error {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, int32(len(resp)))
	if err != nil {
		h.genMsg(cn, reqnum, pp["internalerr"], "could not encode message length", err)
		return err
	}
	resp = append(buf.Bytes(), resp...)
	if h.t > 0 {
		c.SetReadDeadline(time.Now().Add(h.t))
	}
	_, err = c.Write(resp)
	if err != nil {
		h.genMsg(cn, reqnum, pp["netwriteerr"], "", err)
		return err
	}
	return err
}
<|endoftext|>


<commit_before>// Copyright 2015 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 ratelimit

import (
	"io"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
)

// Create a bucket that limits the rate at which it calls the wrapped bucket
// using opThrottle, and limits the bandwidth with which it reads from the
// wrapped bucket using egressThrottle.
func NewThrottledBucket(
	opThrottle Throttle,
	egressThrottle Throttle,
	wrapped gcs.Bucket) (b gcs.Bucket) {
	b = &throttledBucket{
		opThrottle:     opThrottle,
		egressThrottle: egressThrottle,
		wrapped:        wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// throttledBucket
////////////////////////////////////////////////////////////////////////

type throttledBucket struct {
	opThrottle     Throttle
	egressThrottle Throttle
	wrapped        gcs.Bucket
}

func (b *throttledBucket) Name() string {
	return b.wrapped.Name()
}

func (b *throttledBucket) NewReader(
	ctx context.Context,
	req *gcs.ReadObjectRequest) (rc io.ReadCloser, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	rc, err = b.wrapped.NewReader(ctx, req)
	if err != nil {
		return
	}

	// Wrap the result in a throttled layer.
	rc = &readerCloser{
		Reader: ThrottledReader(ctx, rc, b.egressThrottle),
		Closer: rc,
	}

	return
}

func (b *throttledBucket) CreateObject(
	ctx context.Context,
	req *gcs.CreateObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.CreateObject(ctx, req)

	return
}

func (b *throttledBucket) CopyObject(
	ctx context.Context,
	req *gcs.CopyObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.CopyObject(ctx, req)

	return
}

func (b *throttledBucket) StatObject(
	ctx context.Context,
	req *gcs.StatObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.StatObject(ctx, req)

	return
}

func (b *throttledBucket) ListObjects(
	ctx context.Context,
	req *gcs.ListObjectsRequest) (listing *gcs.Listing, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	listing, err = b.wrapped.ListObjects(ctx, req)

	return
}

func (b *throttledBucket) UpdateObject(
	ctx context.Context,
	req *gcs.UpdateObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.UpdateObject(ctx, req)

	return
}

func (b *throttledBucket) DeleteObject(
	ctx context.Context,
	name string) (err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	err = b.wrapped.DeleteObject(ctx, name)

	return
}

////////////////////////////////////////////////////////////////////////
// readerCloser
////////////////////////////////////////////////////////////////////////

// An io.ReadCloser that forwards read requests to an io.Reader and close
// requests to an io.Closer.
type readerCloser struct {
	Reader io.Reader
	Closer io.Closer
}

func (rc *readerCloser) Read(p []byte) (n int, err error) {
	n, err = rc.Reader.Read(p)
	return
}

func (rc *readerCloser) Close() (err error) {
	err = rc.Closer.Close()
	return
}
<commit_msg>Fixed the build, broken by jacobsa/gcloud@58c4ae5.<commit_after>// Copyright 2015 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 ratelimit

import (
	"io"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
)

// Create a bucket that limits the rate at which it calls the wrapped bucket
// using opThrottle, and limits the bandwidth with which it reads from the
// wrapped bucket using egressThrottle.
func NewThrottledBucket(
	opThrottle Throttle,
	egressThrottle Throttle,
	wrapped gcs.Bucket) (b gcs.Bucket) {
	b = &throttledBucket{
		opThrottle:     opThrottle,
		egressThrottle: egressThrottle,
		wrapped:        wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// throttledBucket
////////////////////////////////////////////////////////////////////////

type throttledBucket struct {
	opThrottle     Throttle
	egressThrottle Throttle
	wrapped        gcs.Bucket
}

func (b *throttledBucket) Name() string {
	return b.wrapped.Name()
}

func (b *throttledBucket) NewReader(
	ctx context.Context,
	req *gcs.ReadObjectRequest) (rc io.ReadCloser, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	rc, err = b.wrapped.NewReader(ctx, req)
	if err != nil {
		return
	}

	// Wrap the result in a throttled layer.
	rc = &readerCloser{
		Reader: ThrottledReader(ctx, rc, b.egressThrottle),
		Closer: rc,
	}

	return
}

func (b *throttledBucket) CreateObject(
	ctx context.Context,
	req *gcs.CreateObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.CreateObject(ctx, req)

	return
}

func (b *throttledBucket) CopyObject(
	ctx context.Context,
	req *gcs.CopyObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.CopyObject(ctx, req)

	return
}

func (b *throttledBucket) ComposeObjects(
	ctx context.Context,
	req *gcs.ComposeObjectsRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.ComposeObjects(ctx, req)

	return
}

func (b *throttledBucket) StatObject(
	ctx context.Context,
	req *gcs.StatObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.StatObject(ctx, req)

	return
}

func (b *throttledBucket) ListObjects(
	ctx context.Context,
	req *gcs.ListObjectsRequest) (listing *gcs.Listing, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	listing, err = b.wrapped.ListObjects(ctx, req)

	return
}

func (b *throttledBucket) UpdateObject(
	ctx context.Context,
	req *gcs.UpdateObjectRequest) (o *gcs.Object, err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	o, err = b.wrapped.UpdateObject(ctx, req)

	return
}

func (b *throttledBucket) DeleteObject(
	ctx context.Context,
	name string) (err error) {
	// Wait for permission to call through.
	err = b.opThrottle.Wait(ctx, 1)
	if err != nil {
		return
	}

	// Call through.
	err = b.wrapped.DeleteObject(ctx, name)

	return
}

////////////////////////////////////////////////////////////////////////
// readerCloser
////////////////////////////////////////////////////////////////////////

// An io.ReadCloser that forwards read requests to an io.Reader and close
// requests to an io.Closer.
type readerCloser struct {
	Reader io.Reader
	Closer io.Closer
}

func (rc *readerCloser) Read(p []byte) (n int, err error) {
	n, err = rc.Reader.Read(p)
	return
}

func (rc *readerCloser) Close() (err error) {
	err = rc.Closer.Close()
	return
}
<|endoftext|>


<commit_before>package api

import (
	"fmt"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/mdata"
)

// ReqMap is a map of requests of data,
// it has single requests for which no pre-normalization effort will be performed, and
// requests that can be pre-normalized together to the same resolution, bundled by their PNGroup
type ReqMap struct {
	single   []models.Req
	pngroups map[models.PNGroup][]models.Req
	cnt      uint32
}

func NewReqMap() *ReqMap {
	return &ReqMap{
		pngroups: make(map[models.PNGroup][]models.Req),
	}
}

// Add adds a models.Req to the ReqMap
func (r *ReqMap) Add(req models.Req) {
	r.cnt++
	if req.PNGroup == 0 {
		r.single = append(r.single, req)
		return
	}
	r.pngroups[req.PNGroup] = append(r.pngroups[req.PNGroup], req)
}

// Dump provides a human readable string representation of the ReqsMap
func (r ReqMap) Dump() string {
	out := fmt.Sprintf("ReqsMap (%d entries):\n", r.cnt)
	out += "  Groups:\n"
	for i, reqs := range r.pngroups {
		out += fmt.Sprintf("    * group %d:", i)
		for _, r := range reqs {
			out += "      " + r.DebugString() + "\n"
		}
	}
	out += "  Single:\n"
	for _, r := range r.single {
		out += "    " + r.DebugString() + "\n"
	}
	return out
}

// Requests indexed by their retention ID
type ReqsByRet [][]models.Req

// OutInterval returns the outinterval of the ReqsByRet
// this assumes that all requests have been planned to a consistent out interval, of course.
func (rbr ReqsByRet) OutInterval() uint32 {
	for _, reqs := range rbr {
		if len(reqs) != 0 {
			return reqs[0].OutInterval
		}
	}
	return 0
}

func (rbr ReqsByRet) HasData() bool {
	for _, reqs := range rbr {
		if len(reqs) != 0 {
			return true
		}
	}
	return false
}

func (rbr ReqsByRet) Len() int {
	var cnt int
	for _, reqs := range rbr {
		cnt += len(reqs)
	}
	return cnt
}

// GroupData embodies a PNGroup broken down by whether requests are MDP-optimizable, and by retention
type GroupData struct {
	mdpyes ReqsByRet // MDP-optimizable requests
	mdpno  ReqsByRet // not MDP-optimizable reqs
}

func NewGroupData() GroupData {
	return GroupData{
		mdpyes: make([][]models.Req, mdata.Schemas.Len()),
		mdpno:  make([][]models.Req, mdata.Schemas.Len()),
	}
}

func (gd GroupData) Len() int {
	return gd.mdpno.Len() + gd.mdpyes.Len()
}

// ReqsPlan holds requests that have been planned, broken down by PNGroup and MDP-optimizability
type ReqsPlan struct {
	pngroups map[models.PNGroup]GroupData
	single   GroupData
	cnt      uint32
}

// NewReqsPlan generates a ReqsPlan based on the provided ReqMap.
func NewReqsPlan(reqs ReqMap) ReqsPlan {
	rp := ReqsPlan{
		pngroups: make(map[models.PNGroup]GroupData),
		single:   NewGroupData(),
		cnt:      reqs.cnt,
	}
	for group, groupReqs := range reqs.pngroups {
		data := NewGroupData()
		for _, req := range groupReqs {
			if req.MaxPoints > 0 {
				data.mdpyes[req.SchemaId] = append(data.mdpyes[req.SchemaId], req)
			} else {
				data.mdpno[req.SchemaId] = append(data.mdpno[req.SchemaId], req)
			}
		}
		rp.pngroups[group] = data
	}
	for _, req := range reqs.single {
		if req.MaxPoints > 0 {
			rp.single.mdpyes[req.SchemaId] = append(rp.single.mdpyes[req.SchemaId], req)
		} else {
			rp.single.mdpno[req.SchemaId] = append(rp.single.mdpno[req.SchemaId], req)
		}
	}
	return rp
}

// PointsFetch returns how many points this plan will fetch when executed
func (rp ReqsPlan) PointsFetch() uint32 {
	var cnt uint32
	for _, rbr := range rp.single.mdpyes {
		for _, req := range rbr {
			cnt += req.PointsFetch()
		}
	}
	for _, rbr := range rp.single.mdpno {
		for _, req := range rbr {
			cnt += req.PointsFetch()
		}
	}
	for _, data := range rp.pngroups {
		for _, rbr := range data.mdpyes {
			for _, req := range rbr {
				cnt += req.PointsFetch()
			}
		}
		for _, rbr := range data.mdpno {
			for _, req := range rbr {
				cnt += req.PointsFetch()
			}
		}
	}
	return cnt
}

// Dump provides a human readable string representation of the ReqsPlan
func (rp ReqsPlan) Dump() string {
	out := fmt.Sprintf("ReqsPlan (%d entries):\n", rp.cnt)
	out += "  # Groups:\n"
	for i, data := range rp.pngroups {
		out += fmt.Sprintf("    ## group %d\n", i)
		out += "      ### MDP-yes:\n"
		for schemaID, reqs := range data.mdpyes {
			for _, req := range reqs {
				out += fmt.Sprintf("        [%d] %s\n", schemaID, req.DebugString())
			}
		}
		out += "      ### MDP-no:\n"
		for schemaID, reqs := range data.mdpno {
			for _, req := range reqs {
				out += fmt.Sprintf("        [%d] %s\n", schemaID, req.DebugString())
			}
		}
	}
	out += "  # Single\n"
	out += "   ## MDP-yes:\n"
	for schemaID, reqs := range rp.single.mdpyes {
		for _, req := range reqs {
			out += fmt.Sprintf("    [%d] %s\n", schemaID, req.DebugString())
		}
	}
	out += "    ## MDP-no:\n"
	for schemaID, reqs := range rp.single.mdpno {
		for _, req := range reqs {
			out += fmt.Sprintf("    [%d] %s\n", schemaID, req.DebugString())
		}
	}
	return out
}

// PointsReturn estimates the amount of points that will be returned for this request
// best effort: not aware of summarize(), aggregation functions, runtime normalization. but does account for runtime consolidation
func (rp ReqsPlan) PointsReturn(planMDP uint32) uint32 {
	var cnt uint32
	for _, rbr := range rp.single.mdpyes {
		for _, req := range rbr {
			cnt += req.PointsReturn(planMDP)
		}
	}
	for _, rbr := range rp.single.mdpno {
		for _, req := range rbr {
			cnt += req.PointsReturn(planMDP)
		}
	}
	for _, data := range rp.pngroups {
		for _, rbr := range data.mdpyes {
			for _, req := range rbr {
				cnt += req.PointsReturn(planMDP)
			}
		}
		for _, rbr := range data.mdpno {
			for _, req := range rbr {
				cnt += req.PointsReturn(planMDP)
			}
		}
	}
	return cnt
}

// List returns the requests contained within the plan as a slice
func (rp ReqsPlan) List() []models.Req {
	l := make([]models.Req, 0, rp.cnt)
	for _, reqs := range rp.single.mdpno {
		l = append(l, reqs...)
	}
	for _, reqs := range rp.single.mdpyes {
		l = append(l, reqs...)
	}
	for _, data := range rp.pngroups {
		for _, reqs := range data.mdpno {
			l = append(l, reqs...)
		}
		for _, reqs := range data.mdpyes {
			l = append(l, reqs...)
		}
	}
	return l
}
<commit_msg>typo: GroupData also used for singles<commit_after>package api

import (
	"fmt"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/mdata"
)

// ReqMap is a map of requests of data,
// it has single requests for which no pre-normalization effort will be performed, and
// requests that can be pre-normalized together to the same resolution, bundled by their PNGroup
type ReqMap struct {
	single   []models.Req
	pngroups map[models.PNGroup][]models.Req
	cnt      uint32
}

func NewReqMap() *ReqMap {
	return &ReqMap{
		pngroups: make(map[models.PNGroup][]models.Req),
	}
}

// Add adds a models.Req to the ReqMap
func (r *ReqMap) Add(req models.Req) {
	r.cnt++
	if req.PNGroup == 0 {
		r.single = append(r.single, req)
		return
	}
	r.pngroups[req.PNGroup] = append(r.pngroups[req.PNGroup], req)
}

// Dump provides a human readable string representation of the ReqsMap
func (r ReqMap) Dump() string {
	out := fmt.Sprintf("ReqsMap (%d entries):\n", r.cnt)
	out += "  Groups:\n"
	for i, reqs := range r.pngroups {
		out += fmt.Sprintf("    * group %d:", i)
		for _, r := range reqs {
			out += "      " + r.DebugString() + "\n"
		}
	}
	out += "  Single:\n"
	for _, r := range r.single {
		out += "    " + r.DebugString() + "\n"
	}
	return out
}

// Requests indexed by their retention ID
type ReqsByRet [][]models.Req

// OutInterval returns the outinterval of the ReqsByRet
// this assumes that all requests have been planned to a consistent out interval, of course.
func (rbr ReqsByRet) OutInterval() uint32 {
	for _, reqs := range rbr {
		if len(reqs) != 0 {
			return reqs[0].OutInterval
		}
	}
	return 0
}

func (rbr ReqsByRet) HasData() bool {
	for _, reqs := range rbr {
		if len(reqs) != 0 {
			return true
		}
	}
	return false
}

func (rbr ReqsByRet) Len() int {
	var cnt int
	for _, reqs := range rbr {
		cnt += len(reqs)
	}
	return cnt
}

// GroupData embodies a set of requests broken down by whether requests are MDP-optimizable, and by retention
type GroupData struct {
	mdpyes ReqsByRet // MDP-optimizable requests
	mdpno  ReqsByRet // not MDP-optimizable reqs
}

func NewGroupData() GroupData {
	return GroupData{
		mdpyes: make([][]models.Req, mdata.Schemas.Len()),
		mdpno:  make([][]models.Req, mdata.Schemas.Len()),
	}
}

func (gd GroupData) Len() int {
	return gd.mdpno.Len() + gd.mdpyes.Len()
}

// ReqsPlan holds requests that have been planned, broken down by PNGroup and MDP-optimizability
type ReqsPlan struct {
	pngroups map[models.PNGroup]GroupData
	single   GroupData
	cnt      uint32
}

// NewReqsPlan generates a ReqsPlan based on the provided ReqMap.
func NewReqsPlan(reqs ReqMap) ReqsPlan {
	rp := ReqsPlan{
		pngroups: make(map[models.PNGroup]GroupData),
		single:   NewGroupData(),
		cnt:      reqs.cnt,
	}
	for group, groupReqs := range reqs.pngroups {
		data := NewGroupData()
		for _, req := range groupReqs {
			if req.MaxPoints > 0 {
				data.mdpyes[req.SchemaId] = append(data.mdpyes[req.SchemaId], req)
			} else {
				data.mdpno[req.SchemaId] = append(data.mdpno[req.SchemaId], req)
			}
		}
		rp.pngroups[group] = data
	}
	for _, req := range reqs.single {
		if req.MaxPoints > 0 {
			rp.single.mdpyes[req.SchemaId] = append(rp.single.mdpyes[req.SchemaId], req)
		} else {
			rp.single.mdpno[req.SchemaId] = append(rp.single.mdpno[req.SchemaId], req)
		}
	}
	return rp
}

// PointsFetch returns how many points this plan will fetch when executed
func (rp ReqsPlan) PointsFetch() uint32 {
	var cnt uint32
	for _, rbr := range rp.single.mdpyes {
		for _, req := range rbr {
			cnt += req.PointsFetch()
		}
	}
	for _, rbr := range rp.single.mdpno {
		for _, req := range rbr {
			cnt += req.PointsFetch()
		}
	}
	for _, data := range rp.pngroups {
		for _, rbr := range data.mdpyes {
			for _, req := range rbr {
				cnt += req.PointsFetch()
			}
		}
		for _, rbr := range data.mdpno {
			for _, req := range rbr {
				cnt += req.PointsFetch()
			}
		}
	}
	return cnt
}

// Dump provides a human readable string representation of the ReqsPlan
func (rp ReqsPlan) Dump() string {
	out := fmt.Sprintf("ReqsPlan (%d entries):\n", rp.cnt)
	out += "  # Groups:\n"
	for i, data := range rp.pngroups {
		out += fmt.Sprintf("    ## group %d\n", i)
		out += "      ### MDP-yes:\n"
		for schemaID, reqs := range data.mdpyes {
			for _, req := range reqs {
				out += fmt.Sprintf("        [%d] %s\n", schemaID, req.DebugString())
			}
		}
		out += "      ### MDP-no:\n"
		for schemaID, reqs := range data.mdpno {
			for _, req := range reqs {
				out += fmt.Sprintf("        [%d] %s\n", schemaID, req.DebugString())
			}
		}
	}
	out += "  # Single\n"
	out += "   ## MDP-yes:\n"
	for schemaID, reqs := range rp.single.mdpyes {
		for _, req := range reqs {
			out += fmt.Sprintf("    [%d] %s\n", schemaID, req.DebugString())
		}
	}
	out += "    ## MDP-no:\n"
	for schemaID, reqs := range rp.single.mdpno {
		for _, req := range reqs {
			out += fmt.Sprintf("    [%d] %s\n", schemaID, req.DebugString())
		}
	}
	return out
}

// PointsReturn estimates the amount of points that will be returned for this request
// best effort: not aware of summarize(), aggregation functions, runtime normalization. but does account for runtime consolidation
func (rp ReqsPlan) PointsReturn(planMDP uint32) uint32 {
	var cnt uint32
	for _, rbr := range rp.single.mdpyes {
		for _, req := range rbr {
			cnt += req.PointsReturn(planMDP)
		}
	}
	for _, rbr := range rp.single.mdpno {
		for _, req := range rbr {
			cnt += req.PointsReturn(planMDP)
		}
	}
	for _, data := range rp.pngroups {
		for _, rbr := range data.mdpyes {
			for _, req := range rbr {
				cnt += req.PointsReturn(planMDP)
			}
		}
		for _, rbr := range data.mdpno {
			for _, req := range rbr {
				cnt += req.PointsReturn(planMDP)
			}
		}
	}
	return cnt
}

// List returns the requests contained within the plan as a slice
func (rp ReqsPlan) List() []models.Req {
	l := make([]models.Req, 0, rp.cnt)
	for _, reqs := range rp.single.mdpno {
		l = append(l, reqs...)
	}
	for _, reqs := range rp.single.mdpyes {
		l = append(l, reqs...)
	}
	for _, data := range rp.pngroups {
		for _, reqs := range data.mdpno {
			l = append(l, reqs...)
		}
		for _, reqs := range data.mdpyes {
			l = append(l, reqs...)
		}
	}
	return l
}
<|endoftext|>