<commit_before>package shadowsocks

import (
	"bytes"
	"crypto/rand"
	"errors"
	"io"

	"v2ray.com/core/common/alloc"
	"v2ray.com/core/common/crypto"
	v2io "v2ray.com/core/common/io"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
)

const (
	Version                  = 1
	RequestOptionOneTimeAuth = protocol.RequestOption(101)

	AddrTypeIPv4   = 1
	AddrTypeIPv6   = 4
	AddrTypeDomain = 3
)

func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHeader, v2io.Reader, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	buffer := alloc.NewLocalBuffer(256)
	defer buffer.Release()

	ivLen := account.Cipher.IVSize()
	_, err = io.ReadFull(reader, buffer.Value[:ivLen])
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IV: " + err.Error())
	}

	iv := append([]byte(nil), buffer.Value[:ivLen]...)

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to initialize decoding stream: " + err.Error())
	}
	reader = crypto.NewCryptionReader(stream, reader)

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandTCP,
	}

	lenBuffer := 1
	_, err = io.ReadFull(reader, buffer.Value[:1])
	if err != nil {
		return nil, nil, errors.New("Sahdowsocks|TCP: Failed to read address type: " + err.Error())
	}

	addrType := (buffer.Value[0] & 0x0F)
	if (buffer.Value[0] & 0x10) == 0x10 {
		request.Option |= RequestOptionOneTimeAuth
	}

	switch addrType {
	case AddrTypeIPv4:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+4])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IPv4 address: " + err.Error())
		}
		request.Address = v2net.IPAddress(buffer.Value[lenBuffer : lenBuffer+4])
		lenBuffer += 4
	case AddrTypeIPv6:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+16])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IPv6 address: " + err.Error())
		}
		request.Address = v2net.IPAddress(buffer.Value[lenBuffer : lenBuffer+16])
		lenBuffer += 16
	case AddrTypeDomain:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+1])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read domain lenth: " + err.Error())
		}
		domainLength := int(buffer.Value[lenBuffer])
		lenBuffer++
		_, err = io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+domainLength])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read domain: " + err.Error())
		}
		request.Address = v2net.DomainAddress(string(buffer.Value[lenBuffer : lenBuffer+domainLength]))
		lenBuffer += domainLength
	default:
		return nil, nil, errors.New("Shadowsocks|TCP: Unknown address type.")
	}

	_, err = io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+2])
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to read port: " + err.Error())
	}

	request.Port = v2net.PortFromBytes(buffer.Value[lenBuffer : lenBuffer+2])
	lenBuffer += 2

	if request.Option.Has(RequestOptionOneTimeAuth) {
		authBytes := buffer.Value[lenBuffer : lenBuffer+AuthSize]
		_, err = io.ReadFull(reader, authBytes)
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read OTA: " + err.Error())
		}

		actualAuth := authenticator.Authenticate(nil, buffer.Value[0:lenBuffer])
		if !bytes.Equal(actualAuth, authBytes) {
			return nil, nil, errors.New("Shadowsocks|TCP: Invalid OTA")
		}
	}

	var chunkReader v2io.Reader
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkReader = NewChunkReader(reader, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkReader = v2io.NewAdaptiveReader(reader)
	}

	return request, chunkReader, nil
}

func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Writer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	rand.Read(iv)
	_, err = writer.Write(iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write IV: " + err.Error())
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	writer = crypto.NewCryptionWriter(stream, writer)

	header := alloc.NewLocalBuffer(512).Clear()

	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		header.AppendBytes(AddrTypeIPv4)
		header.Append([]byte(request.Address.IP()))
	case AddrTypeIPv6:
		header.AppendBytes(AddrTypeIPv6)
		header.Append([]byte(request.Address.IP()))
	case AddrTypeDomain:
		header.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		header.Append([]byte(request.Address.Domain()))
	default:
		return nil, errors.New("Shadowsocks|TCP: Unsupported address type. ")
	}

	header.AppendUint16(uint16(request.Port))

	if request.Option.Has(RequestOptionOneTimeAuth) {
		header.Value[0] |= 0x10

		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		header.Value = authenticator.Authenticate(header.Value, header.Value)
	}

	_, err = writer.Write(header.Value)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write header: " + err.Error())
	}

	var chunkWriter v2io.Writer
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkWriter = NewChunkWriter(writer, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkWriter = v2io.NewAdaptiveWriter(writer)
	}

	return chunkWriter, nil
}

func ReadTCPResponse(user *protocol.User, reader io.Reader) (v2io.Reader, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	_, err = io.ReadFull(reader, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to read IV: " + err.Error())
	}

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to initialize decoding stream: " + err.Error())
	}
	return v2io.NewAdaptiveReader(crypto.NewCryptionReader(stream, reader)), nil
}

func WriteTCPResponse(request *protocol.RequestHeader, writer io.Writer) (v2io.Writer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	rand.Read(iv)
	_, err = writer.Write(iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write IV: " + err.Error())
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	return v2io.NewAdaptiveWriter(crypto.NewCryptionWriter(stream, writer)), nil
}

func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*alloc.Buffer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|UDP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	buffer := alloc.NewLocalBuffer(2048)
	ivLen := account.Cipher.IVSize()
	buffer.Slice(0, ivLen)
	rand.Read(buffer.Value)
	iv := buffer.Value

	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		buffer.AppendBytes(AddrTypeIPv4)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyIPv6:
		buffer.AppendBytes(AddrTypeIPv6)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyDomain:
		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		buffer.Append([]byte(request.Address.Domain()))
	default:
		return nil, errors.New("Shadowsocks|UDP: Unsupported address type. ")
	}

	buffer.AppendUint16(uint16(request.Port))
	buffer.Append(payload.Value)

	if request.Option.Has(RequestOptionOneTimeAuth) {
		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		buffer.Value[ivLen] |= 0x10

		buffer.Value = authenticator.Authenticate(buffer.Value, buffer.Value[ivLen:])
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	stream.XORKeyStream(buffer.Value[ivLen:], buffer.Value[ivLen:])
	return buffer, nil
}

func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.RequestHeader, *alloc.Buffer, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|UDP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	ivLen := account.Cipher.IVSize()
	iv := payload.Value[:ivLen]
	payload.SliceFrom(ivLen)

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|UDP: Failed to initialize decoding stream: " + err.Error())
	}
	stream.XORKeyStream(payload.Value, payload.Value)

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandUDP,
	}

	addrType := (payload.Value[0] & 0x0F)
	if (payload.Value[0] & 0x10) == 0x10 {
		request.Option |= RequestOptionOneTimeAuth
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		payloadLen := payload.Len() - AuthSize
		authBytes := payload.Value[payloadLen:]

		actualAuth := authenticator.Authenticate(nil, payload.Value[0:payloadLen])
		if !bytes.Equal(actualAuth, authBytes) {
			return nil, nil, errors.New("Shadowsocks|UDP: Invalid OTA.")
		}

		payload.Slice(0, payloadLen)
	}

	payload.SliceFrom(1)

	switch addrType {
	case AddrTypeIPv4:
		request.Address = v2net.IPAddress(payload.Value[:4])
		payload.SliceFrom(4)
	case AddrTypeIPv6:
		request.Address = v2net.IPAddress(payload.Value[:16])
		payload.SliceFrom(16)
	case AddrTypeDomain:
		domainLength := int(payload.Value[0])
		request.Address = v2net.DomainAddress(string(payload.Value[1 : 1+domainLength]))
		payload.SliceFrom(1 + domainLength)
	default:
		return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type")
	}

	request.Port = v2net.PortFromBytes(payload.Value[:2])
	payload.SliceFrom(2)

	return request, payload, nil
}

type UDPReader struct {
	Reader io.Reader
	User   *protocol.User
}

func (this *UDPReader) Read() (*alloc.Buffer, error) {
	buffer := alloc.NewLocalBuffer(2048)
	nBytes, err := this.Reader.Read(buffer.Value)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	buffer.Slice(0, nBytes)
	_, payload, err := DecodeUDPPacket(this.User, buffer)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	return payload, nil
}

func (this *UDPReader) Release() {
}

type UDPWriter struct {
	Writer  io.Writer
	Request *protocol.RequestHeader
}

func (this *UDPWriter) Write(buffer *alloc.Buffer) error {
	payload, err := EncodeUDPPacket(this.Request, buffer)
	if err != nil {
		return err
	}
	_, err = this.Writer.Write(payload.Value)
	payload.Release()
	return err
}

func (this *UDPWriter) Release() {

}
<commit_msg>fix tcp encoding<commit_after>package shadowsocks

import (
	"bytes"
	"crypto/rand"
	"errors"
	"io"

	"v2ray.com/core/common/alloc"
	"v2ray.com/core/common/crypto"
	v2io "v2ray.com/core/common/io"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
)

const (
	Version                  = 1
	RequestOptionOneTimeAuth = protocol.RequestOption(101)

	AddrTypeIPv4   = 1
	AddrTypeIPv6   = 4
	AddrTypeDomain = 3
)

func ReadTCPSession(user *protocol.User, reader io.Reader) (*protocol.RequestHeader, v2io.Reader, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	buffer := alloc.NewLocalBuffer(256)
	defer buffer.Release()

	ivLen := account.Cipher.IVSize()
	_, err = io.ReadFull(reader, buffer.Value[:ivLen])
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IV: " + err.Error())
	}

	iv := append([]byte(nil), buffer.Value[:ivLen]...)

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to initialize decoding stream: " + err.Error())
	}
	reader = crypto.NewCryptionReader(stream, reader)

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandTCP,
	}

	lenBuffer := 1
	_, err = io.ReadFull(reader, buffer.Value[:1])
	if err != nil {
		return nil, nil, errors.New("Sahdowsocks|TCP: Failed to read address type: " + err.Error())
	}

	addrType := (buffer.Value[0] & 0x0F)
	if (buffer.Value[0] & 0x10) == 0x10 {
		request.Option |= RequestOptionOneTimeAuth
	}

	switch addrType {
	case AddrTypeIPv4:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+4])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IPv4 address: " + err.Error())
		}
		request.Address = v2net.IPAddress(buffer.Value[lenBuffer : lenBuffer+4])
		lenBuffer += 4
	case AddrTypeIPv6:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+16])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read IPv6 address: " + err.Error())
		}
		request.Address = v2net.IPAddress(buffer.Value[lenBuffer : lenBuffer+16])
		lenBuffer += 16
	case AddrTypeDomain:
		_, err := io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+1])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read domain lenth: " + err.Error())
		}
		domainLength := int(buffer.Value[lenBuffer])
		lenBuffer++
		_, err = io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+domainLength])
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read domain: " + err.Error())
		}
		request.Address = v2net.DomainAddress(string(buffer.Value[lenBuffer : lenBuffer+domainLength]))
		lenBuffer += domainLength
	default:
		return nil, nil, errors.New("Shadowsocks|TCP: Unknown address type.")
	}

	_, err = io.ReadFull(reader, buffer.Value[lenBuffer:lenBuffer+2])
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|TCP: Failed to read port: " + err.Error())
	}

	request.Port = v2net.PortFromBytes(buffer.Value[lenBuffer : lenBuffer+2])
	lenBuffer += 2

	if request.Option.Has(RequestOptionOneTimeAuth) {
		authBytes := buffer.Value[lenBuffer : lenBuffer+AuthSize]
		_, err = io.ReadFull(reader, authBytes)
		if err != nil {
			return nil, nil, errors.New("Shadowsocks|TCP: Failed to read OTA: " + err.Error())
		}

		actualAuth := authenticator.Authenticate(nil, buffer.Value[0:lenBuffer])
		if !bytes.Equal(actualAuth, authBytes) {
			return nil, nil, errors.New("Shadowsocks|TCP: Invalid OTA")
		}
	}

	var chunkReader v2io.Reader
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkReader = NewChunkReader(reader, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkReader = v2io.NewAdaptiveReader(reader)
	}

	return request, chunkReader, nil
}

func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (v2io.Writer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	rand.Read(iv)
	_, err = writer.Write(iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write IV: " + err.Error())
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	writer = crypto.NewCryptionWriter(stream, writer)

	header := alloc.NewLocalBuffer(512).Clear()

	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		header.AppendBytes(AddrTypeIPv4)
		header.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyIPv6:
		header.AppendBytes(AddrTypeIPv6)
		header.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyDomain:
		header.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		header.Append([]byte(request.Address.Domain()))
	default:
		return nil, errors.New("Shadowsocks|TCP: Unsupported address type. ")
	}

	header.AppendUint16(uint16(request.Port))

	if request.Option.Has(RequestOptionOneTimeAuth) {
		header.Value[0] |= 0x10

		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		header.Value = authenticator.Authenticate(header.Value, header.Value)
	}

	_, err = writer.Write(header.Value)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write header: " + err.Error())
	}

	var chunkWriter v2io.Writer
	if request.Option.Has(RequestOptionOneTimeAuth) {
		chunkWriter = NewChunkWriter(writer, NewAuthenticator(ChunkKeyGenerator(iv)))
	} else {
		chunkWriter = v2io.NewAdaptiveWriter(writer)
	}

	return chunkWriter, nil
}

func ReadTCPResponse(user *protocol.User, reader io.Reader) (v2io.Reader, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	_, err = io.ReadFull(reader, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to read IV: " + err.Error())
	}

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to initialize decoding stream: " + err.Error())
	}
	return v2io.NewAdaptiveReader(crypto.NewCryptionReader(stream, reader)), nil
}

func WriteTCPResponse(request *protocol.RequestHeader, writer io.Writer) (v2io.Writer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	iv := make([]byte, account.Cipher.IVSize())
	rand.Read(iv)
	_, err = writer.Write(iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to write IV: " + err.Error())
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	return v2io.NewAdaptiveWriter(crypto.NewCryptionWriter(stream, writer)), nil
}

func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*alloc.Buffer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|UDP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	buffer := alloc.NewLocalBuffer(2048)
	ivLen := account.Cipher.IVSize()
	buffer.Slice(0, ivLen)
	rand.Read(buffer.Value)
	iv := buffer.Value

	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		buffer.AppendBytes(AddrTypeIPv4)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyIPv6:
		buffer.AppendBytes(AddrTypeIPv6)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyDomain:
		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		buffer.Append([]byte(request.Address.Domain()))
	default:
		return nil, errors.New("Shadowsocks|UDP: Unsupported address type. ")
	}

	buffer.AppendUint16(uint16(request.Port))
	buffer.Append(payload.Value)

	if request.Option.Has(RequestOptionOneTimeAuth) {
		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		buffer.Value[ivLen] |= 0x10

		buffer.Value = authenticator.Authenticate(buffer.Value, buffer.Value[ivLen:])
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.New("Shadowsocks|TCP: Failed to create encoding stream: " + err.Error())
	}

	stream.XORKeyStream(buffer.Value[ivLen:], buffer.Value[ivLen:])
	return buffer, nil
}

func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.RequestHeader, *alloc.Buffer, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|UDP: Failed to parse account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	ivLen := account.Cipher.IVSize()
	iv := payload.Value[:ivLen]
	payload.SliceFrom(ivLen)

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, nil, errors.New("Shadowsocks|UDP: Failed to initialize decoding stream: " + err.Error())
	}
	stream.XORKeyStream(payload.Value, payload.Value)

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandUDP,
	}

	addrType := (payload.Value[0] & 0x0F)
	if (payload.Value[0] & 0x10) == 0x10 {
		request.Option |= RequestOptionOneTimeAuth
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		payloadLen := payload.Len() - AuthSize
		authBytes := payload.Value[payloadLen:]

		actualAuth := authenticator.Authenticate(nil, payload.Value[0:payloadLen])
		if !bytes.Equal(actualAuth, authBytes) {
			return nil, nil, errors.New("Shadowsocks|UDP: Invalid OTA.")
		}

		payload.Slice(0, payloadLen)
	}

	payload.SliceFrom(1)

	switch addrType {
	case AddrTypeIPv4:
		request.Address = v2net.IPAddress(payload.Value[:4])
		payload.SliceFrom(4)
	case AddrTypeIPv6:
		request.Address = v2net.IPAddress(payload.Value[:16])
		payload.SliceFrom(16)
	case AddrTypeDomain:
		domainLength := int(payload.Value[0])
		request.Address = v2net.DomainAddress(string(payload.Value[1 : 1+domainLength]))
		payload.SliceFrom(1 + domainLength)
	default:
		return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type")
	}

	request.Port = v2net.PortFromBytes(payload.Value[:2])
	payload.SliceFrom(2)

	return request, payload, nil
}

type UDPReader struct {
	Reader io.Reader
	User   *protocol.User
}

func (this *UDPReader) Read() (*alloc.Buffer, error) {
	buffer := alloc.NewLocalBuffer(2048)
	nBytes, err := this.Reader.Read(buffer.Value)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	buffer.Slice(0, nBytes)
	_, payload, err := DecodeUDPPacket(this.User, buffer)
	if err != nil {
		buffer.Release()
		return nil, err
	}
	return payload, nil
}

func (this *UDPReader) Release() {
}

type UDPWriter struct {
	Writer  io.Writer
	Request *protocol.RequestHeader
}

func (this *UDPWriter) Write(buffer *alloc.Buffer) error {
	payload, err := EncodeUDPPacket(this.Request, buffer)
	if err != nil {
		return err
	}
	_, err = this.Writer.Write(payload.Value)
	payload.Release()
	return err
}

func (this *UDPWriter) Release() {

}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "os"
  "flag"
  "strings"
  "gopkg.in/ini.v1"
  "sync"
  "os/exec"
  "time"
  "golang.org/x/net/context"
  "encoding/json"
  "os/signal"
  "syscall"
  "strconv"
  "io/ioutil"
  "net/http"

  "github.com/prometheus/client_golang/prometheus"
  "github.com/prometheus/common/log"
  "github.com/prometheus/client_golang/prometheus/promhttp"
  "github.com/prometheus/common/version"
)

const (
  namespace = "phpfpm"
)

type FpmPoolMetrics struct {
  StartTime int `json:"start time"`
  StartSince int `json:"start since"`
  AcceptedConn int `json:"accepted conn"`
  ListenQueue int `json:"listen queue"`
  MaxListenQueue int `json:"max listen queue"`
  ListenQueueLen int `json:"listen queue len"`
  IdleProcesses int `json:"idle processes"`
  ActiveProcesses int `json:"active processes"`
  TotalProcesses int `json:"total processes"`
  MaxActiveProcesses int `json:"max active processes"`
  MaxChildrenReached int `json:"max children reached"`
  SlowRequests int `json:"slow requests"`
}

type PhpFpmPool struct {
  Name string
  Endpoint string
  StatusUri string
  LastMetrics FpmPoolMetrics
  mu sync.RWMutex
}

type PhpFpmPoolExporter struct {
  poolsToMonitor []*PhpFpmPool
  listenQueue, listenQueueLen, idleProcesses, activeProcesses, totalProcesses *prometheus.GaugeVec
  startSince, acceptedConn, maxListenQueue, maxActiveProcesses, maxChildrenReached, slowRequests *prometheus.CounterVec
}

func (e *PhpFpmPoolExporter) resetMetrics() {
  e.listenQueue.Reset()
  e.listenQueueLen.Reset()
  e.idleProcesses.Reset()
  e.activeProcesses.Reset()
  e.totalProcesses.Reset()
  e.startSince.Reset()
  e.acceptedConn.Reset()
  e.maxListenQueue.Reset()
  e.maxActiveProcesses.Reset()
  e.maxChildrenReached.Reset()
  e.slowRequests.Reset()
}

func (p *PhpFpmPool) GetSyncedCopy() PhpFpmPool {
  p.mu.Lock()
  pfp := p
  p.mu.Unlock()

  return *pfp
}

func (p *PhpFpmPool) PushSyncedLastMetrics(fpm *FpmPoolMetrics) {
  p.mu.Lock()
  p.LastMetrics = *fpm
  p.mu.Unlock()
}

func (p *PhpFpmPool) GetSyncedLastMetricsCopy() FpmPoolMetrics {
  p.mu.Lock()
  lm := &(p).LastMetrics
  p.mu.Unlock()

  return *lm
}

func (e *PhpFpmPoolExporter) Describe(ch chan<- *prometheus.Desc) {
  e.listenQueue.Describe(ch)
  e.listenQueueLen.Describe(ch)
  e.idleProcesses.Describe(ch)
  e.activeProcesses.Describe(ch)
  e.totalProcesses.Describe(ch)
  e.startSince.Describe(ch)
  e.acceptedConn.Describe(ch)
  e.maxListenQueue.Describe(ch)
  e.maxActiveProcesses.Describe(ch)
  e.maxChildrenReached.Describe(ch)
  e.slowRequests.Describe(ch)
}

func (e *PhpFpmPoolExporter) Collect(ch chan<- prometheus.Metric) {
  e.resetMetrics()
  for _, p := range e.poolsToMonitor {
    lastMetric := p.GetSyncedLastMetricsCopy()

    (e.listenQueue.WithLabelValues(p.Name)).Set(float64(lastMetric.ListenQueue))
    (e.listenQueueLen.WithLabelValues(p.Name)).Set(float64(lastMetric.ListenQueueLen))
    (e.idleProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.IdleProcesses))
    (e.activeProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.ActiveProcesses))
    (e.totalProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.TotalProcesses))
    (e.startSince.WithLabelValues(p.Name)).Set(float64(lastMetric.StartSince))
    (e.acceptedConn.WithLabelValues(p.Name)).Set(float64(lastMetric.AcceptedConn))
    (e.maxListenQueue.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxListenQueue))
    (e.maxActiveProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxActiveProcesses))
    (e.maxChildrenReached.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxChildrenReached))
    (e.slowRequests.WithLabelValues(p.Name)).Set(float64(lastMetric.SlowRequests))

    (e.listenQueue.WithLabelValues(p.Name)).Collect(ch)
    (e.listenQueueLen.WithLabelValues(p.Name)).Collect(ch)
    (e.idleProcesses.WithLabelValues(p.Name)).Collect(ch)
    (e.activeProcesses.WithLabelValues(p.Name)).Collect(ch)
    (e.totalProcesses.WithLabelValues(p.Name)).Collect(ch)
    (e.startSince.WithLabelValues(p.Name)).Collect(ch)
    (e.acceptedConn.WithLabelValues(p.Name)).Collect(ch)
    (e.maxListenQueue.WithLabelValues(p.Name)).Collect(ch)
    (e.maxActiveProcesses.WithLabelValues(p.Name)).Collect(ch)
    (e.maxChildrenReached.WithLabelValues(p.Name)).Collect(ch)
    (e.slowRequests.WithLabelValues(p.Name)).Collect(ch)
  }
}

func NewPhpFpmPoolExporter(pools []*PhpFpmPool) (*PhpFpmPoolExporter) {
  poolLabelNames := []string{"pool_name"}

  return &PhpFpmPoolExporter {
    poolsToMonitor: pools,
    startSince: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "start_since",
        Help: "Number of seconds since FPM has started",
      },
      poolLabelNames,
    ),
    acceptedConn: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "accepted_conn",
        Help: "The number of requests accepted by the pool",
      },
      poolLabelNames,
    ),
    listenQueue: prometheus.NewGaugeVec(
      prometheus.GaugeOpts{
        Namespace: namespace,
        Name: "listen_queue",
        Help: "The number of requests in the queue of pending connections",
      },
      poolLabelNames,
    ),
    maxListenQueue: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "max_listen_queue",
        Help: "The maximum number of requests in the queue of pending connections since FPM has started",
      },
      poolLabelNames,
    ),
    listenQueueLen: prometheus.NewGaugeVec(
      prometheus.GaugeOpts{
        Namespace: namespace,
        Name: "listen_queue_len",
        Help: "The size of the socket queue of pending connections",
      },
      poolLabelNames,
    ),
    idleProcesses: prometheus.NewGaugeVec(
      prometheus.GaugeOpts{
        Namespace: namespace,
        Name: "idle_processes",
        Help: "The number of idle processes",
      },
      poolLabelNames,
    ),
    activeProcesses: prometheus.NewGaugeVec(
      prometheus.GaugeOpts{
        Namespace: namespace,
        Name: "active_processes",
        Help: "The number of active processes",
      },
      poolLabelNames,
    ),
    totalProcesses: prometheus.NewGaugeVec(
      prometheus.GaugeOpts{
        Namespace: namespace,
        Name: "total_processes",
        Help: "The number of idle + active processes",
      },
      poolLabelNames,
    ),
    maxActiveProcesses: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "max_active_processes",
        Help: "The maximum number of active processes since FPM has started",
      },
      poolLabelNames,
    ),
    maxChildrenReached: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "max_children_reached",
        Help: "The number of times, the process limit has been reached, when pm tries to start more children (works only for pm 'dynamic' and 'ondemand')",
      },
      poolLabelNames,
    ),
    slowRequests: prometheus.NewCounterVec(
      prometheus.CounterOpts{
        Namespace: namespace,
        Name: "slow_requests",
        Help: "The number of requests that exceeded your request_slowlog_timeout value",
      },
      poolLabelNames,
    ),
  }
}

func GetFilesIn(dirPath string) []string {
  var poolFiles []string

  if strings.HasSuffix(dirPath, "/") {
    dirPath = strings.TrimRight(dirPath, "/")
  }

  dir, err := os.Open(dirPath)

  if err != nil {
    fmt.Errorf("%s", err)
    return nil
  }
  defer dir.Close()

  filesInfo, err := dir.Readdir(-1)

  if err != nil {
    fmt.Errorf("%s", err)
    return nil
  }

  for i := 0; i < len(filesInfo); i++ {
    if (filesInfo[i].Mode().IsRegular()) {
      poolFiles = append(poolFiles, dirPath + "/" + filesInfo[i].Name())
    }
  }

  return poolFiles
}

func PollFpmStatusMetrics(p *PhpFpmPool, pollInterval int, pollTimeout int, mustQuit chan bool, done chan bool) {
  poolCpy := p.GetSyncedCopy()

  env := os.Environ()
  env = append(env, fmt.Sprintf("SCRIPT_NAME=%s",  poolCpy.StatusUri))
  env = append(env, fmt.Sprintf("SCRIPT_FILENAME=%s",  poolCpy.StatusUri))
  env = append(env, "QUERY_STRING=json")
  env = append(env, "REQUEST_METHOD=GET")

  var data []byte
  var err error
  var mts FpmPoolMetrics
  var strData []string

  for i := 0; i < 1; {
    ctx := context.TODO()
    ctxWithCancel, cancel := context.WithTimeout(ctx, time.Duration(pollTimeout*int(time.Second)))
    defer cancel()

    cmd := exec.CommandContext(ctxWithCancel, "/usr/bin/cgi-fcgi", "-bind", "-connect", poolCpy.Endpoint)
    cmd.Env = env
    data, err = cmd.Output()
    if (err != nil) {
      log.Errorln(err.Error())
      continue
    } else {

      strData = strings.SplitAfter(string(data), "\r\n\r\n")
      if (len(strData) < 2) {
        log.Errorln("Unexpected cgi-fcgi response")
        continue
      }

      err := json.Unmarshal([]byte(strData[1]), &mts)
      if err != nil {
        log.Errorln(err.Error())
        continue
      }

      p.PushSyncedLastMetrics(&mts)
    }

    time.Sleep(time.Duration(pollInterval*int(time.Second)))
    select {
      case <-mustQuit:
        i = 1
        log.Infoln("Goroutine received signal asking to quit")
        done <- true
        return
      default:
        continue
    }
  }
  return
}

func main() {
  var (
    listenAddress = flag.String("web.listen-address", ":9101", "Address to listen on for web interface and telemetry.")
    metricsPath = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
    phpfpmPidFile = flag.String("phpfpm.pid-file", "/var/run/php5-fpm.pid", "Path to phpfpm's pid file.")
    configDir = flag.String("config", "/etc/php5/fpm/pool.d/", "Pools conf dir")
    pollInterval = flag.Int("poll_interval", 10, "Poll interval in seconds")
    pollTimeout = flag.Int("poll_timeout", 2, "Poll timeout in seconds")
    showVersion = flag.Bool("version", false, "Print version information.")
  )

  flag.Parse()

  if *showVersion {
    fmt.Fprintln(os.Stdout, version.Print("phpfpm_prometheus_exporter"))
    os.Exit(0)
  }

  log.Infoln("Starting phpfpm_exporter", version.Info())
  log.Infoln("Build context", version.BuildContext())

  //TODO:
  // - Test requirement: cgi-fcgi

  if *phpfpmPidFile != "" {
    procExporter := prometheus.NewProcessCollectorPIDFn(
      func() (int, error) {
        content, err := ioutil.ReadFile(*phpfpmPidFile)
        if err != nil {
          return 0, fmt.Errorf("Can't read pid file: %s", err)
        }
        value, err := strconv.Atoi(strings.TrimSpace(string(content)))
        if err != nil {
          return 0, fmt.Errorf("Can't parse pid file: %s", err)
        }
        return value, nil
      }, namespace)
      prometheus.MustRegister(procExporter)
  }

  sigs := make(chan os.Signal)
  mustQuit := make(chan bool)
  done := make(chan bool)

  signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

  phpFpmPools := []*PhpFpmPool{}

  confFiles := GetFilesIn(*configDir)
  cfg := ini.Empty()

  for _, cf := range confFiles {
    log.Infoln("We will parse: ", cf)
    err := cfg.Append(cf)

    if err != nil {
      fmt.Errorf("%f")
    }
  }

  sections := cfg.SectionStrings()
  sectionsCount := 0

  for _, sect := range sections {
    statusKey, err := cfg.Section(sect).GetKey("pm.status_path")

    if err != nil {
      continue
    }

    listenKey, err := cfg.Section(sect).GetKey("listen")

    if err != nil {
      continue
    }

    pool := PhpFpmPool{Name: sect, Endpoint: listenKey.String(), StatusUri: statusKey.String()}

    sectionsCount++

    go PollFpmStatusMetrics(&pool, *pollInterval, *pollTimeout, mustQuit, done)

    phpFpmPools = append(phpFpmPools, &pool)
  }

  log.Infoln("We will monitor ", sectionsCount, " phpfpm pool(s)")

  phpFpmExporter := NewPhpFpmPoolExporter(phpFpmPools)

  prometheus.MustRegister(phpFpmExporter)
  prometheus.MustRegister(version.NewCollector("phpfpm_exporter"))

  log.Infoln("Listening on", *listenAddress)
  http.Handle(*metricsPath, promhttp.Handler())
  http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    w.Write([]byte(`
      <title>PhpFpm Exporter</title>

PhpFpm Exporter

<a href='` + *metricsPath + `'>Metrics

      `))
  })
  //log.Fatal(http.ListenAndServe(*listenAddress, nil))
  go http.ListenAndServe(*listenAddress, nil)

  log.Infoln("Awaiting quit signal")

  <-sigs

  for j:=0; j<sectionsCount; j++ {
    mustQuit <- true
  }

  log.Infoln("Awaiting all done signals")

  for j:=0; j<sectionsCount; j++ {
    <-done
  }
  close(mustQuit)
  close(done);
  log.Infoln("Clean shutdown!")
}
<commit_msg>Add cgi-fcgi bin configurability, format code<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"golang.org/x/net/context"
	"gopkg.in/ini.v1"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

const (
	namespace = "phpfpm"
)

type FpmPoolMetrics struct {
	StartTime          int `json:"start time"`
	StartSince         int `json:"start since"`
	AcceptedConn       int `json:"accepted conn"`
	ListenQueue        int `json:"listen queue"`
	MaxListenQueue     int `json:"max listen queue"`
	ListenQueueLen     int `json:"listen queue len"`
	IdleProcesses      int `json:"idle processes"`
	ActiveProcesses    int `json:"active processes"`
	TotalProcesses     int `json:"total processes"`
	MaxActiveProcesses int `json:"max active processes"`
	MaxChildrenReached int `json:"max children reached"`
	SlowRequests       int `json:"slow requests"`
}

type PhpFpmPool struct {
	Name        string
	Endpoint    string
	StatusUri   string
	lastMetrics FpmPoolMetrics
	mu          sync.RWMutex
}

type PhpFpmPoolExporter struct {
	poolsToMonitor                                                                                []*PhpFpmPool
	listenQueue, listenQueueLen, idleProcesses, activeProcesses, totalProcesses                    *prometheus.GaugeVec
	startSince, acceptedConn, maxListenQueue, maxActiveProcesses, maxChildrenReached, slowRequests *prometheus.CounterVec
}

func (e *PhpFpmPoolExporter) resetMetrics() {
	e.listenQueue.Reset()
	e.listenQueueLen.Reset()
	e.idleProcesses.Reset()
	e.activeProcesses.Reset()
	e.totalProcesses.Reset()
	e.startSince.Reset()
	e.acceptedConn.Reset()
	e.maxListenQueue.Reset()
	e.maxActiveProcesses.Reset()
	e.maxChildrenReached.Reset()
	e.slowRequests.Reset()
}

func (p *PhpFpmPool) GetSyncedCopy() PhpFpmPool {
	p.mu.Lock()
	pfp := p
	p.mu.Unlock()

	return *pfp
}

func (p *PhpFpmPool) PushSyncedLastMetrics(fpm *FpmPoolMetrics) {
	p.mu.Lock()
	p.lastMetrics = *fpm
	p.mu.Unlock()
}

func (p *PhpFpmPool) GetSyncedLastMetricsCopy() FpmPoolMetrics {
	p.mu.Lock()
	lm := &(p).lastMetrics
	p.mu.Unlock()

	return *lm
}

func (e *PhpFpmPoolExporter) Describe(ch chan<- *prometheus.Desc) {
	e.listenQueue.Describe(ch)
	e.listenQueueLen.Describe(ch)
	e.idleProcesses.Describe(ch)
	e.activeProcesses.Describe(ch)
	e.totalProcesses.Describe(ch)
	e.startSince.Describe(ch)
	e.acceptedConn.Describe(ch)
	e.maxListenQueue.Describe(ch)
	e.maxActiveProcesses.Describe(ch)
	e.maxChildrenReached.Describe(ch)
	e.slowRequests.Describe(ch)
}

func (e *PhpFpmPoolExporter) Collect(ch chan<- prometheus.Metric) {
	e.resetMetrics()
	for _, p := range e.poolsToMonitor {
		lastMetric := p.GetSyncedLastMetricsCopy()

		(e.listenQueue.WithLabelValues(p.Name)).Set(float64(lastMetric.ListenQueue))
		(e.listenQueueLen.WithLabelValues(p.Name)).Set(float64(lastMetric.ListenQueueLen))
		(e.idleProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.IdleProcesses))
		(e.activeProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.ActiveProcesses))
		(e.totalProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.TotalProcesses))
		(e.startSince.WithLabelValues(p.Name)).Set(float64(lastMetric.StartSince))
		(e.acceptedConn.WithLabelValues(p.Name)).Set(float64(lastMetric.AcceptedConn))
		(e.maxListenQueue.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxListenQueue))
		(e.maxActiveProcesses.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxActiveProcesses))
		(e.maxChildrenReached.WithLabelValues(p.Name)).Set(float64(lastMetric.MaxChildrenReached))
		(e.slowRequests.WithLabelValues(p.Name)).Set(float64(lastMetric.SlowRequests))

		(e.listenQueue.WithLabelValues(p.Name)).Collect(ch)
		(e.listenQueueLen.WithLabelValues(p.Name)).Collect(ch)
		(e.idleProcesses.WithLabelValues(p.Name)).Collect(ch)
		(e.activeProcesses.WithLabelValues(p.Name)).Collect(ch)
		(e.totalProcesses.WithLabelValues(p.Name)).Collect(ch)
		(e.startSince.WithLabelValues(p.Name)).Collect(ch)
		(e.acceptedConn.WithLabelValues(p.Name)).Collect(ch)
		(e.maxListenQueue.WithLabelValues(p.Name)).Collect(ch)
		(e.maxActiveProcesses.WithLabelValues(p.Name)).Collect(ch)
		(e.maxChildrenReached.WithLabelValues(p.Name)).Collect(ch)
		(e.slowRequests.WithLabelValues(p.Name)).Collect(ch)
	}
}

func NewPhpFpmPoolExporter(pools []*PhpFpmPool) *PhpFpmPoolExporter {
	poolLabelNames := []string{"pool_name"}

	return &PhpFpmPoolExporter{
		poolsToMonitor: pools,
		startSince: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "start_since",
				Help:      "Number of seconds since FPM has started",
			},
			poolLabelNames,
		),
		acceptedConn: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "accepted_conn",
				Help:      "The number of requests accepted by the pool",
			},
			poolLabelNames,
		),
		listenQueue: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "listen_queue",
				Help:      "The number of requests in the queue of pending connections",
			},
			poolLabelNames,
		),
		maxListenQueue: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "max_listen_queue",
				Help:      "The maximum number of requests in the queue of pending connections since FPM has started",
			},
			poolLabelNames,
		),
		listenQueueLen: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "listen_queue_len",
				Help:      "The size of the socket queue of pending connections",
			},
			poolLabelNames,
		),
		idleProcesses: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "idle_processes",
				Help:      "The number of idle processes",
			},
			poolLabelNames,
		),
		activeProcesses: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "active_processes",
				Help:      "The number of active processes",
			},
			poolLabelNames,
		),
		totalProcesses: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "total_processes",
				Help:      "The number of idle + active processes",
			},
			poolLabelNames,
		),
		maxActiveProcesses: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "max_active_processes",
				Help:      "The maximum number of active processes since FPM has started",
			},
			poolLabelNames,
		),
		maxChildrenReached: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "max_children_reached",
				Help:      "The number of times, the process limit has been reached, when pm tries to start more children (works only for pm 'dynamic' and 'ondemand')",
			},
			poolLabelNames,
		),
		slowRequests: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "slow_requests",
				Help:      "The number of requests that exceeded your request_slowlog_timeout value",
			},
			poolLabelNames,
		),
	}
}

func GetFilesIn(dirPath string) []string {
	var poolFiles []string

	if strings.HasSuffix(dirPath, "/") {
		dirPath = strings.TrimRight(dirPath, "/")
	}

	dir, err := os.Open(dirPath)

	if err != nil {
		fmt.Errorf("%s", err)
		return nil
	}
	defer dir.Close()

	filesInfo, err := dir.Readdir(-1)

	if err != nil {
		fmt.Errorf("%s", err)
		return nil
	}

	for i := 0; i < len(filesInfo); i++ {
		if filesInfo[i].Mode().IsRegular() {
			poolFiles = append(poolFiles, dirPath+"/"+filesInfo[i].Name())
		}
	}

	return poolFiles
}

func PollFpmStatusMetrics(p *PhpFpmPool, pollInterval int, pollTimeout int, cgiFastCgiPath string, cgiFastCgiLdLibPath string, mustQuit chan bool, done chan bool) {
	poolCpy := p.GetSyncedCopy()

	env := os.Environ()

	if cgiFastCgiLdLibPath != "" {
		env = append(env, fmt.Sprintf("LD_LIBRARY_PATH=%s", cgiFastCgiLdLibPath))
	}

	env = append(env, fmt.Sprintf("SCRIPT_NAME=%s", poolCpy.StatusUri))
	env = append(env, fmt.Sprintf("SCRIPT_FILENAME=%s", poolCpy.StatusUri))
	env = append(env, "QUERY_STRING=json")
	env = append(env, "REQUEST_METHOD=GET")

	var data []byte
	var err error
	var mts FpmPoolMetrics
	var strData []string

	for i := 0; i < 1; {
		ctx := context.TODO()
		ctxWithCancel, cancel := context.WithTimeout(ctx, time.Duration(pollTimeout*int(time.Second)))
		defer cancel()

		cmd := exec.CommandContext(ctxWithCancel, cgiFastCgiPath, "-bind", "-connect", poolCpy.Endpoint)
		cmd.Env = env
		data, err = cmd.Output()
		if err != nil {
			log.Errorln(err.Error())
			continue
		} else {

			strData = strings.SplitAfter(string(data), "\r\n\r\n")
			if len(strData) < 2 {
				log.Errorln("Unexpected cgi-fcgi response")
				continue
			}

			err := json.Unmarshal([]byte(strData[1]), &mts)
			if err != nil {
				log.Errorln(err.Error())
				continue
			}

			p.PushSyncedLastMetrics(&mts)
		}

		time.Sleep(time.Duration(pollInterval * int(time.Second)))
		select {
		case <-mustQuit:
			i = 1
			log.Infoln("Goroutine received signal asking to quit")
			done <- true
			return
		default:
			continue
		}
	}
	return
}

func main() {
	var (
		listenAddress       = flag.String("web.listen-address", ":9101", "Address to listen on for web interface and telemetry.")
		metricsPath         = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
		phpfpmPidFile       = flag.String("phpfpm.pid-file", "/var/run/php5-fpm.pid", "Path to phpfpm's pid file.")
		configDir           = flag.String("phpfpm.config", "/etc/php5/fpm/pool.d/", "Pools conf dir")
		pollInterval        = flag.Int("phpfpm.poll-interval", 10, "Poll interval in seconds")
		pollTimeout         = flag.Int("phpfpm.poll-timeout", 2, "Poll timeout in seconds")
		cgiFastCgiPath      = flag.String("cgi-fcgi.path", "/usr/bin/cgi-fcgi", "cgi-fcgi program path")
		cgiFastCgiLdLibPath = flag.String("cgi-fcgi.ld-library-path", "", "LD_LIBRARY_PATH value to run cgi-fcgi")
		showVersion         = flag.Bool("version", false, "Print version information.")
	)

	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("phpfpm_prometheus_exporter"))
		os.Exit(0)
	}

	log.Infoln("Starting phpfpm_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	//TODO:
	// - Test requirement: cgi-fcgi

	if *phpfpmPidFile != "" {
		procExporter := prometheus.NewProcessCollectorPIDFn(
			func() (int, error) {
				content, err := ioutil.ReadFile(*phpfpmPidFile)
				if err != nil {
					return 0, fmt.Errorf("Can't read pid file: %s", err)
				}
				value, err := strconv.Atoi(strings.TrimSpace(string(content)))
				if err != nil {
					return 0, fmt.Errorf("Can't parse pid file: %s", err)
				}
				return value, nil
			}, namespace)
		prometheus.MustRegister(procExporter)
	}

	sigs := make(chan os.Signal)
	mustQuit := make(chan bool)
	done := make(chan bool)

	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	phpFpmPools := []*PhpFpmPool{}

	confFiles := GetFilesIn(*configDir)
	cfg := ini.Empty()

	for _, cf := range confFiles {
		log.Infoln("We will parse: ", cf)
		err := cfg.Append(cf)

		if err != nil {
			fmt.Errorf("%f")
		}
	}

	sections := cfg.SectionStrings()
	sectionsCount := 0

	for _, sect := range sections {
		statusKey, err := cfg.Section(sect).GetKey("pm.status_path")

		if err != nil {
			continue
		}

		listenKey, err := cfg.Section(sect).GetKey("listen")

		if err != nil {
			continue
		}

		pool := PhpFpmPool{Name: sect, Endpoint: listenKey.String(), StatusUri: statusKey.String()}

		sectionsCount++

		go PollFpmStatusMetrics(&pool, *pollInterval, *pollTimeout, *cgiFastCgiPath, *cgiFastCgiLdLibPath, mustQuit, done)

		phpFpmPools = append(phpFpmPools, &pool)
	}

	log.Infoln("We will monitor ", sectionsCount, " phpfpm pool(s)")

	phpFpmExporter := NewPhpFpmPoolExporter(phpFpmPools)

	prometheus.MustRegister(phpFpmExporter)
	prometheus.MustRegister(version.NewCollector("phpfpm_exporter"))

	log.Infoln("Listening on", *listenAddress)
	http.Handle(*metricsPath, promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
      <title>PhpFpm Exporter</title>

PhpFpm Exporter

<a href='` + *metricsPath + `'>Metrics

      `))
	})
	//log.Fatal(http.ListenAndServe(*listenAddress, nil))
	go http.ListenAndServe(*listenAddress, nil)

	log.Infoln("Awaiting quit signal")

	<-sigs

	for j := 0; j < sectionsCount; j++ {
		mustQuit <- true
	}

	log.Infoln("Awaiting all done signals")

	for j := 0; j < sectionsCount; j++ {
		<-done
	}
	close(mustQuit)
	close(done)
	log.Infoln("Clean shutdown!")
}
<|endoftext|>


<commit_before>/*
Copyright 2013 The Camlistore 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 android contains code specific to running the Camlistore client
// code as a child process on Android. This removes ugly API from the
// client package itself.
package android

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"runtime/debug"
	"strconv"
	"sync"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/osutil"
	"camlistore.org/pkg/schema"
)

// TODO(mpl): distinguish CAMPUT, CAMGET, etc
var androidOutput, _ = strconv.ParseBool(os.Getenv("CAMPUT_ANDROID_OUTPUT"))

// IsChild reports whether this process is running as an Android
// child process and should report its output in the form that the
// Android uploader app expects.
func IsChild() bool {
	memOnce.Do(startMemGoroutine)
	return androidOutput
}

var androidOutMu sync.Mutex

func Printf(format string, args ...interface{}) {
	androidOutMu.Lock()
	defer androidOutMu.Unlock()
	fmt.Printf(format, args...)
}

var detectOnce sync.Once
var onAndroidCache bool

func dirExists(f string) bool {
	fi, err := os.Stat(f)
	return err == nil && fi.IsDir()
}

func initOnAndroid() {
	// Good enough heuristic. Suggestions welcome.
	onAndroidCache = dirExists("/data/data") && dirExists("/system/etc")
}

func OnAndroid() bool {
	detectOnce.Do(initOnAndroid)
	return onAndroidCache
}

var pingRx = regexp.MustCompile(`\((.+?)\)`)

type namedInt struct {
	name string
	sync.Mutex
	val int64
}

func (ni *namedInt) Incr(delta int64) {
	ni.Lock()
	ni.val += delta
	nv := ni.val
	ni.Unlock()
	Printf("STAT %s %d\n", ni.name, nv)
}

func (ni *namedInt) Set(v int64) {
	ni.Lock()
	if v == ni.val {
		ni.Unlock()
		return
	}
	ni.val = v
	ni.Unlock()
	Printf("STAT %s %d\n", ni.name, v)
}

var (
	statDNSStart       = &namedInt{name: "dns_start"}
	statDNSDone        = &namedInt{name: "dns_done"}
	statTCPStart       = &namedInt{name: "tcp_start"}
	statTCPStarted     = &namedInt{name: "tcp_started"}
	statTCPFail        = &namedInt{name: "tcp_fail"}
	statTCPDone        = &namedInt{name: "tcp_done"}
	statTCPWrites      = &namedInt{name: "tcp_write_byte"}
	statTCPWrote       = &namedInt{name: "tcp_wrote_byte"}
	statTCPReads       = &namedInt{name: "tcp_read_byte"}
	statHTTPStart      = &namedInt{name: "http_start"}
	statHTTPResHeaders = &namedInt{name: "http_res_headers"}
	statBlobUploaded   = &namedInt{name: "blob_uploaded"}
	statBlobExisted    = &namedInt{name: "blob_existed"}
	statFileUploaded   = &namedInt{name: "file_uploaded"}
	statFileExisted    = &namedInt{name: "file_existed"}
	statMemReleased    = &namedInt{name: "mem_heap_released"}
	statMemAlloc       = &namedInt{name: "mem_alloc"}
	statMemRSS         = &namedInt{name: "mem_rss"}
)

type statTrackingConn struct {
	net.Conn
}

func (c statTrackingConn) Write(p []byte) (n int, err error) {
	statTCPWrites.Incr(int64(len(p)))
	n, err = c.Conn.Write(p)
	statTCPWrote.Incr(int64(n))
	return
}

func (c statTrackingConn) Read(p []byte) (n int, err error) {
	n, err = c.Conn.Read(p)
	statTCPReads.Incr(int64(n))
	return
}

var (
	dnsMu    sync.Mutex
	dnsCache = make(map[string]string)
)

func androidLookupHost(host string) string {
	dnsMu.Lock()
	v, ok := dnsCache[host]
	dnsMu.Unlock()
	if ok {
		return v
	}
	statDNSStart.Incr(1)
	defer statDNSDone.Incr(1)

	// Android has no "dig" or "host" tool, so use "ping -c 1". Ghetto.
	// $ ping -c 1 google.com
	// PING google.com (74.125.224.64) 56(84) bytes of data.
	c := make(chan string, 1)
	go func() {
		cmd := exec.Command("/system/bin/ping", "-c", "1", host)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panic(err)
		}
		defer stdout.Close()
		if err := cmd.Start(); err != nil {
			log.Printf("Error resolving %q with ping: %v", host, err)
			c <- host
			return
		}
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
		}()
		br := bufio.NewReader(stdout)
		line, err := br.ReadString('\n')
		if err != nil {
			log.Printf("Failed to resolve %q with ping", host)
			c <- host
			return
		}
		if m := pingRx.FindStringSubmatch(line); m != nil {
			ip := m[1]
			dnsMu.Lock()
			dnsCache[host] = ip
			dnsMu.Unlock()
			c <- ip
			return
		}
		log.Printf("Failed to resolve %q with ping", host)
		c <- host
	}()
	return <-c
	return v
}

type StatsTransport struct {
	Rt http.RoundTripper
}

func (t StatsTransport) RoundTrip(req *http.Request) (res *http.Response, err error) {
	statHTTPStart.Incr(1)
	res, err = t.Rt.RoundTrip(req)
	statHTTPResHeaders.Incr(1)
	// TODO: track per-response code stats, and also track when body done being read.
	return
}

func Dial(network, addr string) (net.Conn, error) {
	// Temporary laziness hack, avoiding doing a
	// cross-compiled Android cgo build.
	// Without cgo, package net uses
	// /etc/resolv.conf (not available on
	// Android).  We really want a cgo binary to
	// use Android's DNS cache, but it's kinda
	// hard/impossible to cross-compile for now.
	statTCPStart.Incr(1)
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		statTCPFail.Incr(1)
		return nil, fmt.Errorf("couldn't split %q", addr)
	}
	if OnAndroid() {
		// Only do the Android DNS lookups when actually
		// running on a device. We also run in "Android mode"
		// (IsChild) in tests and interactive debugging.
		host = androidLookupHost(host)
	}
	c, err := net.Dial(network, net.JoinHostPort(host, port))
	if err != nil {
		statTCPFail.Incr(1)
		return nil, err
	}
	statTCPStarted.Incr(1)
	return statTrackingConn{Conn: c}, err
}

func TLSConfig() (*tls.Config, error) {
	if !OnAndroid() {
		return nil, nil
	}
	certDir := "/system/etc/security/cacerts"
	fi, err := os.Stat(certDir)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("%q not a dir", certDir)
	}
	pool := x509.NewCertPool()
	cfg := &tls.Config{RootCAs: pool}

	f, err := os.Open(certDir)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	names, _ := f.Readdirnames(-1)
	for _, name := range names {
		pem, err := ioutil.ReadFile(filepath.Join(certDir, name))
		if err != nil {
			return nil, err
		}
		pool.AppendCertsFromPEM(pem)
	}
	return cfg, nil
}

// NoteFileUploaded is a hook for camput to report that a file
// was uploaded.  TODO: move this to pkg/client/android probably.
func NoteFileUploaded(fullPath string, uploaded bool) {
	if !IsChild() {
		return
	}
	if uploaded {
		statFileUploaded.Incr(1)
	} else {
		statFileExisted.Incr(1)
	}
	Printf("FILE_UPLOADED %s\n", fullPath)
}

// androidStatusReceiver is a blobserver.StatReceiver wrapper that
// reports the full filename path and size of uploaded blobs.
// The android app wrapping camput watches stdout for this, for progress bars.
type StatusReceiver struct {
	Sr   blobserver.StatReceiver
	Path string
}

func (asr StatusReceiver) noteChunkOnServer(sb blob.SizedRef) {
	Printf("CHUNK_UPLOADED %d %s %s\n", sb.Size, sb.Ref, asr.Path)
}

func (asr StatusReceiver) ReceiveBlob(blob blob.Ref, source io.Reader) (blob.SizedRef, error) {
	// Sniff the first 1KB of it and don't print the stats if it looks like it was just a schema
	// blob.  We won't update the progress bar for that yet.
	var buf [1024]byte
	contents := buf[:0]
	sb, err := asr.Sr.ReceiveBlob(blob, io.TeeReader(source, writeUntilSliceFull{&contents}))
	if err == nil && !schema.LikelySchemaBlob(contents) {
		statBlobUploaded.Incr(1)
		asr.noteChunkOnServer(sb)
	}
	return sb, err
}

func (asr StatusReceiver) StatBlobs(dest chan<- blob.SizedRef, blobs []blob.Ref) error {
	midc := make(chan blob.SizedRef)
	errc := make(chan error, 1)
	go func() {
		err := asr.Sr.StatBlobs(midc, blobs)
		errc <- err
		close(midc)
	}()
	for sb := range midc {
		asr.noteChunkOnServer(sb)
		statBlobExisted.Incr(1)
		dest <- sb
	}
	return <-errc
}

type writeUntilSliceFull struct {
	s *[]byte
}

func (w writeUntilSliceFull) Write(p []byte) (n int, err error) {
	s := *w.s
	l := len(s)
	growBy := cap(s) - l
	if growBy > len(p) {
		growBy = len(p)
	}
	s = s[0 : l+growBy]
	copy(s[l:], p)
	*w.s = s
	return len(p), nil
}

var memOnce sync.Once

func startMemGoroutine() {
	if !androidOutput {
		return
	}
	go func() {
		var ms runtime.MemStats
		n := 0
		for {
			runtime.ReadMemStats(&ms)
			statMemReleased.Set(int64(ms.HeapReleased))
			statMemAlloc.Set(int64(ms.Alloc))
			statMemRSS.Set(osutil.MemUsage())
			time.Sleep(1 * time.Second)
			n++
			if n%5 == 0 {
				debug.FreeOSMemory()
			}
		}
	}()
}
<commit_msg>client/android: track connection closes too<commit_after>/*
Copyright 2013 The Camlistore 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 android contains code specific to running the Camlistore client
// code as a child process on Android. This removes ugly API from the
// client package itself.
package android

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"runtime/debug"
	"strconv"
	"sync"
	"time"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/osutil"
	"camlistore.org/pkg/schema"
)

// TODO(mpl): distinguish CAMPUT, CAMGET, etc
var androidOutput, _ = strconv.ParseBool(os.Getenv("CAMPUT_ANDROID_OUTPUT"))

// IsChild reports whether this process is running as an Android
// child process and should report its output in the form that the
// Android uploader app expects.
func IsChild() bool {
	memOnce.Do(startMemGoroutine)
	return androidOutput
}

var androidOutMu sync.Mutex

func Printf(format string, args ...interface{}) {
	androidOutMu.Lock()
	defer androidOutMu.Unlock()
	fmt.Printf(format, args...)
}

var detectOnce sync.Once
var onAndroidCache bool

func dirExists(f string) bool {
	fi, err := os.Stat(f)
	return err == nil && fi.IsDir()
}

func initOnAndroid() {
	// Good enough heuristic. Suggestions welcome.
	onAndroidCache = dirExists("/data/data") && dirExists("/system/etc")
}

func OnAndroid() bool {
	detectOnce.Do(initOnAndroid)
	return onAndroidCache
}

var pingRx = regexp.MustCompile(`\((.+?)\)`)

type namedInt struct {
	name string
	sync.Mutex
	val int64
}

func (ni *namedInt) Incr(delta int64) {
	ni.Lock()
	ni.val += delta
	nv := ni.val
	ni.Unlock()
	Printf("STAT %s %d\n", ni.name, nv)
}

func (ni *namedInt) Set(v int64) {
	ni.Lock()
	if v == ni.val {
		ni.Unlock()
		return
	}
	ni.val = v
	ni.Unlock()
	Printf("STAT %s %d\n", ni.name, v)
}

var (
	statDNSStart       = &namedInt{name: "dns_start"}
	statDNSDone        = &namedInt{name: "dns_done"}
	statTCPStart       = &namedInt{name: "tcp_start"}
	statTCPStarted     = &namedInt{name: "tcp_started"}
	statTCPFail        = &namedInt{name: "tcp_fail"}
	statTCPDone        = &namedInt{name: "tcp_done_because_close"}
	statTCPDoneRead    = &namedInt{name: "tcp_done_because_read"}
	statTCPWrites      = &namedInt{name: "tcp_write_byte"}
	statTCPWrote       = &namedInt{name: "tcp_wrote_byte"}
	statTCPReads       = &namedInt{name: "tcp_read_byte"}
	statHTTPStart      = &namedInt{name: "http_start"}
	statHTTPResHeaders = &namedInt{name: "http_res_headers"}
	statBlobUploaded   = &namedInt{name: "blob_uploaded"}
	statBlobExisted    = &namedInt{name: "blob_existed"}
	statFileUploaded   = &namedInt{name: "file_uploaded"}
	statFileExisted    = &namedInt{name: "file_existed"}
	statMemReleased    = &namedInt{name: "mem_heap_released"}
	statMemAlloc       = &namedInt{name: "mem_alloc"}
	statMemRSS         = &namedInt{name: "mem_rss"}
)

type statTrackingConn struct {
	net.Conn
	once sync.Once // guards close stats
}

func (c *statTrackingConn) Write(p []byte) (n int, err error) {
	statTCPWrites.Incr(int64(len(p)))
	n, err = c.Conn.Write(p)
	statTCPWrote.Incr(int64(n))
	return
}

func (c *statTrackingConn) Read(p []byte) (n int, err error) {
	n, err = c.Conn.Read(p)
	statTCPReads.Incr(int64(n))
	if err != nil {
		c.once.Do(func() {
			statTCPDoneRead.Incr(1)
		})
	}
	return
}

func (c *statTrackingConn) Close() error {
	c.once.Do(func() {
		statTCPDone.Incr(1)
	})
	return nil
}

var (
	dnsMu    sync.Mutex
	dnsCache = make(map[string]string)
)

func androidLookupHost(host string) string {
	dnsMu.Lock()
	v, ok := dnsCache[host]
	dnsMu.Unlock()
	if ok {
		return v
	}
	statDNSStart.Incr(1)
	defer statDNSDone.Incr(1)

	// Android has no "dig" or "host" tool, so use "ping -c 1". Ghetto.
	// $ ping -c 1 google.com
	// PING google.com (74.125.224.64) 56(84) bytes of data.
	c := make(chan string, 1)
	go func() {
		cmd := exec.Command("/system/bin/ping", "-c", "1", host)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panic(err)
		}
		defer stdout.Close()
		if err := cmd.Start(); err != nil {
			log.Printf("Error resolving %q with ping: %v", host, err)
			c <- host
			return
		}
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
		}()
		br := bufio.NewReader(stdout)
		line, err := br.ReadString('\n')
		if err != nil {
			log.Printf("Failed to resolve %q with ping", host)
			c <- host
			return
		}
		if m := pingRx.FindStringSubmatch(line); m != nil {
			ip := m[1]
			dnsMu.Lock()
			dnsCache[host] = ip
			dnsMu.Unlock()
			c <- ip
			return
		}
		log.Printf("Failed to resolve %q with ping", host)
		c <- host
	}()
	return <-c
	return v
}

type StatsTransport struct {
	Rt http.RoundTripper
}

func (t StatsTransport) RoundTrip(req *http.Request) (res *http.Response, err error) {
	statHTTPStart.Incr(1)
	res, err = t.Rt.RoundTrip(req)
	statHTTPResHeaders.Incr(1)
	// TODO: track per-response code stats, and also track when body done being read.
	return
}

func Dial(network, addr string) (net.Conn, error) {
	// Temporary laziness hack, avoiding doing a
	// cross-compiled Android cgo build.
	// Without cgo, package net uses
	// /etc/resolv.conf (not available on
	// Android).  We really want a cgo binary to
	// use Android's DNS cache, but it's kinda
	// hard/impossible to cross-compile for now.
	statTCPStart.Incr(1)
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		statTCPFail.Incr(1)
		return nil, fmt.Errorf("couldn't split %q", addr)
	}
	if OnAndroid() {
		// Only do the Android DNS lookups when actually
		// running on a device. We also run in "Android mode"
		// (IsChild) in tests and interactive debugging.
		host = androidLookupHost(host)
	}
	c, err := net.Dial(network, net.JoinHostPort(host, port))
	if err != nil {
		statTCPFail.Incr(1)
		return nil, err
	}
	statTCPStarted.Incr(1)
	return &statTrackingConn{Conn: c}, err
}

func TLSConfig() (*tls.Config, error) {
	if !OnAndroid() {
		return nil, nil
	}
	certDir := "/system/etc/security/cacerts"
	fi, err := os.Stat(certDir)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("%q not a dir", certDir)
	}
	pool := x509.NewCertPool()
	cfg := &tls.Config{RootCAs: pool}

	f, err := os.Open(certDir)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	names, _ := f.Readdirnames(-1)
	for _, name := range names {
		pem, err := ioutil.ReadFile(filepath.Join(certDir, name))
		if err != nil {
			return nil, err
		}
		pool.AppendCertsFromPEM(pem)
	}
	return cfg, nil
}

// NoteFileUploaded is a hook for camput to report that a file
// was uploaded.  TODO: move this to pkg/client/android probably.
func NoteFileUploaded(fullPath string, uploaded bool) {
	if !IsChild() {
		return
	}
	if uploaded {
		statFileUploaded.Incr(1)
	} else {
		statFileExisted.Incr(1)
	}
	Printf("FILE_UPLOADED %s\n", fullPath)
}

// androidStatusReceiver is a blobserver.StatReceiver wrapper that
// reports the full filename path and size of uploaded blobs.
// The android app wrapping camput watches stdout for this, for progress bars.
type StatusReceiver struct {
	Sr   blobserver.StatReceiver
	Path string
}

func (asr StatusReceiver) noteChunkOnServer(sb blob.SizedRef) {
	Printf("CHUNK_UPLOADED %d %s %s\n", sb.Size, sb.Ref, asr.Path)
}

func (asr StatusReceiver) ReceiveBlob(blob blob.Ref, source io.Reader) (blob.SizedRef, error) {
	// Sniff the first 1KB of it and don't print the stats if it looks like it was just a schema
	// blob.  We won't update the progress bar for that yet.
	var buf [1024]byte
	contents := buf[:0]
	sb, err := asr.Sr.ReceiveBlob(blob, io.TeeReader(source, writeUntilSliceFull{&contents}))
	if err == nil && !schema.LikelySchemaBlob(contents) {
		statBlobUploaded.Incr(1)
		asr.noteChunkOnServer(sb)
	}
	return sb, err
}

func (asr StatusReceiver) StatBlobs(dest chan<- blob.SizedRef, blobs []blob.Ref) error {
	midc := make(chan blob.SizedRef)
	errc := make(chan error, 1)
	go func() {
		err := asr.Sr.StatBlobs(midc, blobs)
		errc <- err
		close(midc)
	}()
	for sb := range midc {
		asr.noteChunkOnServer(sb)
		statBlobExisted.Incr(1)
		dest <- sb
	}
	return <-errc
}

type writeUntilSliceFull struct {
	s *[]byte
}

func (w writeUntilSliceFull) Write(p []byte) (n int, err error) {
	s := *w.s
	l := len(s)
	growBy := cap(s) - l
	if growBy > len(p) {
		growBy = len(p)
	}
	s = s[0 : l+growBy]
	copy(s[l:], p)
	*w.s = s
	return len(p), nil
}

var memOnce sync.Once

func startMemGoroutine() {
	if !androidOutput {
		return
	}
	go func() {
		var ms runtime.MemStats
		n := 0
		for {
			runtime.ReadMemStats(&ms)
			statMemReleased.Set(int64(ms.HeapReleased))
			statMemAlloc.Set(int64(ms.Alloc))
			statMemRSS.Set(osutil.MemUsage())
			time.Sleep(1 * time.Second)
			n++
			if n%5 == 0 {
				debug.FreeOSMemory()
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright 2017 Cisco Systems, Inc.
//
// 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 controller

import (
	"fmt"

	"github.com/noironetworks/aci-containers/pkg/apicapi"

	"sort"
	"strings"
)

func (cont *AciController) SubnetChanged(obj apicapi.ApicObject, aciVrfDn string) {
	subnetDn := obj.GetAttrStr("dn")
	subnetIp := obj.GetAttrStr("ip")
	cont.log.Debug("SubnetChanged for dn: ", subnetDn)
	if _, ok := cont.apicConn.CachedSubnetDns[subnetDn]; !ok {
		cont.UpdateSubnetDnCache(subnetDn, subnetIp, aciVrfDn)
		cont.scheduleRdConfig()
	}
}

func (cont *AciController) SubnetDeleted(dn string) {
	cont.log.Debug("Before deleting dn: ", dn, " CachedSubnets.subnetDns Map:  ", cont.apicConn.CachedSubnetDns)
	if _, ok := cont.apicConn.CachedSubnetDns[dn]; ok {
		delete(cont.apicConn.CachedSubnetDns, dn)
		cont.scheduleRdConfig()
	} else {
		cont.log.Debug("This shouldn't happen!")
	}
	cont.log.Debug("After Delete: CachedSubnets.subnetDns Map:  ", cont.apicConn.CachedSubnetDns)
}

func (cont *AciController) UpdateSubnetDnCache(subnetDn string, subnetIp string, aciVrfDn string) {
	cont.log.Debug("aciVrfDn: ", aciVrfDn, "; Processing SubnetDn: ", subnetDn)
	subnetDelimiter := "/subnet"
	subnetParentDn := strings.Split(subnetDn, subnetDelimiter)[0]
	var inCache = false
	for k := range cont.apicConn.CachedSubnetDns {
		if subnetParentDn == strings.Split(k, subnetDelimiter)[0] {
			cont.log.Debug("subnetParentDn is in cachedSubnets:", subnetParentDn)
			cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
			inCache = true
			cont.log.Debug("cachedSubnetDns Map:  ", cont.apicConn.CachedSubnetDns)
			continue
		}
	}
	if !inCache {
		cont.BuildSubnetDnCache(subnetDn, aciVrfDn)
		cont.scheduleRdConfig()
	}
}

func (cont *AciController) BuildSubnetDnCache(dn string, aciVrfDn string) {
	cont.log.Debug("aciVrfDn: ", aciVrfDn, "; Processing dn: ", dn)
	var vrfBdDns []string
	var vrfEpgDns []string

	// Get all the BDs related to aciVrf and then all EPGs related to those BDs
	bdRsDelimiter := "/rsctx"
	epgRsDelimiter := "/rsbd"
	subnetDelimiter := "/subnet"
	bdFilter := fmt.Sprintf("query-target-filter=and(wcard(fvRsCtx.tDn,\"%s\"))", aciVrfDn)
	bdArgs := []string{
		"rsp-prop-include=config-only",
		"query-target=subtree",
		bdFilter,
	}

	epgArgs := []string{
		"query-target=subtree",
		"target-subtree-class=fvRsBd",
	}

	subnetArgs := []string{
		"rsp-prop-include=config-only",
	}

	bdUri := fmt.Sprintf("/api/node/class/fvBD.json?%s", strings.Join(bdArgs, "&"))
	epgUri := fmt.Sprintf("/api/node/class/fvAEPg.json?%s", strings.Join(epgArgs, "&"))
	SubnetUri := fmt.Sprintf("/api/node/class/fvSubnet.json?%s", strings.Join(subnetArgs, "&"))

	apicresp, err := cont.apicConn.GetApicResponse(bdUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			dn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			vrfBdDns = append(vrfBdDns, strings.Split(dn, bdRsDelimiter)[0])
		}
	}
	sort.Strings(vrfBdDns)
	cont.log.Debug("aciVrfBdDns: ", vrfBdDns)

	apicresp, err = cont.apicConn.GetApicResponse(epgUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			epgParentDn, ok := body.Attributes["tDn"].(string)
			cont.log.Debug("epgParentDn: ", epgParentDn)
			if !ok || !cont.contains(vrfBdDns, epgParentDn) {
				continue
			}
			epgDn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			vrfEpgDns = append(vrfEpgDns, strings.Split(epgDn, epgRsDelimiter)[0])
		}
	}
	sort.Strings(vrfEpgDns)
	cont.log.Debug("aciVrfEpgDns: ", vrfEpgDns)
	apicresp, err = cont.apicConn.GetApicResponse(SubnetUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			subnetDn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			subnetParentDn := strings.Split(subnetDn, subnetDelimiter)[0]
			subnetIp, _ := body.Attributes["ip"].(string)
			cont.log.Debug("subnetDn: ", subnetDn, " subnetParentDn: ", subnetParentDn, " subnetIp: ", subnetIp)
			if cont.contains(vrfBdDns, subnetParentDn) || cont.contains(vrfEpgDns, subnetParentDn) {
				cont.indexMutex.Lock()
				cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
				cont.indexMutex.Unlock()
			}
		}
	}
	cont.log.Debug("cachedSubnetsDns Map:  ", cont.apicConn.CachedSubnetDns)
	return
}

func (cont *AciController) contains(s []string, searchterm string) bool {
	i := sort.SearchStrings(s, searchterm)
	return i < len(s) && s[i] == searchterm
}
<commit_msg>Added fn to update SubnetDnCache for Subnetdn<commit_after>// Copyright 2017 Cisco Systems, Inc.
//
// 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 controller

import (
	"fmt"

	"github.com/noironetworks/aci-containers/pkg/apicapi"

	"sort"
	"strings"
)

func (cont *AciController) SubnetChanged(obj apicapi.ApicObject, aciVrfDn string) {
	subnetDn := obj.GetAttrStr("dn")
	subnetIp := obj.GetAttrStr("ip")
	cont.log.Debug("SubnetChanged for dn: ", subnetDn)
	if _, ok := cont.apicConn.CachedSubnetDns[subnetDn]; !ok {
		cont.UpdateSubnetDnCache(subnetDn, subnetIp, aciVrfDn)
		cont.scheduleRdConfig()
	}
}

func (cont *AciController) SubnetDeleted(dn string) {
	cont.log.Debug("Before deleting dn: ", dn, " CachedSubnets.subnetDns Map:  ", cont.apicConn.CachedSubnetDns)
	if _, ok := cont.apicConn.CachedSubnetDns[dn]; ok {
		delete(cont.apicConn.CachedSubnetDns, dn)
		cont.scheduleRdConfig()
	} else {
		cont.log.Debug("This shouldn't happen!")
	}
	cont.log.Debug("After Delete: CachedSubnets.subnetDns Map:  ", cont.apicConn.CachedSubnetDns)
}

func (cont *AciController) UpdateSubnetDnCache(subnetDn string, subnetIp string, aciVrfDn string) {
	cont.log.Debug("aciVrfDn: ", aciVrfDn, "; Processing SubnetDn: ", subnetDn)
	subnetDelimiter := "/subnet"
	subnetParentDn := strings.Split(subnetDn, subnetDelimiter)[0]
	var inCache = false
	for k := range cont.apicConn.CachedSubnetDns {
		if subnetParentDn == strings.Split(k, subnetDelimiter)[0] {
			cont.log.Debug("subnetParentDn is in cachedSubnets:", subnetParentDn)
			cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
			inCache = true
			cont.log.Debug("cachedSubnetDns Map:  ", cont.apicConn.CachedSubnetDns)
			continue
		}
	}
	if !inCache {
		cont.UpdateSubnetDnCacheForDn(subnetDn, subnetIp)
		cont.scheduleRdConfig()
	}
}

func (cont *AciController) isBdPresentInVrf(bdDn string) bool {
	bdsubnetParentArgs := []string{
		"query-target=children&target-subtree-class=fvRsCtx",
	}
	bdsubnetParentDnUri := "/api/node/mo/" + bdDn + ".json?" + strings.Join(bdsubnetParentArgs, "&")
	apicresp, err := cont.apicConn.GetApicResponse(bdsubnetParentDnUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return false
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			tDn, ok := body.Attributes["tDn"].(string)
			if !ok {
				continue
			}
			if tDn == cont.config.AciVrfDn {
				return true
			}
		}
	}
	return false
}

func (cont *AciController) UpdateSubnetDnCacheForDn(subnetDn string, subnetIp string) {
	cont.log.Debug("Updating SubnetDnCache for dn: ", subnetDn, "ip: ", subnetIp)
	subnetDelimiter := "/subnet"
	subnetParentDn := strings.Split(subnetDn, subnetDelimiter)[0]
	subnetParentArgs := []string{
		"query-target=self",
	}
	subnetParentDnUri := "/api/node/mo/" + subnetParentDn + ".json?" + strings.Join(subnetParentArgs, "&")
	apicresp, err := cont.apicConn.GetApicResponse(subnetParentDnUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for class := range obj {
			if class == "fvBD" {
				if cont.isBdPresentInVrf(subnetParentDn) {
					cont.indexMutex.Lock()
					cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
					cont.log.Debug("fvBD: Adding to CachedSubnetDns: ", subnetDn)
					cont.indexMutex.Unlock()
				}
				//if objClass is not fvBD, it will be fvAEPg
			} else {
				epgsubnetParentArgs := []string{
					"query-target=children&target-subtree-class=fvRsBd",
				}
				epgsubnetParentDnUri := "/api/node/mo/" + subnetParentDn + ".json?" + strings.Join(epgsubnetParentArgs, "&")
				apicepgresp, epgerr := cont.apicConn.GetApicResponse(epgsubnetParentDnUri)
				if epgerr != nil {
					cont.log.Debugf("Failed to get APIC response, err: %v", epgerr)
					return
				}
				for _, epgobj := range apicepgresp.Imdata {
					for _, epgbody := range epgobj {
						epgParentDn, ok := epgbody.Attributes["tDn"].(string)
						if !ok {
							continue
						}
						if cont.isBdPresentInVrf(epgParentDn) {
							cont.indexMutex.Lock()
							cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
							cont.log.Debug("fvAEPg: Adding to CachedSubnetDns: ", subnetDn)
							cont.indexMutex.Unlock()
						}
					}
				}
			}
		}
	}

}

func (cont *AciController) BuildSubnetDnCache(dn string, aciVrfDn string) {
	cont.log.Debug("aciVrfDn: ", aciVrfDn, "; Processing dn: ", dn)
	var vrfBdDns []string
	var vrfEpgDns []string

	// Get all the BDs related to aciVrf and then all EPGs related to those BDs
	bdRsDelimiter := "/rsctx"
	epgRsDelimiter := "/rsbd"
	subnetDelimiter := "/subnet"
	bdFilter := fmt.Sprintf("query-target-filter=and(wcard(fvRsCtx.tDn,\"%s\"))", aciVrfDn)
	bdArgs := []string{
		"rsp-prop-include=config-only",
		"query-target=subtree",
		bdFilter,
	}

	epgArgs := []string{
		"query-target=subtree",
		"target-subtree-class=fvRsBd",
	}

	subnetArgs := []string{
		"rsp-prop-include=config-only",
	}

	bdUri := fmt.Sprintf("/api/node/class/fvBD.json?%s", strings.Join(bdArgs, "&"))
	epgUri := fmt.Sprintf("/api/node/class/fvAEPg.json?%s", strings.Join(epgArgs, "&"))
	SubnetUri := fmt.Sprintf("/api/node/class/fvSubnet.json?%s", strings.Join(subnetArgs, "&"))

	apicresp, err := cont.apicConn.GetApicResponse(bdUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			dn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			vrfBdDns = append(vrfBdDns, strings.Split(dn, bdRsDelimiter)[0])
		}
	}
	sort.Strings(vrfBdDns)
	cont.log.Debug("aciVrfBdDns: ", vrfBdDns)

	apicresp, err = cont.apicConn.GetApicResponse(epgUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			epgParentDn, ok := body.Attributes["tDn"].(string)
			cont.log.Debug("epgParentDn: ", epgParentDn)
			if !ok || !cont.contains(vrfBdDns, epgParentDn) {
				continue
			}
			epgDn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			vrfEpgDns = append(vrfEpgDns, strings.Split(epgDn, epgRsDelimiter)[0])
		}
	}
	sort.Strings(vrfEpgDns)
	cont.log.Debug("aciVrfEpgDns: ", vrfEpgDns)
	apicresp, err = cont.apicConn.GetApicResponse(SubnetUri)
	if err != nil {
		cont.log.Debugf("Failed to get APIC response, err: %v", err)
		return
	}
	for _, obj := range apicresp.Imdata {
		for _, body := range obj {
			subnetDn, ok := body.Attributes["dn"].(string)
			if !ok {
				continue
			}
			subnetParentDn := strings.Split(subnetDn, subnetDelimiter)[0]
			subnetIp, _ := body.Attributes["ip"].(string)
			cont.log.Debug("subnetDn: ", subnetDn, " subnetParentDn: ", subnetParentDn, " subnetIp: ", subnetIp)
			if cont.contains(vrfBdDns, subnetParentDn) || cont.contains(vrfEpgDns, subnetParentDn) {
				cont.indexMutex.Lock()
				cont.apicConn.CachedSubnetDns[subnetDn] = subnetIp
				cont.indexMutex.Unlock()
			}
		}
	}
	cont.log.Debug("cachedSubnetsDns Map:  ", cont.apicConn.CachedSubnetDns)
	return
}

func (cont *AciController) contains(s []string, searchterm string) bool {
	i := sort.SearchStrings(s, searchterm)
	return i < len(s) && s[i] == searchterm
}
<|endoftext|>