<commit_before>package flv

import (
	"os"
	"fmt"
	"errors"
	"bytes"
)

type Header struct {
	Version      uint16
	Body         []byte
}

type Frame interface{
	ToBinary() ([]byte, error)
}

type CFrame struct {
	Stream      uint32
	Dts         uint32
	Type        TagType
	Flavor      Flavor
	Position    int64
	Body        []byte
	PrevTagSize uint32
}

func (f CFrame) ToBinary() ([]byte, error) {
	d := make([]byte, 0)
	buf := bytes.NewBuffer(d)
	buf.WriteByte(byte(f.Type))
	bl := uint32(len(f.Body))
	buf.Write([]byte{byte(bl>>16), byte((bl>>8)&0xFF), byte(bl&0xFF)})
	buf.Write([]byte{byte((f.Dts>>16) & 0xFF), byte((f.Dts>>8) & 0xFF), byte(f.Dts & 0xFF)})
	buf.WriteByte(byte((f.Dts >> 24) & 0xFF))
	buf.Write(f.Body)
	buf.Write([]byte{byte((f.PrevTagSize>>24) & 0xFF), byte((f.PrevTagSize>>16) & 0xFF), byte((f.PrevTagSize>>8) & 0xFF), byte(f.PrevTagSize & 0xFF)})
	return buf.Bytes(), nil
}

type VideoFrame struct {
	CFrame
	CodecId     VideoCodec
	Width       uint16
	Height      uint16
}

func (f VideoFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type AudioFrame struct {
	CFrame
	CodecId     AudioCodec
    Rate        uint32
    BitSize     AudioSize
	Channels    AudioType
}

func (f AudioFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type MetaFrame struct {
	CFrame
}

func (f MetaFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type FlvReader struct {
	inFile     *os.File
	width      uint16
	height     uint16
}

func NewReader(inFile *os.File) (*FlvReader) {
	return &FlvReader{
		inFile: inFile,
		width: 0,
		height: 0,
	}
}

type FlvWriter struct {
	outFile     *os.File
}

func NewWriter(outFile *os.File) (*FlvWriter) {
	return &FlvWriter{
		outFile: outFile,
	}
}

func (frReader *FlvReader) ReadHeader() (*Header, error) {
	header := make([]byte, HEADER_LENGTH+4)
	_, err := frReader.inFile.Read(header)
	if err != nil {
		return nil, err
	}

	sig := header[0:3]
	if bytes.Compare(sig, []byte(SIG)) != 0 {
		return nil, fmt.Errorf("bad file format")
	}
	version := (uint16(header[3]) << 8) | (uint16(header[4]) << 0)
	//skip := header[4:5]
	//offset := header[5:9]
	//next_id := header[9:13]

	return &Header{Version: version, Body: header}, nil
}

func (frWriter *FlvWriter) WriteHeader(header *Header) (error) {
	_, err := frWriter.outFile.Write(header.Body)
	if err != nil {
		return err
	}
	return nil
}

func (frReader *FlvReader) ReadFrame() (fr Frame, e error) {

	var n int
	var err error

	curPos, _ := frReader.inFile.Seek(0, os.SEEK_CUR)

	tagHeaderB := make([]byte, TAG_HEADER_LENGTH)
	n, err = frReader.inFile.Read(tagHeaderB)
	if TagSize(n) != TAG_HEADER_LENGTH {
		return nil, fmt.Errorf("bad record")
	}
	if err != nil {
		return nil, err
	}

	tagType := TagType(tagHeaderB[0])
	bodyLen := (uint32(tagHeaderB[1]) << 16) | (uint32(tagHeaderB[2]) << 8) | (uint32(tagHeaderB[3]) << 0)
	ts := (uint32(tagHeaderB[4]) << 16) | (uint32(tagHeaderB[5]) << 8) | (uint32(tagHeaderB[6]) << 0)
	tsExt := uint32(tagHeaderB[7])
	stream := (uint32(tagHeaderB[8]) << 16) | (uint32(tagHeaderB[9]) << 8) | (uint32(tagHeaderB[10]) << 0)

	var dts uint32
	dts = (tsExt << 24) | ts

	bodyBuf := make([]byte, bodyLen)
	n, err = frReader.inFile.Read(bodyBuf)
	if err != nil {
		return nil, err
	}

	prevTagSizeB := make([]byte, PREV_TAG_SIZE_LENGTH)
	n, err = frReader.inFile.Read(prevTagSizeB)
	if err != nil {
		return nil, err
	}
	prevTagSize := (uint32(prevTagSizeB[0]) << 24) | (uint32(prevTagSizeB[1]) << 16) | (uint32(prevTagSizeB[2]) << 8) | (uint32(prevTagSizeB[3]) << 0)

	pFrame := CFrame{
		Stream: stream,
		Dts: dts,
		Type: tagType,
		Position: curPos,
		Body: bodyBuf,
		PrevTagSize: prevTagSize,
	}

	var resFrame Frame

	switch tagType {
	case TAG_TYPE_META:
		pFrame.Flavor = METADATA
		resFrame = MetaFrame{pFrame}
	case TAG_TYPE_VIDEO:
		vft := VideoFrameType(uint8(bodyBuf[0]) >> 4)
		codecId := VideoCodec(uint8(bodyBuf[0]) & 0x0F)
		switch vft {
		case VIDEO_FRAME_TYPE_KEYFRAME:
			pFrame.Flavor = KEYFRAME
			switch codecId {
			case VIDEO_CODEC_ON2VP6:
				hHelper := (uint16(bodyBuf[1]) >> 4) & 0x0F
				wHelper := uint16(bodyBuf[1]) & 0x0F
				w := uint16(bodyBuf[5])
				h := uint16(bodyBuf[6])

				frReader.width = w*16 - wHelper
				frReader.height = h*16 - hHelper
			}
		default:
			pFrame.Flavor = FRAME
		}
		resFrame = VideoFrame{CFrame: pFrame, CodecId: codecId, Width: frReader.width, Height: frReader.height}
	case TAG_TYPE_AUDIO:
		pFrame.Flavor = FRAME
		codecId := AudioCodec(uint8(bodyBuf[0]) >> 4)
		rate := audioRate(AudioRate((uint8(bodyBuf[0]) >> 2) & 0x03))
		bitSize := AudioSize((uint8(bodyBuf[0]) >> 1) & 0x01)
		channels := AudioType(uint8(bodyBuf[0]) & 0x01)
		resFrame = AudioFrame{CFrame: pFrame, CodecId: codecId, Rate: rate, BitSize: bitSize, Channels: channels}
	}

	return resFrame, nil
}

func audioRate(ar AudioRate) (uint32) {
	var ret uint32
	switch ar {
	case AUDIO_RATE_5_5:
		ret = 5500
	case AUDIO_RATE_11:
		ret = 11025
	case AUDIO_RATE_22:
		ret = 22050
	case AUDIO_RATE_44:
		ret = 44056
	}
	return ret
}

func (frWriter *FlvWriter) WriteFrame(fr Frame) (e error) {

	var n int
	var err error

	var tagType TagType
	var body []byte
	var dts uint32
	var stream uint32

	switch fr.(type) {
	case VideoFrame:
		tfr := fr.(VideoFrame)
		tagType = tfr.Type
		body = tfr.Body
		dts = tfr.Dts
		stream = tfr.Stream
	case AudioFrame:
		tfr := fr.(AudioFrame)
		tagType = tfr.Type
		body = tfr.Body
		dts = tfr.Dts
		stream = tfr.Stream
	case MetaFrame:
		tfr := fr.(MetaFrame)
		tagType = tfr.Type
		body = tfr.Body
		dts = tfr.Dts
		stream = tfr.Stream
	}

	tagHeaderB := make([]byte, TAG_HEADER_LENGTH)

	tagHeaderB[0] = byte(tagType)
	bodyLen := len(body)
	tagHeaderB[1] = byte((bodyLen >> 16) & 0xFF)
	tagHeaderB[2] = byte((bodyLen >> 8) & 0xFF)
	tagHeaderB[3] = byte((bodyLen >> 0) & 0xFF)

	tagHeaderB[4] = byte((dts >> 16) & 0xFF)
	tagHeaderB[5] = byte((dts >> 8) & 0xFF)
	tagHeaderB[6] = byte((dts >> 0) & 0xFF)

	tagHeaderB[7] = byte((dts >> 24) & 0xFF)

	tagHeaderB[8] = byte((stream >> 16) & 0xFF)
	tagHeaderB[9] = byte((stream >> 8) & 0xFF)
	tagHeaderB[10] = byte((stream >> 0) & 0xFF)

	n, err = frWriter.outFile.Write(tagHeaderB)
	if err != nil {
		return err
	}
	if n != len(tagHeaderB) {
		return errors.New("bad header write")
	}

	n, err = frWriter.outFile.Write(body)
	if err != nil {
		return err
	}
	if n != len(body) {
		return errors.New("bad body write")
	}

	prevTagSizeB := make([]byte, PREV_TAG_SIZE_LENGTH)

	prevTagSize := uint32(uint32(TAG_HEADER_LENGTH) + uint32(len(body)))

	prevTagSizeB[0] = byte((prevTagSize >> 24) & 0xFF)
	prevTagSizeB[1] = byte((prevTagSize >> 16) & 0xFF)
	prevTagSizeB[2] = byte((prevTagSize >> 8) & 0xFF)
	prevTagSizeB[3] = byte((prevTagSize >> 0) & 0xFF)

	n, err = frWriter.outFile.Write(prevTagSizeB)
	if err != nil {
		return err
	}
	if n != len(prevTagSizeB) {
		return errors.New("bad prev tag write")
	}
	return nil
}
<commit_msg>Change interface for frame writer<commit_after>package flv

import (
	"os"
	"fmt"
	"bytes"
	"io"
)

type Header struct {
	Version      uint16
	Body         []byte
}

type Frame interface{
	WriteFrame(io.Writer) error
}

type CFrame struct {
	Stream      uint32
	Dts         uint32
	Type        TagType
	Flavor      Flavor
	Position    int64
	Body        []byte
	PrevTagSize uint32
}

func (f CFrame) WriteFrame(w io.Writer) (error) {
	bl := uint32(len(f.Body))
	var err error
	err = writeType(w, f.Type); if err != nil {return err}
	err = writeBodyLength(w, bl)
	err = writeDts(w, f.Dts)
	err = writeStream(w, f.Stream)
	err = writeBody(w, f.Body)
	prevTagSize := bl + uint32(TAG_HEADER_LENGTH)
	err = writePrevTagSize(w, prevTagSize)
	return nil
}

func writeType(w io.Writer, t TagType) error {
	_, err := w.Write([]byte{byte(t)})
	return err
}

func writeBodyLength(w io.Writer, bl uint32) error {
	_, err := w.Write([]byte{byte(bl>>16), byte((bl>>8)&0xFF), byte(bl&0xFF)})
	return err
}

func writeDts(w io.Writer, dts uint32) error {
	_, err := w.Write([]byte{byte((dts>>16) & 0xFF), byte((dts>>8) & 0xFF), byte(dts & 0xFF)})
	if err != nil {
		return err
	}
	_, err = w.Write([]byte{byte((dts >> 24) & 0xFF)})
	return err
}

func writeStream(w io.Writer, stream uint32) error {
	_, err := w.Write([]byte{byte(stream>>16), byte((stream>>8)&0xFF), byte(stream&0xFF)})
	return err
}

func writeBody(w io.Writer, body []byte) error {
	_, err := w.Write(body)
	return err
}

func writePrevTagSize(w io.Writer, prevTagSize uint32) error {
	_, err := w.Write([]byte{byte((prevTagSize>>24) & 0xFF), byte((prevTagSize>>16) & 0xFF), byte((prevTagSize>>8) & 0xFF), byte(prevTagSize & 0xFF)})
	return err
}

type VideoFrame struct {
	CFrame
	CodecId     VideoCodec
	Width       uint16
	Height      uint16
}

func (f VideoFrame) WriteFrame(w io.Writer) (error) {
	return f.CFrame.WriteFrame(w)
}

type AudioFrame struct {
	CFrame
	CodecId     AudioCodec
    Rate        uint32
    BitSize     AudioSize
	Channels    AudioType
}

func (f AudioFrame) WriteFrame(w io.Writer) (error) {
	return f.CFrame.WriteFrame(w)
}

type MetaFrame struct {
	CFrame
}

func (f MetaFrame) WriteFrame(w io.Writer) (error) {
	return f.CFrame.WriteFrame(w)
}

type FlvReader struct {
	inFile     *os.File
	width      uint16
	height     uint16
}

func NewReader(inFile *os.File) (*FlvReader) {
	return &FlvReader{
		inFile: inFile,
		width: 0,
		height: 0,
	}
}

type FlvWriter struct {
	outFile     *os.File
}

func NewWriter(outFile *os.File) (*FlvWriter) {
	return &FlvWriter{
		outFile: outFile,
	}
}

func (frReader *FlvReader) ReadHeader() (*Header, error) {
	header := make([]byte, HEADER_LENGTH+4)
	_, err := frReader.inFile.Read(header)
	if err != nil {
		return nil, err
	}

	sig := header[0:3]
	if bytes.Compare(sig, []byte(SIG)) != 0 {
		return nil, fmt.Errorf("bad file format")
	}
	version := (uint16(header[3]) << 8) | (uint16(header[4]) << 0)
	//skip := header[4:5]
	//offset := header[5:9]
	//next_id := header[9:13]

	return &Header{Version: version, Body: header}, nil
}

func (frWriter *FlvWriter) WriteHeader(header *Header) (error) {
	_, err := frWriter.outFile.Write(header.Body)
	if err != nil {
		return err
	}
	return nil
}

func (frReader *FlvReader) ReadFrame() (fr Frame, e error) {

	var n int
	var err error

	curPos, _ := frReader.inFile.Seek(0, os.SEEK_CUR)

	tagHeaderB := make([]byte, TAG_HEADER_LENGTH)
	n, err = frReader.inFile.Read(tagHeaderB)
	if TagSize(n) != TAG_HEADER_LENGTH {
		return nil, fmt.Errorf("bad record")
	}
	if err != nil {
		return nil, err
	}

	tagType := TagType(tagHeaderB[0])
	bodyLen := (uint32(tagHeaderB[1]) << 16) | (uint32(tagHeaderB[2]) << 8) | (uint32(tagHeaderB[3]) << 0)
	ts := (uint32(tagHeaderB[4]) << 16) | (uint32(tagHeaderB[5]) << 8) | (uint32(tagHeaderB[6]) << 0)
	tsExt := uint32(tagHeaderB[7])
	stream := (uint32(tagHeaderB[8]) << 16) | (uint32(tagHeaderB[9]) << 8) | (uint32(tagHeaderB[10]) << 0)

	var dts uint32
	dts = (tsExt << 24) | ts

	bodyBuf := make([]byte, bodyLen)
	n, err = frReader.inFile.Read(bodyBuf)
	if err != nil {
		return nil, err
	}

	prevTagSizeB := make([]byte, PREV_TAG_SIZE_LENGTH)
	n, err = frReader.inFile.Read(prevTagSizeB)
	if err != nil {
		return nil, err
	}
	prevTagSize := (uint32(prevTagSizeB[0]) << 24) | (uint32(prevTagSizeB[1]) << 16) | (uint32(prevTagSizeB[2]) << 8) | (uint32(prevTagSizeB[3]) << 0)

	pFrame := CFrame{
		Stream: stream,
		Dts: dts,
		Type: tagType,
		Position: curPos,
		Body: bodyBuf,
		PrevTagSize: prevTagSize,
	}

	var resFrame Frame

	switch tagType {
	case TAG_TYPE_META:
		pFrame.Flavor = METADATA
		resFrame = MetaFrame{pFrame}
	case TAG_TYPE_VIDEO:
		vft := VideoFrameType(uint8(bodyBuf[0]) >> 4)
		codecId := VideoCodec(uint8(bodyBuf[0]) & 0x0F)
		switch vft {
		case VIDEO_FRAME_TYPE_KEYFRAME:
			pFrame.Flavor = KEYFRAME
			switch codecId {
			case VIDEO_CODEC_ON2VP6:
				hHelper := (uint16(bodyBuf[1]) >> 4) & 0x0F
				wHelper := uint16(bodyBuf[1]) & 0x0F
				w := uint16(bodyBuf[5])
				h := uint16(bodyBuf[6])

				frReader.width = w*16 - wHelper
				frReader.height = h*16 - hHelper
			}
		default:
			pFrame.Flavor = FRAME
		}
		resFrame = VideoFrame{CFrame: pFrame, CodecId: codecId, Width: frReader.width, Height: frReader.height}
	case TAG_TYPE_AUDIO:
		pFrame.Flavor = FRAME
		codecId := AudioCodec(uint8(bodyBuf[0]) >> 4)
		rate := audioRate(AudioRate((uint8(bodyBuf[0]) >> 2) & 0x03))
		bitSize := AudioSize((uint8(bodyBuf[0]) >> 1) & 0x01)
		channels := AudioType(uint8(bodyBuf[0]) & 0x01)
		resFrame = AudioFrame{CFrame: pFrame, CodecId: codecId, Rate: rate, BitSize: bitSize, Channels: channels}
	}

	return resFrame, nil
}

func audioRate(ar AudioRate) (uint32) {
	var ret uint32
	switch ar {
	case AUDIO_RATE_5_5:
		ret = 5500
	case AUDIO_RATE_11:
		ret = 11025
	case AUDIO_RATE_22:
		ret = 22050
	case AUDIO_RATE_44:
		ret = 44056
	}
	return ret
}

func (frWriter *FlvWriter) WriteFrame(fr Frame) (e error) {
	return fr.WriteFrame(frWriter.outFile)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/

package main

import (
	"crypto/dsa"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"time"
)

type CertificatePrincipal struct {
	Email string `json:"email"`
}

func NewCertificatePrincipal(email string) *CertificatePrincipal {
	return &CertificatePrincipal{Email: email}
}

type CertificatePublicKey struct {
	Algorithm string `json:"algorithm"`
	y         string `json:"y"`
	p         string `json:"p"`
	q         string `json:"q"`
	g         string `json:"g"`
}

func NewCertificatePublicKey(priv dsa.PrivateKey) *CertificatePublicKey {
	return &CertificatePublicKey{Algorithm: "DS"}
}

type CertificatePayload struct {
	Principal CertificatePrincipal `json:"principal"`
	PublicKey CertificatePublicKey `json:"public-key"`
	Issuer    string               `json:"iss"`
	IssuedAt  int64                `json:"iat"`
	Audience  string               `json:"aud"`
	ExpiresAt int64                `json:"exp"`
}

func NewCertificatePayload(priv dsa.PrivateKey, email string, issuer string, issuedAt time.Time, expiresAt time.Time) *CertificatePayload {
	return &CertificatePayload{
		Principal: *NewCertificatePrincipal(email),
		PublicKey: *NewCertificatePublicKey(priv),
		Issuer:    issuer,
		IssuedAt:  issuedAt.UnixNano() / int64(time.Millisecond),
		ExpiresAt: expiresAt.UnixNano() / int64(time.Millisecond),
	}
}

type WebTokenHeader struct {
	Algorithm string `json:"alg"`
}

func (p *CertificatePayload) Encode(priv dsa.PrivateKey) (string, error) {
	// Encode the header
	header := WebTokenHeader{Algorithm: "DS128"} // TODO: Create a function to extract this from the dsa.PrivateKey
	headerData, err := json.Marshal(header)
	if err != nil {
		return "", err
	}
	encodedHeader := base64.URLEncoding.EncodeToString(headerData)

	// Encode the certificate
	payloadData, err := json.Marshal(p)
	if err != nil {
		return "", err
	}
	encodedPayload := base64.URLEncoding.EncodeToString(payloadData)

	// The message is simply header plus payload
	message := encodedHeader + "." + encodedPayload

	// Get the SHA1 hash over the message
	messageHash := sha1.New().Sum([]byte(message))

	// Calculate the signature over the message
	r, s, err := dsa.Sign(rand.Reader, &priv, messageHash)
	if err != nil {
		return "", err
	}
	signatureData := append(r.Bytes(), s.Bytes()...)
	encodedSignature := base64.URLEncoding.EncodeToString(signatureData)

	token := message + "." + encodedSignature
	// Prepend

	return token, nil
}

func CreateCertificate(key dsa.PrivateKey, email string, issuer string, issuedAt time.Time, expiresAt time.Time, signingKey dsa.PrivateKey) (string, error) {
	certificatePayload := NewCertificatePayload(key, email, issuer, issuedAt, expiresAt)
	return certificatePayload.Encode(signingKey)
}

//

type AssertionPayload struct {
	Issuer    string `json:"iss"`
	IssuedAt  int64  `json:"iat"`
	ExpiresAt int64  `json:"exp"`
	Audience  string `json:"aud"`
}

func NewAssertionPayload(issuer string, issuedAt time.Time, expiresAt time.Time, audience string) (*AssertionPayload, error) {
	return &AssertionPayload{
		Issuer:    issuer,
		IssuedAt:  issuedAt.UnixNano() / int64(time.Millisecond),
		ExpiresAt: expiresAt.UnixNano() / int64(time.Millisecond),
		Audience:  audience,
	}, nil
}

func (p *AssertionPayload) Encode(signingKey dsa.PrivateKey) (string, error) {
	// Encode the header
	header := WebTokenHeader{Algorithm: "DS128"} // TODO: Create a function to extract this from the dsa.PrivateKey
	headerData, err := json.Marshal(header)
	if err != nil {
		return "", err
	}
	encodedHeader := base64.URLEncoding.EncodeToString(headerData)

	// Encode the certificate
	payloadData, err := json.Marshal(p)
	if err != nil {
		return "", err
	}
	encodedPayload := base64.URLEncoding.EncodeToString(payloadData)

	// The message is simply header plus payload
	message := encodedHeader + "." + encodedPayload

	// Get the SHA1 hash over the message
	messageHash := sha1.New().Sum([]byte(message))

	// Calculate the signature over the message
	r, s, err := dsa.Sign(rand.Reader, &signingKey, messageHash)
	if err != nil {
		return "", err
	}
	signatureData := append(r.Bytes(), s.Bytes()...)
	encodedSignature := base64.URLEncoding.EncodeToString(signatureData)

	token := message + "." + encodedSignature
	// Prepend

	return token, nil
}

func CreateAssertion(signingKey dsa.PrivateKey, certificate string, audience string, issuer string, issuedAt time.Time, expiresAt time.Time) (string, error) {
	assertionPayload, err := NewAssertionPayload(issuer, issuedAt, expiresAt, audience)
	if err != nil {
		return "", err
	}
	assertion, err := assertionPayload.Encode(signingKey)
	if err != nil {
		return "", err
	}
	return certificate + "~" + assertion, nil
}
<commit_msg>Actually fill in the public key<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/

package main

import (
	"crypto/dsa"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"time"
)

type CertificatePrincipal struct {
	Email string `json:"email"`
}

func NewCertificatePrincipal(email string) *CertificatePrincipal {
	return &CertificatePrincipal{Email: email}
}

type CertificatePublicKey struct {
	Algorithm string `json:"algorithm"`
	Y         string `json:"y"`
	P         string `json:"p"`
	Q         string `json:"q"`
	G         string `json:"g"`
}

func NewCertificatePublicKey(key dsa.PrivateKey) *CertificatePublicKey {
	return &CertificatePublicKey{
		Algorithm: "DS",
		Y:         fmt.Sprintf("%x", key.PublicKey.Y),
		P:         fmt.Sprintf("%x", key.PublicKey.Parameters.P),
		Q:         fmt.Sprintf("%x", key.PublicKey.Parameters.Q),
		G:         fmt.Sprintf("%x", key.PublicKey.Parameters.G),
	}
}

type CertificatePayload struct {
	Principal CertificatePrincipal `json:"principal"`
	PublicKey CertificatePublicKey `json:"public-key"`
	Issuer    string               `json:"iss"`
	IssuedAt  int64                `json:"iat"`
	Audience  string               `json:"aud"`
	ExpiresAt int64                `json:"exp"`
}

func NewCertificatePayload(priv dsa.PrivateKey, email string, issuer string, issuedAt time.Time, expiresAt time.Time) *CertificatePayload {
	return &CertificatePayload{
		Principal: *NewCertificatePrincipal(email),
		PublicKey: *NewCertificatePublicKey(priv),
		Issuer:    issuer,
		IssuedAt:  issuedAt.UnixNano() / int64(time.Millisecond),
		ExpiresAt: expiresAt.UnixNano() / int64(time.Millisecond),
	}
}

type WebTokenHeader struct {
	Algorithm string `json:"alg"`
}

func (p *CertificatePayload) Encode(priv dsa.PrivateKey) (string, error) {
	// Encode the header
	header := WebTokenHeader{Algorithm: "DS128"} // TODO: Create a function to extract this from the dsa.PrivateKey
	headerData, err := json.Marshal(header)
	if err != nil {
		return "", err
	}
	encodedHeader := base64.URLEncoding.EncodeToString(headerData)

	// Encode the certificate
	payloadData, err := json.Marshal(p)
	if err != nil {
		return "", err
	}
	encodedPayload := base64.URLEncoding.EncodeToString(payloadData)

	// The message is simply header plus payload
	message := encodedHeader + "." + encodedPayload

	// Get the SHA1 hash over the message
	messageHash := sha1.New().Sum([]byte(message))

	// Calculate the signature over the message
	r, s, err := dsa.Sign(rand.Reader, &priv, messageHash)
	if err != nil {
		return "", err
	}
	signatureData := append(r.Bytes(), s.Bytes()...)
	encodedSignature := base64.URLEncoding.EncodeToString(signatureData)

	token := message + "." + encodedSignature
	// Prepend

	return token, nil
}

func CreateCertificate(key dsa.PrivateKey, email string, issuer string, issuedAt time.Time, expiresAt time.Time, signingKey dsa.PrivateKey) (string, error) {
	certificatePayload := NewCertificatePayload(key, email, issuer, issuedAt, expiresAt)
	return certificatePayload.Encode(signingKey)
}

//

type AssertionPayload struct {
	Issuer    string `json:"iss"`
	IssuedAt  int64  `json:"iat"`
	ExpiresAt int64  `json:"exp"`
	Audience  string `json:"aud"`
}

func NewAssertionPayload(issuer string, issuedAt time.Time, expiresAt time.Time, audience string) (*AssertionPayload, error) {
	return &AssertionPayload{
		Issuer:    issuer,
		IssuedAt:  issuedAt.UnixNano() / int64(time.Millisecond),
		ExpiresAt: expiresAt.UnixNano() / int64(time.Millisecond),
		Audience:  audience,
	}, nil
}

func (p *AssertionPayload) Encode(signingKey dsa.PrivateKey) (string, error) {
	// Encode the header
	header := WebTokenHeader{Algorithm: "DS128"} // TODO: Create a function to extract this from the dsa.PrivateKey
	headerData, err := json.Marshal(header)
	if err != nil {
		return "", err
	}
	encodedHeader := base64.URLEncoding.EncodeToString(headerData)

	// Encode the certificate
	payloadData, err := json.Marshal(p)
	if err != nil {
		return "", err
	}
	encodedPayload := base64.URLEncoding.EncodeToString(payloadData)

	// The message is simply header plus payload
	message := encodedHeader + "." + encodedPayload

	// Get the SHA1 hash over the message
	messageHash := sha1.New().Sum([]byte(message))

	// Calculate the signature over the message
	r, s, err := dsa.Sign(rand.Reader, &signingKey, messageHash)
	if err != nil {
		return "", err
	}
	signatureData := append(r.Bytes(), s.Bytes()...)
	encodedSignature := base64.URLEncoding.EncodeToString(signatureData)

	token := message + "." + encodedSignature
	// Prepend

	return token, nil
}

func CreateAssertion(signingKey dsa.PrivateKey, certificate string, audience string, issuer string, issuedAt time.Time, expiresAt time.Time) (string, error) {
	assertionPayload, err := NewAssertionPayload(issuer, issuedAt, expiresAt, audience)
	if err != nil {
		return "", err
	}
	assertion, err := assertionPayload.Encode(signingKey)
	if err != nil {
		return "", err
	}
	return certificate + "~" + assertion, nil
}
<|endoftext|>


<commit_before>package lexicon

import (
	"errors"
	"sync"

	"github.com/PreetamJinka/orderedlist"
)

var ErrKeyNotPresent = errors.New("lexicon: key not present")

type KeyValue struct {
	Key   interface{}
	Value interface{}
}

// Lexicon is an ordered key-value store.
type Lexicon struct {
	list    *orderedlist.OrderedList
	hashmap map[interface{}]interface{}
	mutex   *sync.Mutex
	compare func(a, b interface{}) int
}

// New returns an initialized lexicon.
func New(compare func(a, b interface{}) int) *Lexicon {
	return &Lexicon{
		list:    orderedlist.New(compare),
		hashmap: make(map[interface{}]interface{}),
		mutex:   &sync.Mutex{},
	}
}

// Set sets a key to a value.
func (lex *Lexicon) Set(key, value interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()

	_, present := lex.hashmap[key]

	if !present {
		lex.hashmap[key] = value
		lex.list.Insert(key)
	} else {
		lex.hashmap[key] = value
	}
}

// SetMany sets multiple key-value pairs.
func (lex *Lexicon) SetMany(kv map[interface{}]interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()
	for key := range kv {
		_, present := lex.hashmap[key]

		if !present {
			lex.hashmap[key] = kv[key]
			lex.list.Insert(key)
		} else {
			lex.hashmap[key] = kv[key]
		}
	}
}

// Get returns a value at the given key.
func (lex *Lexicon) Get(key interface{}) (interface{}, error) {
	val, present := lex.hashmap[key]

	if !present {
		return "", ErrKeyNotPresent
	}

	return val, nil
}

// Remove deletes a key-value pair.
func (lex *Lexicon) Remove(key interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()
	delete(lex.hashmap, key)
	lex.list.Remove(key)
}

// ClearRange removes a range of key-value pairs.
// The range is from [start, end).
func (lex *Lexicon) ClearRange(start, end interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()

	keys := lex.list.GetRange(start, end)

	for _, key := range keys {
		delete(lex.hashmap, key)
		lex.list.Remove(key)
	}
}

// GetRange returns a slice of KeyValue structs.
// The range is from [start, end).
func (lex *Lexicon) GetRange(start, end interface{}) (kv []KeyValue) {
	kv = make([]KeyValue, 0)

	lex.mutex.Lock()
	keys := lex.list.GetRange(start, end)
	lex.mutex.Unlock()

	for _, key := range keys {
		kv = append(kv, KeyValue{
			Key:   key,
			Value: lex.hashmap[key],
		})
	}
	return
}
<commit_msg>Set a slice capacity<commit_after>package lexicon

import (
	"errors"
	"sync"

	"github.com/PreetamJinka/orderedlist"
)

var ErrKeyNotPresent = errors.New("lexicon: key not present")

type KeyValue struct {
	Key   interface{}
	Value interface{}
}

// Lexicon is an ordered key-value store.
type Lexicon struct {
	list    *orderedlist.OrderedList
	hashmap map[interface{}]interface{}
	mutex   *sync.Mutex
	compare func(a, b interface{}) int
}

// New returns an initialized lexicon.
func New(compare func(a, b interface{}) int) *Lexicon {
	return &Lexicon{
		list:    orderedlist.New(compare),
		hashmap: make(map[interface{}]interface{}),
		mutex:   &sync.Mutex{},
	}
}

// Set sets a key to a value.
func (lex *Lexicon) Set(key, value interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()

	_, present := lex.hashmap[key]

	if !present {
		lex.hashmap[key] = value
		lex.list.Insert(key)
	} else {
		lex.hashmap[key] = value
	}
}

// SetMany sets multiple key-value pairs.
func (lex *Lexicon) SetMany(kv map[interface{}]interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()
	for key := range kv {
		_, present := lex.hashmap[key]

		if !present {
			lex.hashmap[key] = kv[key]
			lex.list.Insert(key)
		} else {
			lex.hashmap[key] = kv[key]
		}
	}
}

// Get returns a value at the given key.
func (lex *Lexicon) Get(key interface{}) (interface{}, error) {
	val, present := lex.hashmap[key]

	if !present {
		return "", ErrKeyNotPresent
	}

	return val, nil
}

// Remove deletes a key-value pair.
func (lex *Lexicon) Remove(key interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()
	delete(lex.hashmap, key)
	lex.list.Remove(key)
}

// ClearRange removes a range of key-value pairs.
// The range is from [start, end).
func (lex *Lexicon) ClearRange(start, end interface{}) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()

	keys := lex.list.GetRange(start, end)

	for _, key := range keys {
		delete(lex.hashmap, key)
		lex.list.Remove(key)
	}
}

// GetRange returns a slice of KeyValue structs.
// The range is from [start, end).
func (lex *Lexicon) GetRange(start, end interface{}) (kv []KeyValue) {
	lex.mutex.Lock()
	defer lex.mutex.Unlock()

	keys := lex.list.GetRange(start, end)
	kv = make([]KeyValue, 0, len(keys))

	for _, key := range keys {
		kv = append(kv, KeyValue{
			Key:   key,
			Value: lex.hashmap[key],
		})
	}
	return
}
<|endoftext|>


<commit_before>package protocol

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/binary"
	"fmt"
)

const (
	CONFIGTX_SIZE = 83

	BLOCK_SIZE_ID     = 1
	DIFF_INTERVAL_ID  = 2
	FEE_MINIMUM_ID    = 3
	BLOCK_INTERVAL_ID = 4
	BLOCK_REWARD_ID   = 5

	MIN_BLOCK_SIZE = 1000      //1KB
	MAX_BLOCK_SIZE = 100000000 //100MB

	MIN_DIFF_INTERVAL = 30 //amount in seconds
	MAX_DIFF_INTERVAL = 9223372036854775807

	MIN_FEE_MINIMUM = 0
	MAX_FEE_MINIMUM = 9223372036854775807

	MIN_BLOCK_INTERVAL = 30    //30 seconds
	MAX_BLOCK_INTERVAL = 86400 //24 hours

	MIN_BLOCK_REWARD = 0
	MAX_BLOCK_REWARD = 1152921504606846976 //2^60
)

type ConfigTx struct {
	Header  uint8
	Id      uint8
	Payload uint64
	Fee     uint64
	TxCnt	uint8
	Sig     [64]byte
}

func ConstrConfigTx(header uint8, id uint8, payload uint64, fee uint64, txCnt uint8, rootPrivKey *ecdsa.PrivateKey) (tx *ConfigTx, err error) {

	tx = new(ConfigTx)
	tx.Header = header
	tx.Id = id
	tx.Payload = payload
	tx.Fee = fee
	tx.TxCnt = txCnt

	txHash := tx.Hash()

	r, s, err := ecdsa.Sign(rand.Reader, rootPrivKey, txHash[:])

	if err != nil {
		return nil, err
	}

	copy(tx.Sig[32-len(r.Bytes()):32], r.Bytes())
	copy(tx.Sig[64-len(s.Bytes()):], s.Bytes())

	return tx, nil
}

func (tx *ConfigTx) Hash() (hash [32]byte) {

	if tx == nil {
		return [32]byte{}
	}

	txHash := struct {
		Header  uint8
		Id      uint8
		Payload uint64
		Fee     uint64
		TxCnt	uint8
	}{
		tx.Header,
		tx.Id,
		tx.Payload,
		tx.Fee,
		tx.TxCnt,
	}
	return serializeHashContent(txHash)
}

func (tx *ConfigTx) Encode() (encodedTx []byte) {

	if tx == nil {
		return nil
	}

	var buf bytes.Buffer
	var payloadBuf [8]byte
	var feeBuf [8]byte

	binary.Write(&buf, binary.BigEndian, tx.Payload)
	copy(payloadBuf[:], buf.Bytes())
	buf.Reset()
	binary.Write(&buf, binary.BigEndian, tx.Fee)
	copy(feeBuf[:], buf.Bytes())
	buf.Reset()

	encodedTx = make([]byte, CONFIGTX_SIZE)
	encodedTx[0] = tx.Header
	encodedTx[1] = tx.Id
	copy(encodedTx[2:10], payloadBuf[:])
	copy(encodedTx[10:18], feeBuf[:])
	encodedTx[18] = byte(tx.TxCnt)
	copy(encodedTx[19:83], tx.Sig[:])

	return encodedTx
}

func (*ConfigTx) Decode(encodedTx []byte) (tx *ConfigTx) {

	if len(encodedTx) != CONFIGTX_SIZE {
		return nil
	}

	tx = new(ConfigTx)
	tx.Header = encodedTx[0]
	tx.Id = encodedTx[1]
	tx.Payload = binary.BigEndian.Uint64(encodedTx[2:10])
	tx.Fee = binary.BigEndian.Uint64(encodedTx[10:18])
	tx.TxCnt = uint8(encodedTx[18])
	copy(tx.Sig[:], encodedTx[19:83])

	return tx
}

func (tx *ConfigTx) TxFee() uint64 { return tx.Fee }
func (tx *ConfigTx) Size() uint64   { return CONFIGTX_SIZE }

func (tx ConfigTx) String() string {
	return fmt.Sprintf(
		"\n"+
			"Id: %v\n"+
			"Payload: %v\n"+
			"Fee: %v\n"+
			"TxCnt: %v\n",
		tx.Id,
		tx.Payload,
		tx.Fee,
		tx.TxCnt,
	)
}
<commit_msg>Making header of configTx consistent with other txs.<commit_after>package protocol

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/binary"
	"fmt"
)

const (
	CONFIGTX_SIZE = 83

	BLOCK_SIZE_ID     = 1
	DIFF_INTERVAL_ID  = 2
	FEE_MINIMUM_ID    = 3
	BLOCK_INTERVAL_ID = 4
	BLOCK_REWARD_ID   = 5

	MIN_BLOCK_SIZE = 1000      //1KB
	MAX_BLOCK_SIZE = 100000000 //100MB

	MIN_DIFF_INTERVAL = 30 //amount in seconds
	MAX_DIFF_INTERVAL = 9223372036854775807

	MIN_FEE_MINIMUM = 0
	MAX_FEE_MINIMUM = 9223372036854775807

	MIN_BLOCK_INTERVAL = 30    //30 seconds
	MAX_BLOCK_INTERVAL = 86400 //24 hours

	MIN_BLOCK_REWARD = 0
	MAX_BLOCK_REWARD = 1152921504606846976 //2^60
)

type ConfigTx struct {
	Header  byte
	Id      uint8
	Payload uint64
	Fee     uint64
	TxCnt	uint8
	Sig     [64]byte
}

func ConstrConfigTx(header byte, id uint8, payload uint64, fee uint64, txCnt uint8, rootPrivKey *ecdsa.PrivateKey) (tx *ConfigTx, err error) {

	tx = new(ConfigTx)
	tx.Header = header
	tx.Id = id
	tx.Payload = payload
	tx.Fee = fee
	tx.TxCnt = txCnt

	txHash := tx.Hash()

	r, s, err := ecdsa.Sign(rand.Reader, rootPrivKey, txHash[:])

	if err != nil {
		return nil, err
	}

	copy(tx.Sig[32-len(r.Bytes()):32], r.Bytes())
	copy(tx.Sig[64-len(s.Bytes()):], s.Bytes())

	return tx, nil
}

func (tx *ConfigTx) Hash() (hash [32]byte) {

	if tx == nil {
		return [32]byte{}
	}

	txHash := struct {
		Header  byte
		Id      uint8
		Payload uint64
		Fee     uint64
		TxCnt	uint8
	}{
		tx.Header,
		tx.Id,
		tx.Payload,
		tx.Fee,
		tx.TxCnt,
	}
	return serializeHashContent(txHash)
}

func (tx *ConfigTx) Encode() (encodedTx []byte) {

	if tx == nil {
		return nil
	}

	var buf bytes.Buffer
	var payloadBuf [8]byte
	var feeBuf [8]byte

	binary.Write(&buf, binary.BigEndian, tx.Payload)
	copy(payloadBuf[:], buf.Bytes())
	buf.Reset()
	binary.Write(&buf, binary.BigEndian, tx.Fee)
	copy(feeBuf[:], buf.Bytes())
	buf.Reset()

	encodedTx = make([]byte, CONFIGTX_SIZE)
	encodedTx[0] = tx.Header
	encodedTx[1] = tx.Id
	copy(encodedTx[2:10], payloadBuf[:])
	copy(encodedTx[10:18], feeBuf[:])
	encodedTx[18] = byte(tx.TxCnt)
	copy(encodedTx[19:83], tx.Sig[:])

	return encodedTx
}

func (*ConfigTx) Decode(encodedTx []byte) (tx *ConfigTx) {

	if len(encodedTx) != CONFIGTX_SIZE {
		return nil
	}

	tx = new(ConfigTx)
	tx.Header = encodedTx[0]
	tx.Id = encodedTx[1]
	tx.Payload = binary.BigEndian.Uint64(encodedTx[2:10])
	tx.Fee = binary.BigEndian.Uint64(encodedTx[10:18])
	tx.TxCnt = uint8(encodedTx[18])
	copy(tx.Sig[:], encodedTx[19:83])

	return tx
}

func (tx *ConfigTx) TxFee() uint64 { return tx.Fee }
func (tx *ConfigTx) Size() uint64   { return CONFIGTX_SIZE }

func (tx ConfigTx) String() string {
	return fmt.Sprintf(
		"\n"+
			"Id: %v\n"+
			"Payload: %v\n"+
			"Fee: %v\n"+
			"TxCnt: %v\n",
		tx.Id,
		tx.Payload,
		tx.Fee,
		tx.TxCnt,
	)
}
<|endoftext|>


<commit_before>package main // import "github.com/jmccarty3/kube2consul"

import (
  "github.com/golang/glog"
  //"k8s.io/kubernetes/pkg/api"
  kapi "k8s.io/kubernetes/pkg/api"
  kclient "k8s.io/kubernetes/pkg/client/unversioned"
  kfields "k8s.io/kubernetes/pkg/fields"
	klabels "k8s.io/kubernetes/pkg/labels"
)

//TODO: Chang to store Node pointer. Add getName, getReadyStatus accessors
type KubeNode struct {
  name string
  readyStatus bool;
  serviceIDS map[string]string;
  address string
}

type KubeBookKeeper interface {
  AddNode(*kapi.Node);
  RemoveNode(*kapi.Node);
  UpdateNode(*kapi.Node);
  SyncNodes();
  AddService(*kapi.Service)
  RemoveService(*kapi.Service)
  UpdateService(*kapi.Service)
}

type ClientBookKeeper struct {
  KubeBookKeeper
  client *kclient.Client
  nodes map[string]*KubeNode;
  services map[string]*kapi.Service
  consulQueue chan<- ConsulWork
}

func BuildServiceBaseID(nodeName string, service *kapi.Service) string {
  return nodeName + "-" + service.Name
}

func NewKubeNode() *KubeNode {
  return &KubeNode {
    name: "",
    readyStatus: false,
    serviceIDS: make(map[string]string),
  }
}

func NewClientBookKeeper(client *kclient.Client) *ClientBookKeeper {
  return &ClientBookKeeper {
    client: client,
    nodes: make(map[string]*KubeNode),
    services: make(map[string]*kapi.Service),
  }
}

func RunBookKeeper(workQue <-chan KubeWork, consulQueue chan<- ConsulWork, apiClient *kclient.Client) {

  client := NewClientBookKeeper(apiClient)
  client.consulQueue = consulQueue

  for work := range workQue {
    switch work.Action {
    case KubeWorkAddNode:
      client.AddNode(work.Node)
    case KubeWorkRemoveNode:
      client.RemoveNode(work.Node)
    case KubeWorkAddService:
      client.AddService(work.Service)
    case KubeWorkRemoveService:
      client.RemoveService(work.Service)
    case KubeWorkUpdateService:
      client.UpdateService(work.Service)
    case KubeWorkSync:
      client.Sync()
    default:
      glog.Info("Unsupported work action: ", work.Action)
    }
  }

  glog.Info("Completed all node work")
}

func (client *ClientBookKeeper) AttachServiceToNode(node *KubeNode, service *kapi.Service) {
  baseID := BuildServiceBaseID(node.name, service)
  client.consulQueue <- ConsulWork {
    Action: ConsulWorkAddDNS,
    Service: service,
    Config: DnsInfo {
      BaseID: baseID,
      IPAddress: node.address,
      },
  }
  glog.V(3).Info("Requesting Addition of services with Base ID: ", baseID)
  node.serviceIDS[service.Name] = baseID
}

func (client *ClientBookKeeper) DetachServiceFromNode(node *KubeNode, service *kapi.Service) {
  if baseID,ok := node.serviceIDS[service.Name]; ok {
    //To Consol -> TODO
    client.consulQueue <- ConsulWork {
      Action: ConsulWorkRemoveDNS,
      Config: DnsInfo {
        BaseID: baseID,
      },
    }

    glog.V(3).Info("Requesting Removal of services with Base ID: ", baseID)
    delete(node.serviceIDS, service.Name)
  }
}

func (client *ClientBookKeeper) AddAllServicesToNode(node *KubeNode) {
  for _,service := range client.services {
    client.AttachServiceToNode(node, service)
  }
}

func (client *ClientBookKeeper) RemoveAllServicesFromNode(node *KubeNode) {
  for _,service := range client.services {
    client.DetachServiceFromNode(node,service)
  }
}

func (client *ClientBookKeeper) AddNode(newNode *kapi.Node) {
  if _, ok := client.nodes[newNode.Name]; ok {
    glog.Error("Attempted to Add existing node ", newNode.Name)
    return
  }

  //Add to Node Collection
  createdNode := NewKubeNode()
  createdNode.readyStatus = nodeReady(newNode)
  createdNode.name = newNode.Name
  createdNode.address = newNode.Status.Addresses[0].Address

  //Send request for Service Addition for node and all serviceIDS (Create Service ID here)
  if createdNode.readyStatus {
    client.AddAllServicesToNode(createdNode)
  }
  client.nodes[newNode.Name] = createdNode
  glog.Info("Added Node: ", newNode.Name)
}

func (client *ClientBookKeeper) RemoveNode(oldNode *kapi.Node) {
  if node,ok := client.nodes[oldNode.Name]; ok {
    //Remove All DNS for node
    client.RemoveAllServicesFromNode(node)
    //Remove Node from Collection
    delete(client.nodes, oldNode.Name)
  } else {
    glog.Error("Attmepted to remove missing node: ", oldNode.Name)
  }

}

func (client *ClientBookKeeper) UpdateNode(updatedNode *kapi.Node) {
  //If now ready -> Service Addtion for node
  //TODO Check it exists
  if nodeReady(updatedNode) {
    client.AddAllServicesToNode(client.nodes[updatedNode.Name])
  } else {
    client.RemoveAllServicesFromNode(client.nodes[updatedNode.Name])
  }
  //Else -> Service Removal for Node
  //UnLock
}

func ContainsNodeName(name string, nodes *kapi.NodeList) bool {
  for _,node := range nodes.Items {
    if node.ObjectMeta.Name == name {
      return true
    }
  }
  return false
}

func (client *ClientBookKeeper) Sync() {
  nodes := client.client.Nodes()
  if nodeList, err := nodes.List(klabels.Everything(), kfields.Everything()); err == nil {
    for name,_ := range client.nodes {
      if !ContainsNodeName(name, nodeList) {
          glog.Errorf("Bookkeeper has node: %s that does not exist in api server", name)
      }
    }
  }
  //Add Remove as needed
  //UnLock
}

func (client *ClientBookKeeper) AddService(service *kapi.Service) {
  //TODO Verify it doesn't exist
  client.services[service.Name] = service
  //Perform All DNS Adds
  for _,node := range client.nodes {
    client.AttachServiceToNode(node, service)
  }

  glog.Info("Added Service: ", service.Name)
}

func  (client *ClientBookKeeper) RemoveService(service *kapi.Service) {
  //TODO Verify it does exist
  //Perform All DNS Removes
  for _,node := range client.nodes {
    client.DetachServiceFromNode(node, service)
  }

  delete(client.services, service.Name)
  glog.Info("Removed Service: ", service.Name)
}

func (client *ClientBookKeeper) UpdateService(service *kapi.Service) {
  client.RemoveService(service)
  client.AddService(service)
}
<commit_msg>Removing nodes by name. Allows kube sync to remove nodes easily<commit_after>package main // import "github.com/jmccarty3/kube2consul"

import (
  "github.com/golang/glog"
  //"k8s.io/kubernetes/pkg/api"
  kapi "k8s.io/kubernetes/pkg/api"
  kclient "k8s.io/kubernetes/pkg/client/unversioned"
  kfields "k8s.io/kubernetes/pkg/fields"
	klabels "k8s.io/kubernetes/pkg/labels"
)

//TODO: Chang to store Node pointer. Add getName, getReadyStatus accessors
type KubeNode struct {
  name string
  readyStatus bool;
  serviceIDS map[string]string;
  address string
}

type KubeBookKeeper interface {
  AddNode(*kapi.Node);
  RemoveNode(*kapi.Node);
  UpdateNode(*kapi.Node);
  SyncNodes();
  AddService(*kapi.Service)
  RemoveService(*kapi.Service)
  UpdateService(*kapi.Service)
}

type ClientBookKeeper struct {
  KubeBookKeeper
  client *kclient.Client
  nodes map[string]*KubeNode;
  services map[string]*kapi.Service
  consulQueue chan<- ConsulWork
}

func BuildServiceBaseID(nodeName string, service *kapi.Service) string {
  return nodeName + "-" + service.Name
}

func NewKubeNode() *KubeNode {
  return &KubeNode {
    name: "",
    readyStatus: false,
    serviceIDS: make(map[string]string),
  }
}

func NewClientBookKeeper(client *kclient.Client) *ClientBookKeeper {
  return &ClientBookKeeper {
    client: client,
    nodes: make(map[string]*KubeNode),
    services: make(map[string]*kapi.Service),
  }
}

func RunBookKeeper(workQue <-chan KubeWork, consulQueue chan<- ConsulWork, apiClient *kclient.Client) {

  client := NewClientBookKeeper(apiClient)
  client.consulQueue = consulQueue

  for work := range workQue {
    switch work.Action {
    case KubeWorkAddNode:
      client.AddNode(work.Node)
    case KubeWorkRemoveNode:
      client.RemoveNode(work.Node.Name)
    case KubeWorkAddService:
      client.AddService(work.Service)
    case KubeWorkRemoveService:
      client.RemoveService(work.Service)
    case KubeWorkUpdateService:
      client.UpdateService(work.Service)
    case KubeWorkSync:
      client.Sync()
    default:
      glog.Info("Unsupported work action: ", work.Action)
    }
  }

  glog.Info("Completed all node work")
}

func (client *ClientBookKeeper) AttachServiceToNode(node *KubeNode, service *kapi.Service) {
  baseID := BuildServiceBaseID(node.name, service)
  client.consulQueue <- ConsulWork {
    Action: ConsulWorkAddDNS,
    Service: service,
    Config: DnsInfo {
      BaseID: baseID,
      IPAddress: node.address,
      },
  }
  glog.V(3).Info("Requesting Addition of services with Base ID: ", baseID)
  node.serviceIDS[service.Name] = baseID
}

func (client *ClientBookKeeper) DetachServiceFromNode(node *KubeNode, service *kapi.Service) {
  if baseID,ok := node.serviceIDS[service.Name]; ok {
    //To Consol -> TODO
    client.consulQueue <- ConsulWork {
      Action: ConsulWorkRemoveDNS,
      Config: DnsInfo {
        BaseID: baseID,
      },
    }

    glog.V(3).Info("Requesting Removal of services with Base ID: ", baseID)
    delete(node.serviceIDS, service.Name)
  }
}

func (client *ClientBookKeeper) AddAllServicesToNode(node *KubeNode) {
  for _,service := range client.services {
    client.AttachServiceToNode(node, service)
  }
}

func (client *ClientBookKeeper) RemoveAllServicesFromNode(node *KubeNode) {
  for _,service := range client.services {
    client.DetachServiceFromNode(node,service)
  }
}

func (client *ClientBookKeeper) AddNode(newNode *kapi.Node) {
  if _, ok := client.nodes[newNode.Name]; ok {
    glog.Error("Attempted to Add existing node ", newNode.Name)
    return
  }

  //Add to Node Collection
  createdNode := NewKubeNode()
  createdNode.readyStatus = nodeReady(newNode)
  createdNode.name = newNode.Name
  createdNode.address = newNode.Status.Addresses[0].Address

  //Send request for Service Addition for node and all serviceIDS (Create Service ID here)
  if createdNode.readyStatus {
    client.AddAllServicesToNode(createdNode)
  }
  client.nodes[newNode.Name] = createdNode
  glog.Info("Added Node: ", newNode.Name)
}

func (client *ClientBookKeeper) RemoveNode(oldNodeName string) {
  if node,ok := client.nodes[oldNodeName]; ok {
    //Remove All DNS for node
    client.RemoveAllServicesFromNode(node)
    //Remove Node from Collection
    delete(client.nodes, oldNodeName)
  } else {
    glog.Error("Attmepted to remove missing node: ", oldNodeName)
  }

}

func (client *ClientBookKeeper) UpdateNode(updatedNode *kapi.Node) {
  //If now ready -> Service Addtion for node
  //TODO Check it exists
  if nodeReady(updatedNode) {
    client.AddAllServicesToNode(client.nodes[updatedNode.Name])
  } else {
    client.RemoveAllServicesFromNode(client.nodes[updatedNode.Name])
  }
  //Else -> Service Removal for Node
  //UnLock
}

func ContainsNodeName(name string, nodes *kapi.NodeList) bool {
  for _,node := range nodes.Items {
    if node.ObjectMeta.Name == name {
      return true
    }
  }
  return false
}

func (client *ClientBookKeeper) Sync() {
  nodes := client.client.Nodes()
  if nodeList, err := nodes.List(klabels.Everything(), kfields.Everything()); err == nil {
    for name,_ := range client.nodes {
      if !ContainsNodeName(name, nodeList) {
          glog.Errorf("Bookkeeper has node: %s that does not exist in api server", name)
          client.RemoveNode(name)
      }
    }
  }
  //Add Remove as needed
  //UnLock
}

func (client *ClientBookKeeper) AddService(service *kapi.Service) {
  //TODO Verify it doesn't exist
  client.services[service.Name] = service
  //Perform All DNS Adds
  for _,node := range client.nodes {
    client.AttachServiceToNode(node, service)
  }

  glog.Info("Added Service: ", service.Name)
}

func  (client *ClientBookKeeper) RemoveService(service *kapi.Service) {
  //TODO Verify it does exist
  //Perform All DNS Removes
  for _,node := range client.nodes {
    client.DetachServiceFromNode(node, service)
  }

  delete(client.services, service.Name)
  glog.Info("Removed Service: ", service.Name)
}

func (client *ClientBookKeeper) UpdateService(service *kapi.Service) {
  client.RemoveService(service)
  client.AddService(service)
}
<|endoftext|>


<commit_before>package provider

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
	"unicode"

	"github.com/BurntSushi/toml"
	"github.com/containous/traefik/autogen"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
)

// Provider defines methods of a provider.
type Provider interface {
	// Provide allows the provider to provide configurations to traefik
	// using the given configuration channel.
	Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error
}

// BaseProvider should be inherited by providers
type BaseProvider struct {
	Watch       bool              `description:"Watch provider"`
	Filename    string            `description:"Override default configuration template. For advanced users :)"`
	Constraints types.Constraints `description:"Filter services by constraint, matching with Traefik tags."`
}

// MatchConstraints must match with EVERY single contraint
// returns first constraint that do not match or nil
func (p *BaseProvider) MatchConstraints(tags []string) (bool, *types.Constraint) {
	// if there is no tags and no contraints, filtering is disabled
	if len(tags) == 0 && len(p.Constraints) == 0 {
		return true, nil
	}

	for _, constraint := range p.Constraints {
		// xor: if ok and constraint.MustMatch are equal, then no tag is currently matching with the constraint
		if ok := constraint.MatchConstraintWithAtLeastOneTag(tags); ok != constraint.MustMatch {
			return false, constraint
		}
	}

	// If no constraint or every constraints matching
	return true, nil
}

// GetConfiguration return the provider configuration using templating
func (p *BaseProvider) GetConfiguration(defaultTemplateFile string, funcMap template.FuncMap, templateObjects interface{}) (*types.Configuration, error) {
	var (
		buf []byte
		err error
	)
	configuration := new(types.Configuration)
	var defaultFuncMap = template.FuncMap{
		"replace":   Replace,
		"tolower":   strings.ToLower,
		"normalize": Normalize,
		"split":     split,
		"contains":  contains,
	}

	for funcID, funcElement := range funcMap {
		defaultFuncMap[funcID] = funcElement
	}

	tmpl := template.New(p.Filename).Funcs(defaultFuncMap)
	if len(p.Filename) > 0 {
		buf, err = ioutil.ReadFile(p.Filename)
		if err != nil {
			return nil, err
		}
	} else {
		buf, err = autogen.Asset(defaultTemplateFile)
		if err != nil {
			return nil, err
		}
	}
	_, err = tmpl.Parse(string(buf))
	if err != nil {
		return nil, err
	}

	var buffer bytes.Buffer
	err = tmpl.Execute(&buffer, templateObjects)
	if err != nil {
		return nil, err
	}

	var renderedTemplate = buffer.String()
	//	log.Debugf("Rendering results of %s:\n%s", defaultTemplateFile, renderedTemplate)
	if _, err := toml.Decode(renderedTemplate, configuration); err != nil {
		return nil, err
	}
	return configuration, nil
}

// Replace is an alias for strings.Replace
func Replace(s1 string, s2 string, s3 string) string {
	return strings.Replace(s3, s1, s2, -1)
}

func contains(substr, s string) bool {
	return strings.Contains(s, substr)
}

func split(sep, s string) []string {
	return strings.Split(s, sep)
}

// Normalize transform a string that work with the rest of traefik
func Normalize(name string) string {
	fargs := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	// get function
	return strings.Join(strings.FieldsFunc(name, fargs), "-")
}

// ReverseStringSlice invert the order of the given slice of string
func ReverseStringSlice(slice *[]string) {
	for i, j := 0, len(*slice)-1; i < j; i, j = i+1, j-1 {
		(*slice)[i], (*slice)[j] = (*slice)[j], (*slice)[i]
	}
}

// ClientTLS holds TLS specific configurations as client
// CA, Cert and Key can be either path or file contents
type ClientTLS struct {
	CA                 string `description:"TLS CA"`
	Cert               string `description:"TLS cert"`
	Key                string `description:"TLS key"`
	InsecureSkipVerify bool   `description:"TLS insecure skip verify"`
}

// CreateTLSConfig creates a TLS config from ClientTLS structures
func (clientTLS *ClientTLS) CreateTLSConfig() (*tls.Config, error) {
	var err error
	if clientTLS == nil {
		log.Warnf("clientTLS is nil")
		return nil, nil
	}
	caPool := x509.NewCertPool()
	if clientTLS.CA != "" {
		var ca []byte
		if _, errCA := os.Stat(clientTLS.CA); errCA == nil {
			ca, err = ioutil.ReadFile(clientTLS.CA)
			if err != nil {
				return nil, fmt.Errorf("Failed to read CA. %s", err)
			}
		} else {
			ca = []byte(clientTLS.CA)
		}
		caPool.AppendCertsFromPEM(ca)
	}

	cert := tls.Certificate{}
	_, errKeyIsFile := os.Stat(clientTLS.Key)

	if _, errCertIsFile := os.Stat(clientTLS.Cert); errCertIsFile == nil {
		if errKeyIsFile == nil {
			cert, err = tls.LoadX509KeyPair(clientTLS.Cert, clientTLS.Key)
			if err != nil {
				return nil, fmt.Errorf("Failed to load TLS keypair: %v", err)
			}
		} else {
			return nil, fmt.Errorf("tls cert is a file, but tls key is not")
		}
	} else {
		if errKeyIsFile != nil {
			cert, err = tls.X509KeyPair([]byte(clientTLS.Cert), []byte(clientTLS.Key))
			if err != nil {
				return nil, fmt.Errorf("Failed to load TLS keypair: %v", err)

			}
		} else {
			return nil, fmt.Errorf("tls key is a file, but tls cert is not")
		}
	}

	TLSConfig := &tls.Config{
		Certificates:       []tls.Certificate{cert},
		RootCAs:            caPool,
		InsecureSkipVerify: clientTLS.InsecureSkipVerify,
	}
	return TLSConfig, nil
}
<commit_msg>feat: Add Trace in Base Provider.<commit_after>package provider

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
	"unicode"

	"github.com/BurntSushi/toml"
	"github.com/containous/traefik/autogen"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
)

// Provider defines methods of a provider.
type Provider interface {
	// Provide allows the provider to provide configurations to traefik
	// using the given configuration channel.
	Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error
}

// BaseProvider should be inherited by providers
type BaseProvider struct {
	Watch       bool              `description:"Watch provider"`
	Filename    string            `description:"Override default configuration template. For advanced users :)"`
	Constraints types.Constraints `description:"Filter services by constraint, matching with Traefik tags."`
	Trace       bool              `description:"Display additional provider logs (if available)."`
}

// MatchConstraints must match with EVERY single contraint
// returns first constraint that do not match or nil
func (p *BaseProvider) MatchConstraints(tags []string) (bool, *types.Constraint) {
	// if there is no tags and no contraints, filtering is disabled
	if len(tags) == 0 && len(p.Constraints) == 0 {
		return true, nil
	}

	for _, constraint := range p.Constraints {
		// xor: if ok and constraint.MustMatch are equal, then no tag is currently matching with the constraint
		if ok := constraint.MatchConstraintWithAtLeastOneTag(tags); ok != constraint.MustMatch {
			return false, constraint
		}
	}

	// If no constraint or every constraints matching
	return true, nil
}

// GetConfiguration return the provider configuration using templating
func (p *BaseProvider) GetConfiguration(defaultTemplateFile string, funcMap template.FuncMap, templateObjects interface{}) (*types.Configuration, error) {
	var (
		buf []byte
		err error
	)
	configuration := new(types.Configuration)
	var defaultFuncMap = template.FuncMap{
		"replace":   Replace,
		"tolower":   strings.ToLower,
		"normalize": Normalize,
		"split":     split,
		"contains":  contains,
	}

	for funcID, funcElement := range funcMap {
		defaultFuncMap[funcID] = funcElement
	}

	tmpl := template.New(p.Filename).Funcs(defaultFuncMap)
	if len(p.Filename) > 0 {
		buf, err = ioutil.ReadFile(p.Filename)
		if err != nil {
			return nil, err
		}
	} else {
		buf, err = autogen.Asset(defaultTemplateFile)
		if err != nil {
			return nil, err
		}
	}
	_, err = tmpl.Parse(string(buf))
	if err != nil {
		return nil, err
	}

	var buffer bytes.Buffer
	err = tmpl.Execute(&buffer, templateObjects)
	if err != nil {
		return nil, err
	}

	var renderedTemplate = buffer.String()
	//	log.Debugf("Rendering results of %s:\n%s", defaultTemplateFile, renderedTemplate)
	if _, err := toml.Decode(renderedTemplate, configuration); err != nil {
		return nil, err
	}
	return configuration, nil
}

// Replace is an alias for strings.Replace
func Replace(s1 string, s2 string, s3 string) string {
	return strings.Replace(s3, s1, s2, -1)
}

func contains(substr, s string) bool {
	return strings.Contains(s, substr)
}

func split(sep, s string) []string {
	return strings.Split(s, sep)
}

// Normalize transform a string that work with the rest of traefik
func Normalize(name string) string {
	fargs := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	// get function
	return strings.Join(strings.FieldsFunc(name, fargs), "-")
}

// ReverseStringSlice invert the order of the given slice of string
func ReverseStringSlice(slice *[]string) {
	for i, j := 0, len(*slice)-1; i < j; i, j = i+1, j-1 {
		(*slice)[i], (*slice)[j] = (*slice)[j], (*slice)[i]
	}
}

// ClientTLS holds TLS specific configurations as client
// CA, Cert and Key can be either path or file contents
type ClientTLS struct {
	CA                 string `description:"TLS CA"`
	Cert               string `description:"TLS cert"`
	Key                string `description:"TLS key"`
	InsecureSkipVerify bool   `description:"TLS insecure skip verify"`
}

// CreateTLSConfig creates a TLS config from ClientTLS structures
func (clientTLS *ClientTLS) CreateTLSConfig() (*tls.Config, error) {
	var err error
	if clientTLS == nil {
		log.Warnf("clientTLS is nil")
		return nil, nil
	}
	caPool := x509.NewCertPool()
	if clientTLS.CA != "" {
		var ca []byte
		if _, errCA := os.Stat(clientTLS.CA); errCA == nil {
			ca, err = ioutil.ReadFile(clientTLS.CA)
			if err != nil {
				return nil, fmt.Errorf("Failed to read CA. %s", err)
			}
		} else {
			ca = []byte(clientTLS.CA)
		}
		caPool.AppendCertsFromPEM(ca)
	}

	cert := tls.Certificate{}
	_, errKeyIsFile := os.Stat(clientTLS.Key)

	if _, errCertIsFile := os.Stat(clientTLS.Cert); errCertIsFile == nil {
		if errKeyIsFile == nil {
			cert, err = tls.LoadX509KeyPair(clientTLS.Cert, clientTLS.Key)
			if err != nil {
				return nil, fmt.Errorf("Failed to load TLS keypair: %v", err)
			}
		} else {
			return nil, fmt.Errorf("tls cert is a file, but tls key is not")
		}
	} else {
		if errKeyIsFile != nil {
			cert, err = tls.X509KeyPair([]byte(clientTLS.Cert), []byte(clientTLS.Key))
			if err != nil {
				return nil, fmt.Errorf("Failed to load TLS keypair: %v", err)

			}
		} else {
			return nil, fmt.Errorf("tls key is a file, but tls cert is not")
		}
	}

	TLSConfig := &tls.Config{
		Certificates:       []tls.Certificate{cert},
		RootCAs:            caPool,
		InsecureSkipVerify: clientTLS.InsecureSkipVerify,
	}
	return TLSConfig, nil
}
<|endoftext|>


<commit_before>package provider

import (
	"github.com/rusenask/keel/types"
)

// Provider - generic provider interface
type Provider interface {
	Submit(event types.Event) error
	GetName() string
}
<commit_msg>provider manager so triggers don't have to care about executing all of them<commit_after>package provider

import (
	"github.com/rusenask/keel/types"

	log "github.com/Sirupsen/logrus"
)

// Provider - generic provider interface
type Provider interface {
	Submit(event types.Event) error
	GetName() string
}

// Providers - available providers
type Providers interface {
	Submit(event types.Event) error
	List() []string // list all providers
}

// New - new providers registry
func New(providers []Provider) *DefaultProviders {
	pvs := make(map[string]Provider)

	for _, p := range providers {
		pvs[p.GetName()] = p
	}

	return &DefaultProviders{
		providers: pvs,
	}
}

type DefaultProviders struct {
	providers map[string]Provider
}

func (p *DefaultProviders) Submit(event types.Event) error {
	for _, provider := range p.providers {
		err := provider.Submit(event)
		if err != nil {
			log.WithFields(log.Fields{
				"error":    err,
				"provider": provider.GetName(),
				"event":    event.Repository,
				"trigger":  event.TriggerName,
			}).Error("provider.DefaultProviders: submit event failed")
		}
	}

	return nil
}

func (p *DefaultProviders) List() []string {
	list := []string{}
	for name := range p.providers {
		list = append(list, name)
	}
	return list
}
<|endoftext|>


<commit_before>package pipelines_test

import (
	"errors"
	"os"

	. "github.com/concourse/atc/pipelines"
	"github.com/concourse/atc/pipelines/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/fake_runner"

	"github.com/concourse/atc/db"
	dbfakes "github.com/concourse/atc/db/fakes"
)

var _ = Describe("Pipelines Syncer", func() {
	var (
		pipelinesDB           *fakes.FakePipelinesDB
		pipelineDB            *dbfakes.FakePipelineDB
		otherPipelineDB       *dbfakes.FakePipelineDB
		pipelineDBFactory     *dbfakes.FakePipelineDBFactory
		pipelineRunnerFactory PipelineRunnerFactory

		fakeRunner         *fake_runner.FakeRunner
		fakeRunnerExitChan chan error
		otherFakeRunner    *fake_runner.FakeRunner

		syncer *Syncer
	)

	BeforeEach(func() {
		pipelinesDB = new(fakes.FakePipelinesDB)
		pipelineDB = new(dbfakes.FakePipelineDB)

		pipelineDBFactory = new(dbfakes.FakePipelineDBFactory)

		fakeRunner = new(fake_runner.FakeRunner)
		otherFakeRunner = new(fake_runner.FakeRunner)

		pipelineRunnerFactory = func(pipelineDBArg db.PipelineDB) ifrit.Runner {
			switch pipelineDBArg {
			case pipelineDB:
				return fakeRunner
			case otherPipelineDB:
				return otherFakeRunner
			default:
				panic("unexpected pipelineDB input received")
				return nil
			}
			return fakeRunner
		}

		pipelineDBFactory.BuildStub = func(pipeline db.SavedPipeline) db.PipelineDB {
			switch pipeline.Name {
			case "pipeline":
				return pipelineDB
			case "other-pipeline":
				return otherPipelineDB
			default:
				panic("unexpected pipeline input received")
				return nil
			}
		}

		fakeRunnerExitChan = make(chan error, 1)
		fakeRunner.RunStub = func(signals <-chan os.Signal, ready chan<- struct{}) error {
			close(ready)
			return <-fakeRunnerExitChan
		}

		pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
			{
				ID: 1,
				Pipeline: db.Pipeline{
					Name: "pipeline",
				},
			},
			{
				ID: 2,
				Pipeline: db.Pipeline{
					Name: "other-pipeline",
				},
			},
		}, nil)

		syncer = NewSyncer(
			lagertest.NewTestLogger("test"),

			pipelinesDB,
			pipelineDBFactory,
			pipelineRunnerFactory,
		)
	})

	JustBeforeEach(func() {
		syncer.Sync()
	})

	It("spawns a new process for each pipeline", func() {
		Ω(fakeRunner.RunCallCount()).Should(Equal(1))
		Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))
	})

	Context("when a pipeline is paused", func() {
		BeforeEach(func() {
			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID: 1,
					Pipeline: db.Pipeline{
						Name: "pipeline",
					},
				},
				{
					ID:     2,
					Paused: true,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)
		})

		It("does not spawn a process for it", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(0))
		})
	})

	Context("when we sync again", func() {
		It("does not spawn any processes again", func() {
			syncer.Sync()
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
		})
	})

	Context("when a pipeline is deleted", func() {
		It("stops the process", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID: 2,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)

			syncer.Sync()

			Ω(fakeRunner.RunCallCount()).Should(Equal(1))

			signals, _ := fakeRunner.RunArgsForCall(0)
			Eventually(signals).Should(Receive(Equal(os.Interrupt)))
		})
	})

	Context("when a pipeline is paused", func() {
		It("stops the process", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID:     1,
					Paused: true,
					Pipeline: db.Pipeline{
						Name: "pipeline",
					},
				},
				{
					ID: 2,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)

			syncer.Sync()

			Ω(fakeRunner.RunCallCount()).Should(Equal(1))

			signals, _ := fakeRunner.RunArgsForCall(0)
			Eventually(signals).Should(Receive(Equal(os.Interrupt)))
		})
	})

	Context("when the pipeline's process exits", func() {
		BeforeEach(func() {
			fakeRunnerExitChan <- nil
		})

		Context("when we sync again", func() {
			It("spawns the process again", func() {
				Ω(fakeRunner.RunCallCount()).Should(Equal(1))
				Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

				fakeRunnerExitChan <- errors.New("disaster")
				syncer.Sync()

				Ω(fakeRunner.RunCallCount()).Should(Equal(2))
			})
		})
	})

	Context("when the call to lookup pipelines errors", func() {
		It("does not spawn any processes", func() {
		})
	})
})
<commit_msg>fix data race in syncer test<commit_after>package pipelines_test

import (
	"errors"
	"os"

	. "github.com/concourse/atc/pipelines"
	"github.com/concourse/atc/pipelines/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/fake_runner"

	"github.com/concourse/atc/db"
	dbfakes "github.com/concourse/atc/db/fakes"
)

var _ = Describe("Pipelines Syncer", func() {
	var (
		pipelinesDB           *fakes.FakePipelinesDB
		pipelineDB            *dbfakes.FakePipelineDB
		otherPipelineDB       *dbfakes.FakePipelineDB
		pipelineDBFactory     *dbfakes.FakePipelineDBFactory
		pipelineRunnerFactory PipelineRunnerFactory

		fakeRunner         *fake_runner.FakeRunner
		fakeRunnerExitChan chan error
		otherFakeRunner    *fake_runner.FakeRunner

		syncer *Syncer
	)

	BeforeEach(func() {
		pipelinesDB = new(fakes.FakePipelinesDB)
		pipelineDB = new(dbfakes.FakePipelineDB)

		pipelineDBFactory = new(dbfakes.FakePipelineDBFactory)

		fakeRunner = new(fake_runner.FakeRunner)
		otherFakeRunner = new(fake_runner.FakeRunner)

		pipelineRunnerFactory = func(pipelineDBArg db.PipelineDB) ifrit.Runner {
			switch pipelineDBArg {
			case pipelineDB:
				return fakeRunner
			case otherPipelineDB:
				return otherFakeRunner
			default:
				panic("unexpected pipelineDB input received")
				return nil
			}
			return fakeRunner
		}

		pipelineDBFactory.BuildStub = func(pipeline db.SavedPipeline) db.PipelineDB {
			switch pipeline.Name {
			case "pipeline":
				return pipelineDB
			case "other-pipeline":
				return otherPipelineDB
			default:
				panic("unexpected pipeline input received")
				return nil
			}
		}

		fakeRunnerExitChan = make(chan error, 1)

		// avoid data race
		exitChan := fakeRunnerExitChan

		fakeRunner.RunStub = func(signals <-chan os.Signal, ready chan<- struct{}) error {
			close(ready)
			return <-exitChan
		}

		pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
			{
				ID: 1,
				Pipeline: db.Pipeline{
					Name: "pipeline",
				},
			},
			{
				ID: 2,
				Pipeline: db.Pipeline{
					Name: "other-pipeline",
				},
			},
		}, nil)

		syncer = NewSyncer(
			lagertest.NewTestLogger("test"),

			pipelinesDB,
			pipelineDBFactory,
			pipelineRunnerFactory,
		)
	})

	JustBeforeEach(func() {
		syncer.Sync()
	})

	It("spawns a new process for each pipeline", func() {
		Ω(fakeRunner.RunCallCount()).Should(Equal(1))
		Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))
	})

	Context("when a pipeline is paused", func() {
		BeforeEach(func() {
			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID: 1,
					Pipeline: db.Pipeline{
						Name: "pipeline",
					},
				},
				{
					ID:     2,
					Paused: true,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)
		})

		It("does not spawn a process for it", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(0))
		})
	})

	Context("when we sync again", func() {
		It("does not spawn any processes again", func() {
			syncer.Sync()
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
		})
	})

	Context("when a pipeline is deleted", func() {
		It("stops the process", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID: 2,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)

			syncer.Sync()

			Ω(fakeRunner.RunCallCount()).Should(Equal(1))

			signals, _ := fakeRunner.RunArgsForCall(0)
			Eventually(signals).Should(Receive(Equal(os.Interrupt)))
		})
	})

	Context("when a pipeline is paused", func() {
		It("stops the process", func() {
			Ω(fakeRunner.RunCallCount()).Should(Equal(1))
			Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

			pipelinesDB.GetAllActivePipelinesReturns([]db.SavedPipeline{
				{
					ID:     1,
					Paused: true,
					Pipeline: db.Pipeline{
						Name: "pipeline",
					},
				},
				{
					ID: 2,
					Pipeline: db.Pipeline{
						Name: "other-pipeline",
					},
				},
			}, nil)

			syncer.Sync()

			Ω(fakeRunner.RunCallCount()).Should(Equal(1))

			signals, _ := fakeRunner.RunArgsForCall(0)
			Eventually(signals).Should(Receive(Equal(os.Interrupt)))
		})
	})

	Context("when the pipeline's process exits", func() {
		BeforeEach(func() {
			fakeRunnerExitChan <- nil
		})

		Context("when we sync again", func() {
			It("spawns the process again", func() {
				Ω(fakeRunner.RunCallCount()).Should(Equal(1))
				Ω(otherFakeRunner.RunCallCount()).Should(Equal(1))

				fakeRunnerExitChan <- errors.New("disaster")
				syncer.Sync()

				Ω(fakeRunner.RunCallCount()).Should(Equal(2))
			})
		})
	})

	Context("when the call to lookup pipelines errors", func() {
		It("does not spawn any processes", func() {
		})
	})
})
<|endoftext|>


<commit_before>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 collector

import (
	"database/sql"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sapcc/limes/pkg/db"
	"github.com/sapcc/limes/pkg/limes"
	"github.com/sapcc/limes/pkg/util"
)

////////////////////////////////////////////////////////////////////////////////
// collector metrics

var scrapeSuccessCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "limes_successful_scrapes",
		Help: "Counter for successful scrape operations per Keystone project.",
	},
	[]string{"cluster", "service"},
)

var scrapeFailedCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "limes_failed_scrapes",
		Help: "Counter for failed scrape operations per Keystone project.",
	},
	[]string{"cluster", "service"},
)

func init() {
	prometheus.MustRegister(scrapeSuccessCounter)
	prometheus.MustRegister(scrapeFailedCounter)
}

////////////////////////////////////////////////////////////////////////////////
// data metrics

var clusterCapacityGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_cluster_capacity",
		Help: "Reported capacity of a Limes resource for an OpenStack cluster.",
	},
	[]string{"cluster", "service", "resource"},
)

var domainQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_domain_quota",
		Help: "Assigned quota of a Limes resource for an OpenStack domain.",
	},
	[]string{"cluster", "domain", "service", "resource"},
)

var projectQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_quota",
		Help: "Assigned quota of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "project", "service", "resource"},
)

var projectUsageGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_usage",
		Help: "Actual usage of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "project", "service", "resource"},
)

var projectBackendQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_backendquota",
		Help: "Actual quota of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "project", "service", "resource"},
)

//DataMetricsCollector is a prometheus.Collector that submits
//quota/usage/backend quota from an OpenStack cluster as Prometheus metrics.
type DataMetricsCollector struct {
	Cluster *limes.ClusterConfiguration
}

//Describe implements the prometheus.Collector interface.
func (c *DataMetricsCollector) Describe(ch chan<- *prometheus.Desc) {
	clusterCapacityGauge.Describe(ch)
	domainQuotaGauge.Describe(ch)
	projectQuotaGauge.Describe(ch)
	projectUsageGauge.Describe(ch)
	projectBackendQuotaGauge.Describe(ch)
}

var clusterMetricsQuery = `
	SELECT cs.type, cr.name, cr.capacity
	  FROM cluster_services cs
	  JOIN cluster_resources cr ON cr.service_id = cs.id
	 WHERE cs.cluster_id = $1
`

var domainMetricsQuery = `
	SELECT d.name, ds.type, dr.name, dr.quota
	  FROM domains d
	  JOIN domain_services ds ON ds.domain_id = d.id
	  JOIN domain_resources dr ON dr.service_id = ds.id
	 WHERE d.cluster_id = $1
`

var projectMetricsQuery = `
	SELECT d.name, p.name, ps.type, pr.name, pr.quota, pr.usage, pr.backend_quota
	  FROM domains d
	  JOIN projects p ON p.domain_id = d.id
	  JOIN project_services ps ON ps.project_id = p.id
	  JOIN project_resources pr ON pr.service_id = ps.id
	 WHERE d.cluster_id = $1
`

//Collect implements the prometheus.Collector interface.
func (c *DataMetricsCollector) Collect(ch chan<- prometheus.Metric) {
	//NOTE: I use NewConstMetric() instead of storing the values in the GaugeVec
	//instances,
	//
	//1. because it is faster.
	//2. because this automatically handles deleted projects/domains correctly.
	//   (Their metrics just disappear when Prometheus scrapes next time.)

	//fetch Descs for all metrics
	descCh := make(chan *prometheus.Desc, 1)
	clusterCapacityGauge.Describe(descCh)
	clusterCapacityDesc := <-descCh
	domainQuotaGauge.Describe(descCh)
	domainQuotaDesc := <-descCh
	projectQuotaGauge.Describe(descCh)
	projectQuotaDesc := <-descCh
	projectUsageGauge.Describe(descCh)
	projectUsageDesc := <-descCh
	projectBackendQuotaGauge.Describe(descCh)
	projectBackendQuotaDesc := <-descCh

	//fetch values
	queryArgs := []interface{}{c.Cluster.ID}
	err := db.ForeachRow(db.DB, clusterMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			serviceType  string
			resourceName string
			capacity     uint64
		)
		err := rows.Scan(&serviceType, &resourceName, &capacity)
		if err != nil {
			return err
		}
		ch <- prometheus.MustNewConstMetric(
			clusterCapacityDesc,
			prometheus.GaugeValue, float64(capacity),
			c.Cluster.ID, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect cluster metrics failed: " + err.Error())
	}

	err = db.ForeachRow(db.DB, domainMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			domainName   string
			serviceType  string
			resourceName string
			quota        uint64
		)
		err := rows.Scan(&domainName, &serviceType, &resourceName, &quota)
		if err != nil {
			return err
		}
		ch <- prometheus.MustNewConstMetric(
			domainQuotaDesc,
			prometheus.GaugeValue, float64(quota),
			c.Cluster.ID, domainName, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect domain metrics failed: " + err.Error())
	}

	err = db.ForeachRow(db.DB, projectMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			domainName   string
			projectName  string
			serviceType  string
			resourceName string
			quota        uint64
			usage        uint64
			backendQuota int64
		)
		err := rows.Scan(&domainName, &projectName, &serviceType, &resourceName, &quota, &usage, &backendQuota)
		if err != nil {
			return err
		}

		ch <- prometheus.MustNewConstMetric(
			projectQuotaDesc,
			prometheus.GaugeValue, float64(quota),
			c.Cluster.ID, domainName, projectName, serviceType, resourceName,
		)
		ch <- prometheus.MustNewConstMetric(
			projectUsageDesc,
			prometheus.GaugeValue, float64(usage),
			c.Cluster.ID, domainName, projectName, serviceType, resourceName,
		)
		ch <- prometheus.MustNewConstMetric(
			projectBackendQuotaDesc,
			prometheus.GaugeValue, float64(backendQuota),
			c.Cluster.ID, domainName, projectName, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect project metrics failed: " + err.Error())
	}
}
<commit_msg>expose project/domain IDs in addition to names in data metrics<commit_after>/*******************************************************************************
*
* Copyright 2017 SAP SE
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You should have received a copy of the License along with this
* program. If not, 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 collector

import (
	"database/sql"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sapcc/limes/pkg/db"
	"github.com/sapcc/limes/pkg/limes"
	"github.com/sapcc/limes/pkg/util"
)

////////////////////////////////////////////////////////////////////////////////
// collector metrics

var scrapeSuccessCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "limes_successful_scrapes",
		Help: "Counter for successful scrape operations per Keystone project.",
	},
	[]string{"cluster", "service"},
)

var scrapeFailedCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "limes_failed_scrapes",
		Help: "Counter for failed scrape operations per Keystone project.",
	},
	[]string{"cluster", "service"},
)

func init() {
	prometheus.MustRegister(scrapeSuccessCounter)
	prometheus.MustRegister(scrapeFailedCounter)
}

////////////////////////////////////////////////////////////////////////////////
// data metrics

var clusterCapacityGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_cluster_capacity",
		Help: "Reported capacity of a Limes resource for an OpenStack cluster.",
	},
	[]string{"cluster", "service", "resource"},
)

var domainQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_domain_quota",
		Help: "Assigned quota of a Limes resource for an OpenStack domain.",
	},
	[]string{"cluster", "domain", "domain_id", "service", "resource"},
)

var projectQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_quota",
		Help: "Assigned quota of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "domain_id", "project", "project_id", "service", "resource"},
)

var projectUsageGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_usage",
		Help: "Actual usage of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "domain_id", "project", "project_id", "service", "resource"},
)

var projectBackendQuotaGauge = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "limes_project_backendquota",
		Help: "Actual quota of a Limes resource for an OpenStack project.",
	},
	[]string{"cluster", "domain", "domain_id", "project", "project_id", "service", "resource"},
)

//DataMetricsCollector is a prometheus.Collector that submits
//quota/usage/backend quota from an OpenStack cluster as Prometheus metrics.
type DataMetricsCollector struct {
	Cluster *limes.ClusterConfiguration
}

//Describe implements the prometheus.Collector interface.
func (c *DataMetricsCollector) Describe(ch chan<- *prometheus.Desc) {
	clusterCapacityGauge.Describe(ch)
	domainQuotaGauge.Describe(ch)
	projectQuotaGauge.Describe(ch)
	projectUsageGauge.Describe(ch)
	projectBackendQuotaGauge.Describe(ch)
}

var clusterMetricsQuery = `
	SELECT cs.type, cr.name, cr.capacity
	  FROM cluster_services cs
	  JOIN cluster_resources cr ON cr.service_id = cs.id
	 WHERE cs.cluster_id = $1
`

var domainMetricsQuery = `
	SELECT d.name, d.uuid, ds.type, dr.name, dr.quota
	  FROM domains d
	  JOIN domain_services ds ON ds.domain_id = d.id
	  JOIN domain_resources dr ON dr.service_id = ds.id
	 WHERE d.cluster_id = $1
`

var projectMetricsQuery = `
	SELECT d.name, d.uuid, p.name, p.uuid, ps.type, pr.name, pr.quota, pr.usage, pr.backend_quota
	  FROM domains d
	  JOIN projects p ON p.domain_id = d.id
	  JOIN project_services ps ON ps.project_id = p.id
	  JOIN project_resources pr ON pr.service_id = ps.id
	 WHERE d.cluster_id = $1
`

//Collect implements the prometheus.Collector interface.
func (c *DataMetricsCollector) Collect(ch chan<- prometheus.Metric) {
	//NOTE: I use NewConstMetric() instead of storing the values in the GaugeVec
	//instances,
	//
	//1. because it is faster.
	//2. because this automatically handles deleted projects/domains correctly.
	//   (Their metrics just disappear when Prometheus scrapes next time.)

	//fetch Descs for all metrics
	descCh := make(chan *prometheus.Desc, 1)
	clusterCapacityGauge.Describe(descCh)
	clusterCapacityDesc := <-descCh
	domainQuotaGauge.Describe(descCh)
	domainQuotaDesc := <-descCh
	projectQuotaGauge.Describe(descCh)
	projectQuotaDesc := <-descCh
	projectUsageGauge.Describe(descCh)
	projectUsageDesc := <-descCh
	projectBackendQuotaGauge.Describe(descCh)
	projectBackendQuotaDesc := <-descCh

	//fetch values
	queryArgs := []interface{}{c.Cluster.ID}
	err := db.ForeachRow(db.DB, clusterMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			serviceType  string
			resourceName string
			capacity     uint64
		)
		err := rows.Scan(&serviceType, &resourceName, &capacity)
		if err != nil {
			return err
		}
		ch <- prometheus.MustNewConstMetric(
			clusterCapacityDesc,
			prometheus.GaugeValue, float64(capacity),
			c.Cluster.ID, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect cluster metrics failed: " + err.Error())
	}

	err = db.ForeachRow(db.DB, domainMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			domainName   string
			domainUUID   string
			serviceType  string
			resourceName string
			quota        uint64
		)
		err := rows.Scan(&domainName, &domainUUID, &serviceType, &resourceName, &quota)
		if err != nil {
			return err
		}
		ch <- prometheus.MustNewConstMetric(
			domainQuotaDesc,
			prometheus.GaugeValue, float64(quota),
			c.Cluster.ID, domainName, domainUUID, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect domain metrics failed: " + err.Error())
	}

	err = db.ForeachRow(db.DB, projectMetricsQuery, queryArgs, func(rows *sql.Rows) error {
		var (
			domainName   string
			domainUUID   string
			projectName  string
			projectUUID  string
			serviceType  string
			resourceName string
			quota        uint64
			usage        uint64
			backendQuota int64
		)
		err := rows.Scan(&domainName, &domainUUID, &projectName, &projectUUID, &serviceType, &resourceName, &quota, &usage, &backendQuota)
		if err != nil {
			return err
		}

		ch <- prometheus.MustNewConstMetric(
			projectQuotaDesc,
			prometheus.GaugeValue, float64(quota),
			c.Cluster.ID, domainName, domainUUID, projectName, projectUUID, serviceType, resourceName,
		)
		ch <- prometheus.MustNewConstMetric(
			projectUsageDesc,
			prometheus.GaugeValue, float64(usage),
			c.Cluster.ID, domainName, domainUUID, projectName, projectUUID, serviceType, resourceName,
		)
		ch <- prometheus.MustNewConstMetric(
			projectBackendQuotaDesc,
			prometheus.GaugeValue, float64(backendQuota),
			c.Cluster.ID, domainName, domainUUID, projectName, projectUUID, serviceType, resourceName,
		)
		return nil
	})
	if err != nil {
		util.LogError("collect project metrics failed: " + err.Error())
	}
}
<|endoftext|>


<commit_before>// Copyright 2022 The etcd 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 flags

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestUint32Value(t *testing.T) {
	cases := []struct {
		name        string
		s           string
		expectedVal uint32
		expectError bool
	}{
		{
			name:        "normal uint32 value",
			s:           "200",
			expectedVal: 200,
		},
		{
			name:        "zero value",
			s:           "0",
			expectedVal: 0,
		},
		{
			name:        "negative int value",
			s:           "-200",
			expectError: true,
		},
		{
			name:        "invalid integer value",
			s:           "invalid",
			expectError: true,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var val uint32Value
			err := val.Set(tc.s)

			if tc.expectError {
				if err == nil {
					t.Errorf("Expected failure on parsing uint32 value from %s", tc.s)
				}
			} else {
				if err != nil {
					t.Errorf("Unexpected error when parsing %s: %v", tc.s, err)
				}
				assert.Equal(t, uint32(val), tc.expectedVal)
			}
		})
	}
}
<commit_msg>add unit test for pkg/flags/Uint32FromFlag in uint32_test.go<commit_after>// Copyright 2022 The etcd 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 flags

import (
	"flag"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestUint32Value(t *testing.T) {
	cases := []struct {
		name        string
		s           string
		expectedVal uint32
		expectError bool
	}{
		{
			name:        "normal uint32 value",
			s:           "200",
			expectedVal: 200,
		},
		{
			name:        "zero value",
			s:           "0",
			expectedVal: 0,
		},
		{
			name:        "negative int value",
			s:           "-200",
			expectError: true,
		},
		{
			name:        "invalid integer value",
			s:           "invalid",
			expectError: true,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var val uint32Value
			err := val.Set(tc.s)

			if tc.expectError {
				if err == nil {
					t.Errorf("Expected failure on parsing uint32 value from %s", tc.s)
				}
			} else {
				if err != nil {
					t.Errorf("Unexpected error when parsing %s: %v", tc.s, err)
				}
				assert.Equal(t, uint32(val), tc.expectedVal)
			}
		})
	}
}

func TestUint32FromFlag(t *testing.T) {
	const flagName = "max-concurrent-streams"

	cases := []struct {
		name        string
		defaultVal  uint32
		arguments   []string
		expectedVal uint32
	}{
		{
			name:        "only default value",
			defaultVal:  15,
			arguments:   []string{},
			expectedVal: 15,
		},
		{
			name:        "argument has different value from the default one",
			defaultVal:  16,
			arguments:   []string{"--max-concurrent-streams", "200"},
			expectedVal: 200,
		},
		{
			name:        "argument has the same value from the default one",
			defaultVal:  105,
			arguments:   []string{"--max-concurrent-streams", "105"},
			expectedVal: 105,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			fs := flag.NewFlagSet("etcd", flag.ContinueOnError)
			fs.Var(NewUint32Value(tc.defaultVal), flagName, "Maximum concurrent streams that each client can open at a time.")
			if err := fs.Parse(tc.arguments); err != nil {
				t.Fatalf("Unexpected error: %v\n", err)
			}
			actualMaxStream := Uint32FromFlag(fs, flagName)
			assert.Equal(t, actualMaxStream, tc.expectedVal)
		})
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package provider

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/clientset_generated/release_1_4"
	"k8s.io/kubernetes/pkg/client/restclient"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
)

const (
	UserAgent = "terraform-kubernetes"

	PollInterval = 10 * time.Second
	PollTimeout  = 10 * time.Minute
)

func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		ResourcesMap: map[string]*schema.Resource{
			"kubernetes_kubeconfig": resourceKubeconfig(),
		},
	}
}

func resourceKubeconfig() *schema.Resource {
	return &schema.Resource{
		Create: CreateKubeconfig,
		Delete: DeleteKubeconfig,
		Read:   ReadKubeconfig,

		Schema: map[string]*schema.Schema{
			"server": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				Description: "Domain name or IP address of the API server",
				ForceNew:    true,
			},
			"configdata": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				Description: "kubeconfig in the serialized JSON format",
				ForceNew:    true,
			},
		},
	}
}

func CreateKubeconfig(d *schema.ResourceData, meta interface{}) error {
	server := d.Get("server").(string)
	clientConfig, err := clientcmd.BuildConfigFromKubeconfigGetter(server, kubeConfigGetter(d))
	if err != nil {
		return fmt.Errorf("couldn't parse the supplied config: %v", err)
	}
	clientset, err := release_1_4.NewForConfig(restclient.AddUserAgent(clientConfig, UserAgent))
	if err != nil {
		return fmt.Errorf("failed to initialize the cluster client: %v", err)
	}

	interval := time.NewTicker(PollInterval)
	defer interval.Stop()
	timeout := time.NewTimer(PollTimeout)
	defer timeout.Stop()

	for {
		select {
		case <-interval.C:
			if allComponentsHealthy(clientset) {
				break
			}
		case <-timeout.C:
			return fmt.Errorf("cluster components never turned healthy")
		}
	}

	configAccess := clientcmd.NewDefaultPathOptions()
	kubeConfig, err := kubeConfigGetter(d)()
	if err != nil {
		return fmt.Errorf("couldn't parse the supplied config: %v", err)
	}

	if err := modifyConfig(configAccess, kubeConfig); err != nil {
		return fmt.Errorf("couldn't update kubeconfig: %v", err)
	}

	return nil
}

func DeleteKubeconfig(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func ReadKubeconfig(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func allComponentsHealthy(clientset *release_1_4.Clientset) bool {
	csList, err := clientset.Core().ComponentStatuses().List(api.ListOptions{})
	if err != nil || len(csList.Items) <= 0 {
		log.Printf("[DEBUG] Listing components failed %s", err)
		return false
	}
	for _, cs := range csList.Items {
		if !(len(cs.Conditions) > 0 && cs.Conditions[0].Type == "Healthy" && cs.Conditions[0].Status == "True") {
			log.Printf("[DEBUG] %s isn't healthy. Conditions: %+v", cs.Name, cs.Conditions)
			return false
		}
	}
	return true
}

func kubeConfigGetter(d *schema.ResourceData) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		kubeConfigStr := d.Get("configdata").(string)
		return clientcmd.Load([]byte(kubeConfigStr))
	}
}

func modifyConfig(configAccess clientcmd.ConfigAccess, suppliedConfig *clientcmdapi.Config) error {
	config, err := configAccess.GetStartingConfig()
	if err != nil {
		return err
	}

	for name, cluster := range suppliedConfig.Clusters {
		initial, ok := config.Clusters[name]
		if !ok {
			initial = clientcmdapi.NewCluster()
		}
		modified := *initial

		if len(cluster.Server) > 0 {
			modified.Server = cluster.Server
		}
		if cluster.InsecureSkipTLSVerify {
			modified.InsecureSkipTLSVerify = cluster.InsecureSkipTLSVerify
			// Specifying insecure mode clears any certificate authority
			if modified.InsecureSkipTLSVerify {
				modified.CertificateAuthority = ""
				modified.CertificateAuthorityData = nil
			}
		}
		if len(cluster.CertificateAuthorityData) > 0 {
			modified.CertificateAuthorityData = cluster.CertificateAuthorityData
			modified.InsecureSkipTLSVerify = false
			modified.CertificateAuthority = ""
		} else if len(cluster.CertificateAuthority) > 0 {
			modified.CertificateAuthority = cluster.CertificateAuthority
			modified.InsecureSkipTLSVerify = false
			modified.CertificateAuthorityData = nil
		}
		config.Clusters[name] = &modified
	}

	if err := clientcmd.ModifyConfig(configAccess, *config, true); err != nil {
		return err
	}

	return nil
}
<commit_msg>Reorganize the code a bit to make it amenable to unit tests.<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package provider

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/clientset_generated/release_1_4"
	"k8s.io/kubernetes/pkg/client/restclient"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
)

const (
	UserAgent = "terraform-kubernetes"

	PollInterval = 10 * time.Second
	PollTimeout  = 10 * time.Minute
)

func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		ResourcesMap: map[string]*schema.Resource{
			"kubernetes_kubeconfig": resourceKubeconfig(),
		},
	}
}

func resourceKubeconfig() *schema.Resource {
	return &schema.Resource{
		Create: CreateKubeconfig,
		Delete: DeleteKubeconfig,
		Read:   ReadKubeconfig,

		Schema: map[string]*schema.Schema{
			"server": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				Description: "Domain name or IP address of the API server",
				ForceNew:    true,
			},
			"configdata": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				Description: "kubeconfig in the serialized JSON format",
				ForceNew:    true,
			},
		},
	}
}

func CreateKubeconfig(d *schema.ResourceData, meta interface{}) error {
	server := d.Get("server").(string)
	clientConfig, err := clientcmd.BuildConfigFromKubeconfigGetter(server, kubeConfigGetter(d))
	if err != nil {
		return fmt.Errorf("couldn't parse the supplied config: %v", err)
	}

	clientset, err := release_1_4.NewForConfig(restclient.AddUserAgent(clientConfig, UserAgent))
	if err != nil {
		return fmt.Errorf("failed to initialize the cluster client: %v", err)
	}

	if !clusterHealthy(clientset, PollInterval, PollTimeout) {
		return fmt.Errorf("cluster components never turned healthy")
	}

	configAccess := clientcmd.NewDefaultPathOptions()
	kubeConfig, err := kubeConfigGetter(d)()
	if err != nil {
		return fmt.Errorf("couldn't parse the supplied config: %v", err)
	}

	if err := modifyConfig(configAccess, kubeConfig); err != nil {
		return fmt.Errorf("couldn't update kubeconfig: %v", err)
	}

	return nil
}

func DeleteKubeconfig(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func ReadKubeconfig(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func clusterHealthy(clientset release_1_4.Interface, pollInterval, pollTimeout time.Duration) bool {
	interval := time.NewTicker(pollInterval)
	defer interval.Stop()
	timeout := time.NewTimer(pollTimeout)
	defer timeout.Stop()

	for {
		select {
		case <-interval.C:
			// // Try re-creating the client until all the components turn healthy
			// // or we timeout. Until the server comes up, the clientset creation
			// // completely fails. Also, do not re-use this clientset outside this
			// // block, just to be safe.
			// healthClientSet, err := release_1_4.NewForConfig(restclient.AddUserAgent(clientConfig, UserAgent))
			if allComponentsHealthy(clientset) {
				log.Printf("[DEBUG] all components are healthy")
				return true
			} else {
				log.Printf("[DEBUG] components aren't healthy")
			}
		case <-timeout.C:
			return false
		}
	}
	// Something went wrong
	return false
}

func allComponentsHealthy(clientset release_1_4.Interface) bool {
	csList, err := clientset.Core().ComponentStatuses().List(api.ListOptions{})
	if err != nil || len(csList.Items) <= 0 {
		log.Printf("[DEBUG] Listing components failed %s", err)
		return false
	}
	for _, cs := range csList.Items {
		if !(len(cs.Conditions) > 0 && cs.Conditions[0].Type == "Healthy" && cs.Conditions[0].Status == "True") {
			log.Printf("[DEBUG] %s isn't healthy. Conditions: %+v", cs.Name, cs.Conditions)
			return false
		}
	}
	return true
}

func kubeConfigGetter(d *schema.ResourceData) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		kubeConfigStr := d.Get("configdata").(string)
		return clientcmd.Load([]byte(kubeConfigStr))
	}
}

func modifyConfig(configAccess clientcmd.ConfigAccess, suppliedConfig *clientcmdapi.Config) error {
	config, err := configAccess.GetStartingConfig()
	if err != nil {
		return err
	}

	for name, cluster := range suppliedConfig.Clusters {
		initial, ok := config.Clusters[name]
		if !ok {
			initial = clientcmdapi.NewCluster()
		}
		modified := *initial

		if len(cluster.Server) > 0 {
			modified.Server = cluster.Server
		}
		if cluster.InsecureSkipTLSVerify {
			modified.InsecureSkipTLSVerify = cluster.InsecureSkipTLSVerify
			// Specifying insecure mode clears any certificate authority
			if modified.InsecureSkipTLSVerify {
				modified.CertificateAuthority = ""
				modified.CertificateAuthorityData = nil
			}
		}
		if len(cluster.CertificateAuthorityData) > 0 {
			modified.CertificateAuthorityData = cluster.CertificateAuthorityData
			modified.InsecureSkipTLSVerify = false
			modified.CertificateAuthority = ""
		} else if len(cluster.CertificateAuthority) > 0 {
			modified.CertificateAuthority = cluster.CertificateAuthority
			modified.InsecureSkipTLSVerify = false
			modified.CertificateAuthorityData = nil
		}
		config.Clusters[name] = &modified
	}

	if err := clientcmd.ModifyConfig(configAccess, *config, true); err != nil {
		return err
	}

	return nil
}
<|endoftext|>