<commit_before>// Copyright 2018 the LinuxBoot 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 visitors

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strconv"
	"unsafe"

	"github.com/linuxboot/fiano/pkg/guid"
	"github.com/linuxboot/fiano/pkg/uefi"
)

// CreateFV creates a firmware volume at given offset
type CreateFV struct {
	AbsOffset uint64
	Size      uint64
	Name      guid.GUID

	found bool
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *CreateFV) Run(f uefi.Firmware) error {
	err := f.Apply(v)
	if err != nil {
		return err
	}
	if !v.found {
		return fmt.Errorf("cannot create FV at %#x (+%#x) no BIOS region found", v.AbsOffset, v.Size)
	}
	return nil
}

// Visit applies the CreateFV visitor to any Firmware type.
func (v *CreateFV) Visit(f uefi.Firmware) error {
	if v.found {
		return nil
	}
	var offset, end uint64
	switch f := f.(type) {
	case *uefi.BIOSRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
			end = uint64(f.FRegion.EndOffset())
		} else {
			end = f.Length
		}
		if v.AbsOffset < offset {
			return fmt.Errorf("cannot create FV at %#x, BIOS region starts at %#x", v.AbsOffset, offset)
		}
		if v.AbsOffset+v.Size >= end {
			return fmt.Errorf("cannot create FV ending at %#x (%#x + %#x), BIOS region ends at %#x", v.AbsOffset+v.Size, v.AbsOffset, v.Size, end)
		}

		// Manually visit children, we want the index
		idx := -1
		var bp *uefi.BIOSPadding
	L:
		for i, e := range f.Elements {
			switch f := e.Value.(type) {
			case *uefi.BIOSPadding:
				bp = f
				bpOffset := offset + f.Offset
				bpEnd := bpOffset + uint64(len(f.Buf()))
				// Warning: This won't work if for whatever reason we have two biospads that are contiguous
				// TODO: join contiguous before !
				if v.AbsOffset >= bpOffset && v.AbsOffset+v.Size <= bpEnd {

					idx = i
					break L
				}
			}
		}

		if idx == -1 {
			return fmt.Errorf("cannot create FV at %#x (+%#x) no matching BIOS Pad found", v.AbsOffset, v.Size)
		}
		v.found = true
		fv, err := createEmptyFirmwareVolume(v.AbsOffset-offset, v.Size, &v.Name)
		if err != nil {
			return err
		}
		return insertFVinBP(f, v.AbsOffset-offset, bp, idx, fv)
	}

	return f.ApplyChildren(v)
}

func insertFVinBP(br *uefi.BIOSRegion, offset uint64, bp *uefi.BIOSPadding, idx int, fv *uefi.FirmwareVolume) error {
	// Copy the Elements before the modified BIOS Pad
	newElements := make([]*uefi.TypedFirmware, idx, len(br.Elements))
	copy(newElements, br.Elements[:idx])

	bpBuf := bp.Buf()
	// keep part of the BIOS Pad before the new fv if needed
	if bp.Offset < offset {
		hbp, err := uefi.NewBIOSPadding(bpBuf[:offset-bp.Offset], bp.Offset)
		if err != nil {
			return err
		}
		newElements = append(newElements, uefi.MakeTyped(hbp))
	}

	// Add the FV
	newElements = append(newElements, uefi.MakeTyped(fv))

	// keep part of the BIOS Pad after the new fv if needed
	tbpStart := offset - bp.Offset + fv.Length
	if tbpStart < uint64(len(bpBuf)) {
		tbp, err := uefi.NewBIOSPadding(bpBuf[tbpStart:], offset+fv.Length)
		if err != nil {
			return err
		}
		newElements = append(newElements, uefi.MakeTyped(tbp))
	}

	// Keep the remaining Elements in the BIOS Region
	if idx+1 < len(br.Elements) {
		newElements = append(newElements, br.Elements[idx+1:]...)
	}
	br.Elements = newElements
	return nil
}

func createEmptyFirmwareVolume(fvOffset, size uint64, name *guid.GUID) (*uefi.FirmwareVolume, error) {
	// TODO: can this be refactored with the code in repack.go and assemble.go ?
	fv := &uefi.FirmwareVolume{} // new Firmware Volume
	// Set up volume header first.
	fv.FileSystemGUID = *uefi.FFS2
	fv.Signature = binary.LittleEndian.Uint32([]byte("_FVH"))
	// TODO: retrieve all details from (all) other fv in BIOS Region
	fv.Attributes = 0x0004FEFF
	fv.Revision = 2
	// Create Blocks
	fv.Blocks = make([]uefi.Block, 2)
	fv.Blocks[0] = uefi.Block{Size: 4096, Count: uint32(size / 4096)}
	fv.Blocks[1] = uefi.Block{}
	// Calculate the HeaderLen field
	fv.HeaderLen = uint16(uefi.FirmwareVolumeFixedHeaderSize + int(unsafe.Sizeof(uefi.Block{}))*len(fv.Blocks))

	fv.DataOffset = uint64(fv.HeaderLen) // unless we add the extended header
	fv.Length = size

	if name != nil {
		// TODO: compute the extended header offset (especially if more than 2 Blocks, this offset will be wrong)
		fv.ExtHeaderOffset = 0x60
		fv.FVName = *name
		fv.ExtHeaderSize = uefi.FirmwareVolumeExtHeaderMinSize
	}
	// Generate binary header.
	header := new(bytes.Buffer)
	err := binary.Write(header, binary.LittleEndian, fv.FirmwareVolumeFixedHeader)
	if err != nil {
		return nil, fmt.Errorf("unable to construct binary header of new firmware volume: got %v", err)
	}
	for _, b := range fv.Blocks {
		err = binary.Write(header, binary.LittleEndian, b)
		if err != nil {
			return nil, fmt.Errorf("unable to construct binary header of new firmware volume: got %v", err)
		}

	}
	buf := header.Bytes()

	// Checksum the header
	sum, err := uefi.Checksum16(buf[:fv.HeaderLen])
	if err != nil {
		return nil, err
	}
	newSum := 0 - sum
	binary.LittleEndian.PutUint16(buf[50:], newSum)

	// Store the header buffer in
	fv.SetBuf(buf)

	if name != nil {
		// Build the ExtHeader
		extHeader := new(bytes.Buffer)
		err = binary.Write(extHeader, binary.LittleEndian, fv.FirmwareVolumeExtHeader)
		if err != nil {
			return nil, fmt.Errorf("unable to construct binary extended header of new firmware volume: got %v", err)
		}

		// The extended header in encapsulated in a PadFile just after the header.
		// The UEFI PI Specification is not clear on that point, however the implementation in tianocore GenFv tools is clear:
		// At [1] `GenerateFvImage` gives the extended header as an argument to `AddPadFile` implemented at [2].
		// [1]: https://github.com/tianocore/edk2/blob/master/BaseTools/Source/C/GenFv/GenFvInternalLib.c#L2772
		// [2]: https://github.com/tianocore/edk2/blob/master/BaseTools/Source/C/GenFv/GenFvInternalLib.c#L563
		extHeaderFile, err := uefi.CreatePadFile(uint64(uefi.FileHeaderMinLength + fv.ExtHeaderSize))
		if err != nil {
			return nil, fmt.Errorf("building ExtHeader %v", err)
		}
		if err := extHeaderFile.ChecksumAndAssemble(extHeader.Bytes()); err != nil {
			return nil, fmt.Errorf("building ExtHeader %v", err)
		}

		// Add the extended header in a Padfile just after the header.
		extHeaderFileBuf := extHeaderFile.Buf()
		if err = fv.InsertFile(fv.DataOffset, extHeaderFileBuf); err != nil {
			return nil, fmt.Errorf("adding ExtHeader %v", err)
		}
		fv.DataOffset += uint64(len(extHeaderFileBuf))
	}
	// Add empty space
	extLen := fv.Length - fv.DataOffset
	emptyBuf := make([]byte, extLen)
	uefi.Erase(emptyBuf, uefi.Attributes.ErasePolarity)

	// Store the buffer in
	fv.SetBuf(append(fv.Buf(), emptyBuf...))

	// Make sure DataOffset is 8 byte aligned at least.
	fv.DataOffset = uefi.Align8(fv.DataOffset)

	// Internal fields
	fv.FVOffset = fvOffset
	fv.FreeSpace = fv.Length - fv.DataOffset

	return fv, nil
}

func init() {
	RegisterCLI("create-fv", "creates a FV given an offset, size and volume GUID (can only replace a BIOS Padding)", 3, func(args []string) (uefi.Visitor, error) {
		offset, err := strconv.ParseUint(args[0], 0, 64)
		if err != nil {
			return nil, err
		}
		size, err := strconv.ParseUint(args[1], 0, 64)
		if err != nil {
			return nil, err
		}
		name, err := guid.Parse(args[2])
		if err != nil {
			return nil, err
		}
		return &CreateFV{
			AbsOffset: offset,
			Size:      size,
			Name:      *name,
		}, nil
	})
}
<commit_msg>fix(utk): Correctly detect the end of the free area<commit_after>// Copyright 2018 the LinuxBoot 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 visitors

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strconv"
	"unsafe"

	"github.com/linuxboot/fiano/pkg/guid"
	"github.com/linuxboot/fiano/pkg/uefi"
)

// CreateFV creates a firmware volume at given offset
type CreateFV struct {
	AbsOffset uint64
	Size      uint64
	Name      guid.GUID

	found bool
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *CreateFV) Run(f uefi.Firmware) error {
	err := f.Apply(v)
	if err != nil {
		return err
	}
	if !v.found {
		return fmt.Errorf("cannot create FV at %#x (+%#x) no BIOS region found", v.AbsOffset, v.Size)
	}
	return nil
}

// Visit applies the CreateFV visitor to any Firmware type.
func (v *CreateFV) Visit(f uefi.Firmware) error {
	if v.found {
		return nil
	}
	var offset, end uint64
	switch f := f.(type) {
	case *uefi.BIOSRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
			end = uint64(f.FRegion.EndOffset())
		} else {
			end = f.Length
		}
		if v.AbsOffset < offset {
			return fmt.Errorf("cannot create FV at %#x, BIOS region starts at %#x", v.AbsOffset, offset)
		}
		if v.AbsOffset+v.Size > end {
			return fmt.Errorf("cannot create FV ending at %#x (%#x + %#x), BIOS region ends at %#x", v.AbsOffset+v.Size, v.AbsOffset, v.Size, end)
		}

		// Manually visit children, we want the index
		idx := -1
		var bp *uefi.BIOSPadding
	L:
		for i, e := range f.Elements {
			switch f := e.Value.(type) {
			case *uefi.BIOSPadding:
				bp = f
				bpOffset := offset + f.Offset
				bpEnd := bpOffset + uint64(len(f.Buf()))
				// Warning: This won't work if for whatever reason we have two biospads that are contiguous
				// TODO: join contiguous before !
				if v.AbsOffset >= bpOffset && v.AbsOffset+v.Size <= bpEnd {

					idx = i
					break L
				}
			}
		}

		if idx == -1 {
			return fmt.Errorf("cannot create FV at %#x (+%#x) no matching BIOS Pad found", v.AbsOffset, v.Size)
		}
		v.found = true
		fv, err := createEmptyFirmwareVolume(v.AbsOffset-offset, v.Size, &v.Name)
		if err != nil {
			return err
		}
		return insertFVinBP(f, v.AbsOffset-offset, bp, idx, fv)
	}

	return f.ApplyChildren(v)
}

func insertFVinBP(br *uefi.BIOSRegion, offset uint64, bp *uefi.BIOSPadding, idx int, fv *uefi.FirmwareVolume) error {
	// Copy the Elements before the modified BIOS Pad
	newElements := make([]*uefi.TypedFirmware, idx, len(br.Elements))
	copy(newElements, br.Elements[:idx])

	bpBuf := bp.Buf()
	// keep part of the BIOS Pad before the new fv if needed
	if bp.Offset < offset {
		hbp, err := uefi.NewBIOSPadding(bpBuf[:offset-bp.Offset], bp.Offset)
		if err != nil {
			return err
		}
		newElements = append(newElements, uefi.MakeTyped(hbp))
	}

	// Add the FV
	newElements = append(newElements, uefi.MakeTyped(fv))

	// keep part of the BIOS Pad after the new fv if needed
	tbpStart := offset - bp.Offset + fv.Length
	if tbpStart < uint64(len(bpBuf)) {
		tbp, err := uefi.NewBIOSPadding(bpBuf[tbpStart:], offset+fv.Length)
		if err != nil {
			return err
		}
		newElements = append(newElements, uefi.MakeTyped(tbp))
	}

	// Keep the remaining Elements in the BIOS Region
	if idx+1 < len(br.Elements) {
		newElements = append(newElements, br.Elements[idx+1:]...)
	}
	br.Elements = newElements
	return nil
}

func createEmptyFirmwareVolume(fvOffset, size uint64, name *guid.GUID) (*uefi.FirmwareVolume, error) {
	// TODO: can this be refactored with the code in repack.go and assemble.go ?
	fv := &uefi.FirmwareVolume{} // new Firmware Volume
	// Set up volume header first.
	fv.FileSystemGUID = *uefi.FFS2
	fv.Signature = binary.LittleEndian.Uint32([]byte("_FVH"))
	// TODO: retrieve all details from (all) other fv in BIOS Region
	fv.Attributes = 0x0004FEFF
	fv.Revision = 2
	// Create Blocks
	fv.Blocks = make([]uefi.Block, 2)
	fv.Blocks[0] = uefi.Block{Size: 4096, Count: uint32(size / 4096)}
	fv.Blocks[1] = uefi.Block{}
	// Calculate the HeaderLen field
	fv.HeaderLen = uint16(uefi.FirmwareVolumeFixedHeaderSize + int(unsafe.Sizeof(uefi.Block{}))*len(fv.Blocks))

	fv.DataOffset = uint64(fv.HeaderLen) // unless we add the extended header
	fv.Length = size

	if name != nil {
		// TODO: compute the extended header offset (especially if more than 2 Blocks, this offset will be wrong)
		fv.ExtHeaderOffset = 0x60
		fv.FVName = *name
		fv.ExtHeaderSize = uefi.FirmwareVolumeExtHeaderMinSize
	}
	// Generate binary header.
	header := new(bytes.Buffer)
	err := binary.Write(header, binary.LittleEndian, fv.FirmwareVolumeFixedHeader)
	if err != nil {
		return nil, fmt.Errorf("unable to construct binary header of new firmware volume: got %v", err)
	}
	for _, b := range fv.Blocks {
		err = binary.Write(header, binary.LittleEndian, b)
		if err != nil {
			return nil, fmt.Errorf("unable to construct binary header of new firmware volume: got %v", err)
		}

	}
	buf := header.Bytes()

	// Checksum the header
	sum, err := uefi.Checksum16(buf[:fv.HeaderLen])
	if err != nil {
		return nil, err
	}
	newSum := 0 - sum
	binary.LittleEndian.PutUint16(buf[50:], newSum)

	// Store the header buffer in
	fv.SetBuf(buf)

	if name != nil {
		// Build the ExtHeader
		extHeader := new(bytes.Buffer)
		err = binary.Write(extHeader, binary.LittleEndian, fv.FirmwareVolumeExtHeader)
		if err != nil {
			return nil, fmt.Errorf("unable to construct binary extended header of new firmware volume: got %v", err)
		}

		// The extended header in encapsulated in a PadFile just after the header.
		// The UEFI PI Specification is not clear on that point, however the implementation in tianocore GenFv tools is clear:
		// At [1] `GenerateFvImage` gives the extended header as an argument to `AddPadFile` implemented at [2].
		// [1]: https://github.com/tianocore/edk2/blob/master/BaseTools/Source/C/GenFv/GenFvInternalLib.c#L2772
		// [2]: https://github.com/tianocore/edk2/blob/master/BaseTools/Source/C/GenFv/GenFvInternalLib.c#L563
		extHeaderFile, err := uefi.CreatePadFile(uint64(uefi.FileHeaderMinLength + fv.ExtHeaderSize))
		if err != nil {
			return nil, fmt.Errorf("building ExtHeader %v", err)
		}
		if err := extHeaderFile.ChecksumAndAssemble(extHeader.Bytes()); err != nil {
			return nil, fmt.Errorf("building ExtHeader %v", err)
		}

		// Add the extended header in a Padfile just after the header.
		extHeaderFileBuf := extHeaderFile.Buf()
		if err = fv.InsertFile(fv.DataOffset, extHeaderFileBuf); err != nil {
			return nil, fmt.Errorf("adding ExtHeader %v", err)
		}
		fv.DataOffset += uint64(len(extHeaderFileBuf))
	}
	// Add empty space
	extLen := fv.Length - fv.DataOffset
	emptyBuf := make([]byte, extLen)
	uefi.Erase(emptyBuf, uefi.Attributes.ErasePolarity)

	// Store the buffer in
	fv.SetBuf(append(fv.Buf(), emptyBuf...))

	// Make sure DataOffset is 8 byte aligned at least.
	fv.DataOffset = uefi.Align8(fv.DataOffset)

	// Internal fields
	fv.FVOffset = fvOffset
	fv.FreeSpace = fv.Length - fv.DataOffset

	return fv, nil
}

func init() {
	RegisterCLI("create-fv", "creates a FV given an offset, size and volume GUID (can only replace a BIOS Padding)", 3, func(args []string) (uefi.Visitor, error) {
		offset, err := strconv.ParseUint(args[0], 0, 64)
		if err != nil {
			return nil, err
		}
		size, err := strconv.ParseUint(args[1], 0, 64)
		if err != nil {
			return nil, err
		}
		name, err := guid.Parse(args[2])
		if err != nil {
			return nil, err
		}
		return &CreateFV{
			AbsOffset: offset,
			Size:      size,
			Name:      *name,
		}, nil
	})
}
<|endoftext|>


<commit_before>package s3manager

import (
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/awslabs/aws-sdk-go/aws/awsutil"
	"github.com/awslabs/aws-sdk-go/service/s3"
)

// The default range of bytes to get at a time when using Download().
var DefaultDownloadPartSize int64 = 1024 * 1024 * 5

// The default number of goroutines to spin up when using Download().
var DefaultDownloadConcurrency = 5

// The default set of options used when opts is nil in Download().
var DefaultDownloadOptions = &DownloadOptions{
	PartSize:    DefaultDownloadPartSize,
	Concurrency: DefaultDownloadConcurrency,
}

// DownloadOptions keeps tracks of extra options to pass to an Download() call.
type DownloadOptions struct {
	// The buffer size (in bytes) to use when buffering data into chunks and
	// sending them as parts to S3. The minimum allowed part size is 5MB, and
	// if this value is set to zero, the DefaultPartSize value will be used.
	PartSize int64

	// The number of goroutines to spin up in parallel when sending parts.
	// If this is set to zero, the DefaultConcurrency value will be used.
	Concurrency int

	// An S3 client to use when performing downloads. Leave this as nil to use
	// a default client.
	S3 *s3.S3
}

func NewDownloader(opts *DownloadOptions) *Downloader {
	if opts == nil {
		opts = DefaultDownloadOptions
	}
	return &Downloader{opts: *opts}
}

type Downloader struct {
	opts DownloadOptions
}

func (d *Downloader) Download(w io.WriterAt, input *s3.GetObjectInput) error {
	impl := downloadimpl{w: w, in: input, opts: &d.opts}
	return impl.download()
}

type downloadimpl struct {
	opts *DownloadOptions
	in   *s3.GetObjectInput
	w    io.WriterAt

	wg sync.WaitGroup
	m  sync.Mutex

	pos        int64
	totalBytes int64
	err        error
}

func (d *downloadimpl) init() {
	d.totalBytes = -1

	if d.opts.Concurrency == 0 {
		d.opts.Concurrency = DefaultDownloadConcurrency
	}

	if d.opts.PartSize == 0 {
		d.opts.PartSize = DefaultDownloadPartSize
	}

	if d.opts.S3 == nil {
		d.opts.S3 = s3.New(nil)
	}
}

func (d *downloadimpl) download() error {
	d.init()

	// Spin up workers
	ch := make(chan dlchunk, d.opts.Concurrency)
	d.wg.Add(d.opts.Concurrency)
	for i := 0; i < cap(ch); i++ {
		go d.downloadPart(ch)
	}

	// Assign work
	for d.geterr() == nil {
		if d.pos != 0 {
			total := d.getTotalBytes()
			if total < 0 {
				time.Sleep(10 * time.Millisecond)
				continue
			} else if d.pos >= total {
				break
			}
		}

		ch <- dlchunk{
			dlchunkcounter: &dlchunkcounter{},
			w:              d.w,
			start:          d.pos,
			size:           d.opts.PartSize,
		}
		d.pos += d.opts.PartSize
	}

	// Wait for completion
	close(ch)
	d.wg.Wait()

	// Return error
	return d.err
}

func (d *downloadimpl) downloadPart(ch chan dlchunk) {
	defer d.wg.Done()

	for {
		chunk, ok := <-ch

		if !ok {
			break
		}

		if d.geterr() == nil {
			in := &s3.GetObjectInput{}
			awsutil.Copy(in, d.in)
			rng := fmt.Sprintf("bytes=%d-%d",
				chunk.start, chunk.start+chunk.size-1)
			in.Range = &rng

			resp, err := d.opts.S3.GetObject(in)
			if err != nil {
				d.seterr(err)
			} else {
				d.setTotalBytes(resp)

				_, err := io.Copy(chunk, resp.Body)
				resp.Body.Close()

				if err != nil {
					d.seterr(err)
				}

			}
		}
	}
}

func (d *downloadimpl) getTotalBytes() int64 {
	d.m.Lock()
	defer d.m.Unlock()

	return d.totalBytes
}

func (d *downloadimpl) setTotalBytes(resp *s3.GetObjectOutput) {
	d.m.Lock()
	defer d.m.Unlock()

	if d.totalBytes >= 0 {
		return
	}

	parts := strings.Split(*resp.ContentRange, "/")
	total, err := strconv.ParseInt(parts[len(parts)-1], 10, 64)
	if err != nil {
		d.err = err
		return
	}

	d.totalBytes = total
}

// geterr is a thread-safe getter for the error object
func (d *downloadimpl) geterr() error {
	d.m.Lock()
	defer d.m.Unlock()

	return d.err
}

// seterr is a thread-safe setter for the error object
func (d *downloadimpl) seterr(e error) {
	d.m.Lock()
	defer d.m.Unlock()

	d.err = e
}

type dlchunk struct {
	*dlchunkcounter
	w     io.WriterAt
	start int64
	size  int64
}

type dlchunkcounter struct {
	cur int64
}

func (c dlchunk) Write(p []byte) (n int, err error) {
	if c.cur >= c.size {
		return 0, io.EOF
	}

	n, err = c.w.WriteAt(p, c.start+c.cur)
	c.cur += int64(n)

	return
}
<commit_msg>Add documentation<commit_after>package s3manager

import (
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/awslabs/aws-sdk-go/aws/awsutil"
	"github.com/awslabs/aws-sdk-go/service/s3"
)

// The default range of bytes to get at a time when using Download().
var DefaultDownloadPartSize int64 = 1024 * 1024 * 5

// The default number of goroutines to spin up when using Download().
var DefaultDownloadConcurrency = 5

// The default set of options used when opts is nil in Download().
var DefaultDownloadOptions = &DownloadOptions{
	PartSize:    DefaultDownloadPartSize,
	Concurrency: DefaultDownloadConcurrency,
}

// DownloadOptions keeps tracks of extra options to pass to an Download() call.
type DownloadOptions struct {
	// The buffer size (in bytes) to use when buffering data into chunks and
	// sending them as parts to S3. The minimum allowed part size is 5MB, and
	// if this value is set to zero, the DefaultPartSize value will be used.
	PartSize int64

	// The number of goroutines to spin up in parallel when sending parts.
	// If this is set to zero, the DefaultConcurrency value will be used.
	Concurrency int

	// An S3 client to use when performing downloads. Leave this as nil to use
	// a default client.
	S3 *s3.S3
}

// NewDownloader creates a new Downloader structure that downloads an object
// from S3 in concurrent chunks. Pass in an optional DownloadOptions struct
// to customize the downloader behavior.
func NewDownloader(opts *DownloadOptions) *Downloader {
	if opts == nil {
		opts = DefaultDownloadOptions
	}
	return &Downloader{opts: opts}
}

// The Downloader structure that calls Download(). It is safe to call Download()
// on this structure for multiple objects and across concurrent goroutines.
type Downloader struct {
	opts *DownloadOptions
}

// Download downloads an object in S3 and writes the payload into w using
// concurrent GET requests.
//
// It is safe to call this method for multiple objects and across concurrent
// goroutines.
func (d *Downloader) Download(w io.WriterAt, input *s3.GetObjectInput) error {
	impl := downloader{w: w, in: input, opts: *d.opts}
	return impl.download()
}

// downloader is the implementation structure used internally by Downloader.
type downloader struct {
	opts DownloadOptions
	in   *s3.GetObjectInput
	w    io.WriterAt

	wg sync.WaitGroup
	m  sync.Mutex

	pos        int64
	totalBytes int64
	err        error
}

// init initializes the downloader with default options.
func (d *downloader) init() {
	d.totalBytes = -1

	if d.opts.Concurrency == 0 {
		d.opts.Concurrency = DefaultDownloadConcurrency
	}

	if d.opts.PartSize == 0 {
		d.opts.PartSize = DefaultDownloadPartSize
	}

	if d.opts.S3 == nil {
		d.opts.S3 = s3.New(nil)
	}
}

// download performs the implementation of the object download across ranged
// GETs.
func (d *downloader) download() error {
	d.init()

	// Spin up workers
	ch := make(chan dlchunk, d.opts.Concurrency)
	d.wg.Add(d.opts.Concurrency)
	for i := 0; i < cap(ch); i++ {
		go d.downloadPart(ch)
	}

	// Assign work
	for d.geterr() == nil {
		if d.pos != 0 {
			// This is not the first chunk, let's wait until we know the total
			// size of the payload so we can see if we have read the entire
			// object.
			total := d.getTotalBytes()

			if total < 0 {
				// Total has not yet been set, so sleep and loop around while
				// waiting for our first worker to resolve this value.
				time.Sleep(10 * time.Millisecond)
				continue
			} else if d.pos >= total {
				break // We're finished queueing chunks
			}
		}

		// Queue the next range of bytes to read.
		ch <- dlchunk{
			dlchunkcounter: &dlchunkcounter{},
			w:              d.w,
			start:          d.pos,
			size:           d.opts.PartSize,
		}
		d.pos += d.opts.PartSize
	}

	// Wait for completion
	close(ch)
	d.wg.Wait()

	// Return error
	return d.err
}

// downloadPart is an individual goroutine worker reading from the ch channel
// and performing a GetObject request on the data with a given byte range.
//
// If this is the first worker, this operation also resolves the total number
// of bytes to be read so that the worker manager knows when it is finished.
func (d *downloader) downloadPart(ch chan dlchunk) {
	defer d.wg.Done()

	for {
		chunk, ok := <-ch

		if !ok {
			break
		}

		if d.geterr() == nil {
			// Get the next byte range of data
			in := &s3.GetObjectInput{}
			awsutil.Copy(in, d.in)
			rng := fmt.Sprintf("bytes=%d-%d",
				chunk.start, chunk.start+chunk.size-1)
			in.Range = &rng

			resp, err := d.opts.S3.GetObject(in)
			if err != nil {
				d.seterr(err)
			} else {
				d.setTotalBytes(resp) // Set total if not yet set.

				_, err := io.Copy(chunk, resp.Body)
				resp.Body.Close()

				if err != nil {
					d.seterr(err)
				}

			}
		}
	}
}

// getTotalBytes is a thread-safe getter for retrieving the total byte status.
func (d *downloader) getTotalBytes() int64 {
	d.m.Lock()
	defer d.m.Unlock()

	return d.totalBytes
}

// getTotalBytes is a thread-safe setter for setting the total byte status.
func (d *downloader) setTotalBytes(resp *s3.GetObjectOutput) {
	d.m.Lock()
	defer d.m.Unlock()

	if d.totalBytes >= 0 {
		return
	}

	parts := strings.Split(*resp.ContentRange, "/")
	total, err := strconv.ParseInt(parts[len(parts)-1], 10, 64)
	if err != nil {
		d.err = err
		return
	}

	d.totalBytes = total
}

// geterr is a thread-safe getter for the error object
func (d *downloader) geterr() error {
	d.m.Lock()
	defer d.m.Unlock()

	return d.err
}

// seterr is a thread-safe setter for the error object
func (d *downloader) seterr(e error) {
	d.m.Lock()
	defer d.m.Unlock()

	d.err = e
}

// dlchunk represents a single chunk of data to write by the worker routine.
// This structure also implements an io.SectionReader style interface for
// io.WriterAt, effectively making it an io.SectionWriter (which does not
// exist).
type dlchunk struct {
	*dlchunkcounter
	w     io.WriterAt
	start int64
	size  int64
}

// dlchunkcounter keeps track of the current position the dlchunk struct is
// writing to.
type dlchunkcounter struct {
	cur int64
}

// Write wraps io.WriterAt for the dlchunk, writing from the dlchunk's start
// position to its end (or EOF).
func (c dlchunk) Write(p []byte) (n int, err error) {
	if c.cur >= c.size {
		return 0, io.EOF
	}

	n, err = c.w.WriteAt(p, c.start+c.cur)
	c.cur += int64(n)

	return
}
<|endoftext|>


<commit_before>package dt

import (
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"log"
	"time"

	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
)

type jsonState json.RawMessage

type Msg struct {
	User         *User
	Input        *Input
	LastResponse *Resp
	Route        string
}

// ResponseMsg is used to pass results from packages to Ava
type RespMsg struct {
	ResponseID uint64
	Sentence   string
}

type Resp struct {
	ID         uint64
	UserID     uint64
	InputID    uint64
	FeedbackID uint64
	Sentence   string
	Route      string
	State      map[string]interface{}
	CreatedAt  time.Time
}

type Feedback struct {
	Id        uint64
	Sentence  string
	Sentiment int
	CreatedAt time.Time
}

const (
	SentimentNegative = -1
	SentimentNeutral  = 0
	SentimentPositive = 1
)

type Input struct {
	ID                uint64
	UserID            uint64
	FlexID            string
	FlexIDType        int
	Sentence          string
	SentenceAnnotated string
	ResponseID        uint64
	StructuredInput   *StructuredInput
}

func (j *jsonState) Scan(value interface{}) error {
	if err := json.Unmarshal(value.([]byte), *j); err != nil {
		log.Println("unmarshal jsonState: ", err)
		return err
	}
	return nil
}

func (j *jsonState) Value() (driver.Value, error) {
	return j, nil
}

func NewInput(si *StructuredInput, uid uint64, fid string, fidT int) *Input {
	in := Input{
		StructuredInput: si,
		UserID:          uid,
		FlexID:          fid,
		FlexIDType:      fidT,
	}
	return &in
}

func (m *Msg) GetLastResponse(db *sqlx.DB) error {
	q := `
		SELECT stateid, route, sentence, userid
		FROM responses
		WHERE userid=$1
		ORDER BY createdat DESC`
	if m.User == nil {
		// TODO move to shared errors
		return errors.New("missing user")
	}
	row := db.QueryRowx(q, m.User.ID)
	var tmp struct {
		Route    string
		Sentence string
		StateID  sql.NullInt64
		UserID   uint64
	}
	if err := row.StructScan(&tmp); err != nil {
		log.Println("structscan row ", err)
		return err
	}
	if !tmp.StateID.Valid {
		return errors.New("invalid stateid")
	}
	var state []byte
	q = `SELECT state FROM states WHERE id=$1`
	err := db.Get(&state, q, tmp.StateID)
	if err != nil {
		return err
	}
	m.LastResponse = &Resp{
		Route:    tmp.Route,
		Sentence: tmp.Sentence,
		UserID:   tmp.UserID,
	}
	err = json.Unmarshal(state, &m.LastResponse.State)
	if err != nil {
		log.Println("unmarshaling state", err)
		return err
	}
	return nil
}

func (m *Msg) NewResponse() *Resp {
	var uid uint64
	if m.User != nil {
		uid = m.User.ID
	}
	res := &Resp{
		UserID:  uid,
		InputID: m.Input.ID,
		Route:   m.Route,
	}
	if m.LastResponse != nil {
		res.State = m.LastResponse.State
	}
	return res
}
<commit_msg>Fix nil LastResponse bug<commit_after>package dt

import (
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"log"
	"time"

	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
)

type jsonState json.RawMessage

type Msg struct {
	User         *User
	Input        *Input
	LastResponse *Resp
	Route        string
}

// ResponseMsg is used to pass results from packages to Ava
type RespMsg struct {
	ResponseID uint64
	Sentence   string
}

type Resp struct {
	ID         uint64
	UserID     uint64
	InputID    uint64
	FeedbackID uint64
	Sentence   string
	Route      string
	State      map[string]interface{}
	CreatedAt  time.Time
}

type Feedback struct {
	Id        uint64
	Sentence  string
	Sentiment int
	CreatedAt time.Time
}

const (
	SentimentNegative = -1
	SentimentNeutral  = 0
	SentimentPositive = 1
)

type Input struct {
	ID                uint64
	UserID            uint64
	FlexID            string
	FlexIDType        int
	Sentence          string
	SentenceAnnotated string
	ResponseID        uint64
	StructuredInput   *StructuredInput
}

func (j *jsonState) Scan(value interface{}) error {
	if err := json.Unmarshal(value.([]byte), *j); err != nil {
		log.Println("unmarshal jsonState: ", err)
		return err
	}
	return nil
}

func (j *jsonState) Value() (driver.Value, error) {
	return j, nil
}

func NewInput(si *StructuredInput, uid uint64, fid string, fidT int) *Input {
	in := Input{
		StructuredInput: si,
		UserID:          uid,
		FlexID:          fid,
		FlexIDType:      fidT,
	}
	return &in
}

func (m *Msg) GetLastResponse(db *sqlx.DB) error {
	if m.User == nil {
		return errors.New("missing user")
	}
	q := `SELECT stateid, route, sentence, userid
	      FROM responses
	      WHERE userid=$1
	      ORDER BY createdat DESC`
	row := db.QueryRowx(q, m.User.ID)
	var tmp struct {
		Route    string
		Sentence string
		StateID  sql.NullInt64
		UserID   uint64
	}
	err := row.StructScan(&tmp)
	if err == sql.ErrNoRows {
		m.LastResponse = &Resp{}
		return nil
	}
	if err != nil {
		log.Println("structscan row ", err)
		return err
	}
	if !tmp.StateID.Valid {
		return errors.New("invalid stateid")
	}
	var state []byte
	q = `SELECT state FROM states WHERE id=$1`
	if err = db.Get(&state, q, tmp.StateID); err != nil {
		return err
	}
	m.LastResponse = &Resp{
		Route:    tmp.Route,
		Sentence: tmp.Sentence,
		UserID:   tmp.UserID,
	}
	if err = json.Unmarshal(state, &m.LastResponse.State); err != nil {
		log.Println("unmarshaling state", err)
		return err
	}
	return nil
}

func (m *Msg) NewResponse() *Resp {
	var uid uint64
	if m.User != nil {
		uid = m.User.ID
	}
	res := &Resp{
		UserID:  uid,
		InputID: m.Input.ID,
		Route:   m.Route,
	}
	if m.LastResponse != nil {
		res.State = m.LastResponse.State
	}
	return res
}
<|endoftext|>


<commit_before>// Package coremain contains the functions for starting CoreDNS.
package coremain

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/mholt/caddy"

	"github.com/coredns/coredns/core/dnsserver"
)

func init() {
	caddy.TrapSignals()
	caddy.DefaultConfigFile = "Corefile"
	caddy.Quiet = true // don't show init stuff from caddy
	setVersion()

	flag.StringVar(&conf, "conf", "", "Corefile to load (default \""+caddy.DefaultConfigFile+"\")")
	flag.StringVar(&cpu, "cpu", "100%", "CPU cap")
	flag.BoolVar(&plugins, "plugins", false, "List installed plugins")
	flag.StringVar(&caddy.PidFile, "pidfile", "", "Path to write pid file")
	flag.BoolVar(&version, "version", false, "Show version")
	flag.BoolVar(&dnsserver.Quiet, "quiet", false, "Quiet mode (no initialization output)")
	flag.BoolVar(&logfile, "log", false, "Log to standard output")

	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(defaultLoader))

	caddy.AppName = coreName
	caddy.AppVersion = CoreVersion
}

// Run is CoreDNS's main() function.
func Run() {
	// Reset flag.CommandLine to get rid of unwanted flags for instance from glog (used in kubernetes).
	// And readd the once we want to keep.
	flag.VisitAll(func(f *flag.Flag) {
		if _, ok := flagsBlacklist[f.Name]; ok {
			return
		}
		flagsToKeep = append(flagsToKeep, f)
	})

	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	for _, f := range flagsToKeep {
		flag.Var(f.Value, f.Name, f.Usage)
	}

	flag.Parse()

	if len(flag.Args()) > 0 {
		mustLogFatal(fmt.Errorf("extra command line arguments: %s", flag.Args()))
	}

	// Set up process log before anything bad happens
	if logfile {
		log.SetOutput(os.Stdout)
	}
	log.SetFlags(log.LstdFlags)

	if version {
		showVersion()
		os.Exit(0)
	}
	if plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	if err := setCPU(cpu); err != nil {
		mustLogFatal(err)
	}

	// Get Corefile input
	corefile, err := caddy.LoadCaddyfile(serverType)
	if err != nil {
		mustLogFatal(err)
	}

	// Start your engines
	instance, err := caddy.Start(corefile)
	if err != nil {
		mustLogFatal(err)
	}

	logVersion()
	if !dnsserver.Quiet {
		showVersion()
	}

	// Execute instantiation events
	caddy.EmitEvent(caddy.InstanceStartupEvent, instance)

	// Twiddle your thumbs
	instance.Wait()
}

// mustLogFatal wraps log.Fatal() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatal(args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatal(args...)
}

// confLoader loads the Caddyfile using the -conf flag.
func confLoader(serverType string) (caddy.Input, error) {
	if conf == "" {
		return nil, nil
	}

	if conf == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin, serverType)
	}

	contents, err := ioutil.ReadFile(conf)
	if err != nil {
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       conf,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Corefile from the current working directory.
func defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// logVersion logs the version that is starting.
func logVersion() {
	log.Print("[INFO] " + versionString())
	log.Print("[INFO] " + releaseString())
}

// showVersion prints the version that is starting.
func showVersion() {
	fmt.Print(versionString())
	fmt.Print(releaseString())
	if devBuild && gitShortStat != "" {
		fmt.Printf("%s\n%s\n", gitShortStat, gitFilesModified)
	}
}

// versionString returns the CoreDNS version as a string.
func versionString() string {
	return fmt.Sprintf("%s-%s\n", caddy.AppName, caddy.AppVersion)
}

// releaseString returns the release information related to CoreDNS version:
// <OS>/<ARCH>, <go version>, <commit>
// e.g.,
// linux/amd64, go1.8.3, a6d2d7b5
func releaseString() string {
	return fmt.Sprintf("%s/%s, %s, %s\n", runtime.GOOS, runtime.GOARCH, runtime.Version(), GitCommit)
}

// setVersion figures out the version information
// based on variables set by -ldflags.
func setVersion() {
	// A development build is one that's not at a tag or has uncommitted changes
	devBuild = gitTag == "" || gitShortStat != ""

	// Only set the appVersion if -ldflags was used
	if gitNearestTag != "" || gitTag != "" {
		if devBuild && gitNearestTag != "" {
			appVersion = fmt.Sprintf("%s (+%s %s)",
				strings.TrimPrefix(gitNearestTag, "v"), GitCommit, buildDate)
		} else if gitTag != "" {
			appVersion = strings.TrimPrefix(gitTag, "v")
		}
	}
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

// Flags that control program flow or startup
var (
	conf    string
	cpu     string
	logfile bool
	version bool
	plugins bool
)

// Build information obtained with the help of -ldflags
var (
	appVersion = "(untracked dev build)" // inferred at startup
	devBuild   = true                    // inferred at startup

	buildDate        string // date -u
	gitTag           string // git describe --exact-match HEAD 2> /dev/null
	gitNearestTag    string // git describe --abbrev=0 --tags HEAD
	gitShortStat     string // git diff-index --shortstat
	gitFilesModified string // git diff-index --name-only HEAD

	// Gitcommit contains the commit where we built CoreDNS from.
	GitCommit string
)

// flagsBlacklist removes flags with these names from our flagset.
var flagsBlacklist = map[string]bool{
	"logtostderr":      true,
	"alsologtostderr":  true,
	"v":                true,
	"stderrthreshold":  true,
	"vmodule":          true,
	"log_backtrace_at": true,
	"log_dir":          true,
}

var flagsToKeep []*flag.Flag
<commit_msg>pkg/log usage in coremain as well (#1751)<commit_after>// Package coremain contains the functions for starting CoreDNS.
package coremain

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/coredns/coredns/core/dnsserver"
	clog "github.com/coredns/coredns/plugin/pkg/log"

	"github.com/mholt/caddy"
)

func init() {
	caddy.TrapSignals()
	caddy.DefaultConfigFile = "Corefile"
	caddy.Quiet = true // don't show init stuff from caddy
	setVersion()

	flag.StringVar(&conf, "conf", "", "Corefile to load (default \""+caddy.DefaultConfigFile+"\")")
	flag.StringVar(&cpu, "cpu", "100%", "CPU cap")
	flag.BoolVar(&plugins, "plugins", false, "List installed plugins")
	flag.StringVar(&caddy.PidFile, "pidfile", "", "Path to write pid file")
	flag.BoolVar(&version, "version", false, "Show version")
	flag.BoolVar(&dnsserver.Quiet, "quiet", false, "Quiet mode (no initialization output)")
	flag.BoolVar(&logfile, "log", false, "Log to standard output")

	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(defaultLoader))

	caddy.AppName = coreName
	caddy.AppVersion = CoreVersion
}

// Run is CoreDNS's main() function.
func Run() {
	// Reset flag.CommandLine to get rid of unwanted flags for instance from glog (used in kubernetes).
	// And readd the once we want to keep.
	flag.VisitAll(func(f *flag.Flag) {
		if _, ok := flagsBlacklist[f.Name]; ok {
			return
		}
		flagsToKeep = append(flagsToKeep, f)
	})

	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	for _, f := range flagsToKeep {
		flag.Var(f.Value, f.Name, f.Usage)
	}

	flag.Parse()

	if len(flag.Args()) > 0 {
		mustLogFatal(fmt.Errorf("extra command line arguments: %s", flag.Args()))
	}

	// Set up process log before anything bad happens
	if logfile {
		log.SetOutput(os.Stdout)
	}
	log.SetFlags(log.LstdFlags)

	if version {
		showVersion()
		os.Exit(0)
	}
	if plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	if err := setCPU(cpu); err != nil {
		mustLogFatal(err)
	}

	// Get Corefile input
	corefile, err := caddy.LoadCaddyfile(serverType)
	if err != nil {
		mustLogFatal(err)
	}

	// Start your engines
	instance, err := caddy.Start(corefile)
	if err != nil {
		mustLogFatal(err)
	}

	logVersion()
	if !dnsserver.Quiet {
		showVersion()
	}

	// Execute instantiation events
	caddy.EmitEvent(caddy.InstanceStartupEvent, instance)

	// Twiddle your thumbs
	instance.Wait()
}

// mustLogFatal wraps log.Fatal() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatal(args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatal(args...)
}

// confLoader loads the Caddyfile using the -conf flag.
func confLoader(serverType string) (caddy.Input, error) {
	if conf == "" {
		return nil, nil
	}

	if conf == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin, serverType)
	}

	contents, err := ioutil.ReadFile(conf)
	if err != nil {
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       conf,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Corefile from the current working directory.
func defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// logVersion logs the version that is starting.
func logVersion() {
	clog.Info(versionString())
	clog.Info(releaseString())
}

// showVersion prints the version that is starting.
func showVersion() {
	fmt.Print(versionString())
	fmt.Print(releaseString())
	if devBuild && gitShortStat != "" {
		fmt.Printf("%s\n%s\n", gitShortStat, gitFilesModified)
	}
}

// versionString returns the CoreDNS version as a string.
func versionString() string {
	return fmt.Sprintf("%s-%s\n", caddy.AppName, caddy.AppVersion)
}

// releaseString returns the release information related to CoreDNS version:
// <OS>/<ARCH>, <go version>, <commit>
// e.g.,
// linux/amd64, go1.8.3, a6d2d7b5
func releaseString() string {
	return fmt.Sprintf("%s/%s, %s, %s\n", runtime.GOOS, runtime.GOARCH, runtime.Version(), GitCommit)
}

// setVersion figures out the version information
// based on variables set by -ldflags.
func setVersion() {
	// A development build is one that's not at a tag or has uncommitted changes
	devBuild = gitTag == "" || gitShortStat != ""

	// Only set the appVersion if -ldflags was used
	if gitNearestTag != "" || gitTag != "" {
		if devBuild && gitNearestTag != "" {
			appVersion = fmt.Sprintf("%s (+%s %s)",
				strings.TrimPrefix(gitNearestTag, "v"), GitCommit, buildDate)
		} else if gitTag != "" {
			appVersion = strings.TrimPrefix(gitTag, "v")
		}
	}
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

// Flags that control program flow or startup
var (
	conf    string
	cpu     string
	logfile bool
	version bool
	plugins bool
)

// Build information obtained with the help of -ldflags
var (
	appVersion = "(untracked dev build)" // inferred at startup
	devBuild   = true                    // inferred at startup

	buildDate        string // date -u
	gitTag           string // git describe --exact-match HEAD 2> /dev/null
	gitNearestTag    string // git describe --abbrev=0 --tags HEAD
	gitShortStat     string // git diff-index --shortstat
	gitFilesModified string // git diff-index --name-only HEAD

	// Gitcommit contains the commit where we built CoreDNS from.
	GitCommit string
)

// flagsBlacklist removes flags with these names from our flagset.
var flagsBlacklist = map[string]bool{
	"logtostderr":      true,
	"alsologtostderr":  true,
	"v":                true,
	"stderrthreshold":  true,
	"vmodule":          true,
	"log_backtrace_at": true,
	"log_dir":          true,
}

var flagsToKeep []*flag.Flag
<|endoftext|>


<commit_before>package github

// TODO:
// - testing

import (
	"fmt"
	"regexp"
	"time"

	"github.com/ajm188/slack"
	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
	ghAuth "golang.org/x/oauth2/github" // have to rename so we don't have 2 "github"s
)

var (
	ClientID     string
	ClientSecret string
	AccessToken  string
	RedirectURL  string
	Scopes       []string
)

func Config() *oauth2.Config {
	return &oauth2.Config{
		ClientID:     ClientID,
		ClientSecret: ClientSecret,
		Endpoint:     ghAuth.Endpoint,
		RedirectURL:  RedirectURL,
		Scopes:       Scopes,
	}
}

func Token() *oauth2.Token {
	var noExpire time.Time // this sets noExpire to the zero Time value
	return &oauth2.Token{
		AccessToken:  AccessToken,
		TokenType:    "", // uhhh
		RefreshToken: "",
		Expiry:       noExpire,
	}
}

func OpenIssue(bot *slack.Bot) {
	repoRe := regexp.MustCompile("issue me ([^/ ]+)/([^/ ]+)")
	argsRe := regexp.MustCompile("(\".*?[^\\\\]\")")
	oauthClient := Config().Client(oauth2.NoContext, Token())
	ghClient := github.NewClient(oauthClient)
	issues := ghClient.Issues

	handler := func(b *slack.Bot, event map[string]interface{}) (*slack.Message, slack.Status) {
		text := event["text"].(string)
		owner, repo, err := extractOwnerAndRepo(text, repoRe)
		if err != nil {
			return nil, slack.Continue
		}
		issueRequest, err := extractIssueArgs(text, argsRe)
		if err != nil {
			return nil, slack.Continue
		}
		issue, _, err := issues.Create(owner, repo, issueRequest)
		if err != nil {
			return nil, slack.Continue
		}

		user := event["user"].(string)
		channel := event["channel"].(string)
		message := fmt.Sprintf(
			"I created that issue for you. You can view it here: %s",
			*issue.HTMLURL,
		)
		return bot.Mention(user, channel, message), slack.Continue
	}

	bot.RespondRegexp(repoRe, handler)
}

func extractOwnerAndRepo(text string, re *regexp.Regexp) (string, string, error) {
	m := re.FindStringSubmatch(text)
	if m == nil || len(m) < 3 {
		return "", "", &repoError{text}
	}
	return m[1], m[2], nil
}

func removeQuotes(s string) string {
	return s[1 : len(s)-1]
}

func extractIssueArgs(text string, re *regexp.Regexp) (*github.IssueRequest, error) {
	match := re.FindAllString(text, -1)
	m := make([]string, len(match))
	for i, v := range match {
		m[i] = removeQuotes(v)
	}
	if m == nil || len(m) == 0 {
		return nil, &issueError{text}
	}
	var title, body, assignee *string
	title = &m[0]
	if len(m) >= 2 {
		body = &m[1]
	}
	if len(m) >= 3 {
		assignee = &m[2]
	}
	issueState := "open"
	request := github.IssueRequest{
		Title:     title,
		Body:      body,
		Labels:    nil,
		Assignee:  assignee,
		State:     &issueState,
		Milestone: nil,
	}
	return &request, nil
}
<commit_msg>fix function call argument order<commit_after>package github

// TODO:
// - testing

import (
	"fmt"
	"regexp"
	"time"

	"github.com/ajm188/slack"
	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
	ghAuth "golang.org/x/oauth2/github" // have to rename so we don't have 2 "github"s
)

var (
	ClientID     string
	ClientSecret string
	AccessToken  string
	RedirectURL  string
	Scopes       []string
)

func Config() *oauth2.Config {
	return &oauth2.Config{
		ClientID:     ClientID,
		ClientSecret: ClientSecret,
		Endpoint:     ghAuth.Endpoint,
		RedirectURL:  RedirectURL,
		Scopes:       Scopes,
	}
}

func Token() *oauth2.Token {
	var noExpire time.Time // this sets noExpire to the zero Time value
	return &oauth2.Token{
		AccessToken:  AccessToken,
		TokenType:    "", // uhhh
		RefreshToken: "",
		Expiry:       noExpire,
	}
}

func OpenIssue(bot *slack.Bot) {
	repoRe := regexp.MustCompile("issue me ([^/ ]+)/([^/ ]+)")
	argsRe := regexp.MustCompile("(\".*?[^\\\\]\")")
	oauthClient := Config().Client(oauth2.NoContext, Token())
	ghClient := github.NewClient(oauthClient)
	issues := ghClient.Issues

	handler := func(b *slack.Bot, event map[string]interface{}) (*slack.Message, slack.Status) {
		text := event["text"].(string)
		owner, repo, err := extractOwnerAndRepo(text, repoRe)
		if err != nil {
			return nil, slack.Continue
		}
		issueRequest, err := extractIssueArgs(text, argsRe)
		if err != nil {
			return nil, slack.Continue
		}
		issue, _, err := issues.Create(owner, repo, issueRequest)
		if err != nil {
			return nil, slack.Continue
		}

		user := event["user"].(string)
		channel := event["channel"].(string)
		message := fmt.Sprintf(
			"I created that issue for you. You can view it here: %s",
			*issue.HTMLURL,
		)
		return bot.Mention(user, message, channel), slack.Continue
	}

	bot.RespondRegexp(repoRe, handler)
}

func extractOwnerAndRepo(text string, re *regexp.Regexp) (string, string, error) {
	m := re.FindStringSubmatch(text)
	if m == nil || len(m) < 3 {
		return "", "", &repoError{text}
	}
	return m[1], m[2], nil
}

func removeQuotes(s string) string {
	return s[1 : len(s)-1]
}

func extractIssueArgs(text string, re *regexp.Regexp) (*github.IssueRequest, error) {
	match := re.FindAllString(text, -1)
	m := make([]string, len(match))
	for i, v := range match {
		m[i] = removeQuotes(v)
	}
	if m == nil || len(m) == 0 {
		return nil, &issueError{text}
	}
	var title, body, assignee *string
	title = &m[0]
	if len(m) >= 2 {
		body = &m[1]
	}
	if len(m) >= 3 {
		assignee = &m[2]
	}
	issueState := "open"
	request := github.IssueRequest{
		Title:     title,
		Body:      body,
		Labels:    nil,
		Assignee:  assignee,
		State:     &issueState,
		Milestone: nil,
	}
	return &request, nil
}
<|endoftext|>


<commit_before>package golang

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/chuckpreslar/emission"
	"github.com/go-errors/errors"

	"github.com/markelog/eclectica/pkg"
	"github.com/markelog/eclectica/variables"
	"github.com/markelog/eclectica/versions"
)

var (
	VersionsLink   = "https://golang.org/dl"
	DownloadLink   = "https://storage.googleapis.com/golang"
	versionPattern = "\\d+\\.\\d+(?:\\.\\d+)?(?:(alpha|beta|rc)(?:\\d*)?)?"

	bins = []string{"go", "godoc", "gofmt"}
	dots = []string{".go-version"}
)

type Golang struct {
	Version string
	Emitter *emission.Emitter
	pkg.Base
}

func New(version string, emitter *emission.Emitter) *Golang {
	return &Golang{
		Version: version,
		Emitter: emitter,
	}
}

func (golang Golang) Events() *emission.Emitter {
	return golang.Emitter
}

func (golang Golang) PostInstall() error {
	return dealWithShell()
}

func (golang Golang) Environment() (result []string, err error) {
	result = append(result, "GOROOT="+variables.Path("go", golang.Version))

	// Go versions lower then 1.7 do not have default `GOPATH` environment variable.
	// Starting from 1.7 `GOPATH` is now set to `~/go` path (see `go help gopath` for more)
	// We do the same if for other versions as a default, but only if user didn't set themselves
	if os.Getenv("GOPATH") == "" {
		result = append(result, "GOPATH="+filepath.Join(os.Getenv("HOME"), "go"))
	}

	return
}

func (golang Golang) Info() map[string]string {
	result := make(map[string]string)

	platform, _ := getPlatform()

	version := versions.Unsemverify(golang.Version)
	result["version"] = version
	result["unarchive-filename"] = "go"
	result["filename"] = fmt.Sprintf("go%s.%s", version, platform)
	result["url"] = fmt.Sprintf("%s/%s.tar.gz", DownloadLink, result["filename"])

	return result
}

func (rust Golang) Bins() []string {
	return bins
}

func (rust Golang) Dots() []string {
	return dots
}

func (golang Golang) ListRemote() ([]string, error) {
	doc, err := goquery.NewDocument(VersionsLink)

	if err != nil {
		if _, ok := err.(net.Error); ok {
			return nil, errors.New(variables.ConnectionError)
		}

		return nil, err
	}

	result := []string{}
	rVersion := regexp.MustCompile(versionPattern)

	doc.Find("#archive tr:first-of-type td:first-of-type.filename a").Each(func(i int, node *goquery.Selection) {
		text := node.Text()

		if strings.Contains(text, "bootstrap") {
			return
		}

		// We not checking for duplicates, since it just might create more errors
		version := rVersion.FindAllStringSubmatch(text, 1)[0][0]
		result = append(result, version)
	})

	return result, nil
}

func getPlatform() (string, error) {
	if runtime.GOOS == "linux" {
		return "linux-amd64", nil
	}

	if runtime.GOOS == "darwin" {
		return "darwin-amd64", nil
	}

	return "", errors.New("Not supported envionment")
}
<commit_msg>Change selector for golang lang<commit_after>package golang

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/chuckpreslar/emission"
	"github.com/go-errors/errors"

	"github.com/markelog/eclectica/pkg"
	"github.com/markelog/eclectica/variables"
	"github.com/markelog/eclectica/versions"
)

var (
	VersionsLink   = "https://golang.org/dl"
	DownloadLink   = "https://storage.googleapis.com/golang"
	versionPattern = "\\d+\\.\\d+(?:\\.\\d+)?(?:(alpha|beta|rc)(?:\\d*)?)?"

	bins = []string{"go", "godoc", "gofmt"}
	dots = []string{".go-version"}
)

type Golang struct {
	Version string
	Emitter *emission.Emitter
	pkg.Base
}

func New(version string, emitter *emission.Emitter) *Golang {
	return &Golang{
		Version: version,
		Emitter: emitter,
	}
}

func (golang Golang) Events() *emission.Emitter {
	return golang.Emitter
}

func (golang Golang) PostInstall() error {
	return dealWithShell()
}

func (golang Golang) Environment() (result []string, err error) {
	result = append(result, "GOROOT="+variables.Path("go", golang.Version))

	// Go versions lower then 1.7 do not have default `GOPATH` environment variable.
	// Starting from 1.7 `GOPATH` is now set to `~/go` path (see `go help gopath` for more)
	// We do the same if for other versions as a default, but only if user didn't set themselves
	if os.Getenv("GOPATH") == "" {
		result = append(result, "GOPATH="+filepath.Join(os.Getenv("HOME"), "go"))
	}

	return
}

func (golang Golang) Info() map[string]string {
	result := make(map[string]string)

	platform, _ := getPlatform()

	version := versions.Unsemverify(golang.Version)
	result["version"] = version
	result["unarchive-filename"] = "go"
	result["filename"] = fmt.Sprintf("go%s.%s", version, platform)
	result["url"] = fmt.Sprintf("%s/%s.tar.gz", DownloadLink, result["filename"])

	return result
}

func (rust Golang) Bins() []string {
	return bins
}

func (rust Golang) Dots() []string {
	return dots
}

func (golang Golang) ListRemote() ([]string, error) {
	doc, err := goquery.NewDocument(VersionsLink)

	if err != nil {
		if _, ok := err.(net.Error); ok {
			return nil, errors.New(variables.ConnectionError)
		}

		return nil, err
	}

	result := []string{}
	selector := "#archive tr:first-of-type td:first-of-type.filename a"
	rVersion := regexp.MustCompile(versionPattern)

	doc.Find(selector).Each(func(i int, node *goquery.Selection) {
		text := node.Text()

		if strings.Contains(text, "bootstrap") {
			return
		}

		// We not checking for duplicates, since it just might create more errors
		version := rVersion.FindAllStringSubmatch(text, 1)[0][0]
		result = append(result, version)
	})

	return result, nil
}

func getPlatform() (string, error) {
	if runtime.GOOS == "linux" {
		return "linux-amd64", nil
	}

	if runtime.GOOS == "darwin" {
		return "darwin-amd64", nil
	}

	return "", errors.New("Not supported envionment")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/bitly/go-nsq"
	"log"
	"net/http"
	"encoding/json"
	"os"
	"bytes"
	"github.com/BurntSushi/toml"
	"gitlab.socrate.vsct.fr/dt/haaas"
)

var (
	ip = flag.String("ip", "4.3.2.1", "Node ip address")
	configFile = flag.String("config", "haaas.conf", "Configuration file")
	versionFlag = flag.Bool("version", false, "Print current version")
	config = nsq.NewConfig()
	properties haaas.Config
	daemon *haaas.Daemon
	producer *nsq.Producer

)

func main() {

	flag.Parse()

	if *versionFlag{
		println(haaas.AppVersion)
		os.Exit(0)
	}

	if _, err := toml.DecodeFile(*configFile, &properties); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	properties.IpAddr = *ip
	len := len(properties.HapHome)
	if properties.HapHome[len - 1] == '/' {
		properties.HapHome = properties.HapHome[:len - 1 ]
	}

	daemon = haaas.NewDaemon(&properties)

	log.Printf("Starting haaasd (%s) with id %v", properties.Status, properties.NodeId())

	startTryUpdateConsumer()
	startUpdateConsumer()

	producer, _ = nsq.NewProducer(properties.ProducerAddr, config)

	starRestAPI()
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func startTryUpdateConsumer() {
	tryUpdateConsumer, _ := nsq.NewConsumer(fmt.Sprintf("try_update_%s", properties.ClusterId), properties.NodeId(), config)

	tryUpdateConsumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		defer message.Finish()

		isSlave, err := daemon.IsSlave()
		if err != nil {
			return err
		}

		if isSlave {
			log.Printf("Receive try_update event %s", message.Body)
			data, err := bodyToDatas(message.Body)
			check(err)

			hap := haaas.NewHaproxy(&properties, data.Application, data.Platform)
			err = hap.ApplyConfiguration(data)
			if err == nil {
				commitTryUpdate(data)
			}else {
				log.Fatal(err)
			}
		}

		return nil
	}))

	err := tryUpdateConsumer.ConnectToNSQLookupd(properties.LookupdAddr)
	if err != nil {
		log.Panic("Could not connect")
	}
}

func commitTryUpdate(data haaas.EventMessage) {
	publishMessage("update_", data)
}

func startUpdateConsumer() {
	updateConsumer, err := nsq.NewConsumer("update_" + properties.ClusterId, properties.NodeId(), config)
	check(err)

	updateConsumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		defer message.Finish()
		isMaster, err := daemon.IsMaster()
		if err != nil {
			return err
		}

		if isMaster {
			log.Printf("Receive update event %s", message.Body)
			data, err := bodyToDatas(message.Body)
			check(err)

			hap := haaas.NewHaproxy(&properties, data.Application, data.Platform)
			err = hap.ApplyConfiguration(data)

			if err == nil {
				commitUpdate(data)
			}else {
				log.Fatal(err)
			}
		}

		return nil
	}))

	err = updateConsumer.ConnectToNSQLookupd(properties.LookupdAddr)
	if err != nil {
		log.Panic("Could not connect")
	}
}

func commitUpdate(data haaas.EventMessage) {
	publishMessage("updated_", map[string]string{"application" : data.Application, "platform": data.Platform, "correlationid" : data.Correlationid})
}

// Unmarshal json to EventMessage
func bodyToDatas(jsonStream []byte) (haaas.EventMessage, error) {
	dec := json.NewDecoder(bytes.NewReader(jsonStream))
	var message haaas.EventMessage
	dec.Decode(&message)
	return message, nil
}

// Start web server
// endpoints:
//	/real-ip : get the haaasd ip address as given by the -ip command parameter
func starRestAPI() {
	http.HandleFunc("/real-ip", func(writer http.ResponseWriter, request *http.Request) {
		log.Printf("GET /real-ip")
		fmt.Fprintf(writer, "%s\n", *ip)
	})
	log.Printf("Start listening on port %d", properties.Port)
	if err := http.ListenAndServe(fmt.Sprintf(":%d", properties.Port), nil); err != nil {
		log.Fatal("ListenAndServe:", err)
	}
}

func publishMessage(topic_prefix string, data interface{}) error {
//	producer, _ := nsq.NewProducer("floradora:50150", config)
	jsonMsg, _ := json.Marshal(data)
	topic := topic_prefix + properties.ClusterId
	log.Printf("Publish to %s : %s", topic, jsonMsg)
	return producer.Publish(topic, []byte(jsonMsg))
}

<commit_msg>Update listener<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/bitly/go-nsq"
	"log"
	"net/http"
	"encoding/json"
	"os"
	"bytes"
	"github.com/BurntSushi/toml"
	"gitlab.socrate.vsct.fr/dt/haaas"
	"net"
)

var (
	ip = flag.String("ip", "4.3.2.1", "Node ip address")
	configFile = flag.String("config", "haaas.conf", "Configuration file")
	versionFlag = flag.Bool("version", false, "Print current version")
	config = nsq.NewConfig()
	properties haaas.Config
	daemon *haaas.Daemon
	producer *nsq.Producer

)

func main() {

	flag.Parse()

	if *versionFlag{
		println(haaas.AppVersion)
		os.Exit(0)
	}

	if _, err := toml.DecodeFile(*configFile, &properties); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	properties.IpAddr = *ip
	len := len(properties.HapHome)
	if properties.HapHome[len - 1] == '/' {
		properties.HapHome = properties.HapHome[:len - 1 ]
	}

	daemon = haaas.NewDaemon(&properties)

	log.Printf("Starting haaasd (%s) with id %v", properties.Status, properties.NodeId())

	startTryUpdateConsumer()
	startUpdateConsumer()

	producer, _ = nsq.NewProducer(properties.ProducerAddr, config)

	starRestAPI()
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func startTryUpdateConsumer() {
	tryUpdateConsumer, _ := nsq.NewConsumer(fmt.Sprintf("try_update_%s", properties.ClusterId), properties.NodeId(), config)

	tryUpdateConsumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		defer message.Finish()

		isSlave, err := daemon.IsSlave()
		if err != nil {
			return err
		}

		if isSlave {
			log.Printf("Receive try_update event %s", message.Body)
			data, err := bodyToDatas(message.Body)
			check(err)

			hap := haaas.NewHaproxy(&properties, data.Application, data.Platform)
			err = hap.ApplyConfiguration(data)
			if err == nil {
				commitTryUpdate(data)
			}else {
				log.Fatal(err)
			}
		}

		return nil
	}))

	err := tryUpdateConsumer.ConnectToNSQLookupd(properties.LookupdAddr)
	if err != nil {
		log.Panic("Could not connect")
	}
}

func commitTryUpdate(data haaas.EventMessage) {
	publishMessage("update_", data)
}

func startUpdateConsumer() {
	updateConsumer, err := nsq.NewConsumer("update_" + properties.ClusterId, properties.NodeId(), config)
	check(err)

	updateConsumer.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		defer message.Finish()
		isMaster, err := daemon.IsMaster()
		if err != nil {
			return err
		}

		if isMaster {
			log.Printf("Receive update event %s", message.Body)
			data, err := bodyToDatas(message.Body)
			check(err)

			hap := haaas.NewHaproxy(&properties, data.Application, data.Platform)
			err = hap.ApplyConfiguration(data)

			if err == nil {
				commitUpdate(data)
			}else {
				log.Fatal(err)
			}
		}

		return nil
	}))

	err = updateConsumer.ConnectToNSQLookupd(properties.LookupdAddr)
	if err != nil {
		log.Panic("Could not connect")
	}
}

func commitUpdate(data haaas.EventMessage) {
	publishMessage("updated_", map[string]string{"application" : data.Application, "platform": data.Platform, "correlationid" : data.Correlationid})
}

// Unmarshal json to EventMessage
func bodyToDatas(jsonStream []byte) (haaas.EventMessage, error) {
	dec := json.NewDecoder(bytes.NewReader(jsonStream))
	var message haaas.EventMessage
	dec.Decode(&message)
	return message, nil
}

// Start web server
// endpoints:
//	/real-ip : get the haaasd ip address as given by the -ip command parameter
func starRestAPI() {
	sm := http.NewServeMux()
	sm.HandleFunc("/real-ip", func(writer http.ResponseWriter, request *http.Request) {
		log.Printf("GET /real-ip")
		fmt.Fprintf(writer, "%s\n", *ip)
	})
	log.Printf("Start listening on port %d", properties.Port)
	listener, err := net.Listen("tcp4", fmt.Sprintf(":%d", properties.Port))
	if err != nil {
		log.Fatal(err)
	}
	log.Fatal(http.Serve(listener, sm))
}

func publishMessage(topic_prefix string, data interface{}) error {
	jsonMsg, _ := json.Marshal(data)
	topic := topic_prefix + properties.ClusterId
	log.Printf("Publish to %s : %s", topic, jsonMsg)
	return producer.Publish(topic, []byte(jsonMsg))
}

<|endoftext|>


<commit_before>package handler

import (
	"net/http"

	"github.com/gilcrest/go-api-basic/app"
	"github.com/gilcrest/go-api-basic/controller"
	"github.com/rs/xid"
)

// AppHandler is the struct that serves the application
// and methods for handling all HTTP requests
type AppHandler struct {
	App *app.Application
	controller.StandardResponseFields
}

// AddStandardResponseHeaders middleware is used to add any
// standard HTTP response headers
func (ah *AppHandler) AddStandardResponseHeaders(h http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Add("Content-Type", "application/json")
			h.ServeHTTP(w, r) // call original
		})
}

// AddRequestID middleware is used to add a unique request ID to each request
func (ah *AppHandler) AddRequestID(h http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			// get byte Array representation of guid from xid package (12 bytes)
			id := xid.New()

			// Set the Request ID as xid.ID to
			// controller.RequestID using the constructor
			requestID := controller.NewRequestID(id)

			// Send the RequestID and http.Request to the
			// StandardResponseFields constructor to receive the
			// StandardResponseFields
			ah.StandardResponseFields = controller.NewStandardResponseFields(requestID, r)

			h.ServeHTTP(w, r) // call original
		})
}

// NewAppHandler initializes the AppHandler
func NewAppHandler(app *app.Application) *AppHandler {
	return &AppHandler{App: app}
}
<commit_msg>RequestID to TraceID name changes<commit_after>package handler

import (
	"net/http"

	"github.com/gilcrest/go-api-basic/app"
	"github.com/gilcrest/go-api-basic/controller"
	"github.com/rs/xid"
)

// AppHandler is the struct that serves the application
// and methods for handling all HTTP requests
type AppHandler struct {
	App *app.Application
	controller.StandardResponseFields
}

// AddStandardResponseHeaders middleware is used to add any
// standard HTTP response headers
func (ah *AppHandler) AddStandardResponseHeaders(h http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Add("Content-Type", "application/json")
			h.ServeHTTP(w, r) // call original
		})
}

// SetStandardResponseFields middleware is used to set the fields
// that are part of every response
func (ah *AppHandler) SetStandardResponseFields(h http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			// get byte Array representation of guid from xid package (12 bytes)
			id := xid.New()

			// Send a new TraceID and the http.Request to the
			// StandardResponseFields constructor to set the
			// StandardResponseFields of the AppHandler
			ah.StandardResponseFields = controller.NewStandardResponseFields(controller.NewTraceID(id), r)

			h.ServeHTTP(w, r) // call original
		})
}

// NewAppHandler initializes the AppHandler
func NewAppHandler(app *app.Application) *AppHandler {
	return &AppHandler{App: app}
}

// NewMockAppHandler initializes the AppHandler
func NewMockAppHandler(app *app.Application, r *http.Request) *AppHandler {

	appHandler := &AppHandler{App: app}
	// Send a mocked TraceID and the http.Request to the
	// StandardResponseFields constructor to set the
	// StandardResponseFields of the AppHandler
	appHandler.StandardResponseFields = controller.NewStandardResponseFields(controller.NewMockTraceID(), r)

	return appHandler
}
<|endoftext|>


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

// Package handler provides functions for responding to specific client
// requests by the ePoxy boot server.
package handler

import (
	"fmt"
	"log"
	"net/http"
	"path"

	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/epoxy/template"
)

// Config provides access to Host records.
type Config interface {
	Save(host *storage.Host) error
	Load(name string) (*storage.Host, error)
}

// Env holds data necessary for executing handler functions.
type Env struct {
	// Config provides access to Host records.
	Config Config
	// ServerAddr is the host:port of the public service. Used to generate absolute URLs.
	ServerAddr string
}

// GenerateStage1IPXE creates the stage1 iPXE script for booting machines.
// func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) (int, error) {
func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	// TODO(soltesz):
	// * Verify that the source IP maches the host IP.
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate iPXE script.
	script := template.FormatStage1IPXEScript(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "text/plain; charset=us-ascii")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// GenerateJSONConfig
func (env *Env) GenerateJSONConfig(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]
	// TODO: Verify that the sessionID matches the host.CurrentSessionIDs.Stage2ID.
	// sessionId := mux.Vars(req)["sessionID"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	// TODO(soltesz):
	// * Verify that the source IP maches the host IP.
	// * Save information sent in PostForm, e.g. ssh host key.
	stage := path.Base(req.URL.Path)

	script := template.FormatJSONConfig(host, stage)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// ReceiveReport handles the last step of a boot sequence when the epoxy client reports
// success or failure. In both cases, the session ids are invalidated. In all cases,
// epoxy_client is expected to report the server's public host key.
func (env *Env) ReceiveReport(rw http.ResponseWriter, req *http.Request) {
	// TODO: ParseForm in req.Body()
	// TODO: log or save values where appropriate.
	req.ParseForm()
	log.Println(req.PostForm)
	// TODO: invalidate session ids.
	rw.WriteHeader(http.StatusNoContent)
	return
}
<commit_msg>Update GenerateJSONConfig docstring<commit_after>// Copyright 2016 ePoxy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//////////////////////////////////////////////////////////////////////////////

// Package handler provides functions for responding to specific client
// requests by the ePoxy boot server.
package handler

import (
	"fmt"
	"log"
	"net/http"
	"path"

	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/epoxy/template"
)

// Config provides access to Host records.
type Config interface {
	Save(host *storage.Host) error
	Load(name string) (*storage.Host, error)
}

// Env holds data necessary for executing handler functions.
type Env struct {
	// Config provides access to Host records.
	Config Config
	// ServerAddr is the host:port of the public service. Used to generate absolute URLs.
	ServerAddr string
}

// GenerateStage1IPXE creates the stage1 iPXE script for booting machines.
// func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) (int, error) {
func (env *Env) GenerateStage1IPXE(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	// TODO(soltesz):
	// * Verify that the source IP maches the host IP.
	// * Save information sent in PostForm.

	// Generate new session IDs.
	host.GenerateSessionIDs()

	// Save host record to Datastore to commit session IDs.
	if err := env.Config.Save(host); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate iPXE script.
	script := template.FormatStage1IPXEScript(host, env.ServerAddr)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "text/plain; charset=us-ascii")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// GenerateJSONConfig creates and returns a JSON serialized nextboot.Config
// suitable for responding to stage2 or stage3 requests.
func (env *Env) GenerateJSONConfig(rw http.ResponseWriter, req *http.Request) {
	hostname := mux.Vars(req)["hostname"]
	// TODO: Verify that the sessionID matches the host.CurrentSessionIDs.Stage2ID.
	// sessionId := mux.Vars(req)["sessionID"]

	// Use hostname as key to load record from Datastore.
	host, err := env.Config.Load(hostname)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusNotFound)
		return
	}
	// TODO(soltesz):
	// * Verify that the source IP maches the host IP.
	// * Save information sent in PostForm, e.g. ssh host key.
	stage := path.Base(req.URL.Path)

	script := template.FormatJSONConfig(host, stage)

	// Complete request as successful.
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.WriteHeader(http.StatusOK)
	_, err = fmt.Fprintf(rw, script)
	if err != nil {
		log.Printf("Failed to write response to %q: %v", hostname, err)
	}
	return
}

// ReceiveReport handles the last step of a boot sequence when the epoxy client reports
// success or failure. In both cases, the session ids are invalidated. In all cases,
// epoxy_client is expected to report the server's public host key.
func (env *Env) ReceiveReport(rw http.ResponseWriter, req *http.Request) {
	// TODO: ParseForm in req.Body()
	// TODO: log or save values where appropriate.
	req.ParseForm()
	log.Println(req.PostForm)
	// TODO: invalidate session ids.
	rw.WriteHeader(http.StatusNoContent)
	return
}
<|endoftext|>


<commit_before>// This file is part of graze/golang-service
//
// Copyright (c) 2016 Nature Delivered Ltd. <https://www.graze.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//
// license: https://github.com/graze/golang-service/blob/master/LICENSE
// link:    https://github.com/graze/golang-service

package handlers

import (
    "github.com/gorilla/handlers"
    "fmt"
    "log/syslog"
    "net/http"
    "os"
    "strconv"
)

// rsyslogHandler returns a handlers.CombinedLoggingHandler to output Apache formatter log files to rsyslog
func SyslogHandler(h http.Handler) http.Handler {
    logWriter, err := GetSysLogFromEnv()
    if err != nil {
        panic(err)
    }
    defer logWriter.Close()

    return handlers.CombinedLoggingHandler(logWriter, h)
}

// SyslogConf is a configuration struct to be able to create a SysLog writer
type SyslogConf struct {
    network string
    host, port, application string
    level syslog.Priority
}

// GetSysLog returns a syslog.Writer based on a SyslogConf input
func GetSysLog(conf SyslogConf) (logWriter *syslog.Writer, err error) {
    la := ""
    if conf.host != "" {
        la = conf.host + ":" + conf.port
    }
    logWriter, err = syslog.Dial(
        conf.network,
        la,
        conf.level,
        conf.application)
    if err != nil {
        return nil, err
    }
    return
}

// GetSysLogFromEnv creates a syslog.Writer based on the environment variables supplied
//
// Uses the application environments:
//
//  SYSLOG_NETWORK: The network type of the syslog server (tcp, udp) Leave blank for local syslog
//  SYSLOG_HOST: The host of the syslog server. Leave blank for local syslog
//  SYSLOG_PORT: The port of the syslog server
//  SYSLOG_APPLICATION: The application to report the logs as
//  SYSLOG_LEVEL: The level to limit messages to (default: LEVEL6)
//
// Returns a *syslog.Writer
func GetSysLogFromEnv() (logWriter *syslog.Writer, err error) {
    names := []struct{
        name, env string
        req bool
    }{
        {"network", "SYSLOG_NETWORK", false},
        {"host", "SYSLOG_HOST", false},
        {"port", "SYSLOG_PORT", false},
        {"application", "SYSLOG_APPLICATION", true},
        {"level", "SYSLOG_LEVEL", false},
    }

    env := make(map[string]string)
    for _, v := range names {
        env[v.name] = os.Getenv(v.env)
        if v.req && env[v.name] == "" {
            return nil, fmt.Errorf("Unable to get value for setting: %s", v.env)
        }
    }

    level := syslog.LOG_LOCAL6
    if env["level"] != "" {
        i, err := strconv.Atoi(env["level"])
        if err != nil {
            return nil, err
        }
        level = syslog.Priority(i)
    }

    conf := SyslogConf{env["network"], env["host"], env["port"], env["application"], level}
    return GetSysLog(conf)
}
<commit_msg>set default priority to local6 | notice<commit_after>// This file is part of graze/golang-service
//
// Copyright (c) 2016 Nature Delivered Ltd. <https://www.graze.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//
// license: https://github.com/graze/golang-service/blob/master/LICENSE
// link:    https://github.com/graze/golang-service

package handlers

import (
    "github.com/gorilla/handlers"
    "fmt"
    "log/syslog"
    "net/http"
    "os"
    "strconv"
)

// rsyslogHandler returns a handlers.CombinedLoggingHandler to output Apache formatter log files to rsyslog
func SyslogHandler(h http.Handler) http.Handler {
    logWriter, err := GetSysLogFromEnv()
    if err != nil {
        panic(err)
    }
    defer logWriter.Close()

    return handlers.CombinedLoggingHandler(logWriter, h)
}

// SyslogConf is a configuration struct to be able to create a SysLog writer
type SyslogConf struct {
    network string
    host, port, application string
    level syslog.Priority
}

// GetSysLog returns a syslog.Writer based on a SyslogConf input
func GetSysLog(conf SyslogConf) (logWriter *syslog.Writer, err error) {
    la := ""
    if conf.host != "" {
        la = conf.host + ":" + conf.port
    }
    logWriter, err = syslog.Dial(
        conf.network,
        la,
        conf.level,
        conf.application)
    if err != nil {
        return nil, err
    }
    return
}

// GetSysLogFromEnv creates a syslog.Writer based on the environment variables supplied
//
// Uses the application environments:
//
//  SYSLOG_NETWORK: The network type of the syslog server (tcp, udp) Leave blank for local syslog
//  SYSLOG_HOST: The host of the syslog server. Leave blank for local syslog
//  SYSLOG_PORT: The port of the syslog server
//  SYSLOG_APPLICATION: The application to report the logs as
//  SYSLOG_LEVEL: The level to limit messages to (default: LEVEL6)
//
// Returns a *syslog.Writer
func GetSysLogFromEnv() (logWriter *syslog.Writer, err error) {
    names := []struct{
        name, env string
        req bool
    }{
        {"network", "SYSLOG_NETWORK", false},
        {"host", "SYSLOG_HOST", false},
        {"port", "SYSLOG_PORT", false},
        {"application", "SYSLOG_APPLICATION", true},
        {"level", "SYSLOG_LEVEL", false},
    }

    env := make(map[string]string)
    for _, v := range names {
        env[v.name] = os.Getenv(v.env)
        if v.req && env[v.name] == "" {
            return nil, fmt.Errorf("Unable to get value for setting: %s", v.env)
        }
    }

    level := syslog.LOG_LOCAL6 | syslog.LOG_NOTICE
    if env["level"] != "" {
        i, err := strconv.Atoi(env["level"])
        if err != nil {
            return nil, err
        }
        level = syslog.Priority(i)
    }

    conf := SyslogConf{env["network"], env["host"], env["port"], env["application"], level}
    return GetSysLog(conf)
}
<|endoftext|>