<commit_before>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package localstore

import (
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/store/localstore/engine"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util/bytes"
)

var _ kv.Compactor = (*localstoreCompactor)(nil)

const (
	deleteWorkerCnt = 3
)

var localCompactDefaultPolicy = kv.CompactPolicy{
	SafePoint:       20 * 1000, // in ms
	TriggerInterval: 1 * time.Second,
	BatchDeleteCnt:  100,
}

type localstoreCompactor struct {
	mu              sync.Mutex
	recentKeys      map[string]struct{}
	stopCh          chan struct{}
	delCh           chan kv.EncodedKey
	workerWaitGroup *sync.WaitGroup
	ticker          *time.Ticker
	db              engine.DB
	policy          kv.CompactPolicy
}

func (gc *localstoreCompactor) OnSet(k kv.Key) {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	gc.recentKeys[string(k)] = struct{}{}
}

func (gc *localstoreCompactor) OnGet(k kv.Key) {
	// Do nothing now.
}

func (gc *localstoreCompactor) OnDelete(k kv.Key) {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	gc.recentKeys[string(k)] = struct{}{}
}

func (gc *localstoreCompactor) getAllVersions(key kv.Key) ([]kv.EncodedKey, error) {
	var keys []kv.EncodedKey
	k := key
	for ver := kv.MaxVersion; ver.Ver > 0; ver.Ver-- {
		mvccK, _, err := gc.db.Seek(MvccEncodeVersionKey(key, ver))
		if terror.ErrorEqual(err, engine.ErrNotFound) {
			break
		}
		if err != nil {
			return nil, errors.Trace(err)
		}
		k, ver, err = MvccDecode(mvccK)
		if k.Cmp(key) != 0 {
			break
		}
		if err != nil {
			return nil, errors.Trace(err)
		}
		keys = append(keys, bytes.CloneBytes(mvccK))
	}
	return keys, nil
}

func (gc *localstoreCompactor) deleteWorker() {
	defer gc.workerWaitGroup.Done()
	cnt := 0
	batch := gc.db.NewBatch()
	for {
		select {
		case <-gc.stopCh:
			return
		case key := <-gc.delCh:
			cnt++
			batch.Delete(key)
			// Batch delete.
			if cnt == gc.policy.BatchDeleteCnt {
				err := gc.db.Commit(batch)
				if err != nil {
					log.Error(err)
				}
				batch = gc.db.NewBatch()
				cnt = 0
			}
		}
	}
}

func (gc *localstoreCompactor) checkExpiredKeysWorker() {
	gc.workerWaitGroup.Add(1)
	defer gc.workerWaitGroup.Done()
	for {
		select {
		case <-gc.stopCh:
			log.Info("GC stopped")
			return
		case <-gc.ticker.C:
			gc.mu.Lock()
			m := gc.recentKeys
			if len(m) == 0 {
				gc.mu.Unlock()
				continue
			}
			gc.recentKeys = make(map[string]struct{})
			gc.mu.Unlock()
			log.Info("GC trigger")
			for k := range m {
				err := gc.Compact([]byte(k))
				if err != nil {
					log.Error(err)
				}
			}
		}
	}
}

func (gc *localstoreCompactor) filterExpiredKeys(keys []kv.EncodedKey) []kv.EncodedKey {
	var ret []kv.EncodedKey
	first := true
	currentTS := time.Now().UnixNano() / int64(time.Millisecond)
	// keys are always in descending order.
	for _, k := range keys {
		_, ver, err := MvccDecode(k)
		if err != nil {
			// Should not happen.
			panic(err)
		}
		ts := localVersionToTimestamp(ver)
		// Check timeout keys.
		if currentTS-int64(ts) >= int64(gc.policy.SafePoint) {
			// Skip first version.
			if first {
				first = false
				continue
			}
			ret = append(ret, k)
		}
	}
	return ret
}

func (gc *localstoreCompactor) Compact(k kv.Key) error {
	keys, err := gc.getAllVersions(k)
	if err != nil {
		return errors.Trace(err)
	}
	for _, key := range gc.filterExpiredKeys(keys) {
		// Send timeout key to deleteWorker.
		log.Info("GC send key to deleteWorker", key)
		gc.delCh <- key
	}
	return nil
}

func (gc *localstoreCompactor) Start() {
	// Start workers.
	gc.workerWaitGroup.Add(deleteWorkerCnt)
	for i := 0; i < deleteWorkerCnt; i++ {
		go gc.deleteWorker()
	}

	go gc.checkExpiredKeysWorker()
}

func (gc *localstoreCompactor) Stop() {
	gc.ticker.Stop()
	close(gc.stopCh)
	// Wait for all workers to finish.
	gc.workerWaitGroup.Wait()
}

func newLocalCompactor(policy kv.CompactPolicy, db engine.DB) *localstoreCompactor {
	return &localstoreCompactor{
		recentKeys:      make(map[string]struct{}),
		stopCh:          make(chan struct{}),
		delCh:           make(chan kv.EncodedKey, 100),
		ticker:          time.NewTicker(policy.TriggerInterval),
		policy:          policy,
		db:              db,
		workerWaitGroup: &sync.WaitGroup{},
	}
}
<commit_msg>store/localstore: clean up compactor log.<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package localstore

import (
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/store/localstore/engine"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util/bytes"
)

var _ kv.Compactor = (*localstoreCompactor)(nil)

const (
	deleteWorkerCnt = 3
)

var localCompactDefaultPolicy = kv.CompactPolicy{
	SafePoint:       20 * 1000, // in ms
	TriggerInterval: 10 * time.Second,
	BatchDeleteCnt:  100,
}

type localstoreCompactor struct {
	mu              sync.Mutex
	recentKeys      map[string]struct{}
	stopCh          chan struct{}
	delCh           chan kv.EncodedKey
	workerWaitGroup *sync.WaitGroup
	ticker          *time.Ticker
	db              engine.DB
	policy          kv.CompactPolicy
}

func (gc *localstoreCompactor) OnSet(k kv.Key) {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	gc.recentKeys[string(k)] = struct{}{}
}

func (gc *localstoreCompactor) OnGet(k kv.Key) {
	// Do nothing now.
}

func (gc *localstoreCompactor) OnDelete(k kv.Key) {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	gc.recentKeys[string(k)] = struct{}{}
}

func (gc *localstoreCompactor) getAllVersions(key kv.Key) ([]kv.EncodedKey, error) {
	var keys []kv.EncodedKey
	k := key
	for ver := kv.MaxVersion; ver.Ver > 0; ver.Ver-- {
		mvccK, _, err := gc.db.Seek(MvccEncodeVersionKey(key, ver))
		if terror.ErrorEqual(err, engine.ErrNotFound) {
			break
		}
		if err != nil {
			return nil, errors.Trace(err)
		}
		k, ver, err = MvccDecode(mvccK)
		if k.Cmp(key) != 0 {
			break
		}
		if err != nil {
			return nil, errors.Trace(err)
		}
		keys = append(keys, bytes.CloneBytes(mvccK))
	}
	return keys, nil
}

func (gc *localstoreCompactor) deleteWorker() {
	defer gc.workerWaitGroup.Done()
	cnt := 0
	batch := gc.db.NewBatch()
	for {
		select {
		case <-gc.stopCh:
			return
		case key := <-gc.delCh:
			cnt++
			batch.Delete(key)
			// Batch delete.
			if cnt == gc.policy.BatchDeleteCnt {
				log.Debugf("GC delete commit %d keys", batch.Len())
				err := gc.db.Commit(batch)
				if err != nil {
					log.Error(err)
				}
				batch = gc.db.NewBatch()
				cnt = 0
			}
		}
	}
}

func (gc *localstoreCompactor) checkExpiredKeysWorker() {
	gc.workerWaitGroup.Add(1)
	defer gc.workerWaitGroup.Done()
	for {
		select {
		case <-gc.stopCh:
			log.Debug("GC stopped")
			return
		case <-gc.ticker.C:
			gc.mu.Lock()
			m := gc.recentKeys
			if len(m) == 0 {
				gc.mu.Unlock()
				continue
			}
			gc.recentKeys = make(map[string]struct{})
			gc.mu.Unlock()
			for k := range m {
				err := gc.Compact([]byte(k))
				if err != nil {
					log.Error(err)
				}
			}
		}
	}
}

func (gc *localstoreCompactor) filterExpiredKeys(keys []kv.EncodedKey) []kv.EncodedKey {
	var ret []kv.EncodedKey
	first := true
	currentTS := time.Now().UnixNano() / int64(time.Millisecond)
	// keys are always in descending order.
	for _, k := range keys {
		_, ver, err := MvccDecode(k)
		if err != nil {
			// Should not happen.
			panic(err)
		}
		ts := localVersionToTimestamp(ver)
		// Check timeout keys.
		if currentTS-int64(ts) >= int64(gc.policy.SafePoint) {
			// Skip first version.
			if first {
				first = false
				continue
			}
			ret = append(ret, k)
		}
	}
	return ret
}

func (gc *localstoreCompactor) Compact(k kv.Key) error {
	keys, err := gc.getAllVersions(k)
	if err != nil {
		return errors.Trace(err)
	}
	filteredKeys := gc.filterExpiredKeys(keys)
	log.Debugf("GC send %d keys to delete worker", len(filteredKeys))
	for _, key := range filteredKeys {
		gc.delCh <- key
	}
	return nil
}

func (gc *localstoreCompactor) Start() {
	// Start workers.
	gc.workerWaitGroup.Add(deleteWorkerCnt)
	for i := 0; i < deleteWorkerCnt; i++ {
		go gc.deleteWorker()
	}

	go gc.checkExpiredKeysWorker()
}

func (gc *localstoreCompactor) Stop() {
	gc.ticker.Stop()
	close(gc.stopCh)
	// Wait for all workers to finish.
	gc.workerWaitGroup.Wait()
}

func newLocalCompactor(policy kv.CompactPolicy, db engine.DB) *localstoreCompactor {
	return &localstoreCompactor{
		recentKeys:      make(map[string]struct{}),
		stopCh:          make(chan struct{}),
		delCh:           make(chan kv.EncodedKey, 100),
		ticker:          time.NewTicker(policy.TriggerInterval),
		policy:          policy,
		db:              db,
		workerWaitGroup: &sync.WaitGroup{},
	}
}
<|endoftext|>


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

// Package gif implements a GIF image decoder.
//
// The GIF specification is at http://www.w3.org/Graphics/GIF/spec-gif89a.txt.
package gif

import (
	"bufio"
	"compress/lzw"
	"fmt"
	"image"
	"io"
	"os"
)

// If the io.Reader does not also have ReadByte, then decode will introduce its own buffering.
type reader interface {
	io.Reader
	io.ByteReader
}

// Masks etc.
const (
	// Fields.
	fColorMapFollows = 1 << 7

	// Image fields.
	ifInterlace = 1 << 6

	// Graphic control flags.
	gcTransparentColorSet = 1 << 0
)

// Section indicators.
const (
	sExtension       = 0x21
	sImageDescriptor = 0x2C
	sTrailer         = 0x3B
)

// Extensions.
const (
	eText           = 0x01 // Plain Text
	eGraphicControl = 0xF9 // Graphic Control
	eComment        = 0xFE // Comment
	eApplication    = 0xFF // Application
)

// decoder is the type used to decode a GIF file.
type decoder struct {
	r reader

	// From header.
	vers            string
	width           int
	height          int
	flags           byte
	headerFields    byte
	backgroundIndex byte
	loopCount       int
	delayTime       int

	// Unused from header.
	aspect byte

	// From image descriptor.
	imageFields byte

	// From graphics control.
	transparentIndex byte

	// Computed.
	pixelSize      uint
	globalColorMap image.PalettedColorModel

	// Used when decoding.
	delay []int
	image []*image.Paletted
	tmp   [1024]byte // must be at least 768 so we can read color map
}

// blockReader parses the block structure of GIF image data, which
// comprises (n, (n bytes)) blocks, with 1 <= n <= 255.  It is the
// reader given to the LZW decoder, which is thus immune to the
// blocking.  After the LZW decoder completes, there will be a 0-byte
// block remaining (0, ()), but under normal execution blockReader
// doesn't consume it, so it is handled in decode.
type blockReader struct {
	r     reader
	slice []byte
	tmp   [256]byte
}

func (b *blockReader) Read(p []byte) (int, os.Error) {
	if len(p) == 0 {
		return 0, nil
	}
	if len(b.slice) == 0 {
		blockLen, err := b.r.ReadByte()
		if err != nil {
			return 0, err
		}
		if blockLen == 0 {
			return 0, os.EOF
		}
		b.slice = b.tmp[0:blockLen]
		if _, err = io.ReadFull(b.r, b.slice); err != nil {
			return 0, err
		}
	}
	n := copy(p, b.slice)
	b.slice = b.slice[n:]
	return n, nil
}

// decode reads a GIF image from r and stores the result in d.
func (d *decoder) decode(r io.Reader, configOnly bool) os.Error {
	// Add buffering if r does not provide ReadByte.
	if rr, ok := r.(reader); ok {
		d.r = rr
	} else {
		d.r = bufio.NewReader(r)
	}

	err := d.readHeaderAndScreenDescriptor()
	if err != nil {
		return err
	}
	if configOnly {
		return nil
	}

	if d.headerFields&fColorMapFollows != 0 {
		if d.globalColorMap, err = d.readColorMap(); err != nil {
			return err
		}
	}

Loop:
	for err == nil {
		var c byte
		c, err = d.r.ReadByte()
		if err == os.EOF {
			break
		}
		switch c {
		case sExtension:
			err = d.readExtension()

		case sImageDescriptor:
			var m *image.Paletted
			m, err = d.newImageFromDescriptor()
			if err != nil {
				break
			}
			if d.imageFields&fColorMapFollows != 0 {
				m.Palette, err = d.readColorMap()
				if err != nil {
					break
				}
				// TODO: do we set transparency in this map too? That would be
				// d.setTransparency(m.Palette)
			} else {
				m.Palette = d.globalColorMap
			}
			var litWidth uint8
			litWidth, err = d.r.ReadByte()
			if err != nil {
				return err
			}
			if litWidth < 2 || litWidth > 8 {
				return fmt.Errorf("gif: pixel size in decode out of range: %d", litWidth)
			}
			// A wonderfully Go-like piece of magic.
			lzwr := lzw.NewReader(&blockReader{r: d.r}, lzw.LSB, int(litWidth))
			if _, err = io.ReadFull(lzwr, m.Pix); err != nil {
				break
			}

			// There should be a "0" block remaining; drain that.
			c, err = d.r.ReadByte()
			if err != nil {
				return err
			}
			if c != 0 {
				return os.ErrorString("gif: extra data after image")
			}

			// Undo the interlacing if necessary.
			d.uninterlace(m)

			d.image = append(d.image, m)
			d.delay = append(d.delay, d.delayTime)
			d.delayTime = 0 // TODO: is this correct, or should we hold on to the value?

		case sTrailer:
			break Loop

		default:
			err = fmt.Errorf("gif: unknown block type: 0x%.2x", c)
		}
	}
	if err != nil {
		return err
	}
	if len(d.image) == 0 {
		return io.ErrUnexpectedEOF
	}
	return nil
}

func (d *decoder) readHeaderAndScreenDescriptor() os.Error {
	_, err := io.ReadFull(d.r, d.tmp[0:13])
	if err != nil {
		return err
	}
	d.vers = string(d.tmp[0:6])
	if d.vers != "GIF87a" && d.vers != "GIF89a" {
		return fmt.Errorf("gif: can't recognize format %s", d.vers)
	}
	d.width = int(d.tmp[6]) + int(d.tmp[7])<<8
	d.height = int(d.tmp[8]) + int(d.tmp[9])<<8
	d.headerFields = d.tmp[10]
	d.backgroundIndex = d.tmp[11]
	d.aspect = d.tmp[12]
	d.loopCount = -1
	d.pixelSize = uint(d.headerFields&7) + 1
	return nil
}

func (d *decoder) readColorMap() (image.PalettedColorModel, os.Error) {
	if d.pixelSize > 8 {
		return nil, fmt.Errorf("gif: can't handle %d bits per pixel", d.pixelSize)
	}
	numColors := 1 << d.pixelSize
	numValues := 3 * numColors
	_, err := io.ReadFull(d.r, d.tmp[0:numValues])
	if err != nil {
		return nil, fmt.Errorf("gif: short read on color map: %s", err)
	}
	colorMap := make(image.PalettedColorModel, numColors)
	j := 0
	for i := range colorMap {
		colorMap[i] = image.RGBAColor{d.tmp[j+0], d.tmp[j+1], d.tmp[j+2], 0xFF}
		j += 3
	}
	return colorMap, nil
}

func (d *decoder) readExtension() os.Error {
	extension, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	size := 0
	switch extension {
	case eText:
		size = 13
	case eGraphicControl:
		return d.readGraphicControl()
	case eComment:
		// nothing to do but read the data.
	case eApplication:
		b, err := d.r.ReadByte()
		if err != nil {
			return err
		}
		// The spec requires size be 11, but Adobe sometimes uses 10.
		size = int(b)
	default:
		return fmt.Errorf("gif: unknown extension 0x%.2x", extension)
	}
	if size > 0 {
		if _, err := d.r.Read(d.tmp[0:size]); err != nil {
			return err
		}
	}

	// Application Extension with "NETSCAPE2.0" as string and 1 in data means
	// this extension defines a loop count.
	if extension == eApplication && string(d.tmp[:size]) == "NETSCAPE2.0" {
		n, err := d.readBlock()
		if n == 0 || err != nil {
			return err
		}
		if n == 3 && d.tmp[0] == 1 {
			d.loopCount = int(d.tmp[1]) | int(d.tmp[2])<<8
		}
	}
	for {
		n, err := d.readBlock()
		if n == 0 || err != nil {
			return err
		}
	}
	panic("unreachable")
}

func (d *decoder) readGraphicControl() os.Error {
	if _, err := io.ReadFull(d.r, d.tmp[0:6]); err != nil {
		return fmt.Errorf("gif: can't read graphic control: %s", err)
	}
	d.flags = d.tmp[1]
	d.delayTime = int(d.tmp[2]) | int(d.tmp[3])<<8
	if d.flags&gcTransparentColorSet != 0 {
		d.transparentIndex = d.tmp[4]
		d.setTransparency(d.globalColorMap)
	}
	return nil
}

func (d *decoder) setTransparency(colorMap image.PalettedColorModel) {
	if int(d.transparentIndex) < len(colorMap) {
		colorMap[d.transparentIndex] = image.RGBAColor{}
	}
}

func (d *decoder) newImageFromDescriptor() (*image.Paletted, os.Error) {
	if _, err := io.ReadFull(d.r, d.tmp[0:9]); err != nil {
		return nil, fmt.Errorf("gif: can't read image descriptor: %s", err)
	}
	// TODO: This code (throughout) ignores the top and left values,
	// and assumes (in interlacing, for example) that the images'
	// widths and heights are all the same.
	_ = int(d.tmp[0]) + int(d.tmp[1])<<8 // TODO: honor left value
	_ = int(d.tmp[2]) + int(d.tmp[3])<<8 // TODO: honor top value
	width := int(d.tmp[4]) + int(d.tmp[5])<<8
	height := int(d.tmp[6]) + int(d.tmp[7])<<8
	d.imageFields = d.tmp[8]
	return image.NewPaletted(width, height, nil), nil
}

func (d *decoder) readBlock() (int, os.Error) {
	n, err := d.r.ReadByte()
	if n == 0 || err != nil {
		return 0, err
	}
	return io.ReadFull(d.r, d.tmp[0:n])
}

// interlaceScan defines the ordering for a pass of the interlace algorithm.
type interlaceScan struct {
	skip, start int
}

// interlacing represents the set of scans in an interlaced GIF image.
var interlacing = []interlaceScan{
	{8, 0}, // Group 1 : Every 8th. row, starting with row 0.
	{8, 4}, // Group 2 : Every 8th. row, starting with row 4.
	{4, 2}, // Group 3 : Every 4th. row, starting with row 2.
	{2, 1}, // Group 4 : Every 2nd. row, starting with row 1.
}

func (d *decoder) uninterlace(m *image.Paletted) {
	if d.imageFields&ifInterlace == 0 {
		return
	}
	var nPix []uint8
	dx := d.width
	dy := d.height
	nPix = make([]uint8, dx*dy)
	offset := 0 // steps through the input by sequential scan lines.
	for _, pass := range interlacing {
		nOffset := pass.start * dx // steps through the output as defined by pass.
		for y := pass.start; y < dy; y += pass.skip {
			copy(nPix[nOffset:nOffset+dx], m.Pix[offset:offset+dx])
			offset += dx
			nOffset += dx * pass.skip
		}
	}
	m.Pix = nPix
}

// Decode reads a GIF image from r and returns the first embedded
// image as an image.Image.
func Decode(r io.Reader) (image.Image, os.Error) {
	var d decoder
	if err := d.decode(r, false); err != nil {
		return nil, err
	}
	return d.image[0], nil
}

// GIF represents the possibly multiple images stored in a GIF file.
type GIF struct {
	Image     []*image.Paletted // The successive images.
	Delay     []int             // The successive delay times, one per frame, in 100ths of a second.
	LoopCount int               // The loop count.
}

// DecodeAll reads a GIF image from r and returns the sequential frames
// and timing information.
func DecodeAll(r io.Reader) (*GIF, os.Error) {
	var d decoder
	if err := d.decode(r, false); err != nil {
		return nil, err
	}
	gif := &GIF{
		Image:     d.image,
		LoopCount: d.loopCount,
		Delay:     d.delay,
	}
	return gif, nil
}

// DecodeConfig returns the color model and dimensions of a GIF image without
// decoding the entire image.
func DecodeConfig(r io.Reader) (image.Config, os.Error) {
	var d decoder
	if err := d.decode(r, true); err != nil {
		return image.Config{}, err
	}
	return image.Config{d.globalColorMap, d.width, d.height}, nil
}

func init() {
	image.RegisterFormat("gif", "GIF8?a", Decode, DecodeConfig)
}
<commit_msg>image/gif: fix buglet in graphics extension. need to ReadFull, not Read, to get extension data.<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package gif implements a GIF image decoder.
//
// The GIF specification is at http://www.w3.org/Graphics/GIF/spec-gif89a.txt.
package gif

import (
	"bufio"
	"compress/lzw"
	"fmt"
	"image"
	"io"
	"os"
)

// If the io.Reader does not also have ReadByte, then decode will introduce its own buffering.
type reader interface {
	io.Reader
	io.ByteReader
}

// Masks etc.
const (
	// Fields.
	fColorMapFollows = 1 << 7

	// Image fields.
	ifInterlace = 1 << 6

	// Graphic control flags.
	gcTransparentColorSet = 1 << 0
)

// Section indicators.
const (
	sExtension       = 0x21
	sImageDescriptor = 0x2C
	sTrailer         = 0x3B
)

// Extensions.
const (
	eText           = 0x01 // Plain Text
	eGraphicControl = 0xF9 // Graphic Control
	eComment        = 0xFE // Comment
	eApplication    = 0xFF // Application
)

// decoder is the type used to decode a GIF file.
type decoder struct {
	r reader

	// From header.
	vers            string
	width           int
	height          int
	flags           byte
	headerFields    byte
	backgroundIndex byte
	loopCount       int
	delayTime       int

	// Unused from header.
	aspect byte

	// From image descriptor.
	imageFields byte

	// From graphics control.
	transparentIndex byte

	// Computed.
	pixelSize      uint
	globalColorMap image.PalettedColorModel

	// Used when decoding.
	delay []int
	image []*image.Paletted
	tmp   [1024]byte // must be at least 768 so we can read color map
}

// blockReader parses the block structure of GIF image data, which
// comprises (n, (n bytes)) blocks, with 1 <= n <= 255.  It is the
// reader given to the LZW decoder, which is thus immune to the
// blocking.  After the LZW decoder completes, there will be a 0-byte
// block remaining (0, ()), but under normal execution blockReader
// doesn't consume it, so it is handled in decode.
type blockReader struct {
	r     reader
	slice []byte
	tmp   [256]byte
}

func (b *blockReader) Read(p []byte) (int, os.Error) {
	if len(p) == 0 {
		return 0, nil
	}
	if len(b.slice) == 0 {
		blockLen, err := b.r.ReadByte()
		if err != nil {
			return 0, err
		}
		if blockLen == 0 {
			return 0, os.EOF
		}
		b.slice = b.tmp[0:blockLen]
		if _, err = io.ReadFull(b.r, b.slice); err != nil {
			return 0, err
		}
	}
	n := copy(p, b.slice)
	b.slice = b.slice[n:]
	return n, nil
}

// decode reads a GIF image from r and stores the result in d.
func (d *decoder) decode(r io.Reader, configOnly bool) os.Error {
	// Add buffering if r does not provide ReadByte.
	if rr, ok := r.(reader); ok {
		d.r = rr
	} else {
		d.r = bufio.NewReader(r)
	}

	err := d.readHeaderAndScreenDescriptor()
	if err != nil {
		return err
	}
	if configOnly {
		return nil
	}

	if d.headerFields&fColorMapFollows != 0 {
		if d.globalColorMap, err = d.readColorMap(); err != nil {
			return err
		}
	}

Loop:
	for err == nil {
		var c byte
		c, err = d.r.ReadByte()
		if err == os.EOF {
			break
		}
		switch c {
		case sExtension:
			err = d.readExtension()

		case sImageDescriptor:
			var m *image.Paletted
			m, err = d.newImageFromDescriptor()
			if err != nil {
				break
			}
			if d.imageFields&fColorMapFollows != 0 {
				m.Palette, err = d.readColorMap()
				if err != nil {
					break
				}
				// TODO: do we set transparency in this map too? That would be
				// d.setTransparency(m.Palette)
			} else {
				m.Palette = d.globalColorMap
			}
			var litWidth uint8
			litWidth, err = d.r.ReadByte()
			if err != nil {
				return err
			}
			if litWidth < 2 || litWidth > 8 {
				return fmt.Errorf("gif: pixel size in decode out of range: %d", litWidth)
			}
			// A wonderfully Go-like piece of magic.
			lzwr := lzw.NewReader(&blockReader{r: d.r}, lzw.LSB, int(litWidth))
			if _, err = io.ReadFull(lzwr, m.Pix); err != nil {
				break
			}

			// There should be a "0" block remaining; drain that.
			c, err = d.r.ReadByte()
			if err != nil {
				return err
			}
			if c != 0 {
				return os.ErrorString("gif: extra data after image")
			}

			// Undo the interlacing if necessary.
			d.uninterlace(m)

			d.image = append(d.image, m)
			d.delay = append(d.delay, d.delayTime)
			d.delayTime = 0 // TODO: is this correct, or should we hold on to the value?

		case sTrailer:
			break Loop

		default:
			err = fmt.Errorf("gif: unknown block type: 0x%.2x", c)
		}
	}
	if err != nil {
		return err
	}
	if len(d.image) == 0 {
		return io.ErrUnexpectedEOF
	}
	return nil
}

func (d *decoder) readHeaderAndScreenDescriptor() os.Error {
	_, err := io.ReadFull(d.r, d.tmp[0:13])
	if err != nil {
		return err
	}
	d.vers = string(d.tmp[0:6])
	if d.vers != "GIF87a" && d.vers != "GIF89a" {
		return fmt.Errorf("gif: can't recognize format %s", d.vers)
	}
	d.width = int(d.tmp[6]) + int(d.tmp[7])<<8
	d.height = int(d.tmp[8]) + int(d.tmp[9])<<8
	d.headerFields = d.tmp[10]
	d.backgroundIndex = d.tmp[11]
	d.aspect = d.tmp[12]
	d.loopCount = -1
	d.pixelSize = uint(d.headerFields&7) + 1
	return nil
}

func (d *decoder) readColorMap() (image.PalettedColorModel, os.Error) {
	if d.pixelSize > 8 {
		return nil, fmt.Errorf("gif: can't handle %d bits per pixel", d.pixelSize)
	}
	numColors := 1 << d.pixelSize
	numValues := 3 * numColors
	_, err := io.ReadFull(d.r, d.tmp[0:numValues])
	if err != nil {
		return nil, fmt.Errorf("gif: short read on color map: %s", err)
	}
	colorMap := make(image.PalettedColorModel, numColors)
	j := 0
	for i := range colorMap {
		colorMap[i] = image.RGBAColor{d.tmp[j+0], d.tmp[j+1], d.tmp[j+2], 0xFF}
		j += 3
	}
	return colorMap, nil
}

func (d *decoder) readExtension() os.Error {
	extension, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	size := 0
	switch extension {
	case eText:
		size = 13
	case eGraphicControl:
		return d.readGraphicControl()
	case eComment:
		// nothing to do but read the data.
	case eApplication:
		b, err := d.r.ReadByte()
		if err != nil {
			return err
		}
		// The spec requires size be 11, but Adobe sometimes uses 10.
		size = int(b)
	default:
		return fmt.Errorf("gif: unknown extension 0x%.2x", extension)
	}
	if size > 0 {
		if _, err := io.ReadFull(d.r, d.tmp[0:size]); err != nil {
			return err
		}
	}

	// Application Extension with "NETSCAPE2.0" as string and 1 in data means
	// this extension defines a loop count.
	if extension == eApplication && string(d.tmp[:size]) == "NETSCAPE2.0" {
		n, err := d.readBlock()
		if n == 0 || err != nil {
			return err
		}
		if n == 3 && d.tmp[0] == 1 {
			d.loopCount = int(d.tmp[1]) | int(d.tmp[2])<<8
		}
	}
	for {
		n, err := d.readBlock()
		if n == 0 || err != nil {
			return err
		}
	}
	panic("unreachable")
}

func (d *decoder) readGraphicControl() os.Error {
	if _, err := io.ReadFull(d.r, d.tmp[0:6]); err != nil {
		return fmt.Errorf("gif: can't read graphic control: %s", err)
	}
	d.flags = d.tmp[1]
	d.delayTime = int(d.tmp[2]) | int(d.tmp[3])<<8
	if d.flags&gcTransparentColorSet != 0 {
		d.transparentIndex = d.tmp[4]
		d.setTransparency(d.globalColorMap)
	}
	return nil
}

func (d *decoder) setTransparency(colorMap image.PalettedColorModel) {
	if int(d.transparentIndex) < len(colorMap) {
		colorMap[d.transparentIndex] = image.RGBAColor{}
	}
}

func (d *decoder) newImageFromDescriptor() (*image.Paletted, os.Error) {
	if _, err := io.ReadFull(d.r, d.tmp[0:9]); err != nil {
		return nil, fmt.Errorf("gif: can't read image descriptor: %s", err)
	}
	// TODO: This code (throughout) ignores the top and left values,
	// and assumes (in interlacing, for example) that the images'
	// widths and heights are all the same.
	_ = int(d.tmp[0]) + int(d.tmp[1])<<8 // TODO: honor left value
	_ = int(d.tmp[2]) + int(d.tmp[3])<<8 // TODO: honor top value
	width := int(d.tmp[4]) + int(d.tmp[5])<<8
	height := int(d.tmp[6]) + int(d.tmp[7])<<8
	d.imageFields = d.tmp[8]
	return image.NewPaletted(width, height, nil), nil
}

func (d *decoder) readBlock() (int, os.Error) {
	n, err := d.r.ReadByte()
	if n == 0 || err != nil {
		return 0, err
	}
	return io.ReadFull(d.r, d.tmp[0:n])
}

// interlaceScan defines the ordering for a pass of the interlace algorithm.
type interlaceScan struct {
	skip, start int
}

// interlacing represents the set of scans in an interlaced GIF image.
var interlacing = []interlaceScan{
	{8, 0}, // Group 1 : Every 8th. row, starting with row 0.
	{8, 4}, // Group 2 : Every 8th. row, starting with row 4.
	{4, 2}, // Group 3 : Every 4th. row, starting with row 2.
	{2, 1}, // Group 4 : Every 2nd. row, starting with row 1.
}

func (d *decoder) uninterlace(m *image.Paletted) {
	if d.imageFields&ifInterlace == 0 {
		return
	}
	var nPix []uint8
	dx := d.width
	dy := d.height
	nPix = make([]uint8, dx*dy)
	offset := 0 // steps through the input by sequential scan lines.
	for _, pass := range interlacing {
		nOffset := pass.start * dx // steps through the output as defined by pass.
		for y := pass.start; y < dy; y += pass.skip {
			copy(nPix[nOffset:nOffset+dx], m.Pix[offset:offset+dx])
			offset += dx
			nOffset += dx * pass.skip
		}
	}
	m.Pix = nPix
}

// Decode reads a GIF image from r and returns the first embedded
// image as an image.Image.
func Decode(r io.Reader) (image.Image, os.Error) {
	var d decoder
	if err := d.decode(r, false); err != nil {
		return nil, err
	}
	return d.image[0], nil
}

// GIF represents the possibly multiple images stored in a GIF file.
type GIF struct {
	Image     []*image.Paletted // The successive images.
	Delay     []int             // The successive delay times, one per frame, in 100ths of a second.
	LoopCount int               // The loop count.
}

// DecodeAll reads a GIF image from r and returns the sequential frames
// and timing information.
func DecodeAll(r io.Reader) (*GIF, os.Error) {
	var d decoder
	if err := d.decode(r, false); err != nil {
		return nil, err
	}
	gif := &GIF{
		Image:     d.image,
		LoopCount: d.loopCount,
		Delay:     d.delay,
	}
	return gif, nil
}

// DecodeConfig returns the global color model and dimensions of a GIF image
// without decoding the entire image.
func DecodeConfig(r io.Reader) (image.Config, os.Error) {
	var d decoder
	if err := d.decode(r, true); err != nil {
		return image.Config{}, err
	}
	return image.Config{d.globalColorMap, d.width, d.height}, nil
}

func init() {
	image.RegisterFormat("gif", "GIF8?a", Decode, DecodeConfig)
}
<|endoftext|>


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

package websocket

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
)

func newServerConn(rwc io.ReadWriteCloser, buf *bufio.ReadWriter, req *http.Request) (conn *Conn, err error) {
	config := new(Config)
	var hs serverHandshaker = &hybiServerHandshaker{Config: config}
	code, err := hs.ReadHandshake(buf.Reader, req)
	if err == ErrBadWebSocketVersion {
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		fmt.Fprintf(buf, "Sec-WebSocket-Version: %s\r\n", SupportedProtocolVersion)
		buf.WriteString("\r\n")
		buf.WriteString(err.Error())
		buf.Flush()
		return
	}
	if err != nil {
		hs = &hixie76ServerHandshaker{Config: config}
		code, err = hs.ReadHandshake(buf.Reader, req)
	}
	if err != nil {
		hs = &hixie75ServerHandshaker{Config: config}
		code, err = hs.ReadHandshake(buf.Reader, req)
	}
	if err != nil {
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		buf.WriteString("\r\n")
		buf.WriteString(err.Error())
		buf.Flush()
		return
	}
	config.Protocol = nil

	err = hs.AcceptHandshake(buf.Writer)
	if err != nil {
		code = http.StatusBadRequest
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		buf.WriteString("\r\n")
		buf.Flush()
		return
	}
	conn = hs.NewServerConn(buf, rwc, req)
	return
}

/*
Handler is an interface to a WebSocket.

A trivial example server:

	package main

	import (
		"io"
		"net/http"
		"websocket"
	)

	// Echo the data received on the WebSocket.
	func EchoServer(ws *websocket.Conn) {
		io.Copy(ws, ws);
	}

	func main() {
		http.Handle("/echo", websocket.Handler(EchoServer));
		err := http.ListenAndServe(":12345", nil);
		if err != nil {
			panic("ListenAndServe: " + err.String())
		}
	}
*/
type Handler func(*Conn)

// ServeHTTP implements the http.Handler interface for a Web Socket
func (h Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	rwc, buf, err := w.(http.Hijacker).Hijack()
	if err != nil {
		panic("Hijack failed: " + err.Error())
		return
	}
	// The server should abort the WebSocket connection if it finds
	// the client did not send a handshake that matches with protocol
	// specification.
	defer rwc.Close()
	conn, err := newServerConn(rwc, buf, req)
	if err != nil {
		return
	}
	if conn == nil {
		panic("unepxected nil conn")
	}
	h(conn)
}
<commit_msg>websocket: fix a trivial example server<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package websocket

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
)

func newServerConn(rwc io.ReadWriteCloser, buf *bufio.ReadWriter, req *http.Request) (conn *Conn, err error) {
	config := new(Config)
	var hs serverHandshaker = &hybiServerHandshaker{Config: config}
	code, err := hs.ReadHandshake(buf.Reader, req)
	if err == ErrBadWebSocketVersion {
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		fmt.Fprintf(buf, "Sec-WebSocket-Version: %s\r\n", SupportedProtocolVersion)
		buf.WriteString("\r\n")
		buf.WriteString(err.Error())
		buf.Flush()
		return
	}
	if err != nil {
		hs = &hixie76ServerHandshaker{Config: config}
		code, err = hs.ReadHandshake(buf.Reader, req)
	}
	if err != nil {
		hs = &hixie75ServerHandshaker{Config: config}
		code, err = hs.ReadHandshake(buf.Reader, req)
	}
	if err != nil {
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		buf.WriteString("\r\n")
		buf.WriteString(err.Error())
		buf.Flush()
		return
	}
	config.Protocol = nil

	err = hs.AcceptHandshake(buf.Writer)
	if err != nil {
		code = http.StatusBadRequest
		fmt.Fprintf(buf, "HTTP/1.1 %03d %s\r\n", code, http.StatusText(code))
		buf.WriteString("\r\n")
		buf.Flush()
		return
	}
	conn = hs.NewServerConn(buf, rwc, req)
	return
}

/*
Handler is an interface to a WebSocket.

A trivial example server:

	package main

	import (
		"io"
		"net/http"
		"websocket"
	)

	// Echo the data received on the WebSocket.
	func EchoServer(ws *websocket.Conn) {
		io.Copy(ws, ws);
	}

	func main() {
		http.Handle("/echo", websocket.Handler(EchoServer));
		err := http.ListenAndServe(":12345", nil);
		if err != nil {
			panic("ListenAndServe: " + err.Error())
		}
	}
*/
type Handler func(*Conn)

// ServeHTTP implements the http.Handler interface for a Web Socket
func (h Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	rwc, buf, err := w.(http.Hijacker).Hijack()
	if err != nil {
		panic("Hijack failed: " + err.Error())
		return
	}
	// The server should abort the WebSocket connection if it finds
	// the client did not send a handshake that matches with protocol
	// specification.
	defer rwc.Close()
	conn, err := newServerConn(rwc, buf, req)
	if err != nil {
		return
	}
	if conn == nil {
		panic("unepxected nil conn")
	}
	h(conn)
}
<|endoftext|>


<commit_before>package e2e

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/client/conditions"
	"k8s.io/kubernetes/test/e2e/framework"
	e2elog "k8s.io/kubernetes/test/e2e/framework/log"
)

// getDaemonSetLabelSelector returns labels of daemonset given name and namespace dynamically,
// needed since labels are not same for helm and non-helm deployments.
func getDaemonSetLabelSelector(f *framework.Framework, ns, daemonSetName string) (string, error) {
	ds, err := f.ClientSet.AppsV1().DaemonSets(ns).Get(context.TODO(), daemonSetName, metav1.GetOptions{})
	if err != nil {
		e2elog.Logf("Error getting daemonsets with name %s in namespace %s", daemonSetName, ns)
		return "", err
	}
	s, err := metav1.LabelSelectorAsSelector(ds.Spec.Selector)
	if err != nil {
		e2elog.Logf("Error parsing %s daemonset selector in namespace %s", daemonSetName, ns)
		return "", err
	}
	e2elog.Logf("LabelSelector for %s daemonsets in namespace %s: %s", daemonSetName, ns, s.String())
	return s.String(), nil
}

func waitForDaemonSets(name, ns string, c kubernetes.Interface, t int) error {
	timeout := time.Duration(t) * time.Minute
	start := time.Now()
	e2elog.Logf("Waiting up to %v for all daemonsets in namespace '%s' to start", timeout, ns)

	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		ds, err := c.AppsV1().DaemonSets(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			e2elog.Logf("Error getting daemonsets in namespace: '%s': %v", ns, err)
			if strings.Contains(err.Error(), "not found") {
				return false, nil
			}
			if isRetryableAPIError(err) {
				return false, nil
			}
			return false, err
		}
		dNum := ds.Status.DesiredNumberScheduled
		ready := ds.Status.NumberReady
		e2elog.Logf("%d / %d pods ready in namespace '%s' in daemonset '%s' (%d seconds elapsed)", ready, dNum, ns, ds.ObjectMeta.Name, int(time.Since(start).Seconds()))
		if ready != dNum {
			return false, nil
		}

		return true, nil
	})
}

// Waits for the deployment to complete.

func waitForDeploymentComplete(name, ns string, c kubernetes.Interface, t int) error {
	var (
		deployment *appsv1.Deployment
		reason     string
		err        error
	)
	timeout := time.Duration(t) * time.Minute
	err = wait.PollImmediate(poll, timeout, func() (bool, error) {
		deployment, err = c.AppsV1().Deployments(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			// a StatusError is not marked as 'retryable', but we want to retry anyway
			if isRetryableAPIError(err) || strings.Contains(err.Error(), "etcdserver: request timed out") {
				// hide API-server timeouts, so that PollImmediate() retries
				e2elog.Logf("deployment error: %v", err)
				return false, nil
			}
			return false, err
		}

		// TODO need to check rolling update

		// When the deployment status and its underlying resources reach the
		// desired state, we're done
		if deployment.Status.Replicas == deployment.Status.ReadyReplicas {
			return true, nil
		}
		e2elog.Logf("deployment status: expected replica count %d running replica count %d", deployment.Status.Replicas, deployment.Status.ReadyReplicas)
		reason = fmt.Sprintf("deployment status: %#v", deployment.Status.String())
		return false, nil
	})

	if errors.Is(err, wait.ErrWaitTimeout) {
		err = fmt.Errorf("%s", reason)
	}
	if err != nil {
		return fmt.Errorf("error waiting for deployment %q status to match expectation: %w", name, err)
	}
	return nil
}

func findPodAndContainerName(f *framework.Framework, ns, cn string, opt *metav1.ListOptions) (string, string, error) {
	podList, err := f.PodClientNS(ns).List(context.TODO(), *opt)
	if err != nil {
		return "", "", err
	}

	if len(podList.Items) == 0 {
		return "", "", errors.New("podlist is empty")
	}

	if cn != "" {
		for i := range podList.Items {
			for j := range podList.Items[i].Spec.Containers {
				if podList.Items[i].Spec.Containers[j].Name == cn {
					return podList.Items[i].Name, cn, nil
				}
			}
		}
		return "", "", errors.New("container name not found")
	}
	return podList.Items[0].Name, podList.Items[0].Spec.Containers[0].Name, nil
}

func getCommandInPodOpts(f *framework.Framework, c, ns, cn string, opt *metav1.ListOptions) (framework.ExecOptions, error) {
	cmd := []string{"/bin/sh", "-c", c}
	pName, cName, err := findPodAndContainerName(f, ns, cn, opt)
	if err != nil {
		return framework.ExecOptions{}, err
	}
	return framework.ExecOptions{
		Command:            cmd,
		PodName:            pName,
		Namespace:          ns,
		ContainerName:      cName,
		Stdin:              nil,
		CaptureStdout:      true,
		CaptureStderr:      true,
		PreserveWhitespace: true,
	}, nil
}

// execCommandInDaemonsetPod executes commands inside given container of a daemonset pod on a particular node.
func execCommandInDaemonsetPod(f *framework.Framework, c, daemonsetName, nodeName, containerName, ns string) (string, string, error) {
	selector, err := getDaemonSetLabelSelector(f, ns, daemonsetName)
	if err != nil {
		return "", "", err
	}

	opt := &metav1.ListOptions{
		LabelSelector: selector,
	}
	pods, err := listPods(f, ns, opt)
	if err != nil {
		return "", "", err
	}

	podName := ""
	for i := range pods {
		if pods[i].Spec.NodeName == nodeName {
			podName = pods[i].Name
		}
	}
	if podName == "" {
		return "", "", fmt.Errorf("%s daemonset pod on node %s in namespace %s not found", daemonsetName, nodeName, ns)
	}

	cmd := []string{"/bin/sh", "-c", c}
	podOpt := framework.ExecOptions{
		Command:       cmd,
		Namespace:     ns,
		PodName:       podName,
		ContainerName: containerName,
		CaptureStdout: true,
		CaptureStderr: true,
	}
	return f.ExecWithOptions(podOpt)
}

// listPods returns slice of pods matching given ListOptions and namespace.
func listPods(f *framework.Framework, ns string, opt *metav1.ListOptions) ([]v1.Pod, error) {
	podList, err := f.PodClientNS(ns).List(context.TODO(), *opt)
	if len(podList.Items) == 0 {
		return podList.Items, fmt.Errorf("podlist for label '%s' in namespace %s is empty", opt.LabelSelector, ns)
	}
	return podList.Items, err
}

func execCommandInPod(f *framework.Framework, c, ns string, opt *metav1.ListOptions) (string, string, error) {
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

func execCommandInContainer(f *framework.Framework, c, ns, cn string, opt *metav1.ListOptions) (string, string, error) { //nolint:unparam,lll // cn can be used with different inputs later
	podOpt, err := getCommandInPodOpts(f, c, ns, cn, opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

func execCommandInToolBoxPod(f *framework.Framework, c, ns string) (string, string, error) {
	opt := &metav1.ListOptions{
		LabelSelector: rookToolBoxPodLabel,
	}
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

func execCommandInPodAndAllowFail(f *framework.Framework, c, ns string, opt *metav1.ListOptions) (string, string) {
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", err.Error()
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if err != nil {
		e2elog.Logf("command %s failed: %v", c, err)
	}
	return stdOut, stdErr
}

func loadApp(path string) (*v1.Pod, error) {
	app := v1.Pod{}
	if err := unmarshal(path, &app); err != nil {
		return nil, err
	}
	for i := range app.Spec.Containers {
		app.Spec.Containers[i].ImagePullPolicy = v1.PullIfNotPresent
	}
	return &app, nil
}

func createApp(c kubernetes.Interface, app *v1.Pod, timeout int) error {
	_, err := c.CoreV1().Pods(app.Namespace).Create(context.TODO(), app, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create app: %w", err)
	}
	return waitForPodInRunningState(app.Name, app.Namespace, c, timeout, noError)
}

func createAppErr(c kubernetes.Interface, app *v1.Pod, timeout int, errString string) error {
	_, err := c.CoreV1().Pods(app.Namespace).Create(context.TODO(), app, metav1.CreateOptions{})
	if err != nil {
		return err
	}
	return waitForPodInRunningState(app.Name, app.Namespace, c, timeout, errString)
}

func waitForPodInRunningState(name, ns string, c kubernetes.Interface, t int, expectedError string) error {
	timeout := time.Duration(t) * time.Minute
	start := time.Now()
	e2elog.Logf("Waiting up to %v to be in Running state", name)
	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		pod, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("failed to get app: %w", err)
		}
		switch pod.Status.Phase {
		case v1.PodRunning:
			return true, nil
		case v1.PodFailed, v1.PodSucceeded:
			return false, conditions.ErrPodCompleted
		case v1.PodPending:
			if expectedError != "" {
				events, err := c.CoreV1().Events(ns).List(context.TODO(), metav1.ListOptions{
					FieldSelector: fmt.Sprintf("involvedObject.name=%s", name),
				})
				if err != nil {
					return false, err
				}
				if strings.Contains(events.String(), expectedError) {
					e2elog.Logf("Expected Error %q found successfully", expectedError)
					return true, err
				}
			}
		}
		e2elog.Logf("%s app  is in %s phase expected to be in Running  state (%d seconds elapsed)", name, pod.Status.Phase, int(time.Since(start).Seconds()))
		return false, nil
	})
}

func deletePod(name, ns string, c kubernetes.Interface, t int) error {
	timeout := time.Duration(t) * time.Minute
	err := c.CoreV1().Pods(ns).Delete(context.TODO(), name, metav1.DeleteOptions{})
	if err != nil {
		return fmt.Errorf("failed to delete app: %w", err)
	}
	start := time.Now()
	e2elog.Logf("Waiting for pod %v to be deleted", name)
	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		_, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})

		if apierrs.IsNotFound(err) {
			return true, nil
		}
		e2elog.Logf("%s app  to be deleted (%d seconds elapsed)", name, int(time.Since(start).Seconds()))
		if err != nil {
			return false, fmt.Errorf("failed to get app: %w", err)
		}
		return false, nil
	})
}

func deletePodWithLabel(label, ns string, skipNotFound bool) error { //nolint:unparam // skipNotFound can be used with different inputs later
	_, err := framework.RunKubectl(ns, "delete", "po", "-l", label, fmt.Sprintf("--ignore-not-found=%t", skipNotFound))
	if err != nil {
		e2elog.Logf("failed to delete pod %v", err)
	}
	return err
}

// calculateSHA512sum returns the sha512sum of a file inside a pod.
func calculateSHA512sum(f *framework.Framework, app *v1.Pod, filePath string, opt *metav1.ListOptions) (string, error) {
	cmd := fmt.Sprintf("sha512sum %s", filePath)
	sha512sumOut, stdErr, err := execCommandInPod(f, cmd, app.Namespace, opt)
	if err != nil {
		return "", err
	}
	if stdErr != "" {
		return "", fmt.Errorf("error: sha512sum could not be calculated %v", stdErr)
	}
	// extract checksum from sha512sum output.
	checkSum := strings.Split(sha512sumOut, "")[0]
	e2elog.Logf("Calculated checksum  %s", checkSum)
	return checkSum, nil
}
<commit_msg>e2e: reformat long lines in this package to 120 chars<commit_after>package e2e

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/client/conditions"
	"k8s.io/kubernetes/test/e2e/framework"
	e2elog "k8s.io/kubernetes/test/e2e/framework/log"
)

// getDaemonSetLabelSelector returns labels of daemonset given name and namespace dynamically,
// needed since labels are not same for helm and non-helm deployments.
func getDaemonSetLabelSelector(f *framework.Framework, ns, daemonSetName string) (string, error) {
	ds, err := f.ClientSet.AppsV1().DaemonSets(ns).Get(context.TODO(), daemonSetName, metav1.GetOptions{})
	if err != nil {
		e2elog.Logf("Error getting daemonsets with name %s in namespace %s", daemonSetName, ns)
		return "", err
	}
	s, err := metav1.LabelSelectorAsSelector(ds.Spec.Selector)
	if err != nil {
		e2elog.Logf("Error parsing %s daemonset selector in namespace %s", daemonSetName, ns)
		return "", err
	}
	e2elog.Logf("LabelSelector for %s daemonsets in namespace %s: %s", daemonSetName, ns, s.String())
	return s.String(), nil
}

func waitForDaemonSets(name, ns string, c kubernetes.Interface, t int) error {
	timeout := time.Duration(t) * time.Minute
	start := time.Now()
	e2elog.Logf("Waiting up to %v for all daemonsets in namespace '%s' to start", timeout, ns)

	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		ds, err := c.AppsV1().DaemonSets(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			e2elog.Logf("Error getting daemonsets in namespace: '%s': %v", ns, err)
			if strings.Contains(err.Error(), "not found") {
				return false, nil
			}
			if isRetryableAPIError(err) {
				return false, nil
			}
			return false, err
		}
		dNum := ds.Status.DesiredNumberScheduled
		ready := ds.Status.NumberReady
		e2elog.Logf(
			"%d / %d pods ready in namespace '%s' in daemonset '%s' (%d seconds elapsed)",
			ready,
			dNum,
			ns,
			ds.ObjectMeta.Name,
			int(time.Since(start).Seconds()))
		if ready != dNum {
			return false, nil
		}

		return true, nil
	})
}

// Waits for the deployment to complete.

func waitForDeploymentComplete(name, ns string, c kubernetes.Interface, t int) error {
	var (
		deployment *appsv1.Deployment
		reason     string
		err        error
	)
	timeout := time.Duration(t) * time.Minute
	err = wait.PollImmediate(poll, timeout, func() (bool, error) {
		deployment, err = c.AppsV1().Deployments(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			// a StatusError is not marked as 'retryable', but we want to retry anyway
			if isRetryableAPIError(err) || strings.Contains(err.Error(), "etcdserver: request timed out") {
				// hide API-server timeouts, so that PollImmediate() retries
				e2elog.Logf("deployment error: %v", err)
				return false, nil
			}
			return false, err
		}

		// TODO need to check rolling update

		// When the deployment status and its underlying resources reach the
		// desired state, we're done
		if deployment.Status.Replicas == deployment.Status.ReadyReplicas {
			return true, nil
		}
		e2elog.Logf(
			"deployment status: expected replica count %d running replica count %d",
			deployment.Status.Replicas,
			deployment.Status.ReadyReplicas)
		reason = fmt.Sprintf("deployment status: %#v", deployment.Status.String())
		return false, nil
	})

	if errors.Is(err, wait.ErrWaitTimeout) {
		err = fmt.Errorf("%s", reason)
	}
	if err != nil {
		return fmt.Errorf("error waiting for deployment %q status to match expectation: %w", name, err)
	}
	return nil
}

func findPodAndContainerName(f *framework.Framework, ns, cn string, opt *metav1.ListOptions) (string, string, error) {
	podList, err := f.PodClientNS(ns).List(context.TODO(), *opt)
	if err != nil {
		return "", "", err
	}

	if len(podList.Items) == 0 {
		return "", "", errors.New("podlist is empty")
	}

	if cn != "" {
		for i := range podList.Items {
			for j := range podList.Items[i].Spec.Containers {
				if podList.Items[i].Spec.Containers[j].Name == cn {
					return podList.Items[i].Name, cn, nil
				}
			}
		}
		return "", "", errors.New("container name not found")
	}
	return podList.Items[0].Name, podList.Items[0].Spec.Containers[0].Name, nil
}

func getCommandInPodOpts(
	f *framework.Framework,
	c, ns, cn string,
	opt *metav1.ListOptions) (framework.ExecOptions, error) {
	cmd := []string{"/bin/sh", "-c", c}
	pName, cName, err := findPodAndContainerName(f, ns, cn, opt)
	if err != nil {
		return framework.ExecOptions{}, err
	}
	return framework.ExecOptions{
		Command:            cmd,
		PodName:            pName,
		Namespace:          ns,
		ContainerName:      cName,
		Stdin:              nil,
		CaptureStdout:      true,
		CaptureStderr:      true,
		PreserveWhitespace: true,
	}, nil
}

// execCommandInDaemonsetPod executes commands inside given container of a daemonset pod on a particular node.
func execCommandInDaemonsetPod(
	f *framework.Framework,
	c, daemonsetName, nodeName, containerName, ns string) (string, string, error) {
	selector, err := getDaemonSetLabelSelector(f, ns, daemonsetName)
	if err != nil {
		return "", "", err
	}

	opt := &metav1.ListOptions{
		LabelSelector: selector,
	}
	pods, err := listPods(f, ns, opt)
	if err != nil {
		return "", "", err
	}

	podName := ""
	for i := range pods {
		if pods[i].Spec.NodeName == nodeName {
			podName = pods[i].Name
		}
	}
	if podName == "" {
		return "", "", fmt.Errorf("%s daemonset pod on node %s in namespace %s not found", daemonsetName, nodeName, ns)
	}

	cmd := []string{"/bin/sh", "-c", c}
	podOpt := framework.ExecOptions{
		Command:       cmd,
		Namespace:     ns,
		PodName:       podName,
		ContainerName: containerName,
		CaptureStdout: true,
		CaptureStderr: true,
	}
	return f.ExecWithOptions(podOpt)
}

// listPods returns slice of pods matching given ListOptions and namespace.
func listPods(f *framework.Framework, ns string, opt *metav1.ListOptions) ([]v1.Pod, error) {
	podList, err := f.PodClientNS(ns).List(context.TODO(), *opt)
	if len(podList.Items) == 0 {
		return podList.Items, fmt.Errorf("podlist for label '%s' in namespace %s is empty", opt.LabelSelector, ns)
	}
	return podList.Items, err
}

func execCommandInPod(f *framework.Framework, c, ns string, opt *metav1.ListOptions) (string, string, error) {
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

//nolint:unparam // cn can be used with different inputs later
func execCommandInContainer(
	f *framework.Framework, c, ns, cn string, opt *metav1.ListOptions) (string, string, error) {
	podOpt, err := getCommandInPodOpts(f, c, ns, cn, opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

func execCommandInToolBoxPod(f *framework.Framework, c, ns string) (string, string, error) {
	opt := &metav1.ListOptions{
		LabelSelector: rookToolBoxPodLabel,
	}
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", "", err
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if stdErr != "" {
		e2elog.Logf("stdErr occurred: %v", stdErr)
	}
	return stdOut, stdErr, err
}

func execCommandInPodAndAllowFail(f *framework.Framework, c, ns string, opt *metav1.ListOptions) (string, string) {
	podOpt, err := getCommandInPodOpts(f, c, ns, "", opt)
	if err != nil {
		return "", err.Error()
	}
	stdOut, stdErr, err := f.ExecWithOptions(podOpt)
	if err != nil {
		e2elog.Logf("command %s failed: %v", c, err)
	}
	return stdOut, stdErr
}

func loadApp(path string) (*v1.Pod, error) {
	app := v1.Pod{}
	if err := unmarshal(path, &app); err != nil {
		return nil, err
	}
	for i := range app.Spec.Containers {
		app.Spec.Containers[i].ImagePullPolicy = v1.PullIfNotPresent
	}
	return &app, nil
}

func createApp(c kubernetes.Interface, app *v1.Pod, timeout int) error {
	_, err := c.CoreV1().Pods(app.Namespace).Create(context.TODO(), app, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create app: %w", err)
	}
	return waitForPodInRunningState(app.Name, app.Namespace, c, timeout, noError)
}

func createAppErr(c kubernetes.Interface, app *v1.Pod, timeout int, errString string) error {
	_, err := c.CoreV1().Pods(app.Namespace).Create(context.TODO(), app, metav1.CreateOptions{})
	if err != nil {
		return err
	}
	return waitForPodInRunningState(app.Name, app.Namespace, c, timeout, errString)
}

func waitForPodInRunningState(name, ns string, c kubernetes.Interface, t int, expectedError string) error {
	timeout := time.Duration(t) * time.Minute
	start := time.Now()
	e2elog.Logf("Waiting up to %v to be in Running state", name)
	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		pod, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("failed to get app: %w", err)
		}
		switch pod.Status.Phase {
		case v1.PodRunning:
			return true, nil
		case v1.PodFailed, v1.PodSucceeded:
			return false, conditions.ErrPodCompleted
		case v1.PodPending:
			if expectedError != "" {
				events, err := c.CoreV1().Events(ns).List(context.TODO(), metav1.ListOptions{
					FieldSelector: fmt.Sprintf("involvedObject.name=%s", name),
				})
				if err != nil {
					return false, err
				}
				if strings.Contains(events.String(), expectedError) {
					e2elog.Logf("Expected Error %q found successfully", expectedError)
					return true, err
				}
			}
		}
		e2elog.Logf(
			"%s app  is in %s phase expected to be in Running  state (%d seconds elapsed)",
			name,
			pod.Status.Phase,
			int(time.Since(start).Seconds()))
		return false, nil
	})
}

func deletePod(name, ns string, c kubernetes.Interface, t int) error {
	timeout := time.Duration(t) * time.Minute
	err := c.CoreV1().Pods(ns).Delete(context.TODO(), name, metav1.DeleteOptions{})
	if err != nil {
		return fmt.Errorf("failed to delete app: %w", err)
	}
	start := time.Now()
	e2elog.Logf("Waiting for pod %v to be deleted", name)
	return wait.PollImmediate(poll, timeout, func() (bool, error) {
		_, err := c.CoreV1().Pods(ns).Get(context.TODO(), name, metav1.GetOptions{})

		if apierrs.IsNotFound(err) {
			return true, nil
		}
		e2elog.Logf("%s app  to be deleted (%d seconds elapsed)", name, int(time.Since(start).Seconds()))
		if err != nil {
			return false, fmt.Errorf("failed to get app: %w", err)
		}
		return false, nil
	})
}

//nolint:unparam // skipNotFound arg can be used with different inputs later
func deletePodWithLabel(label, ns string, skipNotFound bool) error {
	_, err := framework.RunKubectl(ns, "delete", "po", "-l", label, fmt.Sprintf("--ignore-not-found=%t", skipNotFound))
	if err != nil {
		e2elog.Logf("failed to delete pod %v", err)
	}
	return err
}

// calculateSHA512sum returns the sha512sum of a file inside a pod.
func calculateSHA512sum(f *framework.Framework, app *v1.Pod, filePath string, opt *metav1.ListOptions) (string, error) {
	cmd := fmt.Sprintf("sha512sum %s", filePath)
	sha512sumOut, stdErr, err := execCommandInPod(f, cmd, app.Namespace, opt)
	if err != nil {
		return "", err
	}
	if stdErr != "" {
		return "", fmt.Errorf("error: sha512sum could not be calculated %v", stdErr)
	}
	// extract checksum from sha512sum output.
	checkSum := strings.Split(sha512sumOut, "")[0]
	e2elog.Logf("Calculated checksum  %s", checkSum)
	return checkSum, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
)

var (
	keyFile       = flag.String("keys", "", "File containing SSH pubkeys.")
	accountsFile  = flag.String("accounts", "", "File containing account definitions.")
	groupsFile    = flag.String("groups", "", "File containing group definitions.")
	matching      = flag.String("matching", ".*", "Only check hosts matching this regex.")
	doAddMissing  = flag.Bool("add_missing", false, "Add missing keys as needed.")
	doDeleteExtra = flag.Bool("delete_extra", false, "Delete extra keys as needed.")
)

type key struct {
	algorithm   string
	key         string
	description string
}

type account struct {
	account       string
	keyGroupNames []string
}

func parseKeys(b []byte) ([]key, error) {
	var keys []key
	seen := make(map[string]string)
	re := regexp.MustCompile(`(?m)^([^\s#]+)\s+([^\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		if _, found := seen[e[3]]; found {
			return nil, fmt.Errorf("key ID %q listed more than once", e[3])
		}
		if prev, found := seen[e[2]]; found {
			return nil, fmt.Errorf("key %q is same as %q", e[3], prev)
		}
		keys = append(keys, key{
			algorithm:   e[1],
			key:         e[2],
			description: e[3],
		})
		seen[e[2]] = e[3]
		seen[e[3]] = e[3]
	}
	return keys, nil
}

func readKeys(fn string) ([]key, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	return parseKeys(b)
}

func readAccounts(fn string) ([]account, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	var accounts []account
	re := regexp.MustCompile(`(?m)^([^#\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		accounts = append(accounts, account{
			account:       e[1],
			keyGroupNames: strings.Split(e[2], " "),
		})
	}
	return accounts, nil
}

func readGroups(fn string) (map[string][]string, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	groups := map[string][]string{}
	re := regexp.MustCompile(`(?m)^([^#\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		groups[e[1]] = strings.Split(e[2], " ")
	}
	return groups, nil
}

func checkAccount(ctx context.Context, kg map[string]*KeyGroup, account account) ([]key, []string, error) {
	cmd := exec.CommandContext(ctx, "ssh", account.account, "cat .ssh/authorized_keys")
	var buf bytes.Buffer
	cmd.Stdout = &buf
	if err := cmd.Run(); err != nil {
		return nil, nil, err
	}
	var err error
	currentKeys, err := parseKeys(buf.Bytes())
	if err != nil {
		return nil, nil, err
	}
	log.Debugf("Found %d current keys", len(currentKeys))

	// Check for extra keys.
	var extra []key
	for _, ck := range currentKeys {
		found := false
		for _, kn := range account.keyGroupNames {
			a := kg[kn]
			if a.Has(ck.key) {
				found = true
				break
			}
		}
		if !found {
			extra = append(extra, ck)
		}
	}

	// Check for missing keys.
	var missing []string
	checked := make(map[string]bool)
	for _, kgn := range account.keyGroupNames {
		for _, k := range kg[kgn].Keys() {
			if checked[k.key] {
				continue
			}
			checked[k.key] = true
			found := false
			for _, ck := range currentKeys {
				if k.key == ck.key {
					found = true
					break
				}
			}
			if !found {
				missing = append(missing, k.description)
			}
		}
	}
	return extra, missing, nil
}

func deleteExtra(ctx context.Context, acct account, extra []key) error {
	var cmds []string
	for _, key := range extra {
		log.Infof("Deleting %q from %q", key.description, acct.account)
		tmpf := ".ssh/authorized_keys.tmp"
		ak := ".ssh/authorized_keys"
		cmds = append(cmds,
			fmt.Sprintf(`grep -v '^%s %s ' %s > %s`, key.algorithm, key.key, ak, tmpf),
			fmt.Sprintf(`chmod --reference=%s %s`, ak, tmpf),
			fmt.Sprintf(`mv %s %s`, tmpf, ak),
		)
	}
	return exec.CommandContext(ctx, "ssh", acct.account, strings.Join(cmds, " && ")).Run()
}

func addMissing(ctx context.Context, keys []key, acct account, missing []string) error {
	var cmds []string
	for _, name := range missing {
		log.Infof("Adding %q to %q", name, acct.account)
		var k *key
		for _, t := range keys {
			if name == t.description {
				k = &t
				break
			}
		}
		if k == nil {
			log.Fatalf("Internal error: missing key %q, but can't find it", name)
		}

		// TODO: random tmpfile name.
		tmpf := ".ssh/authorized_keys.tmp"
		ak := ".ssh/authorized_keys"
		cmds = append(cmds,
			fmt.Sprintf(`cp %s %s`+
				` && echo "%s %s %s" >> %s`+
				` && mv %s %s`,
				ak, tmpf,
				k.algorithm, k.key, name, tmpf,
				tmpf, ak,
			))
	}
	cmd := exec.CommandContext(ctx, "ssh", acct.account, strings.Join(cmds, " && "))
	if err := cmd.Run(); err != nil {
		return err
	}
	return nil
}

func check(ctx context.Context, keys []key, kg map[string]*KeyGroup, accounts []account) error {
	log.Infof("Checking accounts…")
	re, err := regexp.Compile(*matching)
	if err != nil {
		return err
	}
	for _, account := range accounts {
		if re.FindString(account.account) == "" {
			continue
		}

		log.Infof("Checking %q", account.account)
		ctx2, cancel := context.WithTimeout(ctx, 60*time.Second)
		defer cancel()
		extra, missing, err := checkAccount(ctx2, kg, account)
		if err != nil {
			log.Errorf("Failed to check account %q: %v", account.account, err)
			continue
		}

		if len(extra) > 0 {
			var es []string
			for _, e := range extra {
				es = append(es, e.description)
			}
			log.Warningf("Extra keys: %q", es)
		}
		if len(missing) > 0 {
			log.Warningf("Missing keys: %q", missing)
		}
		if *doAddMissing {
			if err := addMissing(ctx, keys, account, missing); err != nil {
				log.Errorf("Failed to add missing keys %q to %q: %v", missing, account, err)
			}
		}
		if *doDeleteExtra {
			if err := deleteExtra(ctx, account, extra); err != nil {
				log.Errorf("Failed to delete extra keys %q from %q: %v", extra, account, err)
			}
		}
	}
	return nil
}

type KeyGroup struct {
	keys []key
}

func NewKeyGroup() *KeyGroup {
	return &KeyGroup{}
}

func (k *KeyGroup) Has(key string) bool {
	for _, k2 := range k.keys {
		if key == k2.key {
			return true
		}
	}
	return false
}

func (k *KeyGroup) Add(ky ...key) {
	k.keys = append(k.keys, ky...)
}

func (k *KeyGroup) Keys() []key {
	return k.keys
}

func main() {
	flag.Parse()
	if flag.NArg() != 1 {
		log.Fatalf("Command not specified. Supported commands: check, expand")
	}
	command := flag.Arg(0)

	// Read keys.
	keys, err := readKeys(*keyFile)
	if err != nil {
		log.Fatalf("Failed to read keyfile %q: %v", *keyFile, err)
	}
	log.Infof("Read %d keys", len(keys))

	// Make keygroups.
	keyGroups := map[string]*KeyGroup{}
	for _, k := range keys {
		kg := NewKeyGroup()
		kg.Add(k)
		keyGroups[k.description] = kg
	}

	groups, err := readGroups(*groupsFile)
	if err != nil {
		log.Fatalf("Failed to read groups file %q: %v", *groupsFile, err)
	}
	log.Infof("Read %d groups", len(groups))
	for g, ks := range groups {
		kg := NewKeyGroup()
		for _, kn := range ks {
			o, found := keyGroups[kn]
			if !found {
				log.Fatalf("Group %q refers to non-existing key %q", g, kn)
			}
			kg.Add(o.Keys()...)
		}
		keyGroups["@"+g] = kg
	}

	// Read accounts.
	accounts, err := readAccounts(*accountsFile)
	if err != nil {
		log.Fatalf("Failed to read accounts file %q: %v", *accountsFile, err)
	}
	log.Infof("Read %d accounts", len(accounts))

	// Verify that all accounts' keys exist in database.
	for _, a := range accounts {
		for _, kn := range a.keyGroupNames {
			if _, found := keyGroups[kn]; !found {
				log.Fatalf("Unknown key %q for account %q", kn, a.account)
			}
		}
	}
	sort.Slice(accounts, func(i, j int) bool {
		return accounts[i].account < accounts[j].account
	})

	if command == "check" {
		ctx := context.Background()
		if err := check(ctx, keys, keyGroups, accounts); err != nil {
			log.Fatalf("blah: %v", err)
		}
		log.Infof("Done")
	} else if command == "expand" {
		re, err := regexp.Compile(*matching)
		if err != nil {
			log.Printf("Invalid regex %q: %v", *matching, err)
		}
		for _, a := range accounts {
			if re.FindString(a.account) == "" {
				continue
			}
			var keys []string
			seen := make(map[string]bool)
			for _, kgn := range a.keyGroupNames {
				for _, k := range keyGroups[kgn].Keys() {
					if !seen[k.description] {
						keys = append(keys, k.description)
						seen[k.description] = true
					}
				}
			}
			sort.Strings(keys)
			fmt.Printf("%s\n  %s\n", a.account, strings.Join(keys, "\n  "))
		}
	} else {
		log.Fatalf("Invalid command %q", command)
	}
}
<commit_msg>Clean up --add_missing code<commit_after>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
)

var (
	keyFile       = flag.String("keys", "", "File containing SSH pubkeys.")
	accountsFile  = flag.String("accounts", "", "File containing account definitions.")
	groupsFile    = flag.String("groups", "", "File containing group definitions.")
	matching      = flag.String("matching", ".*", "Only check hosts matching this regex.")
	doAddMissing  = flag.Bool("add_missing", false, "Add missing keys as needed.")
	doDeleteExtra = flag.Bool("delete_extra", false, "Delete extra keys as needed.")
)

type key struct {
	algorithm   string
	key         string
	description string
}

type account struct {
	account       string
	keyGroupNames []string
}

func parseKeys(b []byte) ([]key, error) {
	var keys []key
	seen := make(map[string]string)
	re := regexp.MustCompile(`(?m)^([^\s#]+)\s+([^\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		if _, found := seen[e[3]]; found {
			return nil, fmt.Errorf("key ID %q listed more than once", e[3])
		}
		if prev, found := seen[e[2]]; found {
			return nil, fmt.Errorf("key %q is same as %q", e[3], prev)
		}
		keys = append(keys, key{
			algorithm:   e[1],
			key:         e[2],
			description: e[3],
		})
		seen[e[2]] = e[3]
		seen[e[3]] = e[3]
	}
	return keys, nil
}

func readKeys(fn string) ([]key, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	return parseKeys(b)
}

func readAccounts(fn string) ([]account, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	var accounts []account
	re := regexp.MustCompile(`(?m)^([^#\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		accounts = append(accounts, account{
			account:       e[1],
			keyGroupNames: strings.Split(e[2], " "),
		})
	}
	return accounts, nil
}

func readGroups(fn string) (map[string][]string, error) {
	b, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	groups := map[string][]string{}
	re := regexp.MustCompile(`(?m)^([^#\s]+)\s+(.*)$`)
	for _, e := range re.FindAllStringSubmatch(string(b), -1) {
		groups[e[1]] = strings.Split(e[2], " ")
	}
	return groups, nil
}

func checkAccount(ctx context.Context, kg map[string]*KeyGroup, account account) ([]key, []string, error) {
	cmd := exec.CommandContext(ctx, "ssh", account.account, "cat .ssh/authorized_keys")
	var buf bytes.Buffer
	cmd.Stdout = &buf
	if err := cmd.Run(); err != nil {
		return nil, nil, err
	}
	var err error
	currentKeys, err := parseKeys(buf.Bytes())
	if err != nil {
		return nil, nil, err
	}
	log.Debugf("Found %d current keys", len(currentKeys))

	// Check for extra keys.
	var extra []key
	for _, ck := range currentKeys {
		found := false
		for _, kn := range account.keyGroupNames {
			a := kg[kn]
			if a.Has(ck.key) {
				found = true
				break
			}
		}
		if !found {
			extra = append(extra, ck)
		}
	}

	// Check for missing keys.
	var missing []string
	checked := make(map[string]bool)
	for _, kgn := range account.keyGroupNames {
		for _, k := range kg[kgn].Keys() {
			if checked[k.key] {
				continue
			}
			checked[k.key] = true
			found := false
			for _, ck := range currentKeys {
				if k.key == ck.key {
					found = true
					break
				}
			}
			if !found {
				missing = append(missing, k.description)
			}
		}
	}
	return extra, missing, nil
}

func deleteExtra(ctx context.Context, acct account, extra []key) error {
	var cmds []string
	for _, key := range extra {
		log.Infof("Deleting %q from %q", key.description, acct.account)
		tmpf := ".ssh/authorized_keys.tmp"
		ak := ".ssh/authorized_keys"
		cmds = append(cmds,
			fmt.Sprintf(`grep -v '^%s %s ' %s > %s`, key.algorithm, key.key, ak, tmpf),
			fmt.Sprintf(`chmod --reference=%s %s`, ak, tmpf),
			fmt.Sprintf(`mv %s %s`, tmpf, ak),
		)
	}
	return exec.CommandContext(ctx, "ssh", acct.account, strings.Join(cmds, " && ")).Run()
}

func addMissing(ctx context.Context, keys []key, acct account, missing []string) error {
	var cmds []string
	for _, name := range missing {
		log.Infof("Adding %q to %q", name, acct.account)
		var k *key
		for _, t := range keys {
			if name == t.description {
				k = &t
				break
			}
		}
		if k == nil {
			log.Fatalf("Internal error: missing key %q, but can't find it", name)
		}

		// TODO: random tmpfile name.
		tmpf := ".ssh/authorized_keys.tmp"
		ak := ".ssh/authorized_keys"
		cmds = append(cmds,
			fmt.Sprintf(`cp %s %s`, ak, tmpf),
			fmt.Sprintf(`chmod --reference=%s %s`, ak, tmpf),
			fmt.Sprintf(`echo "%s %s %s" >> %s`, k.algorithm, k.key, name, tmpf),
			fmt.Sprintf(`mv %s %s`, tmpf, ak),
		)
	}
	cmd := exec.CommandContext(ctx, "ssh", acct.account, strings.Join(cmds, " && "))
	if err := cmd.Run(); err != nil {
		return err
	}
	return nil
}

func check(ctx context.Context, keys []key, kg map[string]*KeyGroup, accounts []account) error {
	log.Infof("Checking accounts…")
	re, err := regexp.Compile(*matching)
	if err != nil {
		return err
	}
	for _, account := range accounts {
		if re.FindString(account.account) == "" {
			continue
		}

		log.Infof("Checking %q", account.account)
		ctx2, cancel := context.WithTimeout(ctx, 60*time.Second)
		defer cancel()
		extra, missing, err := checkAccount(ctx2, kg, account)
		if err != nil {
			log.Errorf("Failed to check account %q: %v", account.account, err)
			continue
		}

		if len(extra) > 0 {
			var es []string
			for _, e := range extra {
				es = append(es, e.description)
			}
			log.Warningf("Extra keys: %q", es)
		}
		if len(missing) > 0 {
			log.Warningf("Missing keys: %q", missing)
		}
		if *doAddMissing {
			if err := addMissing(ctx, keys, account, missing); err != nil {
				log.Errorf("Failed to add missing keys %q to %q: %v", missing, account, err)
			}
		}
		if *doDeleteExtra {
			if err := deleteExtra(ctx, account, extra); err != nil {
				log.Errorf("Failed to delete extra keys %q from %q: %v", extra, account, err)
			}
		}
	}
	return nil
}

type KeyGroup struct {
	keys []key
}

func NewKeyGroup() *KeyGroup {
	return &KeyGroup{}
}

func (k *KeyGroup) Has(key string) bool {
	for _, k2 := range k.keys {
		if key == k2.key {
			return true
		}
	}
	return false
}

func (k *KeyGroup) Add(ky ...key) {
	k.keys = append(k.keys, ky...)
}

func (k *KeyGroup) Keys() []key {
	return k.keys
}

func main() {
	flag.Parse()
	if flag.NArg() != 1 {
		log.Fatalf("Command not specified. Supported commands: check, expand")
	}
	command := flag.Arg(0)

	// Read keys.
	keys, err := readKeys(*keyFile)
	if err != nil {
		log.Fatalf("Failed to read keyfile %q: %v", *keyFile, err)
	}
	log.Infof("Read %d keys", len(keys))

	// Make keygroups.
	keyGroups := map[string]*KeyGroup{}
	for _, k := range keys {
		kg := NewKeyGroup()
		kg.Add(k)
		keyGroups[k.description] = kg
	}

	groups, err := readGroups(*groupsFile)
	if err != nil {
		log.Fatalf("Failed to read groups file %q: %v", *groupsFile, err)
	}
	log.Infof("Read %d groups", len(groups))
	for g, ks := range groups {
		kg := NewKeyGroup()
		for _, kn := range ks {
			o, found := keyGroups[kn]
			if !found {
				log.Fatalf("Group %q refers to non-existing key %q", g, kn)
			}
			kg.Add(o.Keys()...)
		}
		keyGroups["@"+g] = kg
	}

	// Read accounts.
	accounts, err := readAccounts(*accountsFile)
	if err != nil {
		log.Fatalf("Failed to read accounts file %q: %v", *accountsFile, err)
	}
	log.Infof("Read %d accounts", len(accounts))

	// Verify that all accounts' keys exist in database.
	for _, a := range accounts {
		for _, kn := range a.keyGroupNames {
			if _, found := keyGroups[kn]; !found {
				log.Fatalf("Unknown key %q for account %q", kn, a.account)
			}
		}
	}
	sort.Slice(accounts, func(i, j int) bool {
		return accounts[i].account < accounts[j].account
	})

	if command == "check" {
		ctx := context.Background()
		if err := check(ctx, keys, keyGroups, accounts); err != nil {
			log.Fatalf("blah: %v", err)
		}
		log.Infof("Done")
	} else if command == "expand" {
		re, err := regexp.Compile(*matching)
		if err != nil {
			log.Printf("Invalid regex %q: %v", *matching, err)
		}
		for _, a := range accounts {
			if re.FindString(a.account) == "" {
				continue
			}
			var keys []string
			seen := make(map[string]bool)
			for _, kgn := range a.keyGroupNames {
				for _, k := range keyGroups[kgn].Keys() {
					if !seen[k.description] {
						keys = append(keys, k.description)
						seen[k.description] = true
					}
				}
			}
			sort.Strings(keys)
			fmt.Printf("%s\n  %s\n", a.account, strings.Join(keys, "\n  "))
		}
	} else {
		log.Fatalf("Invalid command %q", command)
	}
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////
// Package: easygen
// Purpose: Easy to use universal code/text generator
// Authors: Tong Sun (c) 2015-17, All rights reserved
////////////////////////////////////////////////////////////////////////////

/*

Package easygen is an easy to use universal code/text generator library.

It can be used as a text or html generator for arbitrary purposes with arbitrary data and templates.

It can be used as a code generator, or anything that is structurally repetitive. Some command line parameter handling code generator are provided as examples, including the Go's built-in flag package, and the viper & cobra package.

Many examples have been provided to showcase its functionality, and different ways to use it.

*/
package easygen

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	tt "text/template"

	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

// Opts holds the actual values from the command line parameters
var Opts = Options{ExtYaml: ".yaml", ExtTmpl: ".tmpl"}

////////////////////////////////////////////////////////////////////////////
// Function definitions

// Generate2 will produce output according to the given template and driving data files,
// specified via fileNameTempl and fileNames (for data) respectively.
func Generate2(HTML bool, fileNameTempl string, fileNames ...string) (ret string) {
	Opts.TemplateFile = fileNameTempl

	for _, fileName := range fileNames {
		ret += Generate(HTML, fileName)
	}
	Opts.TemplateFile = ""
	return ret
}

// Generate0 will produce output according from driving data without a template file, using the string from strTempl as the template
func Generate0(HTML bool, strTempl string, fileName string) string {
	Opts.TemplateStr = strTempl
	ret := Generate(HTML, fileName)
	Opts.TemplateStr = ""
	return ret
}

// Generate will produce output from the template according to the corresponding driving data, fileName is for both template and data file name
func Generate(HTML bool, fileName string) string {
	var templates []string

	// Allow to use fileName with and without the @Opts.ExtYaml suffix.
	// If both <file>.yaml and <file> exist, prefer <file>.yaml.
	if _, err := os.Stat(fileName + Opts.ExtYaml); err == nil {
		fileName += Opts.ExtYaml
	} else if _, err = os.Stat(fileName); os.IsNotExist(err) {
		fileName += Opts.ExtYaml
	}

	// Allow to use @Opts.TemplateFile without the @Opts.ExtTmpl suffix.
	if len(Opts.TemplateFile) > 0 {
		for _, template := range strings.Split(Opts.TemplateFile, ",") {
			if _, err := os.Stat(template); os.IsNotExist(err) {
				template += Opts.ExtTmpl
			}
			templates = append(templates, template)
		}
	} else if idx := strings.LastIndex(fileName, "."); idx > 0 {
		templates = []string{fileName[:idx] + Opts.ExtTmpl}
	} else {
		templates = []string{fileName + Opts.ExtTmpl}
	}

	source, err := ioutil.ReadFile(fileName)
	checkError(err)

	m := make(map[interface{}]interface{})

	err = yaml.Unmarshal(source, &m)
	checkError(err)

	env := make(map[string]string)
	for _, e := range os.Environ() {
		sep := strings.Index(e, "=")
		env[e[0:sep]] = e[sep+1:]
	}
	m["ENV"] = env
	//fmt.Printf("] %+v\n", m)

	buf := new(bytes.Buffer)
	for _, template := range templates {
		t, err := ParseFiles(HTML, template)
		checkError(err)

		err = t.Execute(buf, m)
		checkError(err)
	}
	return buf.String()
}

// ParseFiles wraps parsing text or HTML template files into a single
// function, dictated by the first parameter "HTML".
// By Matt Harden @gmail.com
func ParseFiles(HTML bool, filenames ...string) (Template, error) {
	var tname string

	if len(Opts.TemplateStr) > 0 {
		tname = "TT"
	} else if len(filenames) == 0 {
		return nil, fmt.Errorf("ParseFiles called without template filename")
	} else {
		tname = filepath.Base(filenames[0])
	}

	textTemplate := tt.New(tname).Funcs(tt.FuncMap{
		"eqf":      strings.EqualFold,
		"split":    strings.Fields,
		"minus1":   minus1,
		"dateI":    dateI,
		"year4":    year4,
		"replace":  replace,
		"replacec": replacec,
	})

	if len(Opts.TemplateStr) > 0 {
		return textTemplate.Parse(Opts.TemplateStr)
	}
	return textTemplate.ParseFiles(filenames...)
}

////////////////////////////////////////////////////////////////////////////
// Version 2 Function definitions

// Process will process the standard easygen input: the `fileName` is for both template and data file names, and produce output from the template according to the corresponding driving data.
func Process(t Template, wr io.Writer, fileName string) error {
	return Process2(t, wr, fileName, fileName)
}

// Process2 will process the case that both template and data file names are given, and produce output according to the given template and driving data files,
// specified via fileNameTempl and fileNames (for data) respectively.
func Process2(t Template, wr io.Writer, fileNameTempl string, fileNames ...string) error {
	fileName := fileNames[0]

	source, err := ioutil.ReadFile(fileName + Opts.ExtYaml)
	checkError(err)

	m := make(map[interface{}]interface{})

	err = yaml.Unmarshal(source, &m)
	checkError(err)

	env := make(map[string]string)
	for _, e := range os.Environ() {
		sep := strings.Index(e, "=")
		env[e[0:sep]] = e[sep+1:]
	}
	m["ENV"] = env
	//fmt.Printf("] %+v\n", m)

	// template file name
	fileNameT := fileNameTempl
	if len(Opts.TemplateFile) > 0 {
		fileNameT = Opts.TemplateFile
	}
	fileNameT = fileNameT + Opts.ExtTmpl
	tn, err := t.ParseFiles(fileNameT)
	checkError(err)

	return tn.ExecuteTemplate(wr, filepath.Base(fileNameT), m)
}

// Process0 will produce output according to the driving data *without* a template file, using the string from strTempl as the template
func Process0(t Template, wr io.Writer, strTempl string, fileNames ...string) error {
	return nil
}

////////////////////////////////////////////////////////////////////////////
// Support Function definitions

// Exit if error occurs
func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "[%s] Fatal error - %s\n", progname, err)
		os.Exit(1)
	}
}
<commit_msg>- [*] refactor easygen.go, #1<commit_after>////////////////////////////////////////////////////////////////////////////
// Package: easygen
// Purpose: Easy to use universal code/text generator
// Authors: Tong Sun (c) 2015-17, All rights reserved
////////////////////////////////////////////////////////////////////////////

/*

Package easygen is an easy to use universal code/text generator library.

It can be used as a text or html generator for arbitrary purposes with arbitrary data and templates.

It can be used as a code generator, or anything that is structurally repetitive. Some command line parameter handling code generator are provided as examples, including the Go's built-in flag package, and the viper & cobra package.

Many examples have been provided to showcase its functionality, and different ways to use it.

*/
package easygen

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	tt "text/template"

	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

// EgData -- EasyGen driven Data
type EgData map[interface{}]interface{}

// Opts holds the actual values from the command line parameters
var Opts = Options{ExtYaml: ".yaml", ExtTmpl: ".tmpl"}

////////////////////////////////////////////////////////////////////////////
// Function definitions

// Generate2 will produce output according to the given template and driving data files,
// specified via fileNameTempl and fileNames (for data) respectively.
func Generate2(HTML bool, fileNameTempl string, fileNames ...string) (ret string) {
	Opts.TemplateFile = fileNameTempl

	for _, fileName := range fileNames {
		ret += Generate(HTML, fileName)
	}
	Opts.TemplateFile = ""
	return ret
}

// Generate0 will produce output according from driving data without a template file, using the string from strTempl as the template
func Generate0(HTML bool, strTempl string, fileName string) string {
	Opts.TemplateStr = strTempl
	ret := Generate(HTML, fileName)
	Opts.TemplateStr = ""
	return ret
}

// Generate will produce output from the template according to the corresponding driving data, fileName is for both template and data file name
func Generate(HTML bool, fileName string) string {
	var templates []string

	// Allow to use fileName with and without the @Opts.ExtYaml suffix.
	// If both <file>.yaml and <file> exist, prefer <file>.yaml.
	if _, err := os.Stat(fileName + Opts.ExtYaml); err == nil {
		fileName += Opts.ExtYaml
	} else if _, err = os.Stat(fileName); os.IsNotExist(err) {
		fileName += Opts.ExtYaml
	}

	// Allow to use @Opts.TemplateFile without the @Opts.ExtTmpl suffix.
	if len(Opts.TemplateFile) > 0 {
		for _, template := range strings.Split(Opts.TemplateFile, ",") {
			if _, err := os.Stat(template); os.IsNotExist(err) {
				template += Opts.ExtTmpl
			}
			templates = append(templates, template)
		}
	} else if idx := strings.LastIndex(fileName, "."); idx > 0 {
		templates = []string{fileName[:idx] + Opts.ExtTmpl}
	} else {
		templates = []string{fileName + Opts.ExtTmpl}
	}

	m := ReadYamlFile(fileName)

	env := make(map[string]string)
	for _, e := range os.Environ() {
		sep := strings.Index(e, "=")
		env[e[0:sep]] = e[sep+1:]
	}
	m["ENV"] = env
	//fmt.Printf("] %+v\n", m)

	buf := new(bytes.Buffer)
	for _, template := range templates {
		t, err := ParseFiles(HTML, template)
		checkError(err)

		err = t.Execute(buf, m)
		checkError(err)
	}
	return buf.String()
}

// ParseFiles wraps parsing text or HTML template files into a single
// function, dictated by the first parameter "HTML".
// By Matt Harden @gmail.com
func ParseFiles(HTML bool, filenames ...string) (Template, error) {
	var tname string

	if len(Opts.TemplateStr) > 0 {
		tname = "TT"
	} else if len(filenames) == 0 {
		return nil, fmt.Errorf("ParseFiles called without template filename")
	} else {
		tname = filepath.Base(filenames[0])
	}

	textTemplate := tt.New(tname).Funcs(tt.FuncMap{
		"eqf":      strings.EqualFold,
		"split":    strings.Fields,
		"minus1":   minus1,
		"dateI":    dateI,
		"year4":    year4,
		"replace":  replace,
		"replacec": replacec,
	})

	if len(Opts.TemplateStr) > 0 {
		return textTemplate.Parse(Opts.TemplateStr)
	}
	return textTemplate.ParseFiles(filenames...)
}

////////////////////////////////////////////////////////////////////////////
// Version 2 Function definitions

func ReadYamlFile(fileName string) EgData {
	source, err := ioutil.ReadFile(fileName)
	checkError(err)

	m := make(EgData)

	err = yaml.Unmarshal(source, &m)
	checkError(err)

	return m
}

func GetEnv() map[string]string {
	env := make(map[string]string)
	for _, e := range os.Environ() {
		sep := strings.Index(e, "=")
		env[e[0:sep]] = e[sep+1:]
	}
	return env
}

// Process will process the standard easygen input: the `fileName` is for both template and data file names, and produce output from the template according to the corresponding driving data.
func Process(t Template, wr io.Writer, fileName string) error {
	return Process2(t, wr, fileName, fileName)
}

// Process2 will process the case that both template and data file names are given, and produce output according to the given template and driving data files,
// specified via fileNameTempl and fileNames respectively.
// fileNameTempl can be a comma-separated string giving many template files
func Process2(t Template, wr io.Writer, fileNameTempl string, fileNames ...string) error {
}

// Process1 will process the case that both template and data file names are given, and produce output according to the given template and driving data files,
// specified via fileNameTempl and fileName respectively.
// fileNameTempl is not a comma-separated string, but for a single template file.
func Process1(t Template, wr io.Writer, fileNameTempl string, fileName string) error {
	m := ReadYamlFile(fileName + Opts.ExtYaml)
	m["ENV"] = GetEnv()
	//fmt.Printf("] %+v\n", m)

	// template file name
	fileNameT := fileNameTempl
	if len(Opts.TemplateFile) > 0 {
		fileNameT = Opts.TemplateFile
	}
	fileNameT = fileNameT + Opts.ExtTmpl
	tn, err := t.ParseFiles(fileNameT)
	checkError(err)

	return tn.ExecuteTemplate(wr, filepath.Base(fileNameT), m)
}

// Process0 will produce output according to the driving data *without* a template file, using the string from strTempl as the template
func Process0(t Template, wr io.Writer, strTempl string, fileNames ...string) error {
	return nil
}

////////////////////////////////////////////////////////////////////////////
// Support Function definitions

// Exit if error occurs
func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "[%s] Fatal error - %s\n", progname, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package models_test

import (
	"sort"
	"testing"
	"time"

	"github.com/cloudflare/unsee/models"
)

type alertListSortTest struct {
	alert    models.Alert
	position int
}

var alertListSortTests = []alertListSortTest{
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 5, time.UTC)},
		position: 0,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 1, time.UTC)},
		position: 1,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 0, time.UTC)},
		position: 2,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2015, time.March, 10, 0, 0, 0, 0, time.UTC)},
		position: 6,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2016, time.December, 10, 0, 0, 0, 0, time.UTC)},
		position: 4,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 0, time.UTC)},
		position: 3,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2015, time.March, 10, 0, 0, 0, 0, time.UTC)},
		position: 5,
	},
}

func TestUnseeAlertListSort(t *testing.T) {
	al := models.AlertList{}
	for _, testCase := range alertListSortTests {
		testCase.alert.UpdateFingerprints()
		al = append(al, testCase.alert)
	}

	// repeat sort 100 times to ensure we're always sorting same way
	iterations := 100
	failures := 0
	for i := 1; i <= iterations; i++ {
		sort.Sort(al)
		for _, testCase := range alertListSortTests {
			testCase.alert.UpdateFingerprints()
			if al[testCase.position].ContentFingerprint() != testCase.alert.ContentFingerprint() {
				failures++
			}
		}
	}
	if failures > 0 {
		t.Errorf("%d sort failures for %d checks", failures, iterations*len(al))
	}
}

type agFPTest struct {
	ag          models.AlertGroup
	fingerprint string
}

var agFPTests = []agFPTest{
	// empty group fingerprint
	agFPTest{
		ag:          models.AlertGroup{},
		fingerprint: "da39a3ee5e6b4b0d3255bfef95601890afd80709",
	},
	// different Receiver shouldn't change content fingerprint
	agFPTest{
		ag: models.AlertGroup{
			Receiver: "default",
		},
		fingerprint: "da39a3ee5e6b4b0d3255bfef95601890afd80709",
	},
	// different StateCount shouldn't change content fingerprint
	agFPTest{
		ag: models.AlertGroup{
			Receiver:   "default",
			StateCount: map[string]int{"default": 0},
		},
		fingerprint: "da39a3ee5e6b4b0d3255bfef95601890afd80709",
	},
	// different Labels shouldn't change content fingerprint
	agFPTest{
		ag: models.AlertGroup{
			Receiver:   "default",
			Labels:     map[string]string{"foo": "bar"},
			StateCount: map[string]int{"default": 0},
		},
		fingerprint: "da39a3ee5e6b4b0d3255bfef95601890afd80709",
	},
	// different set of alerts should change content fingerprint
	agFPTest{
		ag: models.AlertGroup{
			Receiver: "default",
			Labels:   map[string]string{"foo": "bar"},
			Alerts: models.AlertList{
				models.Alert{
					Labels: map[string]string{"foo1": "bar"},
					State:  models.AlertStateActive,
				},
			},
			StateCount: map[string]int{"default": 0},
		},
		fingerprint: "87e79b3d507d26f21bc5e82b8db19f87ce46c1ad",
	},
}

func TestAlertGroupContentFingerprint(t *testing.T) {
	for _, testCase := range agFPTests {
		alerts := models.AlertList{}
		for _, alert := range testCase.ag.Alerts {
			alert.UpdateFingerprints()
			alerts = append(alerts, alert)
		}
		sort.Sort(alerts)
		testCase.ag.Alerts = alerts
		if testCase.ag.ContentFingerprint() != testCase.fingerprint {
			t.Errorf("Invalid AlertGroup ContentFingerprint(), expected '%s', got '%s', AlertGroup: %v",
				testCase.fingerprint, testCase.ag.ContentFingerprint(), testCase.ag)
		}
	}
}
<commit_msg>Refactor group fingerpint tests<commit_after>package models_test

import (
	"sort"
	"testing"
	"time"

	"github.com/cloudflare/unsee/models"
)

type alertListSortTest struct {
	alert    models.Alert
	position int
}

var alertListSortTests = []alertListSortTest{
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 5, time.UTC)},
		position: 0,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 1, time.UTC)},
		position: 1,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 0, time.UTC)},
		position: 2,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2015, time.March, 10, 0, 0, 0, 0, time.UTC)},
		position: 6,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2016, time.December, 10, 0, 0, 0, 0, time.UTC)},
		position: 4,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2017, time.January, 10, 0, 0, 0, 0, time.UTC)},
		position: 3,
	},
	alertListSortTest{
		alert:    models.Alert{StartsAt: time.Date(2015, time.March, 10, 0, 0, 0, 0, time.UTC)},
		position: 5,
	},
}

func TestUnseeAlertListSort(t *testing.T) {
	al := models.AlertList{}
	for _, testCase := range alertListSortTests {
		testCase.alert.UpdateFingerprints()
		al = append(al, testCase.alert)
	}

	// repeat sort 100 times to ensure we're always sorting same way
	iterations := 100
	failures := 0
	for i := 1; i <= iterations; i++ {
		sort.Sort(al)
		for _, testCase := range alertListSortTests {
			testCase.alert.UpdateFingerprints()
			if al[testCase.position].ContentFingerprint() != testCase.alert.ContentFingerprint() {
				failures++
			}
		}
	}
	if failures > 0 {
		t.Errorf("%d sort failures for %d checks", failures, iterations*len(al))
	}
}

type agFPTest struct {
	name     string            // name of the test
	ag       models.AlertGroup // alert group data
	fpChange bool              // true if fingerprint should change vs previous run
}

var agFPTests = []agFPTest{
	agFPTest{
		name: "empty group fingerprint",
		ag:   models.AlertGroup{},
	},
	agFPTest{
		name: "different Receiver shouldn't change content fingerprint",
		ag: models.AlertGroup{
			Receiver: "default",
		},
		fpChange: false,
	},
	agFPTest{
		name: "different StateCount shouldn't change content fingerprint",
		ag: models.AlertGroup{
			Receiver:   "default",
			StateCount: map[string]int{"default": 0},
		},
		fpChange: false,
	},
	agFPTest{
		name: "different Labels shouldn't change content fingerprint",
		ag: models.AlertGroup{
			Receiver:   "default",
			Labels:     map[string]string{"foo": "bar"},
			StateCount: map[string]int{"default": 0},
		},
		fpChange: false,
	},
	agFPTest{
		name: "different set of alerts should change content fingerprint",
		ag: models.AlertGroup{
			Receiver: "default",
			Labels:   map[string]string{"foo": "bar"},
			Alerts: models.AlertList{
				models.Alert{
					Labels: map[string]string{"foo1": "bar"},
					State:  models.AlertStateActive,
				},
			},
			StateCount: map[string]int{"default": 0},
		},
		fpChange: true,
	},
	agFPTest{
		name: "another different set of alerts should change content fingerprint",
		ag: models.AlertGroup{
			Receiver: "default",
			Labels:   map[string]string{"bar": "foo"},
			Alerts: models.AlertList{
				models.Alert{
					Labels: map[string]string{"bar": "foo"},
					State:  models.AlertStateActive,
				},
			},
			StateCount: map[string]int{"default": 0},
		},
		fpChange: true,
	},
	agFPTest{
		name: "repeating last set of alerts shouldn't change content fingerprint",
		ag: models.AlertGroup{
			Receiver: "default",
			Labels:   map[string]string{"bar": "foo"},
			Alerts: models.AlertList{
				models.Alert{
					Labels: map[string]string{"bar": "foo"},
					State:  models.AlertStateActive,
				},
			},
			StateCount: map[string]int{"default": 0},
		},
		fpChange: false,
	},
}

func TestAlertGroupContentFingerprint(t *testing.T) {
	fps := []string{}
	for i, testCase := range agFPTests {
		t.Run(testCase.name, func(t *testing.T) {
			alerts := models.AlertList{}
			for _, alert := range testCase.ag.Alerts {
				alert.UpdateFingerprints()
				alerts = append(alerts, alert)
			}
			sort.Sort(alerts)
			testCase.ag.Alerts = alerts
			// get alert group fingerprint
			fp := testCase.ag.ContentFingerprint()
			// add it to the list
			fps = append(fps, fp)
			// skip first test case since there's nothing to compare it with
			if i > 0 {
				// check if current test case generated fingerprint that is different
				// from the previous one
				fpChange := fp != fps[i-1]
				// check if we expected a change or not
				if fpChange != testCase.fpChange {
					t.Errorf("Fingerprint changed=%t while expected change=%t, alertgroup: %v",
						fpChange, testCase.fpChange, testCase.ag)
				}
			}
		})
	}
}
<|endoftext|>