<commit_before>/*

Writer implementation.

*/

package bitio

import (
	"bufio"
	"io"
)

// An io.Writer and io.ByteWriter at the same time.
type writerAndByteWriter interface {
	io.Writer
	io.ByteWriter
}

// Writer is the bit writer implementation.
//
// For convenience, it also implements io.WriterCloser and io.ByteWriter.
type Writer struct {
	out       writerAndByteWriter
	wrapperbw *bufio.Writer // wrapper bufio.Writer if the target does not implement io.ByteWriter
	cache     byte          // unwritten bits are stored here
	bits      byte          // number of unwritten bits in cache

	// TryError holds the first error occurred in TryXXX() methods.
	TryError error
}

// NewWriter returns a new Writer using the specified io.Writer as the output.
//
// Must be closed in order to flush cached data.
// If you can't or don't want to close it, flushing data can also be forced
// by calling Align().
func NewWriter(out io.Writer) *Writer {
	w := &Writer{}
	var ok bool
	w.out, ok = out.(writerAndByteWriter)
	if !ok {
		w.wrapperbw = bufio.NewWriter(out)
		w.out = w.wrapperbw
	}
	return w
}

// Write writes len(p) bytes (8 * len(p) bits) to the underlying writer.
//
// Write implements io.Writer, and gives a byte-level interface to the bit stream.
// This will give best performance if the underlying io.Writer is aligned
// to a byte boundary (else all the individual bytes are spread to multiple bytes).
// Byte boundary can be ensured by calling Align().
func (w *Writer) Write(p []byte) (n int, err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	if w.bits == 0 {
		return w.out.Write(p)
	}

	for i, b := range p {
		if err = w.writeUnalignedByte(b); err != nil {
			return i, err
		}
	}

	return len(p), nil
}

// WriteBits writes out the n lowest bits of r.
// Bits of r in positions higher than n are ignored.
//
// For example:
//   err := w.WriteBits(0x1234, 8) // bits higher than the 8th are ignored here
//
// is equivalent to:
//   err := w.WriteBits(0x34, 8)
func (w *Writer) WriteBits(r uint64, n uint8) (err error) {
	// if r would have bits set higher than n-1 (zero indexed),
	// WriteBitsUnsafe's implementation could "corrupt" bits in cache.
	// That is not acceptable. To be on the safe side, mask out higher bits:
	return w.WriteBitsUnsafe((r & (1<<n - 1)), n)
}

// WriteBitsUnsafe writes out the n lowest bits of r.
//
// r cannot have bits set at positions higher than n-1 (zero indexed).
// If r might not satisfy this, a mask must be explicitly applied
// before passing it to WriteBitsUnsafe(), or WriteBits() should be used instead.
// WriteBitsUnsafe() gives slightly better performance than WriteBits() because
// the input r is not masked.
//
// E.g. if you want to write 8 bits:
//   err := w.WriteBitsUnsafe(0x34, 8)        // This is OK,
//                                            // 0x34 has no bits set higher than the 8th
//   err := w.WriteBitsUnsafe(0x1234&0xff, 8) // &0xff masks out bits higher than the 8th
//
// Or:
//   err := w.WriteBits(0x1234, 8) // bits higher than the 8th are ignored here
func (w *Writer) WriteBitsUnsafe(r uint64, n uint8) (err error) {
	// Some optimization, frequent cases
	newbits := w.bits + n
	if newbits < 8 {
		// r fits into cache, no write will occur to out
		w.cache |= byte(r) << (8 - newbits)
		w.bits = newbits
		return nil
	}

	if newbits > 8 {
		// cache will be filled, and there will be more bits to write
		// "Fill cache" and write it out
		free := 8 - w.bits
		err = w.out.WriteByte(w.cache | byte(r>>(n-free)))
		if err != nil {
			return
		}
		n -= free
		// write out whole bytes
		for n >= 8 {
			n -= 8
			// No need to mask r, converting to byte will mask out higher bits
			err = w.out.WriteByte(byte(r >> n))
			if err != nil {
				return
			}
		}
		// Put remaining into cache
		if n > 0 {
			// Note: n < 8 (in case of n=8, 1<<n would overflow byte)
			w.cache, w.bits = (byte(r)&((1<<n)-1))<<(8-n), n
		} else {
			w.cache, w.bits = 0, 0
		}
		return nil
	}

	// cache will be filled exactly with the bits to be written
	bb := w.cache | byte(r)
	w.cache, w.bits = 0, 0
	return w.out.WriteByte(bb)
}

// WriteByte writes 8 bits.
//
// WriteByte implements io.ByteWriter.
func (w *Writer) WriteByte(b byte) (err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	if w.bits == 0 {
		return w.out.WriteByte(b)
	}
	return w.writeUnalignedByte(b)
}

// writeUnalignedByte writes 8 bits which are (may be) unaligned.
func (w *Writer) writeUnalignedByte(b byte) (err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	bits := w.bits
	err = w.out.WriteByte(w.cache | b>>bits)
	if err != nil {
		return
	}
	w.cache = (b & (1<<bits - 1)) << (8 - bits)
	return
}

// WriteBool writes one bit: 1 if param is true, 0 otherwise.
func (w *Writer) WriteBool(b bool) (err error) {
	if w.bits == 7 {
		if b {
			err = w.out.WriteByte(w.cache | 1)
		} else {
			err = w.out.WriteByte(w.cache)
		}
		if err != nil {
			return
		}
		w.cache, w.bits = 0, 0
		return nil
	}

	w.bits++
	if b {
		w.cache |= 1 << (8 - w.bits)
	}
	return nil
}

// Align aligns the bit stream to a byte boundary,
// so next write will start/go into a new byte.
// If there are cached bits, they are first written to the output.
// Returns the number of skipped (unset but still written) bits.
func (w *Writer) Align() (skipped uint8, err error) {
	if w.bits > 0 {
		if err = w.out.WriteByte(w.cache); err != nil {
			return
		}

		skipped = 8 - w.bits
		w.cache, w.bits = 0, 0
	}
	if w.wrapperbw != nil {
		err = w.wrapperbw.Flush()
	}
	return
}

// TryWrite tries to write len(p) bytes (8 * len(p) bits) to the underlying writer.
//
// If there was a previous TryError, it does nothing. Else it calls Write(),
// returns the data it provides and stores the error in the TryError field.
func (w *Writer) TryWrite(p []byte) (n int) {
	if w.TryError == nil {
		n, w.TryError = w.Write(p)
	}
	return
}

// TryWriteBits tries to write out the n lowest bits of r.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBits(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBits(r uint64, n uint8) {
	if w.TryError == nil {
		w.TryError = w.WriteBits(r, n)
	}
}

// TryWriteBitsUnsafe tries to write out the n lowest bits of r.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBitsUnsafe(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBitsUnsafe(r uint64, n uint8) {
	if w.TryError == nil {
		w.TryError = w.WriteBitsUnsafe(r, n)
	}
}

// TryWriteByte tries to write 8 bits.
//
// If there was a previous TryError, it does nothing. Else it calls WriteByte(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteByte(b byte) {
	if w.TryError == nil {
		w.TryError = w.WriteByte(b)
	}
}

// TryWriteBool tries to write one bit: 1 if param is true, 0 otherwise.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBool(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBool(b bool) {
	if w.TryError == nil {
		w.TryError = w.WriteBool(b)
	}
}

// TryAlign tries to align the bit stream to a byte boundary.
//
// If there was a previous TryError, it does nothing. Else it calls Align(),
// returns the data it provides and stores the error in the TryError field.
func (w *Writer) TryAlign() (skipped uint8) {
	if w.TryError == nil {
		skipped, w.TryError = w.Align()
	}
	return
}

// Close closes the bit writer, writes out cached bits.
// It does not close the underlying io.Writer.
//
// Close implements io.Closer.
func (w *Writer) Close() (err error) {
	// Make sure cached bits are flushed:
	if _, err = w.Align(); err != nil {
		return
	}

	return nil
}
<commit_msg>Improve docs<commit_after>/*

Writer implementation.

*/

package bitio

import (
	"bufio"
	"io"
)

// An io.Writer and io.ByteWriter at the same time.
type writerAndByteWriter interface {
	io.Writer
	io.ByteWriter
}

// Writer is the bit writer implementation.
//
// For convenience, it also implements io.WriterCloser and io.ByteWriter.
type Writer struct {
	out       writerAndByteWriter
	wrapperbw *bufio.Writer // wrapper bufio.Writer if the target does not implement io.ByteWriter
	cache     byte          // unwritten bits are stored here
	bits      byte          // number of unwritten bits in cache

	// TryError holds the first error occurred in TryXXX() methods.
	TryError error
}

// NewWriter returns a new Writer using the specified io.Writer as the output.
//
// Must be closed in order to flush cached data.
// If you can't or don't want to close it, flushing data can also be forced
// by calling Align().
func NewWriter(out io.Writer) *Writer {
	w := &Writer{}
	var ok bool
	w.out, ok = out.(writerAndByteWriter)
	if !ok {
		w.wrapperbw = bufio.NewWriter(out)
		w.out = w.wrapperbw
	}
	return w
}

// Write writes len(p) bytes (8 * len(p) bits) to the underlying writer.
//
// Write implements io.Writer, and gives a byte-level interface to the bit stream.
// This will give best performance if the underlying io.Writer is aligned
// to a byte boundary (else all the individual bytes are spread to multiple bytes).
// Byte boundary can be ensured by calling Align().
func (w *Writer) Write(p []byte) (n int, err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	if w.bits == 0 {
		return w.out.Write(p)
	}

	for i, b := range p {
		if err = w.writeUnalignedByte(b); err != nil {
			return i, err
		}
	}

	return len(p), nil
}

// WriteBits writes out the n lowest bits of r.
// Bits of r in positions higher than n are ignored.
//
// For example:
//   err := w.WriteBits(0x1234, 8) // bits higher than the 8th are ignored here
//
// is equivalent to:
//   err := w.WriteBits(0x34, 8)
func (w *Writer) WriteBits(r uint64, n uint8) (err error) {
	// if r would have bits set higher than n-1 (zero indexed),
	// WriteBitsUnsafe's implementation could "corrupt" bits in cache.
	// That is not acceptable. To be on the safe side, mask out higher bits:
	return w.WriteBitsUnsafe((r & (1<<n - 1)), n)
}

// WriteBitsUnsafe writes out the n lowest bits of r.
//
// r must not have bits set at positions higher than n-1 (zero indexed).
// If r might not satisfy this, a mask must be explicitly applied
// before passing it to WriteBitsUnsafe(), or WriteBits() should be used instead.
// WriteBitsUnsafe() gives slightly better performance than WriteBits() because
// the input r is not masked. Calling WriteBitsUnsafe() with an r that does
// not satisfy this is undefined behavior (might corrupt previously written bits).
//
// E.g. if you want to write 8 bits:
//   err := w.WriteBitsUnsafe(0x34, 8)        // This is OK,
//                                            // 0x34 has no bits set higher than the 8th
//   err := w.WriteBitsUnsafe(0x1234&0xff, 8) // &0xff masks out bits higher than the 8th
//
// Or:
//   err := w.WriteBits(0x1234, 8) // bits higher than the 8th are ignored here
func (w *Writer) WriteBitsUnsafe(r uint64, n uint8) (err error) {
	// Some optimization, frequent cases
	newbits := w.bits + n
	if newbits < 8 {
		// r fits into cache, no write will occur to out
		w.cache |= byte(r) << (8 - newbits)
		w.bits = newbits
		return nil
	}

	if newbits > 8 {
		// cache will be filled, and there will be more bits to write
		// "Fill cache" and write it out
		free := 8 - w.bits
		err = w.out.WriteByte(w.cache | byte(r>>(n-free)))
		if err != nil {
			return
		}
		n -= free
		// write out whole bytes
		for n >= 8 {
			n -= 8
			// No need to mask r, converting to byte will mask out higher bits
			err = w.out.WriteByte(byte(r >> n))
			if err != nil {
				return
			}
		}
		// Put remaining into cache
		if n > 0 {
			// Note: n < 8 (in case of n=8, 1<<n would overflow byte)
			w.cache, w.bits = (byte(r)&((1<<n)-1))<<(8-n), n
		} else {
			w.cache, w.bits = 0, 0
		}
		return nil
	}

	// cache will be filled exactly with the bits to be written
	bb := w.cache | byte(r)
	w.cache, w.bits = 0, 0
	return w.out.WriteByte(bb)
}

// WriteByte writes 8 bits.
//
// WriteByte implements io.ByteWriter.
func (w *Writer) WriteByte(b byte) (err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	if w.bits == 0 {
		return w.out.WriteByte(b)
	}
	return w.writeUnalignedByte(b)
}

// writeUnalignedByte writes 8 bits which are (may be) unaligned.
func (w *Writer) writeUnalignedByte(b byte) (err error) {
	// w.bits will be the same after writing 8 bits, so we don't need to update that.
	bits := w.bits
	err = w.out.WriteByte(w.cache | b>>bits)
	if err != nil {
		return
	}
	w.cache = (b & (1<<bits - 1)) << (8 - bits)
	return
}

// WriteBool writes one bit: 1 if param is true, 0 otherwise.
func (w *Writer) WriteBool(b bool) (err error) {
	if w.bits == 7 {
		if b {
			err = w.out.WriteByte(w.cache | 1)
		} else {
			err = w.out.WriteByte(w.cache)
		}
		if err != nil {
			return
		}
		w.cache, w.bits = 0, 0
		return nil
	}

	w.bits++
	if b {
		w.cache |= 1 << (8 - w.bits)
	}
	return nil
}

// Align aligns the bit stream to a byte boundary,
// so next write will start/go into a new byte.
// If there are cached bits, they are first written to the output.
// Returns the number of skipped (unset but still written) bits.
func (w *Writer) Align() (skipped uint8, err error) {
	if w.bits > 0 {
		if err = w.out.WriteByte(w.cache); err != nil {
			return
		}

		skipped = 8 - w.bits
		w.cache, w.bits = 0, 0
	}
	if w.wrapperbw != nil {
		err = w.wrapperbw.Flush()
	}
	return
}

// TryWrite tries to write len(p) bytes (8 * len(p) bits) to the underlying writer.
//
// If there was a previous TryError, it does nothing. Else it calls Write(),
// returns the data it provides and stores the error in the TryError field.
func (w *Writer) TryWrite(p []byte) (n int) {
	if w.TryError == nil {
		n, w.TryError = w.Write(p)
	}
	return
}

// TryWriteBits tries to write out the n lowest bits of r.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBits(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBits(r uint64, n uint8) {
	if w.TryError == nil {
		w.TryError = w.WriteBits(r, n)
	}
}

// TryWriteBitsUnsafe tries to write out the n lowest bits of r.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBitsUnsafe(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBitsUnsafe(r uint64, n uint8) {
	if w.TryError == nil {
		w.TryError = w.WriteBitsUnsafe(r, n)
	}
}

// TryWriteByte tries to write 8 bits.
//
// If there was a previous TryError, it does nothing. Else it calls WriteByte(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteByte(b byte) {
	if w.TryError == nil {
		w.TryError = w.WriteByte(b)
	}
}

// TryWriteBool tries to write one bit: 1 if param is true, 0 otherwise.
//
// If there was a previous TryError, it does nothing. Else it calls WriteBool(),
// and stores the error in the TryError field.
func (w *Writer) TryWriteBool(b bool) {
	if w.TryError == nil {
		w.TryError = w.WriteBool(b)
	}
}

// TryAlign tries to align the bit stream to a byte boundary.
//
// If there was a previous TryError, it does nothing. Else it calls Align(),
// returns the data it provides and stores the error in the TryError field.
func (w *Writer) TryAlign() (skipped uint8) {
	if w.TryError == nil {
		skipped, w.TryError = w.Align()
	}
	return
}

// Close closes the bit writer, writes out cached bits.
// It does not close the underlying io.Writer.
//
// Close implements io.Closer.
func (w *Writer) Close() (err error) {
	// Make sure cached bits are flushed:
	if _, err = w.Align(); err != nil {
		return
	}

	return nil
}
<|endoftext|>


<commit_before>package ltsview

import (
	"fmt"
	"io"

	"github.com/fatih/color"
)

type writer struct {
	io.Writer
}

func (w *writer) WriteSeparator(sep string) (int, error) {
	return fmt.Fprintln(w.Writer, sep)
}

func NewWriter(w io.Writer) Writer {
	return &writer{w}
}

type Writer interface {
	io.Writer
	SeparatorWriter
}

type SeparatorWriter interface {
	WriteSeparator(sep string) (int, error)
}

type ColorableWriter struct {
	io.Writer
}

func (w *ColorableWriter) WriteSeparator(sep string) (int, error) {
	return fmt.Fprintln(w.Writer, sep)
}

func (w *ColorableWriter) Write(p []byte) (int, error) {
	f := color.New(color.FgMagenta).SprintFunc()
	for i := range p {
		str := string(p[i])
		if p[i] == ':' {
			fmt.Fprint(w.Writer, str)
			f = color.New(color.FgGreen).SprintFunc()
			continue
		}
		fmt.Fprint(w.Writer, f(str))
	}
	return len(p), nil
}
<commit_msg>Do not change color if a key has already been printed<commit_after>package ltsview

import (
	"fmt"
	"io"

	"github.com/fatih/color"
)

type writer struct {
	io.Writer
}

func (w *writer) WriteSeparator(sep string) (int, error) {
	return fmt.Fprintln(w.Writer, sep)
}

func NewWriter(w io.Writer) Writer {
	return &writer{w}
}

type Writer interface {
	io.Writer
	SeparatorWriter
}

type SeparatorWriter interface {
	WriteSeparator(sep string) (int, error)
}

type ColorableWriter struct {
	io.Writer
}

func (w *ColorableWriter) WriteSeparator(sep string) (int, error) {
	return fmt.Fprintln(w.Writer, sep)
}

func (w *ColorableWriter) Write(p []byte) (int, error) {
	f := color.New(color.FgMagenta).SprintFunc()
	keyPrinted := false
	for i := range p {
		str := string(p[i])
		if !keyPrinted && p[i] == ':' {
			fmt.Fprint(w.Writer, str)
			keyPrinted = true
			f = color.New(color.FgGreen).SprintFunc()
			continue
		}
		fmt.Fprint(w.Writer, f(str))
	}
	return len(p), nil
}
<|endoftext|>


<commit_before>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webp

import (
	"image"
	"image/color"
	"io"
)

const DefaulQuality = 90

// Options are the encoding parameters.
type Options struct {
	Lossless bool
	Quality  float32 // 0 ~ 100
}

// Encode writes the image m to w in WEBP format.
func Encode(w io.Writer, m image.Image, opt *Options) (err error) {
	var output []byte
	if opt != nil && opt.Lossless {
		switch m := adjustImage(m).(type) {
		case *image.Gray:
			if output, err = EncodeLosslessGray(m); err != nil {
				return
			}
		case *_RGB:
			if output, err = EncodeLosslessRGB(m); err != nil {
				return
			}
		case *image.RGBA:
			if output, err = EncodeLosslessRGBA(m); err != nil {
				return
			}
		default:
			panic("image/webp: Encode, unreachable!")
		}
	} else {
		quality := float32(DefaulQuality)
		if opt != nil {
			quality = opt.Quality
		}
		switch m := adjustImage(m).(type) {
		case *image.Gray:
			if output, err = EncodeGray(m, quality); err != nil {
				return
			}
		case *_RGB:
			if output, err = EncodeRGB(m, quality); err != nil {
				return
			}
		case *image.RGBA:
			if output, err = EncodeRGBA(m, quality); err != nil {
				return
			}
		default:
			panic("image/webp: Encode, unreachable!")
		}
	}
	_, err = w.Write(output)
	return
}

func adjustImage(m image.Image) image.Image {
	if x, ok := m.(Image); ok {
		m = x.BaseType()
	}
	switch m := m.(type) {
	case *image.Gray, *image.RGBA, *_RGB:
		return m
	default:
		b := m.Bounds()
		rgba := image.NewRGBA(b)
		dstColorRGBA64 := &color.RGBA64{}
		dstColor := color.Color(dstColorRGBA64)
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				pr, pg, pb, pa := m.At(x, y).RGBA()
				dstColorRGBA64.R = uint16(pr)
				dstColorRGBA64.G = uint16(pg)
				dstColorRGBA64.B = uint16(pb)
				dstColorRGBA64.A = uint16(pa)
				rgba.Set(x, y, dstColor)
			}
		}
		return rgba
	}
}
<commit_msg>improve user defined format for Encode<commit_after>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webp

import (
	"image"
	"image/color"
	"io"
	"reflect"
)

const DefaulQuality = 90

// Options are the encoding parameters.
type Options struct {
	Lossless bool
	Quality  float32 // 0 ~ 100
}

// Encode writes the image m to w in WEBP format.
func Encode(w io.Writer, m image.Image, opt *Options) (err error) {
	var output []byte
	if opt != nil && opt.Lossless {
		switch m := adjustImage(m).(type) {
		case *image.Gray:
			if output, err = EncodeLosslessGray(m); err != nil {
				return
			}
		case *_RGB:
			if output, err = EncodeLosslessRGB(m); err != nil {
				return
			}
		case *image.RGBA:
			if output, err = EncodeLosslessRGBA(m); err != nil {
				return
			}
		default:
			panic("image/webp: Encode, unreachable!")
		}
	} else {
		quality := float32(DefaulQuality)
		if opt != nil {
			quality = opt.Quality
		}
		switch m := adjustImage(m).(type) {
		case *image.Gray:
			if output, err = EncodeGray(m, quality); err != nil {
				return
			}
		case *_RGB:
			if output, err = EncodeRGB(m, quality); err != nil {
				return
			}
		case *image.RGBA:
			if output, err = EncodeRGBA(m, quality); err != nil {
				return
			}
		default:
			panic("image/webp: Encode, unreachable!")
		}
	}
	_, err = w.Write(output)
	return
}

func adjustImage(m image.Image) image.Image {
	switch m := m.(type) {
	case *image.Gray, *image.RGBA, *_RGB:
		return m
	}

	// try `Image` interface
	if x, ok := m.(Image); ok {
		// try original type
		switch m := x.BaseType().(type) {
		case *image.Gray, *image.RGBA, *_RGB:
			return m
		}
		// create new image with `x.Pix()`
		switch {
		case x.Channels() == 1 && x.Depth() == reflect.Uint8:
			return &image.Gray{
				Pix:    x.Pix(),
				Stride: x.Stride(),
				Rect:   x.Rect(),
			}
		case x.Channels() == 3 && x.Depth() == reflect.Uint8:
			return &_RGB{
				Pix:    x.Pix(),
				Stride: x.Stride(),
				Rect:   x.Rect(),
			}
		case x.Channels() == 4 && x.Depth() == reflect.Uint8:
			return &image.RGBA{
				Pix:    x.Pix(),
				Stride: x.Stride(),
				Rect:   x.Rect(),
			}
		}
	}

	// convert to RGBA
	b := m.Bounds()
	rgba := image.NewRGBA(b)
	dstColorRGBA64 := &color.RGBA64{}
	dstColor := color.Color(dstColorRGBA64)
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			pr, pg, pb, pa := m.At(x, y).RGBA()
			dstColorRGBA64.R = uint16(pr)
			dstColorRGBA64.G = uint16(pg)
			dstColorRGBA64.B = uint16(pb)
			dstColorRGBA64.A = uint16(pa)
			rgba.Set(x, y, dstColor)
		}
	}
	return rgba
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/screensaver"
	"github.com/BurntSushi/xgb/xproto"
	"log"
	"time"
)

type tracker map[window]*track

type track struct {
	Start time.Time
	Spent time.Duration
}

type window struct {
	Class string
	Name  string
}

func (t track) String() string {
	return fmt.Sprint(t.Spent)
}

func (w window) String() string {
	return fmt.Sprintf("%s: %s", w.Class, w.Name)
}

func getClass(b []byte) string {
	i := bytes.IndexByte(b, 0)
	if i == -1 {
		return ""
	}
	return string(b[:i])
}

func asciizToString(b []byte) (s []string) {
	for _, x := range bytes.Split(b, []byte{0}) {
		s = append(s, string(x))
	}
	if len(s) > 0 && s[len(s)-1] == "" {
		s = s[:len(s)-1]
	}
	return
}

func atom(X *xgb.Conn, aname string) *xproto.InternAtomReply {
	a, err := xproto.InternAtom(X, true, uint16(len(aname)), aname).Reply()
	if err != nil {
		log.Fatal("atom", err)
	}
	return a
}

func prop(X *xgb.Conn, w xproto.Window, a *xproto.InternAtomReply) *xproto.GetPropertyReply {
	p, err := xproto.GetProperty(X, false, w, a.Atom, xproto.GetPropertyTypeAny, 0, (1<<32)-1).Reply()
	if err != nil {
		log.Fatal("property", err)
	}
	return p
}

func winName(X *xgb.Conn, root xproto.Window) (window, bool) {
	activeAtom := atom(X, "_NET_ACTIVE_WINDOW")
	netNameAtom := atom(X, "_NET_WM_NAME")
	nameAtom := atom(X, "WM_NAME")
	classAtom := atom(X, "WM_CLASS")

	active := prop(X, root, activeAtom)
	windowId := xproto.Window(xgb.Get32(active.Value))

	/* skip root window */
	if windowId == 0 {
		return window{}, false
	}

	spy(X, windowId)

	name := prop(X, windowId, netNameAtom)
	if string(name.Value) == "" {
		name = prop(X, windowId, nameAtom)
	}
	class := prop(X, windowId, classAtom)

	w := window{
		Class: getClass(class.Value),
		Name:  string(name.Value),
	}

	return w, true
}

func rootWin(X *xgb.Conn) xproto.Window {
	setup := xproto.Setup(X)
	return setup.DefaultScreen(X).Root
}

func spy(X *xgb.Conn, w xproto.Window) {
	xproto.ChangeWindowAttributes(X, w, xproto.CwEventMask,
		[]uint32{xproto.EventMaskPropertyChange})
}

func collect(tracks tracker) {
	var prev *track

	X, err := xgb.NewConn()
	if err != nil {
		log.Fatal("xgb", err)
	}
	defer X.Close()

	err = screensaver.Init(X)
	if err != nil {
		log.Fatal("screensaver", err)
	}

	root := rootWin(X)
	spy(X, root)
	drw := xproto.Drawable(root)
	screensaver.SelectInput(X, drw, screensaver.EventNotifyMask)

	for {
		ev, everr := X.WaitForEvent()
		if everr != nil {
			log.Fatal("wait for event", everr)
		}
		switch event := ev.(type) {
		case xproto.PropertyNotifyEvent:
			if prev != nil {
				prev.Spent += time.Since(prev.Start)
			}
			if win, ok := winName(X, root); ok {
				if _, ok := tracks[win]; !ok {
					tracks[win] = new(track)
				}
				tracks[win].Start = time.Now()
				prev = tracks[win]
			}
		case screensaver.NotifyEvent:
			switch event.State {
			case screensaver.StateOn:
				fmt.Println("screensaver on")
				prev = nil
			}
		}
	}
}

func display(tracks tracker) {
	for {
		var total time.Duration
		classtotal := make(map[string]time.Duration)
		for n, t := range tracks {
			fmt.Println(n, t)
			total += t.Spent
			classtotal[n.Class] += t.Spent
		}
		fmt.Println("")
		for k, v := range classtotal {
			fmt.Println(k, v)
		}
		fmt.Println("Total:", total)
		fmt.Println("")
		time.Sleep(3 * time.Second)
	}
}

func cleanup(tracks tracker) {
	for {
		for k, v := range tracks {
			if time.Since(v.Start).Hours() > 24.0 {
				tracks[k] = nil
			}
		}
		time.Sleep(time.Minute)
	}
}

func main() {
	tracks := make(tracker)
	go collect(tracks)
	go cleanup(tracks)
	display(tracks)
}
<commit_msg>fix delete<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/screensaver"
	"github.com/BurntSushi/xgb/xproto"
	"log"
	"time"
)

type tracker map[window]*track

type track struct {
	Start time.Time
	Spent time.Duration
}

type window struct {
	Class string
	Name  string
}

func (t track) String() string {
	return fmt.Sprint(t.Spent)
}

func (w window) String() string {
	return fmt.Sprintf("%s: %s", w.Class, w.Name)
}

func getClass(b []byte) string {
	i := bytes.IndexByte(b, 0)
	if i == -1 {
		return ""
	}
	return string(b[:i])
}

func asciizToString(b []byte) (s []string) {
	for _, x := range bytes.Split(b, []byte{0}) {
		s = append(s, string(x))
	}
	if len(s) > 0 && s[len(s)-1] == "" {
		s = s[:len(s)-1]
	}
	return
}

func atom(X *xgb.Conn, aname string) *xproto.InternAtomReply {
	a, err := xproto.InternAtom(X, true, uint16(len(aname)), aname).Reply()
	if err != nil {
		log.Fatal("atom", err)
	}
	return a
}

func prop(X *xgb.Conn, w xproto.Window, a *xproto.InternAtomReply) *xproto.GetPropertyReply {
	p, err := xproto.GetProperty(X, false, w, a.Atom, xproto.GetPropertyTypeAny, 0, (1<<32)-1).Reply()
	if err != nil {
		log.Fatal("property", err)
	}
	return p
}

func winName(X *xgb.Conn, root xproto.Window) (window, bool) {
	activeAtom := atom(X, "_NET_ACTIVE_WINDOW")
	netNameAtom := atom(X, "_NET_WM_NAME")
	nameAtom := atom(X, "WM_NAME")
	classAtom := atom(X, "WM_CLASS")

	active := prop(X, root, activeAtom)
	windowId := xproto.Window(xgb.Get32(active.Value))

	/* skip root window */
	if windowId == 0 {
		return window{}, false
	}

	spy(X, windowId)

	name := prop(X, windowId, netNameAtom)
	if string(name.Value) == "" {
		name = prop(X, windowId, nameAtom)
	}
	class := prop(X, windowId, classAtom)

	w := window{
		Class: getClass(class.Value),
		Name:  string(name.Value),
	}

	return w, true
}

func rootWin(X *xgb.Conn) xproto.Window {
	setup := xproto.Setup(X)
	return setup.DefaultScreen(X).Root
}

func spy(X *xgb.Conn, w xproto.Window) {
	xproto.ChangeWindowAttributes(X, w, xproto.CwEventMask,
		[]uint32{xproto.EventMaskPropertyChange})
}

func collect(tracks tracker) {
	var prev *track

	X, err := xgb.NewConn()
	if err != nil {
		log.Fatal("xgb", err)
	}
	defer X.Close()

	err = screensaver.Init(X)
	if err != nil {
		log.Fatal("screensaver", err)
	}

	root := rootWin(X)
	spy(X, root)
	drw := xproto.Drawable(root)
	screensaver.SelectInput(X, drw, screensaver.EventNotifyMask)

	for {
		ev, everr := X.WaitForEvent()
		if everr != nil {
			log.Fatal("wait for event", everr)
		}
		switch event := ev.(type) {
		case xproto.PropertyNotifyEvent:
			if prev != nil {
				prev.Spent += time.Since(prev.Start)
			}
			if win, ok := winName(X, root); ok {
				if _, ok := tracks[win]; !ok {
					tracks[win] = new(track)
				}
				tracks[win].Start = time.Now()
				prev = tracks[win]
			}
		case screensaver.NotifyEvent:
			switch event.State {
			case screensaver.StateOn:
				fmt.Println("screensaver on")
				prev = nil
			}
		}
	}
}

func display(tracks tracker) {
	for {
		var total time.Duration
		classtotal := make(map[string]time.Duration)
		for n, t := range tracks {
			fmt.Println(n, t)
			total += t.Spent
			classtotal[n.Class] += t.Spent
		}
		fmt.Println("")
		for k, v := range classtotal {
			fmt.Println(k, v)
		}
		fmt.Println("Total:", total)
		fmt.Println("")
		time.Sleep(3 * time.Second)
	}
}

func cleanup(tracks tracker) {
	for {
		for k, v := range tracks {
			if time.Since(v.Start).Hours() > 24.0 {
				delete(tracks, k)
			}
		}
		time.Sleep(time.Minute)
	}
}

func main() {
	tracks := make(tracker)
	go collect(tracks)
	go cleanup(tracks)
	display(tracks)
}
<|endoftext|>


<commit_before>//Go beanstalkd client library
//Copyright(2012) Iwan Budi Kusnanto. See LICENSE for detail
package gobeanstalk

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"net"
	"strings"
)

//beanstalkd error
var (
	errOutOfMemory    = errors.New("Out of Memory")
	errInternalError  = errors.New("Internal Error")
	errBadFormat      = errors.New("Bad Format")
	errUnknownCommand = errors.New("Unknown Command")
	errBuried         = errors.New("Buried")
	errExpectedCrlf   = errors.New("Expected CRLF")
	errJobTooBig      = errors.New("Job Too Big")
	errDraining       = errors.New("Draining")
	errDeadlineSoon   = errors.New("Deadline Soon")
	errTimedOut       = errors.New("Timed Out")
	errNotFound       = errors.New("Not Found")
)

//gobeanstalk error
var (
	errInvalidLen	  = errors.New("Invalid Length")
	errUnknown        = errors.New("Unknown Error")
)

//Connection to beanstalkd
type Conn struct {
	conn   net.Conn
	addr   string
	reader *bufio.Reader
}

//create new connection
func NewConn(conn net.Conn, addr string) (*Conn, error) {
	c := new(Conn)
	c.conn = conn
	c.addr = addr
	c.reader = bufio.NewReader(conn)

	return c, nil
}

//A beanstalkd job
type Job struct {
	Id uint64
	Body []byte
}

//Create new job
func NewJob(id uint64, body []byte) *Job {
	j :=  &Job{id, body}
	return j
}

//Connect to beanstalkd server
func Dial(addr string) (*Conn, error) {
	kon, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	c, err := NewConn(kon, addr)

	if err != nil {
		return nil, err
	}

	return c, nil
}

//Watching tube
func (c *Conn) Watch(tubename string) (int, error) {
	cmd := fmt.Sprintf("watch %s\r\n", tubename)

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

/*
Ignore tube.

The "ignore" command is for consumers. It removes the named tube from the
watch list for the current connection
*/
func (c *Conn) Ignore(tubename string) (int, error) {
	//send command and read response string
	cmd := fmt.Sprintf("ignore %s\r\n", tubename)
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	//parse response
	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		if resp == "NOT_IGNORED\r\n" {
			return -1, errors.New("Not Ignored")
		}
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

//Reserve Job
func (c *Conn) Reserve() (*Job, error) {
	//send command and read response
	resp, err := sendGetResp(c, "reserve\r\n")
	if err != nil {
		return nil, err
	}

	//parse response
	var id uint64
	var bodyLen int

	switch {
	case strings.Index(resp, "RESERVED") == 0:
		_, err = fmt.Sscanf(resp, "RESERVED %d %d\r\n", &id, &bodyLen)
		if err != nil {
			return nil, err
		}
	case resp == "DEADLINE_SOON\r\n":
		return nil, errDeadlineSoon
	case resp == "TIMED_OUT\r\n":
		return nil, errTimedOut
	default:
		return nil, parseCommonError(resp)
	}

	//read job body
	body, err := c.readBytes()
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	body = body[:len(body)-2]
	if len(body) != bodyLen {
		return nil, errors.New(fmt.Sprintf("invalid body len = %d/%d", len(body), bodyLen))
	}

	return &Job{id, body}, nil
}

//Delete a job
func (c *Conn) Delete(id uint64) error {
	cmd := fmt.Sprintf("delete %d\r\n", id)
	expected := "DELETED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Use tube

The "use" command is for producers. Subsequent put commands will put jobs into
the tube specified by this command. If no use command has been issued, jobs
will be put into the tube named "default".
*/
func (c *Conn) Use(tubename string) error {
	//check parameter
	if len(tubename) > 200 {
		return errInvalidLen
	}

	cmd := fmt.Sprintf("use %s\r\n", tubename)
	expected := fmt.Sprintf("USING %s\r\n", tubename)
	return sendExpectExact(c, cmd, expected)
}

//Put job
func (c *Conn) Put(data []byte, pri, delay, ttr int) (uint64, error) {
	cmd := fmt.Sprintf("put %d %d %d %d\r\n", pri, delay, ttr, len(data))
	cmd = cmd + string(data) + "\r\n"

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return 0, err
	}

	//parse Put response
	switch {
	case strings.Index(resp, "INSERTED") == 0:
		var id uint64
		_, parseErr := fmt.Sscanf(resp, "INSERTED %d\r\n", &id)
		return id, parseErr
	case strings.Index(resp, "BURIED") == 0:
		var id uint64
		fmt.Sscanf(resp, "BURIED %d\r\n", &id)
		return id, errBuried
	case resp == "EXPECTED_CRLF\r\n":
		return 0, errExpectedCrlf
	case resp == "JOB_TOO_BIG\r\n":
		return 0, errJobTooBig
	case resp == "DRAINING\r\n":
		return 0, errDraining
	default:
		return 0, parseCommonError(resp)
	}
	return 0, errUnknown
}

/*
Release a job.

The release command puts a reserved job back into the ready queue (and marks
its state as "ready") to be run by any client. It is normally used when the job
fails because of a transitory error.
	id is the job id to release.
	pri is a new priority to assign to the job.
	delay is an integer number of seconds to wait before putting the job in
		the ready queue. The job will be in the "delayed" state during this time.
*/
func (c *Conn) Release(id uint64, pri, delay int) error {
	cmd := fmt.Sprintf("release %d %d %d\r\n", id, pri, delay)
	expected := "RELEASED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Bury a job.

The bury command puts a job into the "buried" state. Buried jobs are put into a
FIFO linked list and will not be touched by the server again until a client
kicks them with the "kick" command.
	id is the job id to release.
	pri is a new priority to assign to the job.
*/
func (c *Conn) Bury(id uint64, pri int) error {
	cmd := fmt.Sprintf("bury %d %d\r\n", id, pri)
	expected := "BURIED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Touch a job

The "touch" command allows a worker to request more time to work on a job.
This is useful for jobs that potentially take a long time, but you still want
the benefits of a TTR pulling a job away from an unresponsive worker. A worker
may periodically tell the server that it's still alive and processing a job
(e.g. it may do this on DEADLINE_SOON)
*/
func (c *Conn) Touch(id uint64) error {
	cmd := fmt.Sprintf("touch %d\r\n", id)
	expected := "TOUCHED\r\n"
	return sendExpectExact(c, cmd, expected)
}

//send command and expect some exact response
func sendExpectExact(c *Conn, cmd, expected string) error {
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return err
	}

	if resp != expected {
		return parseCommonError(resp)
	}
	return nil
}

//Send command and read response
func sendGetResp(c *Conn, cmd string) (string, error) {
	_, err := c.conn.Write([]byte(cmd))
	if err != nil {
		return "", err
	}

	//wait for response
	resp, err := c.reader.ReadString('\n')
	if err != nil {
		return "", err
	}
	return resp, nil
}

//read bytes until \n
func (c *Conn) readBytes() ([]byte, error) {
	rsp, err := c.reader.ReadBytes('\n')
	return rsp, err
}

//parse for Common Error
func parseCommonError(str string) error {
	switch str {
	case "BURIED\r\n":
		return errBuried
	case "NOT_FOUND\r\n":
		return errNotFound
	case "OUT_OF_MEMORY\r\n":
		return errOutOfMemory
	case "INTERNAL_ERROR\r\n":
		return errInternalError
	case "BAD_FORMAT\r\n":
		return errBadFormat
	case "UNKNOWN_COMMAND\r\n":
		return errUnknownCommand
	}
	return errUnknown
}

//concat two slices of []byte
func concatSlice(slc1, slc2 []byte) []byte {
	newSlc := make([]byte, len(slc1)+len(slc2))
	copy(newSlc, slc1)
	copy(newSlc[len(slc1):], slc2)
	return newSlc
}
<commit_msg>go fmt<commit_after>//Go beanstalkd client library
//Copyright(2012) Iwan Budi Kusnanto. See LICENSE for detail
package gobeanstalk

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"net"
	"strings"
)

//beanstalkd error
var (
	errOutOfMemory    = errors.New("Out of Memory")
	errInternalError  = errors.New("Internal Error")
	errBadFormat      = errors.New("Bad Format")
	errUnknownCommand = errors.New("Unknown Command")
	errBuried         = errors.New("Buried")
	errExpectedCrlf   = errors.New("Expected CRLF")
	errJobTooBig      = errors.New("Job Too Big")
	errDraining       = errors.New("Draining")
	errDeadlineSoon   = errors.New("Deadline Soon")
	errTimedOut       = errors.New("Timed Out")
	errNotFound       = errors.New("Not Found")
)

//gobeanstalk error
var (
	errInvalidLen = errors.New("Invalid Length")
	errUnknown    = errors.New("Unknown Error")
)

//Connection to beanstalkd
type Conn struct {
	conn   net.Conn
	addr   string
	reader *bufio.Reader
}

//create new connection
func NewConn(conn net.Conn, addr string) (*Conn, error) {
	c := new(Conn)
	c.conn = conn
	c.addr = addr
	c.reader = bufio.NewReader(conn)

	return c, nil
}

//A beanstalkd job
type Job struct {
	Id   uint64
	Body []byte
}

//Create new job
func NewJob(id uint64, body []byte) *Job {
	j := &Job{id, body}
	return j
}

//Connect to beanstalkd server
func Dial(addr string) (*Conn, error) {
	kon, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	c, err := NewConn(kon, addr)

	if err != nil {
		return nil, err
	}

	return c, nil
}

//Watching tube
func (c *Conn) Watch(tubename string) (int, error) {
	cmd := fmt.Sprintf("watch %s\r\n", tubename)

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

/*
Ignore tube.

The "ignore" command is for consumers. It removes the named tube from the
watch list for the current connection
*/
func (c *Conn) Ignore(tubename string) (int, error) {
	//send command and read response string
	cmd := fmt.Sprintf("ignore %s\r\n", tubename)
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	//parse response
	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		if resp == "NOT_IGNORED\r\n" {
			return -1, errors.New("Not Ignored")
		}
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

//Reserve Job
func (c *Conn) Reserve() (*Job, error) {
	//send command and read response
	resp, err := sendGetResp(c, "reserve\r\n")
	if err != nil {
		return nil, err
	}

	//parse response
	var id uint64
	var bodyLen int

	switch {
	case strings.Index(resp, "RESERVED") == 0:
		_, err = fmt.Sscanf(resp, "RESERVED %d %d\r\n", &id, &bodyLen)
		if err != nil {
			return nil, err
		}
	case resp == "DEADLINE_SOON\r\n":
		return nil, errDeadlineSoon
	case resp == "TIMED_OUT\r\n":
		return nil, errTimedOut
	default:
		return nil, parseCommonError(resp)
	}

	//read job body
	body, err := c.readBytes()
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	body = body[:len(body)-2]
	if len(body) != bodyLen {
		return nil, errors.New(fmt.Sprintf("invalid body len = %d/%d", len(body), bodyLen))
	}

	return &Job{id, body}, nil
}

//Delete a job
func (c *Conn) Delete(id uint64) error {
	cmd := fmt.Sprintf("delete %d\r\n", id)
	expected := "DELETED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Use tube

The "use" command is for producers. Subsequent put commands will put jobs into
the tube specified by this command. If no use command has been issued, jobs
will be put into the tube named "default".
*/
func (c *Conn) Use(tubename string) error {
	//check parameter
	if len(tubename) > 200 {
		return errInvalidLen
	}

	cmd := fmt.Sprintf("use %s\r\n", tubename)
	expected := fmt.Sprintf("USING %s\r\n", tubename)
	return sendExpectExact(c, cmd, expected)
}

//Put job
func (c *Conn) Put(data []byte, pri, delay, ttr int) (uint64, error) {
	cmd := fmt.Sprintf("put %d %d %d %d\r\n", pri, delay, ttr, len(data))
	cmd = cmd + string(data) + "\r\n"

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return 0, err
	}

	//parse Put response
	switch {
	case strings.Index(resp, "INSERTED") == 0:
		var id uint64
		_, parseErr := fmt.Sscanf(resp, "INSERTED %d\r\n", &id)
		return id, parseErr
	case strings.Index(resp, "BURIED") == 0:
		var id uint64
		fmt.Sscanf(resp, "BURIED %d\r\n", &id)
		return id, errBuried
	case resp == "EXPECTED_CRLF\r\n":
		return 0, errExpectedCrlf
	case resp == "JOB_TOO_BIG\r\n":
		return 0, errJobTooBig
	case resp == "DRAINING\r\n":
		return 0, errDraining
	default:
		return 0, parseCommonError(resp)
	}
	return 0, errUnknown
}

/*
Release a job.

The release command puts a reserved job back into the ready queue (and marks
its state as "ready") to be run by any client. It is normally used when the job
fails because of a transitory error.
	id is the job id to release.
	pri is a new priority to assign to the job.
	delay is an integer number of seconds to wait before putting the job in
		the ready queue. The job will be in the "delayed" state during this time.
*/
func (c *Conn) Release(id uint64, pri, delay int) error {
	cmd := fmt.Sprintf("release %d %d %d\r\n", id, pri, delay)
	expected := "RELEASED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Bury a job.

The bury command puts a job into the "buried" state. Buried jobs are put into a
FIFO linked list and will not be touched by the server again until a client
kicks them with the "kick" command.
	id is the job id to release.
	pri is a new priority to assign to the job.
*/
func (c *Conn) Bury(id uint64, pri int) error {
	cmd := fmt.Sprintf("bury %d %d\r\n", id, pri)
	expected := "BURIED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Touch a job

The "touch" command allows a worker to request more time to work on a job.
This is useful for jobs that potentially take a long time, but you still want
the benefits of a TTR pulling a job away from an unresponsive worker. A worker
may periodically tell the server that it's still alive and processing a job
(e.g. it may do this on DEADLINE_SOON)
*/
func (c *Conn) Touch(id uint64) error {
	cmd := fmt.Sprintf("touch %d\r\n", id)
	expected := "TOUCHED\r\n"
	return sendExpectExact(c, cmd, expected)
}

//send command and expect some exact response
func sendExpectExact(c *Conn, cmd, expected string) error {
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return err
	}

	if resp != expected {
		return parseCommonError(resp)
	}
	return nil
}

//Send command and read response
func sendGetResp(c *Conn, cmd string) (string, error) {
	_, err := c.conn.Write([]byte(cmd))
	if err != nil {
		return "", err
	}

	//wait for response
	resp, err := c.reader.ReadString('\n')
	if err != nil {
		return "", err
	}
	return resp, nil
}

//read bytes until \n
func (c *Conn) readBytes() ([]byte, error) {
	rsp, err := c.reader.ReadBytes('\n')
	return rsp, err
}

//parse for Common Error
func parseCommonError(str string) error {
	switch str {
	case "BURIED\r\n":
		return errBuried
	case "NOT_FOUND\r\n":
		return errNotFound
	case "OUT_OF_MEMORY\r\n":
		return errOutOfMemory
	case "INTERNAL_ERROR\r\n":
		return errInternalError
	case "BAD_FORMAT\r\n":
		return errBadFormat
	case "UNKNOWN_COMMAND\r\n":
		return errUnknownCommand
	}
	return errUnknown
}

//concat two slices of []byte
func concatSlice(slc1, slc2 []byte) []byte {
	newSlc := make([]byte, len(slc1)+len(slc2))
	copy(newSlc, slc1)
	copy(newSlc[len(slc1):], slc2)
	return newSlc
}
<|endoftext|>


<commit_before>package core

// transfer2go data core module, request implementation
// Author: Valentin Kuznetsov <vkuznet@gmail.com>

import (
	"container/heap"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/vkuznet/transfer2go/utils"
)

// AgentStatus data type
type AgentStatus struct {
	Url       string            `json:"url"`      // agent url
	Name      string            `json:"name"`     // agent name or alias
	TimeStamp int64             `json:"ts"`       // time stamp
	Catalog   string            `json:"catalog"`  // underlying TFC catalog
	Protocol  string            `json:"protocol"` // underlying transfer protocol
	Backend   string            `json:"backend"`  // underlying transfer backend
	Tool      string            `json:"tool"`     // underlying transfer tool, e.g. xrdcp
	ToolOpts  string            `json:"toolopts"` // options for backend tool
	Agents    map[string]string `json:"agents"`   // list of known agents
	Addrs     []string          `json:"addrs"`    // list of all IP addresses
	Metrics   map[string]int64  `json:"metrics"`  // agent metrics
}

// Processor is an object who process' given task
// The logic of the Processor should be implemented.
type Processor struct {
}

// Request interface defines a task process
type Request interface {
	Process(*TransferRequest) error
}

// RequestFunc is a function type that implements the Request interface
type RequestFunc func(*TransferRequest) error

// Decorator wraps a request with extra behavior
type Decorator func(Request) Request

// DefaultProcessor is a default processor instance
var DefaultProcessor = &Processor{}

// String provides string representation of given agent status
func (a *AgentStatus) String() string {
	return fmt.Sprintf("<Agent name=%s url=%s catalog=%s protocol=%s backend=%s tool=%s toolOpts=%s agents=%v addrs=%v metrics(%v)>", a.Name, a.Url, a.Catalog, a.Protocol, a.Backend, a.Tool, a.ToolOpts, a.Agents, a.Addrs, a.Metrics)
}

// Process defines execution process for a given task
func (e *Processor) Process(t *TransferRequest) error {
	return nil
}

// Process is a method of TransferRequest
func (f RequestFunc) Process(t *TransferRequest) error {
	return f(t)
}

// filleTransferRequest creates HTTP request to transfer a given file name
// https://matt.aimonetti.net/posts/2013/07/01/golang-multipart-file-upload-example/
func fileTransferRequest(c CatalogEntry, tr *TransferRequest) (*http.Response, error) {
	file, err := os.Open(c.Pfn)
	defer file.Close()
	if err != nil {
		return nil, err
	}

	pipeOut, pipeIn := io.Pipe()
	writer := multipart.NewWriter(pipeIn)
	// do the request concurrently
	var resp *http.Response
	done := make(chan error)
	go func() {
		// prepare request
		url := fmt.Sprintf("%s/upload", tr.DstUrl)
		req, err := http.NewRequest("POST", url, pipeOut)
		if err != nil {
			done <- err
			return
		}

		req.Header.Set("Content-Type", writer.FormDataContentType())
		req.Header.Set("Pfn", c.Pfn)
		req.Header.Set("Lfn", c.Lfn)
		req.Header.Set("Bytes", fmt.Sprintf("%d", c.Bytes))
		req.Header.Set("Hash", c.Hash)
		req.Header.Set("Src", tr.SrcAlias)
		req.Header.Set("Dst", tr.DstAlias)
		req.Header.Set("Content-Type", writer.FormDataContentType())
		client := utils.HttpClient()
		resp, err = client.Do(req)
		if err != nil {
			done <- err
			return
		}

		done <- nil
	}()

	part, err := writer.CreateFormFile("data", filepath.Base(c.Pfn))
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(part, file)
	if err != nil {
		return nil, err
	}

	err = writer.Close()
	if err != nil {
		return nil, err
	}

	err = pipeIn.Close()
	if err != nil {
		return nil, err
	}

	err = <-done
	if err != nil {
		return nil, err
	}

	return resp, nil
}

// helper function to perform transfer via HTTP protocol
func httpTransfer(c CatalogEntry, t *TransferRequest) (string, error) {
	// create file transfer request
	resp, err := fileTransferRequest(c, t)
	if err != nil {
		return "", err
	}

	var r CatalogEntry
	err = json.NewDecoder(resp.Body).Decode(&r)
	if err != nil {
		return "", err
	}
	return r.Pfn, nil
}

// Store returns a Decorator that stores request
func Store() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			t.Id = time.Now().Unix()
			item := &Item{
				Value:    *t,
				priority: t.Priority,
			}
			fmt.Println(*t)
			err := TFC.InsertRequest(*t)
			if err != nil {
				return err
			} else {
				heap.Push(&RequestQueue, item)
			}
			log.WithFields(log.Fields{
				"Request": t,
			}).Println("Request Saved")
			return r.Process(t)
		})
	}
}

// Delete returns a Decorator that deletes request from heap
func Delete() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			// Delete request from PriorityQueue. The complexity is O(n) where n = heap.Len()
			index := -1
			var err error

			for _, item := range RequestQueue {
				if item.Value.Id == t.Id {
					index = item.index
					break
				}
			}

			if index < RequestQueue.Len() && index >= 0 {
				err = TFC.UpdateRequest(t.Id, "deleted")
				if err != nil {
					t.Status = "error"
					return err
				} else {
					// TODO: May be we need to add lock over here.
					heap.Remove(&RequestQueue, index)
					t.Status = "deleted"
					log.WithFields(log.Fields{
						"Request": t,
					}).Println("Request Deleted")
				}
			} else {
				t.Status = "error"
				err = errors.New("Can't find request in heap")
				return err
			}

			return r.Process(t)
		})
	}
}

// Transfer returns a Decorator that performs request transfers by pull model
func PullTransfer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"Request": t.String(),
			}).Println("Request Transfer")
			// obtain information about source and destination agents
			url := fmt.Sprintf("%s/status", t.DstUrl)
			resp := utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var dstAgent AgentStatus
			err := json.Unmarshal(resp.Data, &dstAgent)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/status", t.SrcUrl)
			resp = utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var srcAgent AgentStatus
			err = json.Unmarshal(resp.Data, &srcAgent)
			if err != nil {
				return err
			}
			// if both are up then send acknowledge message to destination on /pullack url.
			body, err := json.Marshal(t)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/pull", t.DstUrl)
			resp = utils.FetchResponse(url, body)
			// check return status code
			if resp.StatusCode != 200 {
				return fmt.Errorf("Response %s, error=%s", resp.Status, string(resp.Data))
			}
			return r.Process(t)
		})
	}
}

// Transfer returns a Decorator that performs request transfers
func PushTransfer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"Request": t.String(),
			}).Println("Request Transfer")
			var records []CatalogEntry
			// Consider those requests which are failed in previous iteration.
			// If it is nil then request must be passing through first iteration.
			if t.FailedRecords != nil {
				records = t.FailedRecords
			} else {
				records = TFC.Records(*t)
			}
			if len(records) == 0 {
				// file does not exists in TFC, nothing to do, return immediately
				log.WithFields(log.Fields{
					"TransferRequest": t,
				}).Warn("Does not match anything in TFC of this agent\n", t)
				return r.Process(t)
			}
			// obtain information about source and destination agents
			url := fmt.Sprintf("%s/status", t.DstUrl)
			resp := utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var dstAgent AgentStatus
			err := json.Unmarshal(resp.Data, &dstAgent)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/status", t.SrcUrl)
			resp = utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var srcAgent AgentStatus
			err = json.Unmarshal(resp.Data, &srcAgent)
			if err != nil {
				return err
			}

			// TODO: I need to implement bulk transfer for all files in found records
			// so far I loop over them individually and transfer one by one
			var trRecords []CatalogEntry // list of successfully transferred records
			var failedRecords []CatalogEntry
			// Overwrite the previous error status
			t.Status = ""
			for _, rec := range records {

				time0 := time.Now().Unix()

				AgentMetrics.Bytes.Inc(rec.Bytes)

				// if protocol is not given use default one: HTTP
				var rpfn string // remote PFN
				if srcAgent.Protocol == "" || srcAgent.Protocol == "http" {
					log.WithFields(log.Fields{
						"dstAgent": dstAgent.String(),
					}).Println("Transfer via HTTP protocol to", dstAgent.String())
					rpfn, err = httpTransfer(rec, t)
					if err != nil {
						log.WithFields(log.Fields{
							"TransferRequest": t.String(),
							"Record":          rec.String(),
							"Err":             err,
						}).Error("Transfer", rec.String(), t.String(), err)
						t.Status = err.Error()
						failedRecords = append(failedRecords, rec)
						continue // if we fail on single record we continue with others
					}
				} else {
					// construct remote PFN by using destination agent backend and record LFN
					rpfn = fmt.Sprintf("%s%s", dstAgent.Backend, rec.Lfn)
					// perform transfer with the help of backend tool
					var cmd *exec.Cmd
					if srcAgent.ToolOpts == "" {
						cmd = exec.Command(srcAgent.Tool, rec.Pfn, rpfn)
					} else {
						cmd = exec.Command(srcAgent.Tool, srcAgent.ToolOpts, rec.Pfn, rpfn)
					}
					log.WithFields(log.Fields{
						"Command": cmd,
					}).Println("Transfer command")
					err = cmd.Run()
					if err != nil {
						log.WithFields(log.Fields{
							"Tool":         srcAgent.Tool,
							"Tool options": srcAgent.ToolOpts,
							"PFN":          rec.Pfn,
							"Remote PFN":   rpfn,
							"Err":          err,
						}).Error("Transfer")
						t.Status = err.Error()
						failedRecords = append(failedRecords, rec)
						continue // if we fail on single record we continue with others
					}
				}
				r := CatalogEntry{Dataset: rec.Dataset, Block: rec.Block, Lfn: rec.Lfn, Pfn: rpfn, Bytes: rec.Bytes, Hash: rec.Hash, TransferTime: (time.Now().Unix() - time0), Timestamp: time.Now().Unix()}
				trRecords = append(trRecords, r)

				// record how much we transferred
				AgentMetrics.TotalBytes.Inc(r.Bytes) // keep growing
				AgentMetrics.Total.Inc(1)            // keep growing
				AgentMetrics.Bytes.Dec(rec.Bytes)    // decrement since we're done

			}
			// Add entry for remote TFC after transfer is completed
			url = fmt.Sprintf("%s/tfc", t.DstUrl)
			d, e := json.Marshal(trRecords)
			if e != nil {
				return e
			}
			resp = utils.FetchResponse(url, d) // POST request
			if resp.Error != nil {
				return resp.Error
			}
			t.FailedRecords = failedRecords
			return r.Process(t)
		})
	}
}

// Logging returns a Decorator that logs client requests
func Logging(l *log.Logger) Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			l.Println("TransferRequest", t)
			return r.Process(t)
		})
	}
}

// Pause returns a Decorator that pauses request for a given time interval
func Pause(interval time.Duration) Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			if interval > 0 {
				log.WithFields(log.Fields{
					"Request":  t,
					"Interval": interval,
				}).Println("TransferRequest is paused by")
				time.Sleep(interval)
			}
			return r.Process(t)
		})
	}
}

// Tracer returns a Decorator that traces given request
func Tracer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"TransferRequest": t,
			}).Println("Trace")
			return r.Process(t)
		})
	}
}

// Decorate decorates a Request r with all given Decorators
func Decorate(r Request, ds ...Decorator) Request {
	decorated := r
	for _, decorate := range ds {
		decorated = decorate(decorated)
	}
	return decorated
}
<commit_msg>Close body<commit_after>package core

// transfer2go data core module, request implementation
// Author: Valentin Kuznetsov <vkuznet@gmail.com>

import (
	"container/heap"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/vkuznet/transfer2go/utils"
)

// AgentStatus data type
type AgentStatus struct {
	Url       string            `json:"url"`      // agent url
	Name      string            `json:"name"`     // agent name or alias
	TimeStamp int64             `json:"ts"`       // time stamp
	Catalog   string            `json:"catalog"`  // underlying TFC catalog
	Protocol  string            `json:"protocol"` // underlying transfer protocol
	Backend   string            `json:"backend"`  // underlying transfer backend
	Tool      string            `json:"tool"`     // underlying transfer tool, e.g. xrdcp
	ToolOpts  string            `json:"toolopts"` // options for backend tool
	Agents    map[string]string `json:"agents"`   // list of known agents
	Addrs     []string          `json:"addrs"`    // list of all IP addresses
	Metrics   map[string]int64  `json:"metrics"`  // agent metrics
}

// Processor is an object who process' given task
// The logic of the Processor should be implemented.
type Processor struct {
}

// Request interface defines a task process
type Request interface {
	Process(*TransferRequest) error
}

// RequestFunc is a function type that implements the Request interface
type RequestFunc func(*TransferRequest) error

// Decorator wraps a request with extra behavior
type Decorator func(Request) Request

// DefaultProcessor is a default processor instance
var DefaultProcessor = &Processor{}

// String provides string representation of given agent status
func (a *AgentStatus) String() string {
	return fmt.Sprintf("<Agent name=%s url=%s catalog=%s protocol=%s backend=%s tool=%s toolOpts=%s agents=%v addrs=%v metrics(%v)>", a.Name, a.Url, a.Catalog, a.Protocol, a.Backend, a.Tool, a.ToolOpts, a.Agents, a.Addrs, a.Metrics)
}

// Process defines execution process for a given task
func (e *Processor) Process(t *TransferRequest) error {
	return nil
}

// Process is a method of TransferRequest
func (f RequestFunc) Process(t *TransferRequest) error {
	return f(t)
}

// filleTransferRequest creates HTTP request to transfer a given file name
// https://matt.aimonetti.net/posts/2013/07/01/golang-multipart-file-upload-example/
func fileTransferRequest(c CatalogEntry, tr *TransferRequest) (*http.Response, error) {
	file, err := os.Open(c.Pfn)
	defer file.Close()
	if err != nil {
		return nil, err
	}

	pipeOut, pipeIn := io.Pipe()
	writer := multipart.NewWriter(pipeIn)
	// do the request concurrently
	var resp *http.Response
	done := make(chan error)
	go func() {
		// prepare request
		url := fmt.Sprintf("%s/upload", tr.DstUrl)
		req, err := http.NewRequest("POST", url, pipeOut)
		if err != nil {
			done <- err
			return
		}

		req.Header.Set("Content-Type", writer.FormDataContentType())
		req.Header.Set("Pfn", c.Pfn)
		req.Header.Set("Lfn", c.Lfn)
		req.Header.Set("Bytes", fmt.Sprintf("%d", c.Bytes))
		req.Header.Set("Hash", c.Hash)
		req.Header.Set("Src", tr.SrcAlias)
		req.Header.Set("Dst", tr.DstAlias)
		req.Header.Set("Content-Type", writer.FormDataContentType())
		client := utils.HttpClient()
		resp, err = client.Do(req)
		defer resp.Body.Close()
		if err != nil {
			done <- err
			return
		}

		done <- nil
	}()

	part, err := writer.CreateFormFile("data", filepath.Base(c.Pfn))
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(part, file)
	if err != nil {
		return nil, err
	}

	err = writer.Close()
	if err != nil {
		return nil, err
	}

	err = pipeIn.Close()
	if err != nil {
		return nil, err
	}

	err = <-done
	if err != nil {
		return nil, err
	}

	return resp, nil
}

// helper function to perform transfer via HTTP protocol
func httpTransfer(c CatalogEntry, t *TransferRequest) (string, error) {
	// create file transfer request
	resp, err := fileTransferRequest(c, t)
	if err != nil {
		return "", err
	}

	var r CatalogEntry
	err = json.NewDecoder(resp.Body).Decode(&r)
	if err != nil {
		return "", err
	}
	return r.Pfn, nil
}

// Store returns a Decorator that stores request
func Store() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			t.Id = time.Now().Unix()
			item := &Item{
				Value:    *t,
				priority: t.Priority,
			}
			fmt.Println(*t)
			err := TFC.InsertRequest(*t)
			if err != nil {
				return err
			} else {
				heap.Push(&RequestQueue, item)
			}
			log.WithFields(log.Fields{
				"Request": t,
			}).Println("Request Saved")
			return r.Process(t)
		})
	}
}

// Delete returns a Decorator that deletes request from heap
func Delete() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			// Delete request from PriorityQueue. The complexity is O(n) where n = heap.Len()
			index := -1
			var err error

			for _, item := range RequestQueue {
				if item.Value.Id == t.Id {
					index = item.index
					break
				}
			}

			if index < RequestQueue.Len() && index >= 0 {
				err = TFC.UpdateRequest(t.Id, "deleted")
				if err != nil {
					t.Status = "error"
					return err
				} else {
					// TODO: May be we need to add lock over here.
					heap.Remove(&RequestQueue, index)
					t.Status = "deleted"
					log.WithFields(log.Fields{
						"Request": t,
					}).Println("Request Deleted")
				}
			} else {
				t.Status = "error"
				err = errors.New("Can't find request in heap")
				return err
			}

			return r.Process(t)
		})
	}
}

// Transfer returns a Decorator that performs request transfers by pull model
func PullTransfer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"Request": t.String(),
			}).Println("Request Transfer")
			// obtain information about source and destination agents
			url := fmt.Sprintf("%s/status", t.DstUrl)
			resp := utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var dstAgent AgentStatus
			err := json.Unmarshal(resp.Data, &dstAgent)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/status", t.SrcUrl)
			resp = utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var srcAgent AgentStatus
			err = json.Unmarshal(resp.Data, &srcAgent)
			if err != nil {
				return err
			}
			// if both are up then send acknowledge message to destination on /pullack url.
			body, err := json.Marshal(t)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/pull", t.DstUrl)
			resp = utils.FetchResponse(url, body)
			// check return status code
			if resp.StatusCode != 200 {
				return fmt.Errorf("Response %s, error=%s", resp.Status, string(resp.Data))
			}
			return r.Process(t)
		})
	}
}

// Transfer returns a Decorator that performs request transfers
func PushTransfer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"Request": t.String(),
			}).Println("Request Transfer")
			var records []CatalogEntry
			// Consider those requests which are failed in previous iteration.
			// If it is nil then request must be passing through first iteration.
			if t.FailedRecords != nil {
				records = t.FailedRecords
			} else {
				records = TFC.Records(*t)
			}
			if len(records) == 0 {
				// file does not exists in TFC, nothing to do, return immediately
				log.WithFields(log.Fields{
					"TransferRequest": t,
				}).Warn("Does not match anything in TFC of this agent\n", t)
				return r.Process(t)
			}
			// obtain information about source and destination agents
			url := fmt.Sprintf("%s/status", t.DstUrl)
			resp := utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var dstAgent AgentStatus
			err := json.Unmarshal(resp.Data, &dstAgent)
			if err != nil {
				return err
			}
			url = fmt.Sprintf("%s/status", t.SrcUrl)
			resp = utils.FetchResponse(url, []byte{})
			if resp.Error != nil {
				return resp.Error
			}
			var srcAgent AgentStatus
			err = json.Unmarshal(resp.Data, &srcAgent)
			if err != nil {
				return err
			}

			// TODO: I need to implement bulk transfer for all files in found records
			// so far I loop over them individually and transfer one by one
			var trRecords []CatalogEntry // list of successfully transferred records
			var failedRecords []CatalogEntry
			// Overwrite the previous error status
			t.Status = ""
			for _, rec := range records {

				time0 := time.Now().Unix()

				AgentMetrics.Bytes.Inc(rec.Bytes)

				// if protocol is not given use default one: HTTP
				var rpfn string // remote PFN
				if srcAgent.Protocol == "" || srcAgent.Protocol == "http" {
					log.WithFields(log.Fields{
						"dstAgent": dstAgent.String(),
					}).Println("Transfer via HTTP protocol to", dstAgent.String())
					rpfn, err = httpTransfer(rec, t)
					if err != nil {
						log.WithFields(log.Fields{
							"TransferRequest": t.String(),
							"Record":          rec.String(),
							"Err":             err,
						}).Error("Transfer", rec.String(), t.String(), err)
						t.Status = err.Error()
						failedRecords = append(failedRecords, rec)
						continue // if we fail on single record we continue with others
					}
				} else {
					// construct remote PFN by using destination agent backend and record LFN
					rpfn = fmt.Sprintf("%s%s", dstAgent.Backend, rec.Lfn)
					// perform transfer with the help of backend tool
					var cmd *exec.Cmd
					if srcAgent.ToolOpts == "" {
						cmd = exec.Command(srcAgent.Tool, rec.Pfn, rpfn)
					} else {
						cmd = exec.Command(srcAgent.Tool, srcAgent.ToolOpts, rec.Pfn, rpfn)
					}
					log.WithFields(log.Fields{
						"Command": cmd,
					}).Println("Transfer command")
					err = cmd.Run()
					if err != nil {
						log.WithFields(log.Fields{
							"Tool":         srcAgent.Tool,
							"Tool options": srcAgent.ToolOpts,
							"PFN":          rec.Pfn,
							"Remote PFN":   rpfn,
							"Err":          err,
						}).Error("Transfer")
						t.Status = err.Error()
						failedRecords = append(failedRecords, rec)
						continue // if we fail on single record we continue with others
					}
				}
				r := CatalogEntry{Dataset: rec.Dataset, Block: rec.Block, Lfn: rec.Lfn, Pfn: rpfn, Bytes: rec.Bytes, Hash: rec.Hash, TransferTime: (time.Now().Unix() - time0), Timestamp: time.Now().Unix()}
				trRecords = append(trRecords, r)

				// record how much we transferred
				AgentMetrics.TotalBytes.Inc(r.Bytes) // keep growing
				AgentMetrics.Total.Inc(1)            // keep growing
				AgentMetrics.Bytes.Dec(rec.Bytes)    // decrement since we're done

			}
			// Add entry for remote TFC after transfer is completed
			url = fmt.Sprintf("%s/tfc", t.DstUrl)
			d, e := json.Marshal(trRecords)
			if e != nil {
				return e
			}
			resp = utils.FetchResponse(url, d) // POST request
			if resp.Error != nil {
				return resp.Error
			}
			t.FailedRecords = failedRecords
			return r.Process(t)
		})
	}
}

// Logging returns a Decorator that logs client requests
func Logging(l *log.Logger) Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			l.Println("TransferRequest", t)
			return r.Process(t)
		})
	}
}

// Pause returns a Decorator that pauses request for a given time interval
func Pause(interval time.Duration) Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			if interval > 0 {
				log.WithFields(log.Fields{
					"Request":  t,
					"Interval": interval,
				}).Println("TransferRequest is paused by")
				time.Sleep(interval)
			}
			return r.Process(t)
		})
	}
}

// Tracer returns a Decorator that traces given request
func Tracer() Decorator {
	return func(r Request) Request {
		return RequestFunc(func(t *TransferRequest) error {
			log.WithFields(log.Fields{
				"TransferRequest": t,
			}).Println("Trace")
			return r.Process(t)
		})
	}
}

// Decorate decorates a Request r with all given Decorators
func Decorate(r Request, ds ...Decorator) Request {
	decorated := r
	for _, decorate := range ds {
		decorated = decorate(decorated)
	}
	return decorated
}
<|endoftext|>


<commit_before>package gonets

import "testing"

func TestBasicNetDirect(t *testing.T) {
	inp_l := Layer(inputLayer{3})
	sfm := Layer(&hiddenLayer{
		u: []unit{
			unit{0.0, []float64{0.3, 0.4, -0.1}},
			unit{0.1, []float64{-0.1, 0.1, 0.0}},
			unit{-0.3, []float64{-0.3, 0.0, 0.0}},
		},
		a: Softmax,
	})
	sfm.AttachInputs(inp_l)
	o, err := sfm.GetOutput([]float64{3, 2, 1})
	if err != nil {
		t.Fail()
	}
	tgt := []float64{0.7919496226186542, 0.1598918712338877, 0.048158506147458105}
	for i := range tgt {
		if o[i] != tgt[i] {
			t.Fail()
		}
	}
}

func TestBasicNetConstructors(t *testing.T) {
	biases := []float64{0.0, 0.1, -0.3}
	weights := [][]float64{
		{0.3, 0.4, -0.1},
		{-0.1, 0.1, 0.0},
		{-0.3, 0.0, 0.0},
	}
	sfm, err := SoftmaxLayer(InputLayer(3), biases, weights)
	if err != nil {
		t.Fatal(err)
	}
	o, err := sfm.GetOutput([]float64{3, 2, 1})
	if err != nil {
		t.Fatal(err)
	}
	tgt := []float64{0.7919496226186542, 0.1598918712338877, 0.048158506147458105}
	for i := range tgt {
		if o[i] != tgt[i] {
			t.Fatalf("%+v != %+v", o[i], tgt[i])
		}
	}
}
<commit_msg>complex real-world Lasagne (Theano-based ANN package) compatibility test<commit_after>package gonets

import (
	"math"
	"testing"
)

func TestBasicNetDirect(t *testing.T) {
	inp_l := Layer(inputLayer{3})
	sfm := Layer(&hiddenLayer{
		u: []unit{
			unit{0.0, []float64{0.3, 0.4, -0.1}},
			unit{0.1, []float64{-0.1, 0.1, 0.0}},
			unit{-0.3, []float64{-0.3, 0.0, 0.0}},
		},
		a: Softmax,
	})
	sfm.AttachInputs(inp_l)
	o, err := sfm.GetOutput([]float64{3, 2, 1})
	if err != nil {
		t.Fail()
	}
	tgt := []float64{0.7919496226186542, 0.1598918712338877, 0.048158506147458105}
	for i := range tgt {
		if o[i] != tgt[i] {
			t.Fail()
		}
	}
}

func TestBasicNetConstructors(t *testing.T) {
	biases := []float64{0.0, 0.1, -0.3}
	weights := [][]float64{
		{0.3, 0.4, -0.1},
		{-0.1, 0.1, 0.0},
		{-0.3, 0.0, 0.0},
	}
	sfm, err := SoftmaxLayer(InputLayer(3), biases, weights)
	if err != nil {
		t.Fatal(err)
	}
	o, err := sfm.GetOutput([]float64{3, 2, 1})
	if err != nil {
		t.Fatal(err)
	}
	tgt := []float64{0.7919496226186542, 0.1598918712338877, 0.048158506147458105}
	for i := range tgt {
		if o[i] != tgt[i] {
			t.Fatalf("%+v != %+v", o[i], tgt[i])
		}
	}
}

func TestLasagneCompatibility(t *testing.T) {
	w1 := [][]float64{
		[]float64{-0.109595783122, 0.209483981573, 0.122117914097, 0.0716140910105, 0.112168624432, 0.0159258536371, -0.0151195257045, 0.0962302240911, 0.0458310692353, -0.0257266391724, -0.422622765543, -0.0165887479738, -0.218176051145, -0.317418433196, -0.179569247066, -0.0914008534842, -0.191939092518, -0.000108985369519, -0.174754967039, 0.160224194794, -0.0639363286407, -0.129051442081, 0.386961666847, 0.0150403736405, -0.70029036513, 0.460468224536, 0.121088748999, 0.117098491034},
		[]float64{0.211631211691, 0.134569394477, 0.115913647831, 0.0783742465631, 0.0413851115954, 0.0479523717229, 0.0475213332174, 0.122535496511, 0.0647260567506, 0.0886193293738, -0.0252346973127, -6.73335956438e-05, -0.0331124776276, -0.0985513085794, -0.0312408551554, -0.220824313021, -0.0558351858373, -0.0212985755574, 0.0195726728305, 0.053837776281, 0.0262370827495, 0.019485934755, 0.215210739953, 0.0243230441865, 0.0852560049886, 0.139319097071, -1.81127046853, 0.0507853151942},
		[]float64{-0.0835754939928, -0.375922688896, 0.0487980670745, -0.0112866143314, 0.0967032912317, 0.0316108126655, 0.0475973181855, 0.0341371087019, 0.00719269855435, -0.0686237820401, -0.181060805933, -0.0489050970871, -0.0851999611865, -0.0714608505252, -0.193938304972, 0.135735304938, -0.215817260336, -0.0372146245005, 0.102994743277, 0.0554789320264, -0.0611136747927, -0.0526758103667, 0.307219039928, 0.11109501764, -0.478557124179, 0.0555677729173, -1.39206529973, -0.145246469005},
		[]float64{0.729282052301, -1.07478893256, 0.438879343319, 0.0813187884668, 0.145962614996, 0.132920880587, 0.285901713238, 0.167462902015, 0.200951338505, 0.218269255499, -0.429428320235, -0.00972881578782, 0.0930611149157, 0.157340375191, 0.407451641692, 0.0276353807711, -0.0996143652946, -0.0851605921999, -0.0574387955867, 0.229387382637, 0.00450069390171, -0.0250284174177, 0.369870997728, -0.0533345239895, -0.95769297787, 0.603201435829, 0.015390587776, 0.31935992233},
		[]float64{0.0764285076257, -0.0972343313161, 0.0613440747912, 0.140477699738, 0.151804130187, 0.110119209635, 0.177345192999, 0.237278875926, 0.0927115482558, 0.057776602444, 0.0443569202995, -0.12813757285, -0.112509774897, -0.246624945246, -0.537158679623, 0.0431268143683, -0.277826764246, -0.248726000094, 0.0981974202112, -0.0377599799245, -0.00381517759309, 0.0689059526618, 0.7291501788, 0.0321694999756, 0.379698120135, -0.189547451852, 0.0753656864124, -0.469598580609},
	}
	b1 := []float64{1.06529159066, -0.271159499769, -0.360616089225, 0.758574607807, 0.18340392505}
	w2 := [][]float64{
		[]float64{0.617280342307, -0.929100661125, -0.647544732776, 0.621824740939, 1.05812547924},
		[]float64{-0.210913081076, 0.930927754667, 0.895306960904, -0.056629859462, 0.194744500807},
	}
	b2 := []float64{1.03980357614, -1.03980357614}
	l0 := InputLayer(28)
	l1, err := RectifyLayer(l0, b1, w1)
	if err != nil {
		t.Fatal(err)
	}
	l2, err := SoftmaxLayer(l1, b2, w2)
	if err != nil {
		t.Fatal(err)
	}

	inp := []float64{-0.311443752632, -0.33276644248, -0.279749856683, -0.23357894362, -0.388451400757, -0.290115813651, -0.288566611959, -0.486674811383, -0.336321895253, 2.87178876171, -0.0264714918414, -0.470125525019, -0.376509260503, 1.92171443031, -0.327773601422, -0.134484865699, -0.517244623798, -0.670748122659, 1.30312227345, -0.293501652974, -1.11866423067, -0.646212635966, 2.69934924792, 1.3389814651, -1.05340000685, 0.827605976572, 1.79816100342, 0.739421624031}
	target := []float64{9.99812890e-01, 1.87110350e-04}

	out, err := l2.GetOutput(inp)
	if err != nil {
		t.Fatal(err)
	}
	approx := func(a, b float64) bool {
		if math.Abs(a-b) <= 0.01 {
			return true
		}
		return false
	}
	for i := range out {
		if !approx(target[i], out[i]) {
			t.Fatalf("%+v != %+v", out[i], target[i])
		}
	}
}
<|endoftext|>


<commit_before>// Package gitlab provides authentication strategies using Gitlab.
package gitlab

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"

	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/groups"
	"github.com/dexidp/dex/pkg/log"
)

const (
	// read operations of the /api/v4/user endpoint
	scopeUser = "read_user"
	// used to retrieve groups from /oauth/userinfo
	// https://docs.gitlab.com/ee/integration/openid_connect_provider.html
	scopeOpenID = "openid"
)

// Config holds configuration options for gilab logins.
type Config struct {
	BaseURL      string   `json:"baseURL"`
	ClientID     string   `json:"clientID"`
	ClientSecret string   `json:"clientSecret"`
	RedirectURI  string   `json:"redirectURI"`
	Groups       []string `json:"groups"`
	UseLoginAsID bool     `json:"useLoginAsID"`
}

type gitlabUser struct {
	ID       int
	Name     string
	Username string
	State    string
	Email    string
	IsAdmin  bool
}

// Open returns a strategy for logging in through GitLab.
func (c *Config) Open(id string, logger log.Logger) (connector.Connector, error) {
	if c.BaseURL == "" {
		c.BaseURL = "https://gitlab.com"
	}
	return &gitlabConnector{
		baseURL:      c.BaseURL,
		redirectURI:  c.RedirectURI,
		clientID:     c.ClientID,
		clientSecret: c.ClientSecret,
		logger:       logger,
		groups:       c.Groups,
		useLoginAsID: c.UseLoginAsID,
	}, nil
}

type connectorData struct {
	// GitLab's OAuth2 tokens never expire. We don't need a refresh token.
	AccessToken string `json:"accessToken"`
}

var (
	_ connector.CallbackConnector = (*gitlabConnector)(nil)
	_ connector.RefreshConnector  = (*gitlabConnector)(nil)
)

type gitlabConnector struct {
	baseURL      string
	redirectURI  string
	groups       []string
	clientID     string
	clientSecret string
	logger       log.Logger
	httpClient   *http.Client
	// if set to true will use the user's handle rather than their numeric id as the ID
	useLoginAsID bool
}

func (c *gitlabConnector) oauth2Config(scopes connector.Scopes) *oauth2.Config {
	gitlabScopes := []string{scopeUser}
	if scopes.Groups {
		gitlabScopes = []string{scopeUser, scopeOpenID}
	}

	gitlabEndpoint := oauth2.Endpoint{AuthURL: c.baseURL + "/oauth/authorize", TokenURL: c.baseURL + "/oauth/token"}
	return &oauth2.Config{
		ClientID:     c.clientID,
		ClientSecret: c.clientSecret,
		Endpoint:     gitlabEndpoint,
		Scopes:       gitlabScopes,
		RedirectURL:  c.redirectURI,
	}
}

func (c *gitlabConnector) LoginURL(scopes connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", c.redirectURI, callbackURL)
	}
	return c.oauth2Config(scopes).AuthCodeURL(state), nil
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}

func (c *gitlabConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}

	oauth2Config := c.oauth2Config(s)

	ctx := r.Context()
	if c.httpClient != nil {
		ctx = context.WithValue(r.Context(), oauth2.HTTPClient, c.httpClient)
	}

	token, err := oauth2Config.Exchange(ctx, q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("gitlab: failed to get token: %v", err)
	}

	client := oauth2Config.Client(ctx, token)

	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("gitlab: get user: %v", err)
	}

	username := user.Name
	if username == "" {
		username = user.Email
	}
	identity = connector.Identity{
		UserID:        strconv.Itoa(user.ID),
		Username:      username,
		Email:         user.Email,
		EmailVerified: true,
	}
	if c.useLoginAsID {
		identity.UserID = user.Username
	}

	if s.Groups {
		groups, err := c.getGroups(ctx, client, s.Groups, user.Username)
		if err != nil {
			return identity, fmt.Errorf("gitlab: get groups: %v", err)
		}
		identity.Groups = groups
	}

	if s.OfflineAccess {
		data := connectorData{AccessToken: token.AccessToken}
		connData, err := json.Marshal(data)
		if err != nil {
			return identity, fmt.Errorf("marshal connector data: %v", err)
		}
		identity.ConnectorData = connData
	}

	return identity, nil
}

func (c *gitlabConnector) Refresh(ctx context.Context, s connector.Scopes, ident connector.Identity) (connector.Identity, error) {
	if len(ident.ConnectorData) == 0 {
		return ident, errors.New("no upstream access token found")
	}

	var data connectorData
	if err := json.Unmarshal(ident.ConnectorData, &data); err != nil {
		return ident, fmt.Errorf("gitlab: unmarshal access token: %v", err)
	}

	client := c.oauth2Config(s).Client(ctx, &oauth2.Token{AccessToken: data.AccessToken})
	user, err := c.user(ctx, client)
	if err != nil {
		return ident, fmt.Errorf("gitlab: get user: %v", err)
	}

	username := user.Name
	if username == "" {
		username = user.Email
	}
	ident.Username = username
	ident.Email = user.Email

	if s.Groups {
		groups, err := c.getGroups(ctx, client, s.Groups, user.Username)
		if err != nil {
			return ident, fmt.Errorf("gitlab: get groups: %v", err)
		}
		ident.Groups = groups
	}
	return ident, nil
}

// user queries the GitLab API for profile information using the provided client. The HTTP
// client is expected to be constructed by the golang.org/x/oauth2 package, which inserts
// a bearer token as part of the request.
func (c *gitlabConnector) user(ctx context.Context, client *http.Client) (gitlabUser, error) {
	var u gitlabUser
	req, err := http.NewRequest("GET", c.baseURL+"/api/v4/user", nil)
	if err != nil {
		return u, fmt.Errorf("gitlab: new req: %v", err)
	}
	req = req.WithContext(ctx)
	resp, err := client.Do(req)
	if err != nil {
		return u, fmt.Errorf("gitlab: get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return u, fmt.Errorf("gitlab: read body: %v", err)
		}
		return u, fmt.Errorf("%s: %s", resp.Status, body)
	}

	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return u, fmt.Errorf("failed to decode response: %v", err)
	}
	return u, nil
}

type userInfo struct {
	Groups []string
}

// userGroups queries the GitLab API for group membership.
//
// The HTTP passed client is expected to be constructed by the golang.org/x/oauth2 package,
// which inserts a bearer token as part of the request.
func (c *gitlabConnector) userGroups(ctx context.Context, client *http.Client) ([]string, error) {
	req, err := http.NewRequest("GET", c.baseURL+"/oauth/userinfo", nil)
	if err != nil {
		return nil, fmt.Errorf("gitlab: new req: %v", err)
	}
	req = req.WithContext(ctx)
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("gitlab: get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("gitlab: read body: %v", err)
		}
		return nil, fmt.Errorf("%s: %s", resp.Status, body)
	}
	var u userInfo
	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return nil, fmt.Errorf("failed to decode response: %v", err)
	}

	return u.Groups, nil
}

func (c *gitlabConnector) getGroups(ctx context.Context, client *http.Client, groupScope bool, userLogin string) ([]string, error) {
	gitlabGroups, err := c.userGroups(ctx, client)
	if err != nil {
		return nil, err
	}

	if len(c.groups) > 0 {
		filteredGroups := groups.Filter(gitlabGroups, c.groups)
		if len(filteredGroups) == 0 {
			return nil, fmt.Errorf("gitlab: user %q is not in any of the required groups", userLogin)
		}
		return filteredGroups, nil
	} else if groupScope {
		return gitlabGroups, nil
	}

	return nil, nil
}
<commit_msg>gitlab: add groups scope by default when filtering is requested<commit_after>// Package gitlab provides authentication strategies using Gitlab.
package gitlab

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"

	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/groups"
	"github.com/dexidp/dex/pkg/log"
)

const (
	// read operations of the /api/v4/user endpoint
	scopeUser = "read_user"
	// used to retrieve groups from /oauth/userinfo
	// https://docs.gitlab.com/ee/integration/openid_connect_provider.html
	scopeOpenID = "openid"
)

// Config holds configuration options for gilab logins.
type Config struct {
	BaseURL      string   `json:"baseURL"`
	ClientID     string   `json:"clientID"`
	ClientSecret string   `json:"clientSecret"`
	RedirectURI  string   `json:"redirectURI"`
	Groups       []string `json:"groups"`
	UseLoginAsID bool     `json:"useLoginAsID"`
}

type gitlabUser struct {
	ID       int
	Name     string
	Username string
	State    string
	Email    string
	IsAdmin  bool
}

// Open returns a strategy for logging in through GitLab.
func (c *Config) Open(id string, logger log.Logger) (connector.Connector, error) {
	if c.BaseURL == "" {
		c.BaseURL = "https://gitlab.com"
	}
	return &gitlabConnector{
		baseURL:      c.BaseURL,
		redirectURI:  c.RedirectURI,
		clientID:     c.ClientID,
		clientSecret: c.ClientSecret,
		logger:       logger,
		groups:       c.Groups,
		useLoginAsID: c.UseLoginAsID,
	}, nil
}

type connectorData struct {
	// GitLab's OAuth2 tokens never expire. We don't need a refresh token.
	AccessToken string `json:"accessToken"`
}

var (
	_ connector.CallbackConnector = (*gitlabConnector)(nil)
	_ connector.RefreshConnector  = (*gitlabConnector)(nil)
)

type gitlabConnector struct {
	baseURL      string
	redirectURI  string
	groups       []string
	clientID     string
	clientSecret string
	logger       log.Logger
	httpClient   *http.Client
	// if set to true will use the user's handle rather than their numeric id as the ID
	useLoginAsID bool
}

func (c *gitlabConnector) oauth2Config(scopes connector.Scopes) *oauth2.Config {
	gitlabScopes := []string{scopeUser}
	if c.groupsRequired(scopes.Groups) {
		gitlabScopes = []string{scopeUser, scopeOpenID}
	}

	gitlabEndpoint := oauth2.Endpoint{AuthURL: c.baseURL + "/oauth/authorize", TokenURL: c.baseURL + "/oauth/token"}
	return &oauth2.Config{
		ClientID:     c.clientID,
		ClientSecret: c.clientSecret,
		Endpoint:     gitlabEndpoint,
		Scopes:       gitlabScopes,
		RedirectURL:  c.redirectURI,
	}
}

func (c *gitlabConnector) LoginURL(scopes connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", c.redirectURI, callbackURL)
	}
	return c.oauth2Config(scopes).AuthCodeURL(state), nil
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}

func (c *gitlabConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}

	oauth2Config := c.oauth2Config(s)

	ctx := r.Context()
	if c.httpClient != nil {
		ctx = context.WithValue(r.Context(), oauth2.HTTPClient, c.httpClient)
	}

	token, err := oauth2Config.Exchange(ctx, q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("gitlab: failed to get token: %v", err)
	}

	client := oauth2Config.Client(ctx, token)

	user, err := c.user(ctx, client)
	if err != nil {
		return identity, fmt.Errorf("gitlab: get user: %v", err)
	}

	username := user.Name
	if username == "" {
		username = user.Email
	}
	identity = connector.Identity{
		UserID:        strconv.Itoa(user.ID),
		Username:      username,
		Email:         user.Email,
		EmailVerified: true,
	}
	if c.useLoginAsID {
		identity.UserID = user.Username
	}

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, s.Groups, user.Username)
		if err != nil {
			return identity, fmt.Errorf("gitlab: get groups: %v", err)
		}
		identity.Groups = groups
	}

	if s.OfflineAccess {
		data := connectorData{AccessToken: token.AccessToken}
		connData, err := json.Marshal(data)
		if err != nil {
			return identity, fmt.Errorf("marshal connector data: %v", err)
		}
		identity.ConnectorData = connData
	}

	return identity, nil
}

func (c *gitlabConnector) Refresh(ctx context.Context, s connector.Scopes, ident connector.Identity) (connector.Identity, error) {
	if len(ident.ConnectorData) == 0 {
		return ident, errors.New("no upstream access token found")
	}

	var data connectorData
	if err := json.Unmarshal(ident.ConnectorData, &data); err != nil {
		return ident, fmt.Errorf("gitlab: unmarshal access token: %v", err)
	}

	client := c.oauth2Config(s).Client(ctx, &oauth2.Token{AccessToken: data.AccessToken})
	user, err := c.user(ctx, client)
	if err != nil {
		return ident, fmt.Errorf("gitlab: get user: %v", err)
	}

	username := user.Name
	if username == "" {
		username = user.Email
	}
	ident.Username = username
	ident.Email = user.Email

	if c.groupsRequired(s.Groups) {
		groups, err := c.getGroups(ctx, client, s.Groups, user.Username)
		if err != nil {
			return ident, fmt.Errorf("gitlab: get groups: %v", err)
		}
		ident.Groups = groups
	}
	return ident, nil
}

func (c *gitlabConnector) groupsRequired(groupScope bool) bool {
	return len(c.groups) > 0 || groupScope
}

// user queries the GitLab API for profile information using the provided client. The HTTP
// client is expected to be constructed by the golang.org/x/oauth2 package, which inserts
// a bearer token as part of the request.
func (c *gitlabConnector) user(ctx context.Context, client *http.Client) (gitlabUser, error) {
	var u gitlabUser
	req, err := http.NewRequest("GET", c.baseURL+"/api/v4/user", nil)
	if err != nil {
		return u, fmt.Errorf("gitlab: new req: %v", err)
	}
	req = req.WithContext(ctx)
	resp, err := client.Do(req)
	if err != nil {
		return u, fmt.Errorf("gitlab: get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return u, fmt.Errorf("gitlab: read body: %v", err)
		}
		return u, fmt.Errorf("%s: %s", resp.Status, body)
	}

	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return u, fmt.Errorf("failed to decode response: %v", err)
	}
	return u, nil
}

type userInfo struct {
	Groups []string
}

// userGroups queries the GitLab API for group membership.
//
// The HTTP passed client is expected to be constructed by the golang.org/x/oauth2 package,
// which inserts a bearer token as part of the request.
func (c *gitlabConnector) userGroups(ctx context.Context, client *http.Client) ([]string, error) {
	req, err := http.NewRequest("GET", c.baseURL+"/oauth/userinfo", nil)
	if err != nil {
		return nil, fmt.Errorf("gitlab: new req: %v", err)
	}
	req = req.WithContext(ctx)
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("gitlab: get URL %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("gitlab: read body: %v", err)
		}
		return nil, fmt.Errorf("%s: %s", resp.Status, body)
	}
	var u userInfo
	if err := json.NewDecoder(resp.Body).Decode(&u); err != nil {
		return nil, fmt.Errorf("failed to decode response: %v", err)
	}

	return u.Groups, nil
}

func (c *gitlabConnector) getGroups(ctx context.Context, client *http.Client, groupScope bool, userLogin string) ([]string, error) {
	gitlabGroups, err := c.userGroups(ctx, client)
	if err != nil {
		return nil, err
	}

	if len(c.groups) > 0 {
		filteredGroups := groups.Filter(gitlabGroups, c.groups)
		if len(filteredGroups) == 0 {
			return nil, fmt.Errorf("gitlab: user %q is not in any of the required groups", userLogin)
		}
		return filteredGroups, nil
	} else if groupScope {
		return gitlabGroups, nil
	}

	return nil, nil
}
<|endoftext|>