<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.

// +build cgo

package webp

import (
	"errors"
	"image"
	"image/color"
	"io"
	"io/ioutil"
	"os"
)

func LoadConfig(name string, cbuf ...CBuffer) (config image.Config, err error) {
	f, err := os.Open(name)
	if err != nil {
		return image.Config{}, err
	}
	defer f.Close()

	if len(cbuf) == 0 || cbuf[0] == nil {
		cbuf = []CBuffer{NewCBuffer(maxWebpHeaderSize)}
		defer cbuf[0].Close()
	}
	if len(cbuf[0].CData()) == 0 {
		if err = cbuf[0].Resize(maxWebpHeaderSize); err != nil {
			return
		}
	}

	header := cbuf[0].CData()
	if len(header) > maxWebpHeaderSize {
		header = header[:maxWebpHeaderSize]
	}
	n, err := f.Read(header)
	if err != nil && err != io.EOF {
		return
	}
	header = header[:n]
	width, height, _, err := GetInfoEx(header, cbuf[0])
	if err != nil {
		return
	}

	config.Width = width
	config.Height = height
	config.ColorModel = color.RGBAModel
	return
}

func Load(name string, cbuf ...CBuffer) (m image.Image, err error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}
	if fi.Size() > (2 << 30) {
		return nil, errors.New("webp: Load, file size is too large (> 2GB)!")
	}
	if len(cbuf) == 0 || cbuf[0] == nil {
		cbuf = []CBuffer{NewCBuffer(int(fi.Size()))}
		defer cbuf[0].Close()
	}

	if len(cbuf[0].CData()) < int(fi.Size()) {
		if err = cbuf[0].Resize(int(fi.Size())); err != nil {
			return
		}
	}
	data := cbuf[0].CData()
	if n := int(fi.Size()); len(data) > n {
		data = data[:n]
	}

	_, err = f.Read(data)
	if err != nil {
		return nil, err
	}
	m, pix, err := DecodeRGBAEx(data, cbuf[0])
	if err != nil {
		return
	}
	pix.Close()

	return m, nil
}

// DecodeConfig returns the color model and dimensions of a WEBP image without
// decoding the entire image.
func DecodeConfig(r io.Reader) (config image.Config, err error) {
	header := make([]byte, maxWebpHeaderSize)
	n, err := r.Read(header)
	if err != nil && err != io.EOF {
		return
	}
	header, err = header[:n], nil
	width, height, _, err := GetInfo(header)
	if err != nil {
		return
	}
	config.Width = width
	config.Height = height
	config.ColorModel = color.RGBAModel
	return
}

// Decode reads a WEBP image from r and returns it as an image.Image.
func Decode(r io.Reader) (m image.Image, err error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if m, err = DecodeRGBA(data); err != nil {
		return
	}
	return
}

func init() {
	image.RegisterFormat("webp", "RIFF????WEBPVP8", Decode, DecodeConfig)
}
<commit_msg>improve cbuffer support<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.

// +build cgo

package webp

import (
	"errors"
	"image"
	"image/color"
	"io"
	"io/ioutil"
	"os"
)

func LoadConfig(name string, cbuf ...CBuffer) (config image.Config, err error) {
	f, err := os.Open(name)
	if err != nil {
		return image.Config{}, err
	}
	defer f.Close()

	if len(cbuf) == 0 || cbuf[0] == nil {
		cbuf = []CBuffer{NewCBuffer(maxWebpHeaderSize)}
		defer cbuf[0].Close()
	}
	if len(cbuf[0].CData()) < maxWebpHeaderSize {
		if err = cbuf[0].Resize(maxWebpHeaderSize); err != nil {
			return
		}
	}

	header := cbuf[0].CData()
	if len(header) > maxWebpHeaderSize {
		header = header[:maxWebpHeaderSize]
	}
	n, err := f.Read(header)
	if err != nil && err != io.EOF {
		return
	}
	header = header[:n]
	width, height, _, err := GetInfoEx(header, cbuf[0])
	if err != nil {
		return
	}

	config.Width = width
	config.Height = height
	config.ColorModel = color.RGBAModel
	return
}

func Load(name string, cbuf ...CBuffer) (m image.Image, err error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}
	if fi.Size() > (2 << 30) {
		return nil, errors.New("webp: Load, file size is too large (> 2GB)!")
	}
	if len(cbuf) == 0 || cbuf[0] == nil {
		cbuf = []CBuffer{NewCBuffer(int(fi.Size()))}
		defer cbuf[0].Close()
	}

	if len(cbuf[0].CData()) < int(fi.Size()) {
		if err = cbuf[0].Resize(int(fi.Size())); err != nil {
			return
		}
	}
	data := cbuf[0].CData()
	if n := int(fi.Size()); len(data) > n {
		data = data[:n]
	}

	_, err = f.Read(data)
	if err != nil {
		return nil, err
	}
	m, pix, err := DecodeRGBAEx(data, cbuf[0])
	if err != nil {
		return
	}
	pix.Close()

	return m, nil
}

// DecodeConfig returns the color model and dimensions of a WEBP image without
// decoding the entire image.
func DecodeConfig(r io.Reader) (config image.Config, err error) {
	header := make([]byte, maxWebpHeaderSize)
	n, err := r.Read(header)
	if err != nil && err != io.EOF {
		return
	}
	header, err = header[:n], nil
	width, height, _, err := GetInfo(header)
	if err != nil {
		return
	}
	config.Width = width
	config.Height = height
	config.ColorModel = color.RGBAModel
	return
}

// Decode reads a WEBP image from r and returns it as an image.Image.
func Decode(r io.Reader) (m image.Image, err error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if m, err = DecodeRGBA(data); err != nil {
		return
	}
	return
}

func init() {
	image.RegisterFormat("webp", "RIFF????WEBPVP8", Decode, DecodeConfig)
}
<|endoftext|>


<commit_before>package nats

import (
	"errors"
	"bytes"
	"encoding/json"
	"bufio"
	"regexp"
	"strconv"
)

var (
	ErrLineTooLong   = errors.New("reader: line too long")
	ErrUnknownObject = errors.New("reader: unknown object")
	ErrInvalidObject = errors.New("reader: invalid object")
)

var (
	nonSpaceRegexp = regexp.MustCompile("\\S+")
)

type readObject interface {
	read(line []byte, rd *bufio.Reader) (err error)
}

type readMessage struct {
	Subscription   []byte
	SubscriptionId uint
	ReplyTo        []byte
	Payload        []byte
}

func (self *readMessage) read(line []byte, rd *bufio.Reader) (err error) {
	var chunks [][]byte

	chunks = nonSpaceRegexp.FindAll(line, -1)

	if len(chunks) < 4 {
		return ErrInvalidObject
	}

	// Skip +MSG
	idx := 1
	self.Subscription = chunks[idx]

	idx += 1
	sid, err := strconv.ParseUint(string(chunks[idx]), 10, 0)
	if err != nil {
		return err
	}

	self.SubscriptionId = uint(sid)

	if len(chunks) == 5 {
		idx += 1
		self.ReplyTo = chunks[idx]
	}

	idx += 1
	bytes, err := strconv.ParseUint(string(chunks[idx]), 10, 0)
	if err != nil {
		return err
	}

	self.Payload = make([]byte, bytes+2)

	// Read until self.Payload is filled
	var target []byte = self.Payload
	for len(target) > 0 {
		n, err := rd.Read(target)
		if err != nil {
			return err
		}

		target = target[n:]
	}

	// Trim to actual payload size, removing CRLF
	self.Payload = self.Payload[:bytes]

	return
}

type readOk struct {
	// No content
}

func (self *readOk) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readErr struct {
	Payload []byte
}

func (self *readErr) read(line []byte, rd *bufio.Reader) (err error) {
	var index [][]int

	index = nonSpaceRegexp.FindAllIndex(line, 2)

	if len(index) == 2 {
		self.Payload = line[index[1][0]:]
	}

	return
}

type readPing struct {
	// No content
}

func (self *readPing) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readPong struct {
	// No content
}

func (self *readPong) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readInfo struct {
	ServerId     string `json:"server_id"`
	Version      string `json:"version"`
	AuthRequired bool   `json:"auth_required"`
	SslRequired  bool   `json:"ssl_required"`
	MaxPayload   int64  `json:"max_payload"`
}

func (self *readInfo) read(line []byte, rd *bufio.Reader) (err error) {
	var index [][]int

	index = nonSpaceRegexp.FindAllIndex(line, 2)

	if len(index) == 2 {
		err = json.Unmarshal(line[index[1][0]:], &self)
		if err != nil {
			return err
		}
	}

	return
}

func read(rd *bufio.Reader) (readObject, error) {
	var line []byte
	var more bool
	var err error

	line, more, err = rd.ReadLine()

	if err != nil {
		return nil, err
	}

	if more {
		return nil, ErrLineTooLong
	}

	var head []byte = nonSpaceRegexp.Find(line)
	var obj readObject

	switch string(bytes.ToLower(head)) {
	case "msg":
		obj = new(readMessage)
	case "+ok":
		obj = new(readOk)
	case "-err":
		obj = new(readErr)
	case "ping":
		obj = new(readPing)
	case "pong":
		obj = new(readPong)
	case "info":
		obj = new(readInfo)
	default:
		return nil, ErrUnknownObject
	}

	err = obj.read(line, rd)
	if err != nil {
		return nil, err
	}

	return obj, nil
}
<commit_msg>Create slice copies where needed<commit_after>package nats

import (
	"errors"
	"bytes"
	"encoding/json"
	"bufio"
	"regexp"
	"strconv"
)

var (
	ErrLineTooLong   = errors.New("reader: line too long")
	ErrUnknownObject = errors.New("reader: unknown object")
	ErrInvalidObject = errors.New("reader: invalid object")
)

var (
	nonSpaceRegexp = regexp.MustCompile("\\S+")
)

type readObject interface {
	read(line []byte, rd *bufio.Reader) (err error)
}

type readMessage struct {
	Subscription   []byte
	SubscriptionId uint
	ReplyTo        []byte
	Payload        []byte
}

func (self *readMessage) read(line []byte, rd *bufio.Reader) (err error) {
	var chunks [][]byte

	chunks = nonSpaceRegexp.FindAll(line, -1)

	if len(chunks) < 4 {
		return ErrInvalidObject
	}

	// Skip +MSG
	idx := 1
	self.Subscription = make([]byte, len(chunks[idx]))
	copy(self.Subscription, chunks[idx])

	idx += 1
	sid, err := strconv.ParseUint(string(chunks[idx]), 10, 0)
	if err != nil {
		return err
	}

	self.SubscriptionId = uint(sid)

	if len(chunks) == 5 {
		idx += 1
		self.ReplyTo = make([]byte, len(chunks[idx]))
		copy(self.ReplyTo, chunks[idx])
	}

	idx += 1
	bytes, err := strconv.ParseUint(string(chunks[idx]), 10, 0)
	if err != nil {
		return err
	}

	self.Payload = make([]byte, bytes+2)

	// Read until self.Payload is filled
	var target []byte = self.Payload
	for len(target) > 0 {
		n, err := rd.Read(target)
		if err != nil {
			return err
		}

		target = target[n:]
	}

	// Trim to actual payload size, removing CRLF
	self.Payload = self.Payload[:bytes]

	return
}

type readOk struct {
	// No content
}

func (self *readOk) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readErr struct {
	Payload []byte
}

func (self *readErr) read(line []byte, rd *bufio.Reader) (err error) {
	var index [][]int

	index = nonSpaceRegexp.FindAllIndex(line, 2)

	if len(index) == 2 {
		self.Payload = line[index[1][0]:]
	}

	return
}

type readPing struct {
	// No content
}

func (self *readPing) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readPong struct {
	// No content
}

func (self *readPong) read(line []byte, rd *bufio.Reader) (err error) {
	return
}

type readInfo struct {
	ServerId     string `json:"server_id"`
	Version      string `json:"version"`
	AuthRequired bool   `json:"auth_required"`
	SslRequired  bool   `json:"ssl_required"`
	MaxPayload   int64  `json:"max_payload"`
}

func (self *readInfo) read(line []byte, rd *bufio.Reader) (err error) {
	var index [][]int

	index = nonSpaceRegexp.FindAllIndex(line, 2)

	if len(index) == 2 {
		err = json.Unmarshal(line[index[1][0]:], &self)
		if err != nil {
			return err
		}
	}

	return
}

func read(rd *bufio.Reader) (readObject, error) {
	var line []byte
	var more bool
	var err error

	line, more, err = rd.ReadLine()

	if err != nil {
		return nil, err
	}

	if more {
		return nil, ErrLineTooLong
	}

	var head []byte = nonSpaceRegexp.Find(line)
	var obj readObject

	switch string(bytes.ToLower(head)) {
	case "msg":
		obj = new(readMessage)
	case "+ok":
		obj = new(readOk)
	case "-err":
		obj = new(readErr)
	case "ping":
		obj = new(readPing)
	case "pong":
		obj = new(readPong)
	case "info":
		obj = new(readInfo)
	default:
		return nil, ErrUnknownObject
	}

	err = obj.read(line, rd)
	if err != nil {
		return nil, err
	}

	return obj, nil
}
<|endoftext|>


<commit_before>package yenc

import (
	"bufio"
	"bytes"
	"fmt"
	"hash"
	"hash/crc32"
	"io"
	"strconv"
)

const (
	keywordLine = "=y"
	headerLine  = "begin"
	partLine    = "part"
	trailerLine = "end"

	byteOffset    byte = 42
	specialOffset byte = 64
)

var (
	keywordBytes = []byte(keywordLine)
	headerBytes  = []byte(headerLine)
	partBytes    = []byte(partLine)
	trailerBytes = []byte(trailerLine)
)

//Error constants
var (
	ErrBadCRC    = fmt.Errorf("CRC check error")
	ErrWrongSize = fmt.Errorf("size check error")
)

//Reader implements the io.Reader methods for an underlying YENC
//document/stream. It additionally exposes some of the metadata that may be
//useful for consumers.
type Reader struct {
	br         *bufio.Reader
	begun, eof bool

	Length          int
	CRC             hash.Hash32
	Headers, Footer *YENCHeader
}

//NewReader returns a reader from an input reader.
func NewReader(r io.Reader) *Reader {
	return &Reader{
		br:  bufio.NewReader(r),
		CRC: crc32.New(crc32.IEEETable),
	}
}

func (d *Reader) Read(p []byte) (bytesRead int, err error) {
	defer func() {
		d.Length += bytesRead
	}()

	if d.eof {
		return
	}

	var b byte
	var bs []byte

	var lp = len(p)

readLoop:
	for bytesRead < lp {
		b, err = d.br.ReadByte()

		if err != nil {
			return
		}

		switch b {
		case '\r':
			bs, err = d.br.Peek(1)

			if err != nil {
				return
			}
			if len(bs) < 1 {
				return
			}

			if bs[0] == '\n' {
				_, err = d.br.ReadByte()

				if err != nil {
					return
				}

				continue readLoop
			} else {
				break
			}
		case escape:
			b, err = d.br.ReadByte()

			if err != nil {
				return
			}

			if b == 'y' {
				err = d.checkKeywordLine()
				if err != nil {
					return
				}

				continue readLoop
			}

			b -= specialOffset
		}

		if !d.begun {
			continue readLoop
		}

		p[bytesRead] = b - byteOffset
		d.CRC.Write(p[bytesRead : bytesRead+1])
		bytesRead++
	}

	return
}

func (d *Reader) checkKeywordLine() error {
	bs, err := d.br.Peek(5)

	if err != nil {
		return err
	}

	if beginsWith(bs, headerBytes) || beginsWith(bs, partBytes) {
		d.begun = true

		h, err := ReadYENCHeader(d.br)
		if err != nil {
			return err
		}
		d.Headers = h

		return err
	} else if beginsWith(bs, trailerBytes) {
		d.eof = true

		if err = d.checkTrailer(bs); err != nil {
			return err
		}

		if err == nil {
			err = io.EOF
		}

		return err
	}

	return nil
}

func beginsWith(l, c []byte) bool {
	return len(l) >= len(c) && bytes.Equal(c, l[:len(c)])
}

func (d *Reader) checkTrailer(l []byte) error {
	f, err := ReadYENCHeader(d.br)

	if err != nil {
		return err
	}

	d.Footer = f
	preCrc := d.Footer.Get("pcrc32")

	if preCrc == "" {
		return nil
	}

	i, err := strconv.ParseUint(preCrc, 16, 0)
	if err != nil {
		return fmt.Errorf("error parsing uint: %v", err)
	}

	length, err := strconv.Atoi(d.Footer.Get("size"))

	if err != nil && length != d.Length {
		return ErrWrongSize
	}

	sum := d.CRC.Sum32()
	if sum != uint32(i) {
		return ErrBadCRC
	}

	return nil
}
<commit_msg>Faster crc check<commit_after>package yenc

import (
	"bufio"
	"bytes"
	"fmt"
	"hash"
	"hash/crc32"
	"io"
	"strconv"
)

const (
	keywordLine = "=y"
	headerLine  = "begin"
	partLine    = "part"
	trailerLine = "end"

	byteOffset    byte = 42
	specialOffset byte = 64
)

var (
	keywordBytes = []byte(keywordLine)
	headerBytes  = []byte(headerLine)
	partBytes    = []byte(partLine)
	trailerBytes = []byte(trailerLine)
)

//Error constants
var (
	ErrBadCRC    = fmt.Errorf("CRC check error")
	ErrWrongSize = fmt.Errorf("size check error")
)

//Reader implements the io.Reader methods for an underlying YENC
//document/stream. It additionally exposes some of the metadata that may be
//useful for consumers.
type Reader struct {
	br         *bufio.Reader
	begun, eof bool

	Length          int
	CRC             hash.Hash32
	Headers, Footer *YENCHeader
}

//NewReader returns a reader from an input reader.
func NewReader(r io.Reader) *Reader {
	return &Reader{
		br:  bufio.NewReader(r),
		CRC: crc32.New(crc32.IEEETable),
	}
}

func (d *Reader) Read(p []byte) (bytesRead int, err error) {
	defer func() {
		d.Length += bytesRead
	}()

	if d.eof {
		return
	}

	var b byte
	var bs []byte

	var lp = len(p)

readLoop:
	for bytesRead < lp {
		b, err = d.br.ReadByte()

		if err != nil {
			break readLoop
		}

		switch b {
		case '\r':
			bs, err = d.br.Peek(1)

			if err != nil || len(bs) < 1 {
				break readLoop
			}

			if bs[0] == '\n' {
				_, err = d.br.ReadByte()

				if err != nil {
					break readLoop
				}

				continue readLoop
			} else {
				break
			}
		case escape:
			b, err = d.br.ReadByte()

			if err != nil {
				break readLoop
			}

			if b == 'y' {
				//Must write crc before checking bytes. This took me forever to find.
				d.CRC.Write(p[:bytesRead])
				err = d.checkKeywordLine()
				if err != nil {
					return
				}

				continue readLoop
			}

			b -= specialOffset
		}

		if !d.begun {
			continue readLoop
		}

		p[bytesRead] = b - byteOffset
		bytesRead++
	}

	d.CRC.Write(p[:bytesRead])
	return
}

func (d *Reader) checkKeywordLine() error {
	bs, err := d.br.Peek(5)

	if err != nil {
		return err
	}

	if beginsWith(bs, headerBytes) || beginsWith(bs, partBytes) {
		d.begun = true

		h, err := ReadYENCHeader(d.br)
		if err != nil {
			return err
		}
		d.Headers = h

		return err
	} else if beginsWith(bs, trailerBytes) {
		d.eof = true

		if err = d.checkTrailer(bs); err != nil {
			return err
		}

		if err == nil {
			err = io.EOF
		}

		return err
	}

	return nil
}

func beginsWith(l, c []byte) bool {
	return len(l) >= len(c) && bytes.Equal(c, l[:len(c)])
}

func (d *Reader) checkTrailer(l []byte) error {
	f, err := ReadYENCHeader(d.br)

	if err != nil {
		return err
	}

	d.Footer = f
	preCrc := d.Footer.Get("pcrc32")

	if preCrc == "" {
		return nil
	}

	i, err := strconv.ParseUint(preCrc, 16, 0)
	if err != nil {
		return fmt.Errorf("error parsing uint: %v", err)
	}

	length, err := strconv.Atoi(d.Footer.Get("size"))

	if err != nil && length != d.Length {
		return ErrWrongSize
	}

	sum := d.CRC.Sum32()
	if sum != uint32(i) {
		return ErrBadCRC
	}

	return nil
}
<|endoftext|>


<commit_before>package http

import (
	"bytes"
	"fmt"
	"testing"

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

func TestUnmarshalJSON(t *testing.T) {
	input := `{"test":"data","validation":"process"}`
	expected := map[string]interface{}{
		"test":       "data",
		"validation": "process",
	}

	var actual map[string]interface{}
	err := UnmarshalJSON(bytes.NewReader([]byte(input)), &actual)
	if err != nil {
		fmt.Printf("decoding err: %v\n", err)
	}

	assert.EqualValues(t, expected, actual)
}
<commit_msg>added some http tests<commit_after>package http

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

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

type testHandler struct{}

type testObject struct {
	Field string `json:"field"`
}

func TestUnmarshalJSON(t *testing.T) {
	input := `{"test":"data","validation":"process"}`
	expected := map[string]interface{}{
		"test":       "data",
		"validation": "process",
	}

	var actual map[string]interface{}
	err := UnmarshalJSON(bytes.NewReader([]byte(input)), &actual)
	if err != nil {
		fmt.Printf("decoding err: %v\n", err)
	}

	assert.EqualValues(t, expected, actual)
}

func (h *testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	b, _ := json.Marshal(&testObject{
		Field: "response",
	})
	w.Write(b)
}

func TestPostGet(t *testing.T) {
	s := httptest.NewServer(new(testHandler))
	defer s.Close()
	body := testObject{Field: "hello"}
	expect := &testObject{Field: "response"}

	{
		output := new(testObject)
		err := Post(s.URL, body, output)
		assert.NoError(t, err)
		assert.EqualValues(t, expect, output)
	}

	{
		output := new(testObject)
		err := Get(s.URL, output)
		assert.NoError(t, err)
		assert.EqualValues(t, expect, output)
	}
}
<|endoftext|>


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

// Package vulnerability implements functions to validate that the fields of vulnerability entities
// being passed into the API meet our requirements.
package vulnerability

import (
	"errors"
	"fmt"

	"github.com/grafeas/grafeas/go/v1/api/validators/package"
	vulnpb "github.com/grafeas/grafeas/proto/v1/vulnerability_go_proto"
)

// ValidateVulnerability validates that a vulnerability has all its required fields filled in.
func ValidateVulnerability(v *vulnpb.Vulnerability) []error {
	errs := []error{}

	for i, detail := range v.GetDetails() {
		if detail == nil {
			errs = append(errs, fmt.Errorf("details[%d] detail cannot be null", i))
		} else {
			for _, err := range validateVulnerabilityDetail(detail) {
				errs = append(errs, fmt.Errorf("details[%d].%s", i, err))
			}
		}
	}

	return errs
}

func validateVulnerabilityDetail(vd *vulnpb.Vulnerability_Detail) []error {
	errs := []error{}

	if vd.GetCpeUri() == "" {
		errs = append(errs, errors.New("cpe_uri is required"))
	}
	if vd.GetPackage() == "" {
		errs = append(errs, errors.New("package is required"))
	}

	if ver := vd.GetMinAffectedVersion(); ver != nil {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("min_affected_version.%s", err))
		}
	}
	if ver := vd.GetMaxAffectedVersion(); ver != nil {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("max_affected_version.%s", err))
		}
	}
	if fl := vd.GetFixedLocation(); fl != nil {
		for _, err := range validateVulnerabilityLocation(fl) {
			errs = append(errs, fmt.Errorf("fixed_location.%s", err))
		}
	}

	return errs
}

func validateVulnerabilityLocation(vl *vulnpb.VulnerabilityLocation) []error {
	errs := []error{}

	if vl.GetCpeUri() == "" {
		errs = append(errs, errors.New("cpe_uri is required"))
	}
	if vl.GetPackage() == "" {
		errs = append(errs, errors.New("package is required"))
	}
	if ver := vl.GetVersion(); ver == nil {
		errs = append(errs, errors.New("version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
	}

	return errs
}

// ValidateDetails validates that a details has all its required fields filled in.
func ValidateDetails(d *vulnpb.Details) []error {
	errs := []error{}

	if pkgIssue := d.GetPackageIssue(); pkgIssue == nil {
		errs = append(errs, errors.New("package_issue is required"))
	} else if len(pkgIssue) == 0 {
		errs = append(errs, errors.New("package_issue requires at least 1 element"))
	} else {
		for i, p := range pkgIssue {
			if p == nil {
				errs = append(errs, fmt.Errorf("package_issue[%d] package issue cannot be null", i))
			} else {
				for _, err := range validatePackageIssue(p) {
					errs = append(errs, fmt.Errorf("package_issue[%d].%s", i, err))
				}
			}
		}
	}

	return errs
}

func validatePackageIssue(p *vulnpb.PackageIssue) []error {
	errs := []error{}

        if p.GetAffectedCpeUri() == "" {
		errs = append(errs, errors.New("affected cpe_uri is required"))
	}
	if p.GetAffectedPackage() == "" {
		errs = append(errs, errors.New("affected package is required"))
	}
	if ver := p.GetAffectedVersion(); ver == nil {
		errs = append(errs, errors.New("affected version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
	}

        if p.GetFixedCpeUri() == "" {
		errs = append(errs, errors.New("fixed cpe_uri is required"))
	}
	if p.GetFixedPackage() == "" {
		errs = append(errs, errors.New("fixed package is required"))
	}
	if ver := p.GetFixedVersion(); ver == nil {
		errs = append(errs, errors.New("fixed version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
	}
	return errs
}
<commit_msg>fixed vulnerability<commit_after>// Copyright 2018 The Grafeas Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package vulnerability implements functions to validate that the fields of vulnerability entities
// being passed into the API meet our requirements.
package vulnerability

import (
	"errors"
	"fmt"

	"github.com/grafeas/grafeas/go/v1/api/validators/package"
	vulnpb "github.com/grafeas/grafeas/proto/v1/vulnerability_go_proto"
	ppb "google3/third_party/grafeas/proto/v1/package_go_proto"
)

// ValidateVulnerability validates that a vulnerability has all its required fields filled in.
func ValidateVulnerability(v *vulnpb.Vulnerability) []error {
	errs := []error{}

	for i, detail := range v.GetDetails() {
		if detail == nil {
			errs = append(errs, fmt.Errorf("details[%d] detail cannot be null", i))
		} else {
			for _, err := range validateVulnerabilityDetail(detail) {
				errs = append(errs, fmt.Errorf("details[%d].%s", i, err))
			}
		}
	}

	return errs
}

func validateVulnerabilityDetail(vd *vulnpb.Vulnerability_Detail) []error {
	errs := []error{}

	if vd.GetCpeUri() == "" {
		errs = append(errs, errors.New("cpe_uri is required"))
	}
	if vd.GetPackage() == "" {
		errs = append(errs, errors.New("package is required"))
	}

	if ver := vd.GetMinAffectedVersion(); ver != nil {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("min_affected_version.%s", err))
		}
	}
	if ver := vd.GetMaxAffectedVersion(); ver != nil {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("max_affected_version.%s", err))
		}
	}
	if fl := vd.GetFixedLocation(); fl != nil {
		for _, err := range validateVulnerabilityLocation(fl) {
			errs = append(errs, fmt.Errorf("fixed_location.%s", err))
		}
	}

	return errs
}

func validateVulnerabilityLocation(vl *vulnpb.VulnerabilityLocation) []error {
	errs := []error{}

	if vl.GetCpeUri() == "" {
		errs = append(errs, errors.New("cpe_uri is required"))
	}
	if vl.GetPackage() == "" {
		errs = append(errs, errors.New("package is required"))
	}
	if ver := vl.GetVersion(); ver == nil {
		errs = append(errs, errors.New("version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
	}

	return errs
}

// ValidateDetails validates that a details has all its required fields filled in.
func ValidateDetails(d *vulnpb.Details) []error {
	errs := []error{}

	if pkgIssue := d.GetPackageIssue(); pkgIssue == nil {
		errs = append(errs, errors.New("package_issue is required"))
	} else if len(pkgIssue) == 0 {
		errs = append(errs, errors.New("package_issue requires at least 1 element"))
	} else {
		for i, p := range pkgIssue {
			if p == nil {
				errs = append(errs, fmt.Errorf("package_issue[%d] package issue cannot be null", i))
			} else {
				for _, err := range validatePackageIssue(p) {
					errs = append(errs, fmt.Errorf("package_issue[%d].%s", i, err))
				}
			}
		}
	}

	return errs
}

func validatePackageIssue(p *vulnpb.PackageIssue) []error {
	errs := []error{}

	if p.GetAffectedCpeUri() == "" {
		errs = append(errs, errors.New("affected cpe_uri is required"))
	}
	if p.GetAffectedPackage() == "" {
		errs = append(errs, errors.New("affected package is required"))
	}
	if ver := p.GetAffectedVersion(); ver == nil {
		errs = append(errs, errors.New("affected version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
	}

	if ver := p.GetFixedVersion(); ver == nil {
		errs = append(errs, errors.New("fixed version is required"))
	} else {
		for _, err := range pkg.ValidateVersion(ver) {
			errs = append(errs, fmt.Errorf("version.%s", err))
		}
		if ver.Kind == ppb.Version_NORMAL {
			if p.GetFixedCpeUri() == "" {
				errs = append(errs, errors.New("fixed cpe_uri is required when fixed version kind is NORMAL"))
			}
			if p.GetFixedPackage() == "" {
				errs = append(errs, errors.New("fixed package is required when fixed version kind is NORMAL"))
			}
		}
	}
	return errs
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/blevesearch/bleve"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	biblescholar "github.com/turtlemonvh/biblescholar/search"
	"github.com/turtlemonvh/biblescholar/search/server"
)

/*
Project is called `biblescholar`. The command line interface is called `bblsearch`.
*/

var searchCmdV *cobra.Command
var buildCommit string = "UNKNOWN"
var buildBranch string = "UNKNOWN"

func main() {
	RootCmd.Execute()
}

func InitializeConfig() {
	viper.SetConfigName("bblsearch")
	viper.AddConfigPath("/etc/bblsearch/")
	viper.AddConfigPath("$HOME/.bblsearch")
	viper.AddConfigPath(".")

	viper.SetEnvPrefix("BBLSEARCH_")
	viper.AutomaticEnv()
}

func init() {
	RootCmd.AddCommand(indexCmd)
	RootCmd.AddCommand(serverCmd)
	searchCmdV = RootCmd

	// FIXME: Add support for multiple outputs and handling log levels via command line or env variable
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)

	RootCmd.PersistentFlags().StringP(
		"index-path", "i", biblescholar.DefaultIndexName,
		fmt.Sprintf("path to bleve index. Default is: %s", biblescholar.DefaultIndexName),
	)
	RootCmd.PersistentFlags().Bool("debug-logging", false, "turn on debug level logging")
	indexCmd.Flags().StringP("data-dir", "d", "downloads", "directory containing tsv data files to use in indexing")
	serverCmd.Flags().IntP("port", "p", 8000, "port to run server on")
	serverCmd.Flags().Bool("validate-alexa", false, "should the application validate that requests are from the alexa service?")
}

func HandleLogLevel() {
	if viper.GetBool("debug-logging") {
		fmt.Println("Setting log level to debug")
		log.SetLevel(log.DebugLevel)
		log.Debug("Set log level to debug")
	}
}

var RootCmd = &cobra.Command{
	Use:   "bblsearch",
	Short: "bblsearch is a search interface for the Bible",
	Run: func(cmd *cobra.Command, args []string) {
		InitializeConfig()
		fmt.Printf("bblsearch %s (%s)\n", buildBranch, buildCommit)
	},
}

var indexLongDesc = `Searches data-dir for tsv files to add to index.  Files should look like:

<version>  <Book>  <Chapter #>  <Verse #>  <Verse Text>
`
var indexCmd = &cobra.Command{
	Use:   "index",
	Short: "Index from a collection of tsv files",
	Long:  indexLongDesc,
	Run: func(cmd *cobra.Command, args []string) {
		viper.BindPFlag("data-dir", cmd.Flags().Lookup("data-dir"))
		viper.BindPFlag("index-path", cmd.Flags().Lookup("index-path"))
		viper.BindPFlag("debug-logging", cmd.Flags().Lookup("debug-logging"))

		HandleLogLevel()

		index := biblescholar.CreateOrOpenIndex(viper.GetString("index-path"))
		fmt.Println("Adding content to: ", viper.GetString("index-path"))

		_, err := biblescholar.IndexFromTSVs(index, viper.GetString("data-dir"))
		if err != nil {
			log.Fatal(err)
		}
	},
}

var serverCmd = &cobra.Command{
	Use:   "server",
	Short: "Start a server that fields responses to queries from alexa",
	Long:  `Start a server that fields responses to queries from alexa`,
	Run: func(cmd *cobra.Command, args []string) {
		viper.BindPFlag("port", cmd.Flags().Lookup("port"))
		viper.BindPFlag("index-path", cmd.Flags().Lookup("index-path"))
		viper.BindPFlag("debug-logging", cmd.Flags().Lookup("debug-logging"))
		viper.BindPFlag("should-validate-alexa-requests", cmd.Flags().Lookup("validate-alexa"))

		HandleLogLevel()

		// Always text logs, because docker thinks there is a tty
		// https://godoc.org/github.com/sirupsen/logrus#TextFormatter
		log.SetFormatter(&log.TextFormatter{
			DisableColors: true,
			FullTimestamp: true,
		})

		idx, err := bleve.Open(viper.GetString("index-path"))
		if err != nil {
			log.Fatal(err)
		}

		svr := server.ServerConfig{
			Port:                viper.GetInt("port"),
			BuildCommit:         buildCommit,
			BuildBranch:         buildBranch,
			Index:               idx,
			ShouldValidateAlexa: viper.GetBool("should-validate-alexa-requests"),
		}
		svr.StartServer()
	},
}
<commit_msg>Derive binary name for help text from the name of the binary used in the call<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/blevesearch/bleve"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	biblescholar "github.com/turtlemonvh/biblescholar/search"
	"github.com/turtlemonvh/biblescholar/search/server"
)

/*
Project is called `biblescholar`. The command line interface is called `bblsearch`.
*/

var searchCmdV *cobra.Command
var buildCommit string = "UNKNOWN"
var buildBranch string = "UNKNOWN"

func main() {
	RootCmd.Execute()
}

func InitializeConfig() {
	viper.SetConfigName("bblsearch")
	viper.AddConfigPath("/etc/bblsearch/")
	viper.AddConfigPath("$HOME/.bblsearch")
	viper.AddConfigPath(".")

	viper.SetEnvPrefix("BBLSEARCH_")
	viper.AutomaticEnv()
}

func init() {
	RootCmd.AddCommand(indexCmd)
	RootCmd.AddCommand(serverCmd)
	searchCmdV = RootCmd

	// FIXME: Add support for multiple outputs and handling log levels via command line or env variable
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)

	RootCmd.PersistentFlags().StringP(
		"index-path", "i", biblescholar.DefaultIndexName,
		fmt.Sprintf("path to bleve index. Default is: %s", biblescholar.DefaultIndexName),
	)
	RootCmd.PersistentFlags().Bool("debug-logging", false, "turn on debug level logging")
	indexCmd.Flags().StringP("data-dir", "d", "downloads", "directory containing tsv data files to use in indexing")
	serverCmd.Flags().IntP("port", "p", 8000, "port to run server on")
	serverCmd.Flags().Bool("validate-alexa", false, "should the application validate that requests are from the alexa service?")
}

func HandleLogLevel() {
	if viper.GetBool("debug-logging") {
		fmt.Println("Setting log level to debug")
		log.SetLevel(log.DebugLevel)
		log.Debug("Set log level to debug")
	}
}

var RootCmd = &cobra.Command{
	Use:   os.Args[0],
	Short: fmt.Sprintf("%s is a search interface for the Bible", os.Args[0]),
	Run: func(cmd *cobra.Command, args []string) {
		InitializeConfig()
		fmt.Printf("bblsearch %s (%s)\n", buildBranch, buildCommit)
	},
}

var indexLongDesc = `Searches data-dir for tsv files to add to index.  Files should look like:

<version>  <Book>  <Chapter #>  <Verse #>  <Verse Text>
`
var indexCmd = &cobra.Command{
	Use:   "index",
	Short: "Index from a collection of tsv files",
	Long:  indexLongDesc,
	Run: func(cmd *cobra.Command, args []string) {
		viper.BindPFlag("data-dir", cmd.Flags().Lookup("data-dir"))
		viper.BindPFlag("index-path", cmd.Flags().Lookup("index-path"))
		viper.BindPFlag("debug-logging", cmd.Flags().Lookup("debug-logging"))

		HandleLogLevel()

		index := biblescholar.CreateOrOpenIndex(viper.GetString("index-path"))
		fmt.Println("Adding content to: ", viper.GetString("index-path"))

		_, err := biblescholar.IndexFromTSVs(index, viper.GetString("data-dir"))
		if err != nil {
			log.Fatal(err)
		}
	},
}

var serverCmd = &cobra.Command{
	Use:   "server",
	Short: "Start a server that fields responses to queries from alexa",
	Long:  `Start a server that fields responses to queries from alexa`,
	Run: func(cmd *cobra.Command, args []string) {
		viper.BindPFlag("port", cmd.Flags().Lookup("port"))
		viper.BindPFlag("index-path", cmd.Flags().Lookup("index-path"))
		viper.BindPFlag("debug-logging", cmd.Flags().Lookup("debug-logging"))
		viper.BindPFlag("should-validate-alexa-requests", cmd.Flags().Lookup("validate-alexa"))

		HandleLogLevel()

		// Always text logs, because docker thinks there is a tty
		// https://godoc.org/github.com/sirupsen/logrus#TextFormatter
		log.SetFormatter(&log.TextFormatter{
			DisableColors: true,
			FullTimestamp: true,
		})

		idx, err := bleve.Open(viper.GetString("index-path"))
		if err != nil {
			log.Fatal(err)
		}

		svr := server.ServerConfig{
			Port:                viper.GetInt("port"),
			BuildCommit:         buildCommit,
			BuildBranch:         buildBranch,
			Index:               idx,
			ShouldValidateAlexa: viper.GetBool("should-validate-alexa-requests"),
		}
		svr.StartServer()
	},
}
<|endoftext|>


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

package main

import (
	"database/sql"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync/atomic"
	"time"

	"github.com/google/keytransparency/core/keyserver"
	"github.com/google/keytransparency/core/mutator/entry"
	"github.com/google/keytransparency/core/vrf"
	"github.com/google/keytransparency/core/vrf/p256"
	"github.com/google/keytransparency/impl/etcd/queue"
	"github.com/google/keytransparency/impl/google/authentication"
	"github.com/google/keytransparency/impl/sql/appender"
	"github.com/google/keytransparency/impl/sql/commitments"
	"github.com/google/keytransparency/impl/sql/engine"
	"github.com/google/keytransparency/impl/sql/sqlhist"
	"github.com/google/keytransparency/impl/transaction"

	"github.com/google/trillian"
	"github.com/google/trillian/extension/builtin"
	"github.com/google/trillian/server"
	"github.com/google/trillian/util"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/reflection"

	pb "github.com/google/keytransparency/impl/proto/keytransparency_v1_service"
)

var (
	addr          = flag.String("addr", ":8080", "The ip:port combination to listen on")
	serverDBPath  = flag.String("db", "db", "Database connection string")
	etcdEndpoints = flag.String("etcd", "", "Comma delimited list of etcd endpoints")
	mapID         = flag.String("domain", "example.com", "Distinguished name for this key server")
	realm         = flag.String("auth-realm", "registered-users@gmail.com", "Authentication realm for WWW-Authenticate response header")
	vrfPath       = flag.String("vrf", "private_vrf_key.dat", "Path to VRF private key")
	mapLogURL     = flag.String("maplog", "", "URL of CT server for Signed Map Heads")
	keyFile       = flag.String("key", "testdata/server.key", "TLS private key file")
	certFile      = flag.String("cert", "testdata/server.pem", "TLS cert file")
	verbose       = flag.Bool("verbose", false, "Log requests and responses")
)

func newLogServer(db *sql.DB) (*server.TrillianLogRPCServer, error) {
	timesource := &util.SystemTimeSource{}
	registry, err := builtin.NewExtensionRegistry(db, nil)
	if err != nil {
		return nil, err
	}

	return server.NewTrillianLogRPCServer(registry, timesource), nil
}

func openDB() *sql.DB {
	db, err := sql.Open(engine.DriverName, *serverDBPath)
	if err != nil {
		log.Fatalf("sql.Open(): %v", err)
	}
	if err := db.Ping(); err != nil {
		log.Fatalf("db.Ping(): %v", err)
	}
	return db
}

func openEtcd() *clientv3.Client {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   strings.Split(*etcdEndpoints, ","),
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		log.Fatalf("Failed to connect to etcd: %v", err)
	}
	return cli
}

func openVRFKey() vrf.PrivateKey {
	vrfBytes, err := ioutil.ReadFile(*vrfPath)
	if err != nil {
		log.Fatalf("Failed opening VRF private key: %v", err)
	}
	vrfPriv, err := p256.NewVRFSignerFromPEM(vrfBytes)
	if err != nil {
		log.Fatalf("Failed parsing VRF private key: %v", err)
	}
	return vrfPriv
}

func grpcGatewayMux(addr string) (*runtime.ServeMux, error) {
	ctx := context.Background()

	creds, err := credentials.NewClientTLSFromFile(*certFile, "")
	if err != nil {
		return nil, err
	}
	dopts := []grpc.DialOption{grpc.WithTransportCredentials(creds)}

	gwmux := runtime.NewServeMux()
	if err := pb.RegisterKeyTransparencyServiceHandlerFromEndpoint(ctx, gwmux, addr, dopts); err != nil {
		return nil, err
	}

	return gwmux, nil
}

// grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
// connections or otherHandler otherwise. Copied from cockroachdb.
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// This is a partial recreation of gRPC's internal checks.
		// https://github.com/grpc/grpc-go/blob/master/transport/handler_server.go#L62
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

var marshaler = jsonpb.Marshaler{Indent: "  ", OrigName: true}
var requestCounter uint64

// jsonLogger logs the request and response protobufs as json objects.
func jsonLogger(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
	atomic.AddUint64(&requestCounter, 1)
	// Print request.
	pb, ok := req.(proto.Message)
	if !ok {
		log.Printf("req %t, %v, not a proto.Message", req, req)
		return handler(ctx, req)
	}
	s, err := marshaler.MarshalToString(pb)
	if err != nil {
		log.Printf("Failed to marshal %v", pb)
		return handler(ctx, req)
	}
	log.Printf("%v>%v", requestCounter, s)

	resp, err = handler(ctx, req)

	// Print response.
	pb, ok = resp.(proto.Message)
	if !ok {
		log.Printf("req %t, %v, not a proto.Message", req, req)
		return resp, err
	}
	s, err = marshaler.MarshalToString(pb)
	if err != nil {
		log.Printf("Failed to marshal %v", pb)
		return resp, err
	}
	log.Printf("%v<%v", requestCounter, s)

	return resp, err
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	flag.Parse()

	// Open Resources.
	sqldb := openDB()
	defer sqldb.Close()
	etcdCli := openEtcd()
	defer etcdCli.Close()
	factory := transaction.NewFactory(sqldb, etcdCli)

	creds, err := credentials.NewServerTLSFromFile(*certFile, *keyFile)
	if err != nil {
		log.Fatalf("Failed to load server credentials %v", err)
	}
	auth, err := authentication.NewGoogleAuth()
	if err != nil {
		log.Fatalf("Failed to create authentication library instance: %v", err)
	}

	// Create database and helper objects.
	commitments, err := commitments.New(sqldb, *mapID)
	if err != nil {
		log.Fatalf("Failed to create committer: %v", err)
	}
	queue := queue.New(context.Background(), etcdCli, *mapID, factory)
	tree, err := sqlhist.New(context.Background(), sqldb, *mapID, factory)
	if err != nil {
		log.Fatalf("Failed to create SQL history: %v", err)
	}
	sths, err := appender.New(context.Background(), sqldb, *mapID, *mapLogURL, nil)
	if err != nil {
		log.Fatalf("Failed to create appender: %v", err)
	}
	vrfPriv := openVRFKey()
	mutator := entry.New()

	// Create Trillian server.
	logServer, err := newLogServer(sqldb)
	if err != nil {
		log.Fatalf("Failed to create trillian log server: %v", err)
	}

	// Create gRPC server.
	svr := keyserver.New(commitments, queue, tree, sths, vrfPriv, mutator, auth, factory)
	opts := []grpc.ServerOption{grpc.Creds(creds)}
	if *verbose {
		opts = append(opts, grpc.UnaryInterceptor(jsonLogger))
	}
	grpcServer := grpc.NewServer(opts...)
	pb.RegisterKeyTransparencyServiceServer(grpcServer, svr)
	trillian.RegisterTrillianLogServer(grpcServer, logServer)
	reflection.Register(grpcServer)

	// Create HTTP handlers and gRPC gateway.
	gwmux, err := grpcGatewayMux(*addr)
	if err != nil {
		log.Fatalf("Failed setting up REST proxy: %v", err)
	}

	mux := http.NewServeMux()
	// Insert handlers for other http paths here.
	mux.Handle("/", gwmux)

	// Serve HTTP2 server over TLS.
	log.Printf("Listening on %v", *addr)
	if err := http.ListenAndServeTLS(*addr, *certFile, *keyFile,
		grpcHandlerFunc(grpcServer, mux)); err != nil {
		log.Fatalf("ListenAndServeTLS: %v", err)
	}
}
<commit_msg>Follow Trillian HEAD's registry conversion<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"database/sql"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync/atomic"
	"time"

	"github.com/google/keytransparency/core/keyserver"
	"github.com/google/keytransparency/core/mutator/entry"
	"github.com/google/keytransparency/core/vrf"
	"github.com/google/keytransparency/core/vrf/p256"
	"github.com/google/keytransparency/impl/etcd/queue"
	"github.com/google/keytransparency/impl/google/authentication"
	"github.com/google/keytransparency/impl/sql/appender"
	"github.com/google/keytransparency/impl/sql/commitments"
	"github.com/google/keytransparency/impl/sql/engine"
	"github.com/google/keytransparency/impl/sql/sqlhist"
	"github.com/google/keytransparency/impl/transaction"

	"github.com/google/trillian"
	"github.com/google/trillian/crypto/keys"
	"github.com/google/trillian/extension"
	"github.com/google/trillian/server"
	"github.com/google/trillian/storage/mysql"
	"github.com/google/trillian/util"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/reflection"

	pb "github.com/google/keytransparency/impl/proto/keytransparency_v1_service"
)

var (
	addr          = flag.String("addr", ":8080", "The ip:port combination to listen on")
	serverDBPath  = flag.String("db", "db", "Database connection string")
	etcdEndpoints = flag.String("etcd", "", "Comma delimited list of etcd endpoints")
	mapID         = flag.String("domain", "example.com", "Distinguished name for this key server")
	realm         = flag.String("auth-realm", "registered-users@gmail.com", "Authentication realm for WWW-Authenticate response header")
	vrfPath       = flag.String("vrf", "private_vrf_key.dat", "Path to VRF private key")
	mapLogURL     = flag.String("maplog", "", "URL of CT server for Signed Map Heads")
	keyFile       = flag.String("key", "testdata/server.key", "TLS private key file")
	certFile      = flag.String("cert", "testdata/server.pem", "TLS cert file")
	verbose       = flag.Bool("verbose", false, "Log requests and responses")
)

func newLogServer(db *sql.DB) (*server.TrillianLogRPCServer, error) {
	registry := extension.Registry{
		AdminStorage:  mysql.NewAdminStorage(db),
		SignerFactory: keys.PEMSignerFactory{},
		LogStorage:    mysql.NewLogStorage(db),
	}
	timesource := &util.SystemTimeSource{}

	return server.NewTrillianLogRPCServer(registry, timesource), nil
}

func openDB() *sql.DB {
	db, err := sql.Open(engine.DriverName, *serverDBPath)
	if err != nil {
		log.Fatalf("sql.Open(): %v", err)
	}
	if err := db.Ping(); err != nil {
		log.Fatalf("db.Ping(): %v", err)
	}
	return db
}

func openEtcd() *clientv3.Client {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   strings.Split(*etcdEndpoints, ","),
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		log.Fatalf("Failed to connect to etcd: %v", err)
	}
	return cli
}

func openVRFKey() vrf.PrivateKey {
	vrfBytes, err := ioutil.ReadFile(*vrfPath)
	if err != nil {
		log.Fatalf("Failed opening VRF private key: %v", err)
	}
	vrfPriv, err := p256.NewVRFSignerFromPEM(vrfBytes)
	if err != nil {
		log.Fatalf("Failed parsing VRF private key: %v", err)
	}
	return vrfPriv
}

func grpcGatewayMux(addr string) (*runtime.ServeMux, error) {
	ctx := context.Background()

	creds, err := credentials.NewClientTLSFromFile(*certFile, "")
	if err != nil {
		return nil, err
	}
	dopts := []grpc.DialOption{grpc.WithTransportCredentials(creds)}

	gwmux := runtime.NewServeMux()
	if err := pb.RegisterKeyTransparencyServiceHandlerFromEndpoint(ctx, gwmux, addr, dopts); err != nil {
		return nil, err
	}

	return gwmux, nil
}

// grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
// connections or otherHandler otherwise. Copied from cockroachdb.
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// This is a partial recreation of gRPC's internal checks.
		// https://github.com/grpc/grpc-go/blob/master/transport/handler_server.go#L62
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

var marshaler = jsonpb.Marshaler{Indent: "  ", OrigName: true}
var requestCounter uint64

// jsonLogger logs the request and response protobufs as json objects.
func jsonLogger(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
	atomic.AddUint64(&requestCounter, 1)
	// Print request.
	pb, ok := req.(proto.Message)
	if !ok {
		log.Printf("req %t, %v, not a proto.Message", req, req)
		return handler(ctx, req)
	}
	s, err := marshaler.MarshalToString(pb)
	if err != nil {
		log.Printf("Failed to marshal %v", pb)
		return handler(ctx, req)
	}
	log.Printf("%v>%v", requestCounter, s)

	resp, err = handler(ctx, req)

	// Print response.
	pb, ok = resp.(proto.Message)
	if !ok {
		log.Printf("req %t, %v, not a proto.Message", req, req)
		return resp, err
	}
	s, err = marshaler.MarshalToString(pb)
	if err != nil {
		log.Printf("Failed to marshal %v", pb)
		return resp, err
	}
	log.Printf("%v<%v", requestCounter, s)

	return resp, err
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	flag.Parse()

	// Open Resources.
	sqldb := openDB()
	defer sqldb.Close()
	etcdCli := openEtcd()
	defer etcdCli.Close()
	factory := transaction.NewFactory(sqldb, etcdCli)

	creds, err := credentials.NewServerTLSFromFile(*certFile, *keyFile)
	if err != nil {
		log.Fatalf("Failed to load server credentials %v", err)
	}
	auth, err := authentication.NewGoogleAuth()
	if err != nil {
		log.Fatalf("Failed to create authentication library instance: %v", err)
	}

	// Create database and helper objects.
	commitments, err := commitments.New(sqldb, *mapID)
	if err != nil {
		log.Fatalf("Failed to create committer: %v", err)
	}
	queue := queue.New(context.Background(), etcdCli, *mapID, factory)
	tree, err := sqlhist.New(context.Background(), sqldb, *mapID, factory)
	if err != nil {
		log.Fatalf("Failed to create SQL history: %v", err)
	}
	sths, err := appender.New(context.Background(), sqldb, *mapID, *mapLogURL, nil)
	if err != nil {
		log.Fatalf("Failed to create appender: %v", err)
	}
	vrfPriv := openVRFKey()
	mutator := entry.New()

	// Create Trillian server.
	logServer, err := newLogServer(sqldb)
	if err != nil {
		log.Fatalf("Failed to create trillian log server: %v", err)
	}

	// Create gRPC server.
	svr := keyserver.New(commitments, queue, tree, sths, vrfPriv, mutator, auth, factory)
	opts := []grpc.ServerOption{grpc.Creds(creds)}
	if *verbose {
		opts = append(opts, grpc.UnaryInterceptor(jsonLogger))
	}
	grpcServer := grpc.NewServer(opts...)
	pb.RegisterKeyTransparencyServiceServer(grpcServer, svr)
	trillian.RegisterTrillianLogServer(grpcServer, logServer)
	reflection.Register(grpcServer)

	// Create HTTP handlers and gRPC gateway.
	gwmux, err := grpcGatewayMux(*addr)
	if err != nil {
		log.Fatalf("Failed setting up REST proxy: %v", err)
	}

	mux := http.NewServeMux()
	// Insert handlers for other http paths here.
	mux.Handle("/", gwmux)

	// Serve HTTP2 server over TLS.
	log.Printf("Listening on %v", *addr)
	if err := http.ListenAndServeTLS(*addr, *certFile, *keyFile,
		grpcHandlerFunc(grpcServer, mux)); err != nil {
		log.Fatalf("ListenAndServeTLS: %v", err)
	}
}
<|endoftext|>


<commit_before>/*
 Copyright (c) 2013 Jeremy Torres, https://github.com/jeremytorres/jpegextract

 Permission is hereby granted, free of charge, to any person obtaining
 a copy of this software and associated documentation files (the
 "Software"), to deal in the Software without restriction, including
 without limitation the rights to use, copy, modify, merge, publish,
 distribute, sublicense, and/or sell copies of the Software, and to
 permit persons to whom the Software is furnished to do so, subject to
 the following conditions:

 The above copyright notice and this permission notice shall be
 included in all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
 LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

// Command utilizing the rawparser library for extracting jpegs from RAW iles.
// The embedded JPEGs are extracted and optionally:
// (1) rotated via oritenation provided in RawFile EXIF info;
//
// Usage:
//     jpgextract --raws "NEF,CR2" --dest-dir "/path_to/output_dir"
//                --src-dirs "/path_to/source1,/path_to/source2"
//                [--num-routines "8" --quality "80" --rotate]
package main

import (
	"github.com/codegangsta/cli"
	"github.com/jeremytorres/rawparser"
	"log"
	"math"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	// RawTypesKey is the constant representing the command line argument for the RAW file
	// types to be processed.
	RawTypesKey = "raws"

	// SrcDirKey is the constant representing the command line argument for the source directories
	// containing the source RawFiles.
	SrcDirKey = "src-dirs"
	// DestDirKey is a constant representing the command line argument for the destination directory
	// where extracted jpegs shall be placed.
	DestDirKey = "dest-dir"
	// NumRoutinesKey is the constant representing the command line argument for the max number of
	// go routines to be processed.  As this utility is mostly file IO bound, it is important
	// not to use too many routines ("too many files open..." errors can occur from OS).  Therefore,
	// if this value is greater than the logical cores, a warning is logged and the value is set to
	// the number of logical cores.
	NumRoutinesKey = "num-routines"
	// QualityKey is the constant representing the command line argument for the JPEG quality used
	// when processing the embedded JPEG in a RawFile.
	QualityKey = "quality"
	// RotateKey is the constant representing the command line argument indicating rotation of
	// JPEGs should occur based on the EXIF info embedded within the RawFile.  ImageMagic's 'convert'
	// utility is used and is checked at startup for existence.
	RotateKey = "rotate"
	// AppVersionKey is the constant defining the current version of this command utility.
	AppVersionKey = "1.0"
	// ImageMagicConvertBin is the constant representing ImageMagic's 'convert' utility.
	ImageMagicConvertBin = "convert"
)

var (
	destDir, sqlLiteDb     string
	rawFileExts, srcDirs   []string
	numOfRoutines, quality int
	rotate                 bool
	parsers                *rawparser.RawParsers
	// validParserKeys is a slice of RAW file parsers supported by this implementation.
	validParserKeys = []string{
		rawparser.NefParserKey,
		rawparser.Cr2ParserKey,
	}
)

/*
// processFilesConcurrent processes RawFiles concurrently, based on the number goroutines
// specified.
func processFilesConcurrent(rp *RawFileParserPair, c chan<- bool) {

	go func(rp *RawFileParserPair, c chan<- bool) {
		rawfile, err := rp.parser.ProcessFile(&rawparser.RawFileInfo{rp.file, destDir, quality})
		if err != nil {
			log.Printf("Error processing file: '%s' error: %v\n", rp.file, err)
		} else {
			if rotate && rawfile.JpegOrientation != 0.0 {
				// rotate jpeg
				go func(fileName string, radiansCw float64) {
					degrees := radiansCw * (180 / math.Pi)
					log.Printf("Rotating image %f degrees for jpeg: '%s'\n", degrees, fileName)
					cmd := exec.Command(ImageMagicConvertBin, "-rotate", strconv.FormatFloat(degrees, 'f', 2, 64), fileName, fileName)
					err := cmd.Start()
					if err != nil {
						log.Fatal(err)
					}
					err = cmd.Wait()
					if err != nil {
						log.Printf("Command finished with error: %v", err)
					}
				}(rawfile.JpegPath, rawfile.JpegOrientation)
			}
		}
		// signal completion of work
		c <- true
	}(rp, c)
}
*/

func isRawFileExtValid(ext string) bool {
	for _, validExt := range validParserKeys {
		if ext == validExt {
			return true
		}
	}
	return false
}

// processCli parses command line arguments and checks for validity of user-specified
// values.
func processCli() bool {
	var processed = false

	app := cli.NewApp()
	app.Name = "jpgextract"
	app.Usage = "Processes RAW files and extracts JPEGs and optionally create a SQLite database."
	app.Version = AppVersionKey
	app.Author = "Jeremy Torres"
	app.Flags = []cli.Flag{
		cli.StringFlag{RawTypesKey, "", "comma-separated list of RAW file extensions to process.  Current supported " +
			"RAW files: [NEF | CR2]"},
		cli.StringFlag{SrcDirKey, "", "comma-separated list of full paths to the directories containing RawFiles " +
			"(e.g., /path_to/dir1,/path_to/dir2..."},
		cli.StringFlag{DestDirKey, "", "the full path to the directory containing extracted jpegs"},
		cli.IntFlag{NumRoutinesKey, 2, "the number of concurrent files to be processed"},
		cli.IntFlag{QualityKey, 80, "JPEG encoding quality used for extracted jpegs"},
		cli.BoolFlag{RotateKey, "ImageMagic's 'convert' command will be used to rotate jpegs based on EXIF info from RawFile.  'convert' must be in PATH"},
	}
	app.Action = func(c *cli.Context) {
		rawExts := strings.TrimSpace(c.String(RawTypesKey))
		srcDir := strings.TrimSpace(c.String(SrcDirKey))
		destDir = strings.TrimSpace(c.String(DestDirKey))
		numOfRoutines = c.Int(NumRoutinesKey)
		quality = c.Int(QualityKey)
		rotate = c.Bool(RotateKey)

		// src and dest dirs required; remaing args have sane defaults
		if rawExts == "" || srcDir == "" || destDir == "" {
			cli.ShowAppHelp(c)
			os.Exit(1)
		}

		if rotate && !isImagicConvertInPath(ImageMagicConvertBin) {
			log.Fatal("Rotation of jpegs was enables, but ImageMagic's 'convert' is not in path!")
			exitWithErr()
		}

		// verify RAW file extensions
		rawFileExts = strings.Split(rawExts, ",")
		for i, ext := range rawFileExts {
			rawFileExts[i] = strings.ToUpper(strings.TrimSpace(ext))
			if !isRawFileExtValid(rawFileExts[i]) {
				log.Printf("Error: Invalid Raw File extension: %s\n", rawFileExts[i])
				exitWithErr()
			}
		}

		// verify user-provided src dirs
		srcDirs = strings.Split(srcDir, ",")
		//log.Printf("SrcDirs: %s", srcDirs)
		for i, dir := range srcDirs {
			srcDirs[i] = strings.TrimSpace(dir)
			// ensure directory ends with '/'
			if !strings.HasSuffix(dir, "/") {
				srcDirs[i] += "/"
			}
			isValid, errStr := validateUserDir(srcDirs[i])
			if !isValid {
				log.Println("Error: ", errStr)
				exitWithErr()
			}
		}

		// verify user-provided dest dir
		isValid, errStr := validateUserDir(destDir)
		if !isValid {
			log.Println(errStr)
			exitWithErr()
		}
		// ensure directory ends with '/'
		if !strings.HasSuffix(destDir, "/") {
			destDir += "/"
		}

		processed = true
	}

	app.Run(os.Args)

	return processed
}

func initParsers() {
	parsers = rawparser.NewRawParsers()
	cr2Parser, cr2Key := rawparser.NewCr2Parser(isHostLittleEndian())
	parsers.Register(cr2Key, cr2Parser)

	nefParser, nefKey := rawparser.NewNefParser(isHostLittleEndian())
	parsers.Register(nefKey, nefParser)
}

func doProcess() int {
	log.Printf("RawTypes: %v SourceDirs: %v DestinationDir: %s JPEG Quality: %d Rotate Images: %v\n",
		rawFileExts, srcDirs, destDir, quality, rotate)

	var done sync.WaitGroup
	total := 0
	finish := make(chan struct{})

	// process all src dirs
	for _, dir := range srcDirs {

		// process all raw file types
		for i, rawType := range rawFileExts {

			globPattern := dir + "*." + rawType
			files, _ := getFilesForExt(globPattern)
			fileCnt := len(files)

			log.Println("Dir: ", dir)

			parser := parsers.GetParser(rawType)

			if fileCnt > 0 {
				log.Printf("Raw Type: %s ==> Processing '%s' %d files with %d NumCPU:\n",
					rawFileExts[i], dir, len(files), runtime.NumCPU())

				for _, file := range files {
					done.Add(1)

					go func(c chan struct{}) {
						rawfile, err := parser.ProcessFile(&rawparser.RawFileInfo{file, destDir, quality})
						if err != nil {
							log.Printf("Error processing file: '%s' error: %v\n", files, err)
						} else {
							if rotate && rawfile.JpegOrientation != 0.0 {
								// rotate jpeg
								go func(fileName string, radiansCw float64) {
									degrees := radiansCw * (180 / math.Pi)
									log.Printf("Rotating image %f degrees for jpeg: '%s'\n", degrees, fileName)
									cmd := exec.Command(ImageMagicConvertBin, "-rotate", strconv.FormatFloat(degrees, 'f', 2, 64), fileName, fileName)
									err := cmd.Start()
									if err != nil {
										log.Fatal(err)
									}
									err = cmd.Wait()
									if err != nil {
										log.Printf("Command finished with error: %v", err)
									}
								}(rawfile.JpegPath, rawfile.JpegOrientation)
							}
						}

						select {
						case <-c:
						}
						// signal completion of work
						done.Done()

					}(finish)
				}
				total += fileCnt
			}
		}
	}

	close(finish)

	done.Wait()

	return total
}

func setup() {
	success := processCli()
	if success {
		runtime.GOMAXPROCS(numOfRoutines)
		initParsers()
	} else {
		exitWithErr()
	}
}

func main() {
	t0 := time.Now()

	setup()

	cnt := doProcess()

	duration := time.Since(t0)

	log.Printf("jpgextract processed %d files in %02f mintues (%02f seconds)\n",
		cnt, duration.Minutes(), duration.Seconds())
}
<commit_msg>Added file param to go func.<commit_after>/*
 Copyright (c) 2013 Jeremy Torres, https://github.com/jeremytorres/jpegextract

 Permission is hereby granted, free of charge, to any person obtaining
 a copy of this software and associated documentation files (the
 "Software"), to deal in the Software without restriction, including
 without limitation the rights to use, copy, modify, merge, publish,
 distribute, sublicense, and/or sell copies of the Software, and to
 permit persons to whom the Software is furnished to do so, subject to
 the following conditions:

 The above copyright notice and this permission notice shall be
 included in all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
 LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

// Command utilizing the rawparser library for extracting jpegs from RAW iles.
// The embedded JPEGs are extracted and optionally:
// (1) rotated via oritenation provided in RawFile EXIF info;
//
// Usage:
//     jpgextract --raws "NEF,CR2" --dest-dir "/path_to/output_dir"
//                --src-dirs "/path_to/source1,/path_to/source2"
//                [--num-routines "8" --quality "80" --rotate]
package main

import (
	"github.com/codegangsta/cli"
	"github.com/jeremytorres/rawparser"
	"log"
	"math"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	// RawTypesKey is the constant representing the command line argument for the RAW file
	// types to be processed.
	RawTypesKey = "raws"

	// SrcDirKey is the constant representing the command line argument for the source directories
	// containing the source RawFiles.
	SrcDirKey = "src-dirs"
	// DestDirKey is a constant representing the command line argument for the destination directory
	// where extracted jpegs shall be placed.
	DestDirKey = "dest-dir"
	// NumRoutinesKey is the constant representing the command line argument for the max number of
	// go routines to be processed.  As this utility is mostly file IO bound, it is important
	// not to use too many routines ("too many files open..." errors can occur from OS).  Therefore,
	// if this value is greater than the logical cores, a warning is logged and the value is set to
	// the number of logical cores.
	NumRoutinesKey = "num-routines"
	// QualityKey is the constant representing the command line argument for the JPEG quality used
	// when processing the embedded JPEG in a RawFile.
	QualityKey = "quality"
	// RotateKey is the constant representing the command line argument indicating rotation of
	// JPEGs should occur based on the EXIF info embedded within the RawFile.  ImageMagic's 'convert'
	// utility is used and is checked at startup for existence.
	RotateKey = "rotate"
	// AppVersionKey is the constant defining the current version of this command utility.
	AppVersionKey = "1.0"
	// ImageMagicConvertBin is the constant representing ImageMagic's 'convert' utility.
	ImageMagicConvertBin = "convert"
)

var (
	destDir, sqlLiteDb     string
	rawFileExts, srcDirs   []string
	numOfRoutines, quality int
	rotate                 bool
	parsers                *rawparser.RawParsers
	// validParserKeys is a slice of RAW file parsers supported by this implementation.
	validParserKeys = []string{
		rawparser.NefParserKey,
		rawparser.Cr2ParserKey,
	}
)

/*
// processFilesConcurrent processes RawFiles concurrently, based on the number goroutines
// specified.
func processFilesConcurrent(rp *RawFileParserPair, c chan<- bool) {

	go func(rp *RawFileParserPair, c chan<- bool) {
		rawfile, err := rp.parser.ProcessFile(&rawparser.RawFileInfo{rp.file, destDir, quality})
		if err != nil {
			log.Printf("Error processing file: '%s' error: %v\n", rp.file, err)
		} else {
			if rotate && rawfile.JpegOrientation != 0.0 {
				// rotate jpeg
				go func(fileName string, radiansCw float64) {
					degrees := radiansCw * (180 / math.Pi)
					log.Printf("Rotating image %f degrees for jpeg: '%s'\n", degrees, fileName)
					cmd := exec.Command(ImageMagicConvertBin, "-rotate", strconv.FormatFloat(degrees, 'f', 2, 64), fileName, fileName)
					err := cmd.Start()
					if err != nil {
						log.Fatal(err)
					}
					err = cmd.Wait()
					if err != nil {
						log.Printf("Command finished with error: %v", err)
					}
				}(rawfile.JpegPath, rawfile.JpegOrientation)
			}
		}
		// signal completion of work
		c <- true
	}(rp, c)
}
*/

func isRawFileExtValid(ext string) bool {
	for _, validExt := range validParserKeys {
		if ext == validExt {
			return true
		}
	}
	return false
}

// processCli parses command line arguments and checks for validity of user-specified
// values.
func processCli() bool {
	var processed = false

	app := cli.NewApp()
	app.Name = "jpgextract"
	app.Usage = "Processes RAW files and extracts JPEGs and optionally create a SQLite database."
	app.Version = AppVersionKey
	app.Author = "Jeremy Torres"
	app.Flags = []cli.Flag{
		cli.StringFlag{RawTypesKey, "", "comma-separated list of RAW file extensions to process.  Current supported " +
			"RAW files: [NEF | CR2]"},
		cli.StringFlag{SrcDirKey, "", "comma-separated list of full paths to the directories containing RawFiles " +
			"(e.g., /path_to/dir1,/path_to/dir2..."},
		cli.StringFlag{DestDirKey, "", "the full path to the directory containing extracted jpegs"},
		cli.IntFlag{NumRoutinesKey, 2, "the number of concurrent files to be processed"},
		cli.IntFlag{QualityKey, 80, "JPEG encoding quality used for extracted jpegs"},
		cli.BoolFlag{RotateKey, "ImageMagic's 'convert' command will be used to rotate jpegs based on EXIF info from RawFile.  'convert' must be in PATH"},
	}
	app.Action = func(c *cli.Context) {
		rawExts := strings.TrimSpace(c.String(RawTypesKey))
		srcDir := strings.TrimSpace(c.String(SrcDirKey))
		destDir = strings.TrimSpace(c.String(DestDirKey))
		numOfRoutines = c.Int(NumRoutinesKey)
		quality = c.Int(QualityKey)
		rotate = c.Bool(RotateKey)

		// src and dest dirs required; remaing args have sane defaults
		if rawExts == "" || srcDir == "" || destDir == "" {
			cli.ShowAppHelp(c)
			os.Exit(1)
		}

		if rotate && !isImagicConvertInPath(ImageMagicConvertBin) {
			log.Fatal("Rotation of jpegs was enables, but ImageMagic's 'convert' is not in path!")
			exitWithErr()
		}

		// verify RAW file extensions
		rawFileExts = strings.Split(rawExts, ",")
		for i, ext := range rawFileExts {
			rawFileExts[i] = strings.ToUpper(strings.TrimSpace(ext))
			if !isRawFileExtValid(rawFileExts[i]) {
				log.Printf("Error: Invalid Raw File extension: %s\n", rawFileExts[i])
				exitWithErr()
			}
		}

		// verify user-provided src dirs
		srcDirs = strings.Split(srcDir, ",")
		//log.Printf("SrcDirs: %s", srcDirs)
		for i, dir := range srcDirs {
			srcDirs[i] = strings.TrimSpace(dir)
			// ensure directory ends with '/'
			if !strings.HasSuffix(dir, "/") {
				srcDirs[i] += "/"
			}
			isValid, errStr := validateUserDir(srcDirs[i])
			if !isValid {
				log.Println("Error: ", errStr)
				exitWithErr()
			}
		}

		// verify user-provided dest dir
		isValid, errStr := validateUserDir(destDir)
		if !isValid {
			log.Println(errStr)
			exitWithErr()
		}
		// ensure directory ends with '/'
		if !strings.HasSuffix(destDir, "/") {
			destDir += "/"
		}

		processed = true
	}

	app.Run(os.Args)

	return processed
}

func initParsers() {
	parsers = rawparser.NewRawParsers()
	cr2Parser, cr2Key := rawparser.NewCr2Parser(isHostLittleEndian())
	parsers.Register(cr2Key, cr2Parser)

	nefParser, nefKey := rawparser.NewNefParser(isHostLittleEndian())
	parsers.Register(nefKey, nefParser)
}

func doProcess() int {
	log.Printf("RawTypes: %v SourceDirs: %v DestinationDir: %s JPEG Quality: %d Rotate Images: %v\n",
		rawFileExts, srcDirs, destDir, quality, rotate)

	var done sync.WaitGroup
	total := 0
	finish := make(chan struct{})

	// process all src dirs
	for _, dir := range srcDirs {

		// process all raw file types
		for i, rawType := range rawFileExts {

			globPattern := dir + "*." + rawType
			files, _ := getFilesForExt(globPattern)
			fileCnt := len(files)

			parser := parsers.GetParser(rawType)

			if fileCnt > 0 {
				log.Printf("Raw Type: %s ==> Processing '%s' %d files with %d NumCPU:\n",
					rawFileExts[i], dir, len(files), runtime.NumCPU())

				for _, file := range files {
					done.Add(1)

					go func(c chan struct{}, file string) {
						rawfile, err := parser.ProcessFile(&rawparser.RawFileInfo{file, destDir, quality})
						if err != nil {
							log.Printf("Error with file: '%s'.  Error: %v\n", file, err)
						} else {
							if rotate && rawfile.JpegOrientation != 0.0 {
								// rotate jpeg
								go func(fileName string, radiansCw float64) {
									degrees := radiansCw * (180 / math.Pi)
									log.Printf("Rotating image %f degrees for jpeg: '%s'\n", degrees, fileName)
									cmd := exec.Command(ImageMagicConvertBin, "-rotate", strconv.FormatFloat(degrees, 'f', 2, 64), fileName, fileName)
									err := cmd.Start()
									if err != nil {
										log.Fatal(err)
									}
									err = cmd.Wait()
									if err != nil {
										log.Printf("Command finished with error: %v", err)
									}
								}(rawfile.JpegPath, rawfile.JpegOrientation)
							}
						}

						select {
						case <-c:
						}
						// signal completion of work
						done.Done()

					}(finish, file)
				}
				total += fileCnt
			}
		}
	}

	close(finish)

	done.Wait()

	return total
}

func setup() {
	success := processCli()
	if success {
		runtime.GOMAXPROCS(numOfRoutines)
		initParsers()
	} else {
		exitWithErr()
	}
}

func main() {
	t0 := time.Now()

	setup()

	cnt := doProcess()

	duration := time.Since(t0)

	log.Printf("jpgextract processed %d files in %02f mintues (%02f seconds)\n",
		cnt, duration.Minutes(), duration.Seconds())
}
<|endoftext|>


<commit_before>package isolated

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("ssh command", func() {
	Describe("help", func() {
		When("--help flag is provided", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("ssh", "--help")
				Eventually(session).Should(Say(`NAME:`))
				Eventually(session).Should(Say(`ssh - SSH to an application container instance`))
				Eventually(session).Should(Say(`USAGE:`))
				Eventually(session).Should(Say(`cf ssh APP_NAME \[-i INDEX\] \[-c COMMAND\]\.\.\. \[-L \[BIND_ADDRESS:\]PORT:HOST:HOST_PORT\] \[--skip-host-validation\] \[--skip-remote-execution\] \[--disable-pseudo-tty \| --force-pseudo-tty \| --request-pseudo-tty\]`))
				Eventually(session).Should(Say(`--app-instance-index, -i\s+Application instance index \(Default: 0\)`))
				Eventually(session).Should(Say(`--command, -c\s+Command to run\. This flag can be defined more than once\.`))
				Eventually(session).Should(Say(`--disable-pseudo-tty, -T\s+Disable pseudo-tty allocation`))
				Eventually(session).Should(Say(`--force-pseudo-tty\s+Force pseudo-tty allocation`))
				Eventually(session).Should(Say(`-L\s+Local port forward specification\. This flag can be defined more than once\.`))
				Eventually(session).Should(Say(`--request-pseudo-tty, -t\s+Request pseudo-tty allocation`))
				Eventually(session).Should(Say(`--skip-host-validation, -k\s+Skip host key validation`))
				Eventually(session).Should(Say(`--skip-remote-execution, -N\s+Do not execute a remote command`))
				Eventually(session).Should(Say(`SEE ALSO:`))
				Eventually(session).Should(Say(`allow-space-ssh, enable-ssh, space-ssh-allowed, ssh-code, ssh-enabled`))
				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("an application with multiple instances has been pushed", func() {
		var (
			appName          string
			appDirForCleanup string
			domainName       string
			orgName          string
			spaceName        string
		)

		BeforeEach(func() {
			helpers.LogoutCF()

			orgName = helpers.NewOrgName()
			spaceName = helpers.NewSpaceName()

			helpers.SetupCF(orgName, spaceName)
			appName = helpers.PrefixedRandomName("app")
			domainName = helpers.DefaultSharedDomain()
			helpers.WithHelloWorldApp(func(appDir string) {
				manifestContents := []byte(fmt.Sprintf(`
---
applications:
- name: %s
  memory: 128M
  instances: 2
  disk_quota: 128M
  routes:
  - route: %s.%s
`, appName, appName, domainName))
				manifestPath := filepath.Join(appDir, "manifest.yml")
				err := ioutil.WriteFile(manifestPath, manifestContents, 0666)
				Expect(err).ToNot(HaveOccurred())

				// Create manifest
				Eventually(helpers.CF("push", appName, "-p", appDir, "-f", manifestPath, "-b", "staticfile_buildpack", "--random-route")).Should(Exit(0))
				appDirForCleanup = appDir
			})
		})

		AfterEach(func() {
			Eventually(helpers.CF("delete", appName, "-f", "-r")).Should(Exit(0))
			Expect(os.RemoveAll(appDirForCleanup)).NotTo(HaveOccurred())
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app index is specified", func() {
			When("it is negative", func() {
				It("throws an error and informs the user that the app instance index cannot be negative", func() {
					session := helpers.CF("ssh", appName, "-i", "-1")
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Say("The application instance index cannot be negative"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the app index exceeds the last valid index", func() {
				It("throws an error and informs the user that the specified application does not exist", func() {
					session := helpers.CF("ssh", appName, "-i", "2")
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Say("The specified application instance does not exist"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("it is a valid index", func() {
				It("does not throw any error", func() {
					buffer := NewBuffer()
					buffer.Write([]byte("exit\n"))
					session := helpers.CFWithStdin(buffer, "ssh", appName, "-i", "0")
					Eventually(session).Should(Exit(0))
				})
			})
		})
	})
})
<commit_msg>Fix flaky SSH test in ssh_command_test<commit_after>package isolated

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"time"

	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("ssh command", func() {
	Describe("help", func() {
		When("--help flag is provided", func() {
			It("displays command usage to output", func() {
				session := helpers.CF("ssh", "--help")
				Eventually(session).Should(Say(`NAME:`))
				Eventually(session).Should(Say(`ssh - SSH to an application container instance`))
				Eventually(session).Should(Say(`USAGE:`))
				Eventually(session).Should(Say(`cf ssh APP_NAME \[-i INDEX\] \[-c COMMAND\]\.\.\. \[-L \[BIND_ADDRESS:\]PORT:HOST:HOST_PORT\] \[--skip-host-validation\] \[--skip-remote-execution\] \[--disable-pseudo-tty \| --force-pseudo-tty \| --request-pseudo-tty\]`))
				Eventually(session).Should(Say(`--app-instance-index, -i\s+Application instance index \(Default: 0\)`))
				Eventually(session).Should(Say(`--command, -c\s+Command to run\. This flag can be defined more than once\.`))
				Eventually(session).Should(Say(`--disable-pseudo-tty, -T\s+Disable pseudo-tty allocation`))
				Eventually(session).Should(Say(`--force-pseudo-tty\s+Force pseudo-tty allocation`))
				Eventually(session).Should(Say(`-L\s+Local port forward specification\. This flag can be defined more than once\.`))
				Eventually(session).Should(Say(`--request-pseudo-tty, -t\s+Request pseudo-tty allocation`))
				Eventually(session).Should(Say(`--skip-host-validation, -k\s+Skip host key validation`))
				Eventually(session).Should(Say(`--skip-remote-execution, -N\s+Do not execute a remote command`))
				Eventually(session).Should(Say(`SEE ALSO:`))
				Eventually(session).Should(Say(`allow-space-ssh, enable-ssh, space-ssh-allowed, ssh-code, ssh-enabled`))
				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("an application with multiple instances has been pushed", func() {
		var (
			appName          string
			appDirForCleanup string
			domainName       string
			orgName          string
			spaceName        string
		)

		BeforeEach(func() {
			helpers.LogoutCF()

			orgName = helpers.NewOrgName()
			spaceName = helpers.NewSpaceName()

			helpers.SetupCF(orgName, spaceName)
			appName = helpers.PrefixedRandomName("app")
			domainName = helpers.DefaultSharedDomain()
			helpers.WithHelloWorldApp(func(appDir string) {
				manifestContents := []byte(fmt.Sprintf(`
---
applications:
- name: %s
  memory: 128M
  instances: 2
  disk_quota: 128M
  routes:
  - route: %s.%s
`, appName, appName, domainName))
				manifestPath := filepath.Join(appDir, "manifest.yml")
				err := ioutil.WriteFile(manifestPath, manifestContents, 0666)
				Expect(err).ToNot(HaveOccurred())

				// Create manifest
				Eventually(helpers.CF("push", appName, "-p", appDir, "-f", manifestPath, "-b", "staticfile_buildpack", "--random-route")).Should(Exit(0))
				appDirForCleanup = appDir
			})
		})

		AfterEach(func() {
			Eventually(helpers.CF("delete", appName, "-f", "-r")).Should(Exit(0))
			Expect(os.RemoveAll(appDirForCleanup)).NotTo(HaveOccurred())
			helpers.QuickDeleteOrg(orgName)
		})

		When("the app index is specified", func() {
			When("it is negative", func() {
				It("throws an error and informs the user that the app instance index cannot be negative", func() {
					session := helpers.CF("ssh", appName, "-i", "-1")
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Say("The application instance index cannot be negative"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("the app index exceeds the last valid index", func() {
				It("throws an error and informs the user that the specified application does not exist", func() {
					session := helpers.CF("ssh", appName, "-i", "2")
					Eventually(session).Should(Say("FAILED"))
					Eventually(session).Should(Say("The specified application instance does not exist"))
					Eventually(session).Should(Exit(1))
				})
			})

			When("it is a valid index", func() {
				It("does not throw any error", func() {
					buffer := NewBuffer()
					buffer.Write([]byte("exit\n"))

					By("waiting for all instances to be running")
					Eventually(func() bool {
						session := helpers.CF("app", appName)
						Eventually(session).Should(Exit(0))
						output := session.Out.Contents()
						return regexp.MustCompile(`#\d.*running.*\n#\d.*running.*`).Match(output)
					}, 30*time.Second, 2*time.Second).Should(BeTrue())

					session := helpers.CFWithStdin(buffer, "ssh", appName, "-i", "0")
					Eventually(session).Should(Exit(0))
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>/*
Package sshmgr is a goroutine safe manager for SSH clients sharing between ssh/sftp sessions

It makes possible to share and reutilize existing client connections
for the same host `made with the same user and port` between multiple sessions and goroutines.

This is useful when yout application relies on SSH/SFTP for interacting with several
hosts and not spawn multiple connections to the same hosts, saving resources on both sides.

Clients are reference counted per session, and automatically closed/removed from the manager when all dependent sessions are closed.

	Usage:

		package main

		import (
			"github.com/brunotm/sshmgr"
		)

		func main() {
			config := sshmgr.NewConfig("hostA.domain.com", "user", "password", "or_key_file_path")
			sshSession, err := sshmgr.Manager.GetSSHSession(config)
			if err != nil {
				panic(err)
			}
			defer sshSession.Close()

			data, err := sshSession.CombinedOutput("uptime")
			if err != nil {
				panic(err)
			}

			fmt.Printf("%s: %s", config.NetAddr, string(data))
		}
*/
package sshmgr

import (
	"fmt"
	"sync"

	"github.com/brunotm/sshmgr/locker"
	"github.com/pkg/sftp"
)

// Manager is the package default ssh manager
var Manager = NewManager()

// SSHManager manage ssh clients and sessions.
// Clients are reference counted per session and removed from manager when the refcount reaches 0
type SSHManager struct {
	mtx     *sync.RWMutex
	locker  *locker.Locker
	clients map[string]*sshClient
}

// GetSSHSession creates a session from a active managed client or create a new one on demand
func (m *SSHManager) GetSSHSession(config *SSHConfig) (*SSHSession, error) {
	clientName := fmt.Sprintf("%s@%s:%s", config.User, config.NetAddr, config.Port)
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	m.mtx.RLock()
	client, exists := m.clients[clientName]
	m.mtx.RUnlock()

	// Found a active client, try to create a new session from it
	if exists {
		session, err := client.NewSession()
		if err != nil {
			return nil, err
		}
		client.incr()
		return &SSHSession{clientName: clientName, manager: m, Session: session}, nil
	}

	// Create a new client and session
	client, err := newSSHClient(config)
	if err != nil {
		return nil, err
	}

	session, err := client.NewSession()
	if err != nil {
		return nil, err
	}
	client.incr()

	m.mtx.Lock()
	m.clients[clientName] = client
	m.mtx.Unlock()

	return &SSHSession{clientName: clientName, manager: m, Session: session}, nil
}

// GetSFTPSession creates a session from a active managed client or create a new one on demand
func (m *SSHManager) GetSFTPSession(config *SSHConfig) (*SFTPSession, error) {
	clientName := fmt.Sprintf("%s@%s:%s", config.User, config.NetAddr, config.Port)
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	m.mtx.RLock()
	client, exists := m.clients[clientName]
	m.mtx.RUnlock()

	// Found a active client, try to create a new session from it
	if exists {
		session, err := sftp.NewClient(client.Client)
		if err != nil {
			return nil, err
		}
		client.incr()
		return &SFTPSession{clientName: clientName, manager: m, Client: session}, nil
	}

	// Create a new client and session
	client, err := newSSHClient(config)
	if err != nil {
		return nil, err
	}

	session, err := sftp.NewClient(client.Client)
	if err != nil {
		return nil, err
	}
	client.incr()

	m.mtx.Lock()
	m.clients[clientName] = client
	m.mtx.Unlock()

	return &SFTPSession{clientName: clientName, manager: m, Client: session}, nil
}

func (m *SSHManager) notifySessionClose(clientName string) {
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	m.mtx.RLock()
	client, exists := m.clients[clientName]
	m.mtx.RUnlock()

	if !exists {
		// We should never get here
		panic(fmt.Sprintf("Client not found: %s", clientName))
	}

	if client.decr() == 0 {
		defer client.Close()
		m.mtx.Lock()
		delete(m.clients, clientName)
		m.mtx.Unlock()

	}
}

// NewManager creates a new SSHManager
func NewManager() *SSHManager {
	return &SSHManager{&sync.RWMutex{}, locker.New(), map[string]*sshClient{}}
}
<commit_msg>Generalize client add and get from manager<commit_after>/*
Package sshmgr is a goroutine safe manager for SSH clients sharing between ssh/sftp sessions

It makes possible to share and reutilize existing client connections
for the same host `made with the same user and port` between multiple sessions and goroutines.

This is useful when yout application relies on SSH/SFTP for interacting with several
hosts and not spawn multiple connections to the same hosts, saving resources on both sides.

Clients are reference counted per session, and automatically closed/removed from the manager when all dependent sessions are closed.

	Usage:

		package main

		import (
			"github.com/brunotm/sshmgr"
		)

		func main() {
			config := sshmgr.NewConfig("hostA.domain.com", "user", "password", "or_key_file_path")
			sshSession, err := sshmgr.Manager.GetSSHSession(config)
			if err != nil {
				panic(err)
			}
			defer sshSession.Close()

			data, err := sshSession.CombinedOutput("uptime")
			if err != nil {
				panic(err)
			}

			fmt.Printf("%s: %s", config.NetAddr, string(data))
		}
*/
package sshmgr

import (
	"fmt"
	"sync"

	"github.com/brunotm/sshmgr/locker"
	"github.com/pkg/sftp"
)

// Manager is the package default ssh manager
var Manager = NewManager()

// SSHManager manage ssh clients and sessions.
// Clients are reference counted per session and removed from manager when the refcount reaches 0
type SSHManager struct {
	mtx     *sync.RWMutex
	locker  *locker.Locker
	clients map[string]*sshClient
}

// addClient adds a client to manager client map
func (m *SSHManager) addClient(clientName string, client *sshClient) {
	m.mtx.Lock()
	defer m.mtx.Unlock()
	m.clients[clientName] = client
}

// getClient get a existing client from manager client map
func (m *SSHManager) getClient(clientName string) *sshClient {
	m.mtx.RLock()
	defer m.mtx.RUnlock()
	return m.clients[clientName]
}

// GetSSHSession creates a session from a active managed client or create a new one on demand
func (m *SSHManager) GetSSHSession(config *SSHConfig) (*SSHSession, error) {
	clientName := fmt.Sprintf("%s@%s:%s", config.User, config.NetAddr, config.Port)
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	// If a existing client is found try to create a session from it
	client := m.getClient(clientName)
	if client != nil {
		session, err := client.NewSession()
		if err != nil {
			return nil, err
		}
		client.incr()
		return &SSHSession{clientName: clientName, manager: m, Session: session}, nil
	}

	// Create a new client and session
	client, err := newSSHClient(config)
	if err != nil {
		return nil, err
	}

	session, err := client.NewSession()
	if err != nil {
		return nil, err
	}
	client.incr()
	m.addClient(clientName, client)

	return &SSHSession{clientName: clientName, manager: m, Session: session}, nil
}

// GetSFTPSession creates a session from a active managed client or create a new one on demand
func (m *SSHManager) GetSFTPSession(config *SSHConfig) (*SFTPSession, error) {
	clientName := fmt.Sprintf("%s@%s:%s", config.User, config.NetAddr, config.Port)
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	// If a existing client is found try to create a session from it
	client := m.getClient(clientName)
	if client != nil {
		session, err := sftp.NewClient(client.Client)
		if err != nil {
			return nil, err
		}
		client.incr()
		return &SFTPSession{clientName: clientName, manager: m, Client: session}, nil
	}

	// Create a new client and session
	client, err := newSSHClient(config)
	if err != nil {
		return nil, err
	}

	session, err := sftp.NewClient(client.Client)
	if err != nil {
		return nil, err
	}
	client.incr()
	m.addClient(clientName, client)

	return &SFTPSession{clientName: clientName, manager: m, Client: session}, nil
}

func (m *SSHManager) notifySessionClose(clientName string) {
	m.locker.Lock(clientName)
	defer m.locker.Unlock(clientName)

	client := m.getClient(clientName)

	if client == nil {
		// We should never get here
		panic(fmt.Sprintf("Client not found: %s", clientName))
	}

	if client.decr() == 0 {
		defer client.Close()
		m.mtx.Lock()
		delete(m.clients, clientName)
		m.mtx.Unlock()

	}
}

// NewManager creates a new SSHManager
func NewManager() *SSHManager {
	return &SSHManager{&sync.RWMutex{}, locker.New(), map[string]*sshClient{}}
}
<|endoftext|>


<commit_before>package job

import (
	"gonzbee/config"
	"gonzbee/nntp"
	"gonzbee/nzb"
	"gonzbee/yenc"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

type messagejob struct {
	group string
	msgId string
	ch    chan io.ReadCloser
}

func init() {
	go poolHandler()
}

var download = make(chan *messagejob)
var downloadMux = make(chan *messagejob)
var reaper = make(chan int)

func newConnection() error {
	s := config.C.Server.GetAddressStr()
	var err error
	var n *nntp.Conn
	if config.C.Server.TLS {
		n, err = nntp.DialTLS(s)
	} else {
		n, err = nntp.Dial(s)
	}
	if err != nil {
		return err
	}

	err = n.Authenticate(config.C.Server.Username, config.C.Server.Password)
	if err != nil {
		n.Close()
		return err
	}
	log.Println("spun up nntp connection")
	go func() {
		defer n.Close()
		for {
			select {
			case m := <-downloadMux:
				err = n.SwitchGroup(m.group)
				if err != nil {
					panic(err)
				}
				b, err := n.GetMessageReader(m.msgId)
				if err != nil {
					log.Print("Error getting Message ", m.msgId, ": ", err.Error())
				}
				m.ch <- b
			case <-(after(10 * time.Second)):
				reaper <- 1
				return
			}
		}
	}()
	return nil
}

func after(d time.Duration) <-chan time.Time {
	t := time.NewTimer(d)
	return t.C
}

func poolHandler() {
	var number int
	for {
		select {
		case msg := <-download:
			if number < 10 {
				err := newConnection()
				if err == nil {
					number++
				}
			}
			downloadMux <- msg
		case <-reaper:
			number--
		}
	}
}

func (j *job) handle() {
	wg := new(sync.WaitGroup)
	for _, f := range j.n.File {
		ch := make(chan io.ReadCloser)
		wg.Add(1)
		partsLeft := len(f.Segments)
		go func(ret chan io.ReadCloser) {
			m := <-ret
			for m == nil {
				partsLeft--
				m = <-ret
			}
			part, _ := yenc.NewPart(m)
			file, _ := os.Create(filepath.Join(j.dir, part.Filename))
			file.Seek(part.Begin, os.SEEK_SET)
			part.Decode(file)
			m.Close()
			partsLeft--
			for ; partsLeft > 0; partsLeft-- {
				m = <-ret
				if m == nil {
					continue
				}
				part, _ := yenc.NewPart(m)
				file.Seek(part.Begin, os.SEEK_SET)
				part.Decode(file)
				m.Close()
			}
			file.Close()
			wg.Done()
		}(ch)
		for _, seg := range f.Segments {
			msg := &messagejob{
				msgId: seg.MsgId,
				group: f.Groups[0],
				ch:    ch,
			}
			download <- msg
		}
	}
	wg.Wait()
}

func Start(n *nzb.Nzb, name string) {
	incDir := config.C.GetIncompleteDir()
	workDir := filepath.Join(incDir, name)
	os.Mkdir(workDir, 0777)
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
}
<commit_msg>make job a bit better<commit_after>package job

import (
	"gonzbee/config"
	"gonzbee/nntp"
	"gonzbee/nzb"
	"gonzbee/yenc"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

type messagejob struct {
	group string
	msgId string
	ch    chan io.ReadCloser
}

func init() {
	go poolHandler()
}

var download = make(chan *messagejob)
var downloadMux = make(chan *messagejob)
var reaper = make(chan int)

func newConnection() error {
	s := config.C.Server.GetAddressStr()
	var err error
	var n *nntp.Conn
	if config.C.Server.TLS {
		n, err = nntp.DialTLS(s)
	} else {
		n, err = nntp.Dial(s)
	}
	if err != nil {
		return err
	}

	err = n.Authenticate(config.C.Server.Username, config.C.Server.Password)
	if err != nil {
		n.Close()
		return err
	}
	log.Println("spun up nntp connection")
	go func() {
		defer n.Close()
		for {
			select {
			case m := <-downloadMux:
				err = n.SwitchGroup(m.group)
				if err != nil {
					panic(err)
				}
				b, err := n.GetMessageReader(m.msgId)
				if err != nil {
					log.Print("Error getting Message ", m.msgId, ": ", err.Error())
				}
				m.ch <- b
			case <-(after(10 * time.Second)):
				reaper <- 1
				return
			}
		}
	}()
	return nil
}

func after(d time.Duration) <-chan time.Time {
	t := time.NewTimer(d)
	return t.C
}

func poolHandler() {
	var number int
	for {
		select {
		case msg := <-download:
			if number < 10 {
				err := newConnection()
				if err == nil {
					number++
				}
			}
			downloadMux <- msg
		case <-reaper:
			number--
		}
	}
}

func (j *job) handle() {
	wg := new(sync.WaitGroup)
	for _, f := range j.n.File {
		ch := make(chan io.ReadCloser)
		wg.Add(1)
		partsLeft := len(f.Segments)
		go func(ret chan io.ReadCloser) {
			var file *os.File
			var part *yenc.Part
			var err error
			defer wg.Done()
			for ; partsLeft > 0; partsLeft-- {
				m := <-ret
				if m == nil {
					continue
				}
				part, err = yenc.NewPart(m)
				if err != nil {
					log.Print(err.Error())
					m.Close()
					continue
				}
				if file == nil {
					file, err = os.Create(filepath.Join(j.dir, part.Filename))
					if err != nil {
						panic("could not create file: " + err.Error())
					}
					defer file.Close()
				}
				file.Seek(part.Begin, os.SEEK_SET)
				part.Decode(file)
				m.Close()
			}
			if part != nil {
				log.Print("Done Decoding file " + part.Filename)
			} else {
				log.Print("Could not decode entire file")
			}
		}(ch)
		for _, seg := range f.Segments {
			msg := &messagejob{
				msgId: seg.MsgId,
				group: f.Groups[0],
				ch:    ch,
			}
			download <- msg
		}
	}
	wg.Wait()
}

func Start(n *nzb.Nzb, name string) {
	incDir := config.C.GetIncompleteDir()
	workDir := filepath.Join(incDir, name)
	os.Mkdir(workDir, 0777)
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
}
<|endoftext|>