<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/tarm/serial"
	"gopkg.in/readline.v1"
)

var (
	rlInstance *readline.Instance
	conn       *serial.Port
	serIn      = make(chan []byte)
	outBound   = make(chan string)
	progress   = make(chan bool)
	incLevel   = make(chan int)
)

func main() {
	var err error
	rlInstance, err = readline.NewEx(&readline.Config{
		UniqueEditLine: true,
	})
	check(err)
	defer rlInstance.Close()

	tty := "/dev/cu.SLAB_USBtoUART"
	conn, err = serial.OpenPort(&serial.Config{Name: tty, Baud: 115200})
	check(err)

	go serialInput()
	go serialExchange()

	outBound <- ""
	<-progress
	for {
		line, err := rlInstance.Readline()
		if err != nil { // io.EOF, readline.ErrInterrupt
			break
		}
		if strings.HasPrefix(line, "include ") {
			doInclude(line[8:])
		} else {
			outBound <- line
			<-progress
		}
	}
}

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

func serialInput() {
	buf := make([]byte, 128)
	for {
		n, err := conn.Read(buf)
		check(err)
		if n == 0 {
			close(serIn)
			return
		}
		serIn <- buf[:n]
	}
}

func readWithTimeout() []byte {
	select {
	case data := <-serIn:
		return data
	case <-time.After(500 * time.Millisecond):
		return nil
	}
}

func serialExchange() {
	includeDepth := 0
	for {
		select {
		case data := <-serIn:
			if len(data) == 0 {
				return
			}
			print(string(data))
		case line := <-outBound:
			// the task here is to omit "normal" output for included lines,
			// i.e. lines which only generate an echo, a space, and " ok.\n"
			// everything else should be echoed in full, including the input
			including := includeDepth > 0
			if len(line) > 0 {
				serialSend(line)
				prefix, _ := expectEcho(line)
				print(prefix)
				if !including {
					print(line)
				}
			}
			serialSend("\r")
			prefix, _ := expectEcho(" ok.\n")
			if including && prefix != " " {
				print(line)
			}
			print(prefix)
			if !including || prefix != " " {
				print(" ok.\n")
			}
			progress <- true
		case n := <-incLevel:
			includeDepth += n
		}
	}
}

func expectEcho(match string) (string, bool) {
	var collected []byte
	for {
		data := readWithTimeout()
		collected = append(collected, data...)
		if bytes.HasSuffix(collected, []byte(match)) {
			bytesBefore := len(collected) - len(match)
			return string(collected[:bytesBefore]), true
		}
		if len(data) == 0 || bytes.IndexByte(collected, '\n') >= 0 {
			return string(collected), false
		}
	}
}

func serialSend(data string) {
	_, err := conn.Write([]byte(data))
	check(err)
}

func doInclude(fname string) {
	incLevel <- +1
	defer func() { incLevel <- -1 }()

	lineNum := 0
	fmt.Printf("\t>>> include %s\n", fname)
	defer func() {
		fmt.Printf("\t<<<<<<<<<<< %s (%d lines)\n", fname, lineNum)
	}()

	f, err := os.Open(fname)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		lineNum++

		s := strings.TrimLeft(line, " ")
		if s == "" || strings.HasPrefix(s, "\\ ") {
			continue // don't send empty or comment-only lines
		}

		outBound <- line
		<-progress
	}
}
<commit_msg>better matching logic, still fails on large inputs<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/tarm/serial"
	"gopkg.in/readline.v1"
)

var (
	rlInstance *readline.Instance
	conn       *serial.Port
	serIn      = make(chan []byte)
	outBound   = make(chan string)
	progress   = make(chan bool)
	incLevel   = make(chan int)
)

func main() {
	var err error
	rlInstance, err = readline.NewEx(&readline.Config{
		UniqueEditLine: true,
	})
	check(err)
	defer rlInstance.Close()

	tty := "/dev/cu.SLAB_USBtoUART"
	conn, err = serial.OpenPort(&serial.Config{Name: tty, Baud: 115200})
	check(err)

	go serialInput()
	go serialExchange()

	outBound <- ""
	<-progress
	for {
		line, err := rlInstance.Readline()
		if err != nil { // io.EOF, readline.ErrInterrupt
			break
		}
		if strings.HasPrefix(line, "include ") {
			doInclude(line[8:])
		} else {
			outBound <- line
			<-progress
		}
	}
}

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

func serialInput() {
	buf := make([]byte, 128)
	for {
		n, err := conn.Read(buf)
		check(err)
		if n == 0 {
			close(serIn)
			return
		}
		serIn <- buf[:n]
	}
}

func readWithTimeout() []byte {
	select {
	case data := <-serIn:
		return data
	case <-time.After(500 * time.Millisecond):
		return nil
	}
}

func serialExchange() {
	includeDepth := 0
	for {
		select {
		case data := <-serIn:
			if len(data) == 0 {
				return
			}
			print(string(data))
		case line := <-outBound:
			// the task here is to omit "normal" output for included lines,
			// i.e. lines which only generate an echo, a space, and " ok.\n"
			// everything else should be echoed in full, including the input
			including := includeDepth > 0
			prompt := " ok.\n"
			if len(line) > 0 {
				serialSend(line)
				prefix, matched := expectEcho(line)
				print(prefix)
				if matched && !including {
					print(line)
					line = ""
				}
			}
			// now that the echo is done, send a CR and wait for the prompt
			serialSend("\r")
			prefix, matched := expectEcho(prompt)
			if !matched {
				prompt = ""
			}
			if !including || prefix != " " || !matched {
				print(line + prefix + prompt)
			}
			// signal to sender that this request has been processed
			progress <- matched
		case n := <-incLevel:
			includeDepth += n
		}
	}
}

func expectEcho(match string) (string, bool) {
	var collected []byte
	for {
		data := readWithTimeout()
		collected = append(collected, data...)
		if bytes.HasSuffix(collected, []byte(match)) {
			bytesBefore := len(collected) - len(match)
			return string(collected[:bytesBefore]), true
		}
		if len(data) == 0 || bytes.IndexByte(collected, '\n') >= 0 {
			return string(collected), false
		}
	}
}

func serialSend(data string) {
	_, err := conn.Write([]byte(data))
	check(err)
}

func doInclude(fname string) {
	incLevel <- +1
	defer func() { incLevel <- -1 }()

	lineNum := 0
	fmt.Printf("\t>>> include %s\n", fname)
	defer func() {
		fmt.Printf("\t<<<<<<<<<<< %s (%d lines)\n", fname, lineNum)
	}()

	f, err := os.Open(fname)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		lineNum++

		s := strings.TrimLeft(line, " ")
		if s == "" || strings.HasPrefix(s, "\\ ") {
			continue // don't send empty or comment-only lines
		}

		outBound <- line
		<-progress
	}
}
<|endoftext|>


<commit_before>package fifo

import (
	"context"
	"io"
	"os"
	"runtime"
	"sync"
	"syscall"

	"github.com/pkg/errors"
)

type fifo struct {
	flag        int
	opened      chan struct{}
	closed      chan struct{}
	closing     chan struct{}
	err         error
	file        *os.File
	closingOnce sync.Once // close has been called
	closedOnce  sync.Once // fifo is closed
	handle      *handle
}

var leakCheckWg *sync.WaitGroup

// OpenFifo opens a fifo. Returns io.ReadWriteCloser.
// Context can be used to cancel this function until open(2) has not returned.
// Accepted flags:
// - syscall.O_CREAT - create new fifo if one doesn't exist
// - syscall.O_RDONLY - open fifo only from reader side
// - syscall.O_WRONLY - open fifo only from writer side
// - syscall.O_RDWR - open fifo from both sides, never block on syscall level
// - syscall.O_NONBLOCK - return io.ReadWriteCloser even if other side of the
//     fifo isn't open. read/write will be connected after the actual fifo is
//     open or after fifo is closed.
func OpenFifo(ctx context.Context, fn string, flag int, perm os.FileMode) (io.ReadWriteCloser, error) {
	if _, err := os.Stat(fn); err != nil {
		if os.IsNotExist(err) && flag&syscall.O_CREAT != 0 {
			if err := syscall.Mkfifo(fn, uint32(perm&os.ModePerm)); err != nil && !os.IsExist(err) {
				return nil, errors.Wrapf(err, "error creating fifo %v", fn)
			}
		} else {
			return nil, err
		}
	}

	block := flag&syscall.O_NONBLOCK == 0 || flag&syscall.O_RDWR != 0

	flag &= ^syscall.O_CREAT
	flag &= ^syscall.O_NONBLOCK

	h, err := getHandle(fn)
	if err != nil {
		return nil, err
	}

	f := &fifo{
		handle:  h,
		flag:    flag,
		opened:  make(chan struct{}),
		closed:  make(chan struct{}),
		closing: make(chan struct{}),
	}

	wg := leakCheckWg
	if wg != nil {
		wg.Add(2)
	}

	go func() {
		if wg != nil {
			defer wg.Done()
		}
		select {
		case <-ctx.Done():
			f.Close()
		case <-f.opened:
		case <-f.closed:
		}
	}()
	go func() {
		if wg != nil {
			defer wg.Done()
		}
		var file *os.File
		fn, err := h.Path()
		if err == nil {
			file, err = os.OpenFile(fn, flag, 0)
		}
		select {
		case <-f.closing:
			if err == nil {
				select {
				case <-ctx.Done():
					err = ctx.Err()
				default:
					err = errors.Errorf("fifo %v was closed before opening", fn)
				}
				if file != nil {
					file.Close()
				}
			}
		default:
		}
		if err != nil {
			f.closedOnce.Do(func() {
				f.err = err
				close(f.closed)
			})
			return
		}
		f.file = file
		close(f.opened)
	}()
	if block {
		select {
		case <-f.opened:
		case <-f.closed:
			return nil, f.err
		}
	}
	return f, nil
}

// Read from a fifo to a byte array.
func (f *fifo) Read(b []byte) (int, error) {
	if f.flag&syscall.O_WRONLY > 0 {
		return 0, errors.New("reading from write-only fifo")
	}
	select {
	case <-f.opened:
		return f.file.Read(b)
	default:
	}
	select {
	case <-f.opened:
		return f.file.Read(b)
	case <-f.closed:
		return 0, errors.New("reading from a closed fifo")
	}
}

// Write from byte array to a fifo.
func (f *fifo) Write(b []byte) (int, error) {
	if f.flag&(syscall.O_WRONLY|syscall.O_RDWR) == 0 {
		return 0, errors.New("writing to read-only fifo")
	}
	select {
	case <-f.opened:
		return f.file.Write(b)
	default:
	}
	select {
	case <-f.opened:
		return f.file.Write(b)
	case <-f.closed:
		return 0, errors.New("writing to a closed fifo")
	}
}

// Close the fifo. Next reads/writes will error. This method can also be used
// before open(2) has returned and fifo was never opened.
func (f *fifo) Close() error {
	for {
		select {
		case <-f.closed:
			f.handle.Close()
			return f.err
		default:
			select {
			case <-f.opened:
				f.closedOnce.Do(func() {
					f.err = f.file.Close()
					close(f.closed)
				})
			default:
				if f.flag&syscall.O_RDWR != 0 {
					runtime.Gosched()
					break
				}
				f.closingOnce.Do(func() {
					close(f.closing)
				})
				reverseMode := syscall.O_WRONLY
				if f.flag&syscall.O_WRONLY > 0 {
					reverseMode = syscall.O_RDONLY
				}
				fn, err := f.handle.Path()
				// if Close() is called concurrently(shouldn't) it may cause error
				// because handle is closed
				select {
				case <-f.closed:
				default:
					if err != nil {
						// Path has become invalid. We will leak a goroutine.
						// This case should not happen in linux.
						f.closedOnce.Do(func() {
							f.err = err
							close(f.closed)
						})
						<-f.closed
						break
					}
					f, err := os.OpenFile(fn, reverseMode|syscall.O_NONBLOCK, 0)
					if err == nil {
						f.Close()
					}
					runtime.Gosched()
				}
			}
		}
	}
}
<commit_msg>use context from net pkg for go1.6 compat<commit_after>package fifo

import (
	"io"
	"os"
	"runtime"
	"sync"
	"syscall"

	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

type fifo struct {
	flag        int
	opened      chan struct{}
	closed      chan struct{}
	closing     chan struct{}
	err         error
	file        *os.File
	closingOnce sync.Once // close has been called
	closedOnce  sync.Once // fifo is closed
	handle      *handle
}

var leakCheckWg *sync.WaitGroup

// OpenFifo opens a fifo. Returns io.ReadWriteCloser.
// Context can be used to cancel this function until open(2) has not returned.
// Accepted flags:
// - syscall.O_CREAT - create new fifo if one doesn't exist
// - syscall.O_RDONLY - open fifo only from reader side
// - syscall.O_WRONLY - open fifo only from writer side
// - syscall.O_RDWR - open fifo from both sides, never block on syscall level
// - syscall.O_NONBLOCK - return io.ReadWriteCloser even if other side of the
//     fifo isn't open. read/write will be connected after the actual fifo is
//     open or after fifo is closed.
func OpenFifo(ctx context.Context, fn string, flag int, perm os.FileMode) (io.ReadWriteCloser, error) {
	if _, err := os.Stat(fn); err != nil {
		if os.IsNotExist(err) && flag&syscall.O_CREAT != 0 {
			if err := syscall.Mkfifo(fn, uint32(perm&os.ModePerm)); err != nil && !os.IsExist(err) {
				return nil, errors.Wrapf(err, "error creating fifo %v", fn)
			}
		} else {
			return nil, err
		}
	}

	block := flag&syscall.O_NONBLOCK == 0 || flag&syscall.O_RDWR != 0

	flag &= ^syscall.O_CREAT
	flag &= ^syscall.O_NONBLOCK

	h, err := getHandle(fn)
	if err != nil {
		return nil, err
	}

	f := &fifo{
		handle:  h,
		flag:    flag,
		opened:  make(chan struct{}),
		closed:  make(chan struct{}),
		closing: make(chan struct{}),
	}

	wg := leakCheckWg
	if wg != nil {
		wg.Add(2)
	}

	go func() {
		if wg != nil {
			defer wg.Done()
		}
		select {
		case <-ctx.Done():
			f.Close()
		case <-f.opened:
		case <-f.closed:
		}
	}()
	go func() {
		if wg != nil {
			defer wg.Done()
		}
		var file *os.File
		fn, err := h.Path()
		if err == nil {
			file, err = os.OpenFile(fn, flag, 0)
		}
		select {
		case <-f.closing:
			if err == nil {
				select {
				case <-ctx.Done():
					err = ctx.Err()
				default:
					err = errors.Errorf("fifo %v was closed before opening", fn)
				}
				if file != nil {
					file.Close()
				}
			}
		default:
		}
		if err != nil {
			f.closedOnce.Do(func() {
				f.err = err
				close(f.closed)
			})
			return
		}
		f.file = file
		close(f.opened)
	}()
	if block {
		select {
		case <-f.opened:
		case <-f.closed:
			return nil, f.err
		}
	}
	return f, nil
}

// Read from a fifo to a byte array.
func (f *fifo) Read(b []byte) (int, error) {
	if f.flag&syscall.O_WRONLY > 0 {
		return 0, errors.New("reading from write-only fifo")
	}
	select {
	case <-f.opened:
		return f.file.Read(b)
	default:
	}
	select {
	case <-f.opened:
		return f.file.Read(b)
	case <-f.closed:
		return 0, errors.New("reading from a closed fifo")
	}
}

// Write from byte array to a fifo.
func (f *fifo) Write(b []byte) (int, error) {
	if f.flag&(syscall.O_WRONLY|syscall.O_RDWR) == 0 {
		return 0, errors.New("writing to read-only fifo")
	}
	select {
	case <-f.opened:
		return f.file.Write(b)
	default:
	}
	select {
	case <-f.opened:
		return f.file.Write(b)
	case <-f.closed:
		return 0, errors.New("writing to a closed fifo")
	}
}

// Close the fifo. Next reads/writes will error. This method can also be used
// before open(2) has returned and fifo was never opened.
func (f *fifo) Close() error {
	for {
		select {
		case <-f.closed:
			f.handle.Close()
			return f.err
		default:
			select {
			case <-f.opened:
				f.closedOnce.Do(func() {
					f.err = f.file.Close()
					close(f.closed)
				})
			default:
				if f.flag&syscall.O_RDWR != 0 {
					runtime.Gosched()
					break
				}
				f.closingOnce.Do(func() {
					close(f.closing)
				})
				reverseMode := syscall.O_WRONLY
				if f.flag&syscall.O_WRONLY > 0 {
					reverseMode = syscall.O_RDONLY
				}
				fn, err := f.handle.Path()
				// if Close() is called concurrently(shouldn't) it may cause error
				// because handle is closed
				select {
				case <-f.closed:
				default:
					if err != nil {
						// Path has become invalid. We will leak a goroutine.
						// This case should not happen in linux.
						f.closedOnce.Do(func() {
							f.err = err
							close(f.closed)
						})
						<-f.closed
						break
					}
					f, err := os.OpenFile(fn, reverseMode|syscall.O_NONBLOCK, 0)
					if err == nil {
						f.Close()
					}
					runtime.Gosched()
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
 * transport_test.go
 *
 * Copyright 2017 Bill Zissimopoulos
 */
/*
 * This file is part of netchan.
 *
 * It is licensed under the MIT license. The full license text can be found
 * in the License.txt file at the root of this project.
 */

package netchan

import (
	"crypto/tls"
	"net/url"
	"reflect"
	"sync/atomic"
	"testing"
)

type testTransportRecverSender struct {
	id       string
	msg      string
	t        *testing.T
	rcnt     uint32
	done     chan struct{}
	doneflag bool
}

func newTestTransportRecverSender(id string, msg string, t *testing.T) *testTransportRecverSender {
	return &testTransportRecverSender{
		id,
		msg,
		t,
		0,
		make(chan struct{}),
		false,
	}
}

func (self *testTransportRecverSender) transportRecver(link Link) error {
	id, vmsg, err := link.Recv()
	if nil != err {
		if self.doneflag {
			return ErrTransportClosed
		}
		self.t.Errorf("Recver: error %v", err)
		return err
	}
	if id != self.id {
		self.t.Errorf("Recver: id mismatch")
	}
	msg := vmsg.Interface()
	if !reflect.DeepEqual(msg, self.msg) {
		self.t.Errorf("Recver: msg mismatch")
	}

	if 2 == atomic.AddUint32(&self.rcnt, 1) {
		self.doneflag = true
		close(self.done)
	}

	return nil
}

func (self *testTransportRecverSender) transportSender(link Link) error {
	err := link.Send(self.id, reflect.ValueOf(self.msg))
	if nil != err {
		self.t.Errorf("Sender: error %v", err)
		return err
	}

	<-self.done

	return ErrTransportClosed
}

func testTransport(t *testing.T, transport Transport, scheme string) {
	id0 := "NAME"
	msg := "42,43,44,45,46"
	trs := newTestTransportRecverSender(id0, msg, t)
	transport.SetRecver(trs.transportRecver)
	transport.SetSender(trs.transportSender)

	err := transport.Listen()
	if nil != err {
		panic(err)
	}
	err = transport.Listen()
	if nil != err {
		panic(err)
	}

	id, link, err := transport.Connect(&url.URL{
		Scheme: scheme,
		Host:   "127.0.0.1",
		Path:   "/" + id0,
	})
	if nil != err {
		panic(err)
	}

	if id0 != id {
		t.Errorf("Connect: incorrect id: %v", id)
	}

	link.Open()
	link.Open()

	<-trs.done
}

func TestTcpTransport(t *testing.T) {
	marshaler := NewGobMarshaler()
	transport := NewNetTransport(
		marshaler,
		&url.URL{
			Scheme: "tcp",
			Host:   ":25000",
		},
		nil)
	defer transport.Close()

	testTransport(t, transport, "tcp")
}

func TestTlsTransport(t *testing.T) {
	cert, err := tls.X509KeyPair([]byte(tlscert), []byte(tlskey))
	if nil != err {
		panic(err)
	}

	marshaler := NewGobMarshaler()
	transport := NewNetTransportTLS(
		marshaler,
		&url.URL{
			Scheme: "tls",
			Host:   ":25000",
		},
		nil,
		&tls.Config{
			Certificates:       []tls.Certificate{cert},
			InsecureSkipVerify: true,
		})
	defer transport.Close()

	testTransport(t, transport, "tls")
}
<commit_msg>netTransport: RedialTimeout testing<commit_after>/*
 * transport_test.go
 *
 * Copyright 2017 Bill Zissimopoulos
 */
/*
 * This file is part of netchan.
 *
 * It is licensed under the MIT license. The full license text can be found
 * in the License.txt file at the root of this project.
 */

package netchan

import (
	"crypto/tls"
	"net/url"
	"reflect"
	"sync/atomic"
	"testing"
	"time"
)

type testTransportRecverSender struct {
	id       string
	msg      string
	t        *testing.T
	rcnt     uint32
	done     chan struct{}
	doneflag bool
}

func newTestTransportRecverSender(id string, msg string, t *testing.T) *testTransportRecverSender {
	return &testTransportRecverSender{
		id,
		msg,
		t,
		0,
		make(chan struct{}),
		false,
	}
}

func (self *testTransportRecverSender) transportRecver(link Link) error {
	id, vmsg, err := link.Recv()
	if nil != err {
		if self.doneflag {
			return ErrTransportClosed
		}
		self.t.Errorf("Recver: error %v", err)
		return err
	}
	if id != self.id {
		self.t.Errorf("Recver: id mismatch")
	}
	msg := vmsg.Interface()
	if !reflect.DeepEqual(msg, self.msg) {
		self.t.Errorf("Recver: msg mismatch")
	}

	if 2 == atomic.AddUint32(&self.rcnt, 1) {
		self.doneflag = true
		close(self.done)
	}

	return nil
}

func (self *testTransportRecverSender) transportSender(link Link) error {
	err := link.Send(self.id, reflect.ValueOf(self.msg))
	if nil != err {
		self.t.Errorf("Sender: error %v", err)
		return err
	}

	<-self.done

	return ErrTransportClosed
}

func testTransport(t *testing.T, transport Transport, scheme string) {
	id0 := "NAME"
	msg := "42,43,44,45,46"
	trs := newTestTransportRecverSender(id0, msg, t)
	transport.SetRecver(trs.transportRecver)
	transport.SetSender(trs.transportSender)

	err := transport.Listen()
	if nil != err {
		panic(err)
	}
	err = transport.Listen()
	if nil != err {
		panic(err)
	}

	id, link, err := transport.Connect(&url.URL{
		Scheme: scheme,
		Host:   "127.0.0.1",
		Path:   "/" + id0,
	})
	if nil != err {
		panic(err)
	}

	if id0 != id {
		t.Errorf("Connect: incorrect id: %v", id)
	}

	link.Open()
	link.Open()

	<-trs.done
}

func TestTcpTransport(t *testing.T) {
	marshaler := NewGobMarshaler()
	transport := NewNetTransport(
		marshaler,
		&url.URL{
			Scheme: "tcp",
			Host:   ":25000",
		},
		nil)
	defer transport.Close()

	testTransport(t, transport, "tcp")
}

func TestTlsTransport(t *testing.T) {
	cert, err := tls.X509KeyPair([]byte(tlscert), []byte(tlskey))
	if nil != err {
		panic(err)
	}

	marshaler := NewGobMarshaler()
	transport := NewNetTransportTLS(
		marshaler,
		&url.URL{
			Scheme: "tls",
			Host:   ":25000",
		},
		nil,
		&tls.Config{
			Certificates:       []tls.Certificate{cert},
			InsecureSkipVerify: true,
		})
	defer transport.Close()

	testTransport(t, transport, "tls")
}

type testTransportRedialRecverSender struct {
	t    *testing.T
	done chan struct{}
}

func newtestTransportRedialRecverSender(t *testing.T) *testTransportRedialRecverSender {
	return &testTransportRedialRecverSender{
		t,
		make(chan struct{}),
	}
}

func (self *testTransportRedialRecverSender) transportRecver(link Link) error {
	link.Recv()
	return ErrTransportClosed
}

func (self *testTransportRedialRecverSender) transportSender(link Link) error {
	err := link.Send("id", reflect.ValueOf("msg"))
	if nil == err {
		self.t.Error("Sender: expect error, got nil")
	}

	close(self.done)
	return ErrTransportClosed
}

func testTransportRedial(t *testing.T, transport Transport, scheme string) {
	now := time.Now()
	trs := newtestTransportRedialRecverSender(t)
	transport.SetRecver(trs.transportRecver)
	transport.SetSender(trs.transportSender)

	err := transport.Listen()
	if nil != err {
		panic(err)
	}

	_, link, err := transport.Connect(&url.URL{
		Scheme: scheme,
		Host:   "127.0.0.1:25001",
		Path:   "/id",
	})
	if nil != err {
		panic(err)
	}

	link.Open()

	<-trs.done

	if time.Second > time.Now().Sub(now) {
		t.Error("expect 1 second of redial attempts; got less")
	}
}

func TestTcpTransportRedial(t *testing.T) {
	marshaler := NewGobMarshaler()
	transport := NewNetTransport(
		marshaler,
		&url.URL{
			Scheme: "tcp",
			Host:   ":25000",
		},
		&Config{RedialTimeout: time.Second})
	defer transport.Close()

	testTransportRedial(t, transport, "tcp")
}
<|endoftext|>


<commit_before>package goblet

import (
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/url"
	"reflect"
	"strings"

	"github.com/extrame/unmarshall"
	"github.com/sirupsen/logrus"
)

type FormFillFn func(content string) (interface{}, error)
type MultiFormFillFn func(ctx *Context, id string) (interface{}, error)

func (s *Server) AddFillForTypeInForm(typ string, fn FormFillFn) {
	s.filler[typ] = fn
}

func (s *Server) AddFillForTypeInMultiForm(typ string, fn MultiFormFillFn) {
	s.multiFiller[typ] = fn
}

// types that impliment RequestDecoder can unmarshal
// the request body into an apropriate type/struct
type RequestDecoder interface {
	Unmarshal(cx *Context, v interface{}, autofill bool) error
}

// a JSON decoder for request body (just a wrapper to json.Unmarshal)
type JsonRequestDecoder struct{}

func (d *JsonRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) (err error) {
	// read body
	if cx.fill_bts == nil {
		cx.fill_bts, err = ioutil.ReadAll(cx.request.Body)
	}
	if err != nil {
		return err
	}
	return json.Unmarshal(cx.fill_bts, v)
}

// an XML decoder for request body
type XmlRequestDecoder struct{}

func (d *XmlRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) (err error) {
	// read body
	if cx.fill_bts == nil {
		cx.fill_bts, err = ioutil.ReadAll(cx.request.Body)
	}
	if err != nil {
		return err
	}
	if err = xml.Unmarshal(cx.fill_bts, v); err != nil {
		logrus.Errorf("[Fill Error]Request:%s,Err:%s\n", string(cx.fill_bts), err.Error())
	}
	return err
}

// a form-enc decoder for request body
type FormRequestDecoder struct{}

type FileGetter func(string) (multipart.File, *multipart.FileHeader, error)

func (d *FormRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) error {
	if cx.request.Form == nil {
		cx.request.ParseForm()
	}

	var maxlength = 0
	for k, _ := range cx.request.Form {
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	var unmarshaller = unmarshall.Unmarshaller{
		Values: func() map[string][]string {
			return cx.request.Form
		},
		ValuesGetter: func(prefix string) url.Values {
			values := (*map[string][]string)(&cx.request.Form)
			var sub = make(url.Values)
			if values != nil {
				for k, v := range *values {
					if strings.HasPrefix(k, prefix+"[") {
						sub[k] = v
					}
				}
			}
			return sub
		},
		ValueGetter: func(tag string) []string {
			values := (*map[string][]string)(&cx.request.Form)
			if values != nil {
				if results, ok := (*values)[tag]; !ok {
					//get the value of [Tag] from [tag](lower case), it maybe a problem TODO
					return (*values)[strings.ToLower(tag)]
				} else {
					return results
				}
			}
			return []string{}
		},
		Tag:          "goblet",
		MaxLength:    maxlength,
		TagConcatter: concatPrefix,
		AutoFill:     autofill,
	}

	for typ, fn := range cx.Server.filler {
		if unmarshaller.FillForSpecifiledType == nil {
			unmarshaller.FillForSpecifiledType = make(map[string]func(id string) (reflect.Value, error))
		}
		unmarshaller.FillForSpecifiledType[typ] = func(content string) (reflect.Value, error) {
			obj, err := fn(content)
			return reflect.ValueOf(obj), err
		}
	}

	return unmarshaller.Unmarshall(v)
}

// a form-enc decoder for request body
type MultiFormRequestDecoder struct{}

func (d *MultiFormRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) error {
	err := cx.request.ParseMultipartForm(32 << 20)
	if err != nil {
		return err
	}
	values := (map[string][]string)(cx.request.Form)
	if cx.request.MultipartForm == nil {
		return errors.New("MultipartForm is empty")
	}
	var maxlength = 0

	for k, v := range cx.request.MultipartForm.Value {
		values[k] = v
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	for k, _ := range cx.request.MultipartForm.File {
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	var unmarshaller = unmarshall.Unmarshaller{
		Values: func() map[string][]string {
			return values
		},
		MaxLength: maxlength,
		ValueGetter: func(tag string) []string {
			return values[tag]
		},
		ValuesGetter: func(prefix string) url.Values {
			var sub = make(url.Values)
			if values != nil {
				for k, v := range values {
					if strings.HasPrefix(k, prefix+"[") {
						sub[k] = v
					}
				}
			}
			return sub
		},
		TagConcatter: concatPrefix,
		FillForSpecifiledType: map[string]func(id string) (reflect.Value, error){
			"github.com/extrame/goblet.File": func(id string) (reflect.Value, error) {
				var file File
				var err error
				var f multipart.File
				var h *multipart.FileHeader
				if f, h, err = cx.request.FormFile(id); err == nil {
					file.Name = h.Filename
					file.Header = h.Header
					file.rc = f
					return reflect.ValueOf(file), err
				} else {
					return reflect.ValueOf(nil), err
				}
			},
		},
		AutoFill: autofill,
		Tag:      "goblet",
	}

	for typ, fn := range cx.Server.multiFiller {
		unmarshaller.FillForSpecifiledType[typ] = func(id string) (reflect.Value, error) {
			obj, err := fn(cx, id)
			return reflect.ValueOf(obj), err
		}
	}

	return unmarshaller.Unmarshall(v)
}

// map of Content-Type -> RequestDecoders
var decoders map[string]RequestDecoder = map[string]RequestDecoder{
	"application/json":                  new(JsonRequestDecoder),
	"application/xml":                   new(XmlRequestDecoder),
	"text/xml":                          new(XmlRequestDecoder),
	"application/x-www-form-urlencoded": new(FormRequestDecoder),
	"text/plain":                        new(FormRequestDecoder),
	"multipart/form-data":               new(MultiFormRequestDecoder),
}

// goweb.Context Helper function to fill a variable with the contents
// of the request body. The body will be decoded based
// on the content-type and an apropriate RequestDecoder
// automatically selected
// If you want to use md5 function for the specified field, please add
// md5 tag for it. AND the md5 tag must be the last one of the tags, so
// if you have no other tag, please add ',' before md5
func (cx *Context) Fill(v interface{}, fills ...bool) error {
	// get content type
	ct := cx.request.Header.Get("Content-Type")
	// default to urlencoded
	if ct == "" {
		if cx.Server.Basic.DefaultType != "" {
			ct = cx.Server.Basic.DefaultType
		} else {
			ct = "application/x-www-form-urlencoded"
		}
	} else if strings.HasPrefix(ct, "text/plain") && cx.Server.Basic.DefaultType != "" {
		ct = cx.Server.Basic.DefaultType
	}
	autofill := true
	if len(fills) > 0 {
		autofill = fills[0]
	}
	return cx.FillAs(v, autofill, ct)
}

func (cx *Context) FillAs(v interface{}, autofill bool, ct string) error {
	// ignore charset (after ';')
	ct = strings.Split(ct, ";")[0]
	// get request decoder
	decoder, ok := decoders[ct]
	if ok != true {
		return fmt.Errorf("Cannot decode request for %s data", ct)
	}
	// decode
	err := decoder.Unmarshal(cx, v, autofill)
	if err != nil {
		fmt.Println(err)
		return err
	}
	// all clear
	return nil
}

func concatPrefix(prefix, tag string) string {
	return prefix + "[" + tag + "]"
}
<commit_msg>change default fill method when GETE<commit_after>package goblet

import (
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/url"
	"reflect"
	"strings"

	"github.com/extrame/unmarshall"
	"github.com/sirupsen/logrus"
)

type FormFillFn func(content string) (interface{}, error)
type MultiFormFillFn func(ctx *Context, id string) (interface{}, error)

func (s *Server) AddFillForTypeInForm(typ string, fn FormFillFn) {
	s.filler[typ] = fn
}

func (s *Server) AddFillForTypeInMultiForm(typ string, fn MultiFormFillFn) {
	s.multiFiller[typ] = fn
}

// types that impliment RequestDecoder can unmarshal
// the request body into an apropriate type/struct
type RequestDecoder interface {
	Unmarshal(cx *Context, v interface{}, autofill bool) error
}

// a JSON decoder for request body (just a wrapper to json.Unmarshal)
type JsonRequestDecoder struct{}

func (d *JsonRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) (err error) {
	// read body
	if cx.fill_bts == nil {
		cx.fill_bts, err = ioutil.ReadAll(cx.request.Body)
	}
	if err != nil {
		return err
	}
	return json.Unmarshal(cx.fill_bts, v)
}

// an XML decoder for request body
type XmlRequestDecoder struct{}

func (d *XmlRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) (err error) {
	// read body
	if cx.fill_bts == nil {
		cx.fill_bts, err = ioutil.ReadAll(cx.request.Body)
	}
	if err != nil {
		return err
	}
	if err = xml.Unmarshal(cx.fill_bts, v); err != nil {
		logrus.Errorf("[Fill Error]Request:%s,Err:%s\n", string(cx.fill_bts), err.Error())
	}
	return err
}

// a form-enc decoder for request body
type FormRequestDecoder struct{}

type FileGetter func(string) (multipart.File, *multipart.FileHeader, error)

func (d *FormRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) error {
	if cx.request.Form == nil {
		cx.request.ParseForm()
	}

	var maxlength = 0
	for k, _ := range cx.request.Form {
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	var unmarshaller = unmarshall.Unmarshaller{
		Values: func() map[string][]string {
			return cx.request.Form
		},
		ValuesGetter: func(prefix string) url.Values {
			values := (*map[string][]string)(&cx.request.Form)
			var sub = make(url.Values)
			if values != nil {
				for k, v := range *values {
					if strings.HasPrefix(k, prefix+"[") {
						sub[k] = v
					}
				}
			}
			return sub
		},
		ValueGetter: func(tag string) []string {
			values := (*map[string][]string)(&cx.request.Form)
			if values != nil {
				if results, ok := (*values)[tag]; !ok {
					//get the value of [Tag] from [tag](lower case), it maybe a problem TODO
					return (*values)[strings.ToLower(tag)]
				} else {
					return results
				}
			}
			return []string{}
		},
		Tag:          "goblet",
		MaxLength:    maxlength,
		TagConcatter: concatPrefix,
		AutoFill:     autofill,
	}

	for typ, fn := range cx.Server.filler {
		if unmarshaller.FillForSpecifiledType == nil {
			unmarshaller.FillForSpecifiledType = make(map[string]func(id string) (reflect.Value, error))
		}
		unmarshaller.FillForSpecifiledType[typ] = func(content string) (reflect.Value, error) {
			obj, err := fn(content)
			return reflect.ValueOf(obj), err
		}
	}

	return unmarshaller.Unmarshall(v)
}

// a form-enc decoder for request body
type MultiFormRequestDecoder struct{}

func (d *MultiFormRequestDecoder) Unmarshal(cx *Context, v interface{}, autofill bool) error {
	err := cx.request.ParseMultipartForm(32 << 20)
	if err != nil {
		return err
	}
	values := (map[string][]string)(cx.request.Form)
	if cx.request.MultipartForm == nil {
		return errors.New("MultipartForm is empty")
	}
	var maxlength = 0

	for k, v := range cx.request.MultipartForm.Value {
		values[k] = v
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	for k, _ := range cx.request.MultipartForm.File {
		if len(k) > maxlength {
			maxlength = len(k)
		}
	}

	var unmarshaller = unmarshall.Unmarshaller{
		Values: func() map[string][]string {
			return values
		},
		MaxLength: maxlength,
		ValueGetter: func(tag string) []string {
			return values[tag]
		},
		ValuesGetter: func(prefix string) url.Values {
			var sub = make(url.Values)
			if values != nil {
				for k, v := range values {
					if strings.HasPrefix(k, prefix+"[") {
						sub[k] = v
					}
				}
			}
			return sub
		},
		TagConcatter: concatPrefix,
		FillForSpecifiledType: map[string]func(id string) (reflect.Value, error){
			"github.com/extrame/goblet.File": func(id string) (reflect.Value, error) {
				var file File
				var err error
				var f multipart.File
				var h *multipart.FileHeader
				if f, h, err = cx.request.FormFile(id); err == nil {
					file.Name = h.Filename
					file.Header = h.Header
					file.rc = f
					return reflect.ValueOf(file), err
				} else {
					return reflect.ValueOf(nil), err
				}
			},
		},
		AutoFill: autofill,
		Tag:      "goblet",
	}

	for typ, fn := range cx.Server.multiFiller {
		unmarshaller.FillForSpecifiledType[typ] = func(id string) (reflect.Value, error) {
			obj, err := fn(cx, id)
			return reflect.ValueOf(obj), err
		}
	}

	return unmarshaller.Unmarshall(v)
}

// map of Content-Type -> RequestDecoders
var decoders map[string]RequestDecoder = map[string]RequestDecoder{
	"application/json":                  new(JsonRequestDecoder),
	"application/xml":                   new(XmlRequestDecoder),
	"text/xml":                          new(XmlRequestDecoder),
	"application/x-www-form-urlencoded": new(FormRequestDecoder),
	"text/plain":                        new(FormRequestDecoder),
	"multipart/form-data":               new(MultiFormRequestDecoder),
}

// goweb.Context Helper function to fill a variable with the contents
// of the request body. The body will be decoded based
// on the content-type and an apropriate RequestDecoder
// automatically selected
// If you want to use md5 function for the specified field, please add
// md5 tag for it. AND the md5 tag must be the last one of the tags, so
// if you have no other tag, please add ',' before md5
func (cx *Context) Fill(v interface{}, fills ...bool) error {
	// get content type
	ct := cx.request.Header.Get("Content-Type")
	//if method is GET, only form in url is supported
	if cx.request.Method == "GET" {
		ct = "application/x-www-form-urlencoded"
	}
	// default to urlencoded
	if ct == "" {
		if cx.Server.Basic.DefaultType != "" {
			ct = cx.Server.Basic.DefaultType
		} else {
			ct = "application/x-www-form-urlencoded"
		}
	} else if strings.HasPrefix(ct, "text/plain") && cx.Server.Basic.DefaultType != "" {
		ct = cx.Server.Basic.DefaultType
	}

	autofill := true
	if len(fills) > 0 {
		autofill = fills[0]
	}
	return cx.FillAs(v, autofill, ct)
}

func (cx *Context) FillAs(v interface{}, autofill bool, ct string) error {
	// ignore charset (after ';')
	ct = strings.Split(ct, ";")[0]
	// get request decoder
	decoder, ok := decoders[ct]
	if ok != true {
		return fmt.Errorf("Cannot decode request for %s data", ct)
	}
	// decode
	err := decoder.Unmarshal(cx, v, autofill)
	if err != nil {
		fmt.Println(err)
		return err
	}
	// all clear
	return nil
}

func concatPrefix(prefix, tag string) string {
	return prefix + "[" + tag + "]"
}
<|endoftext|>


<commit_before>package storage

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"

	"fmt"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

type DiskLocation struct {
	Directory      string
	MaxVolumeCount int
	volumes        map[VolumeId]*Volume
	sync.RWMutex
}

func NewDiskLocation(dir string, maxVolumeCount int) *DiskLocation {
	location := &DiskLocation{Directory: dir, MaxVolumeCount: maxVolumeCount}
	location.volumes = make(map[VolumeId]*Volume)
	return location
}

func (l *DiskLocation) volumeIdFromPath(dir os.FileInfo) (VolumeId, string, error) {
	name := dir.Name()
	if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
		collection := ""
		base := name[:len(name)-len(".dat")]
		i := strings.LastIndex(base, "_")
		if i > 0 {
			collection, base = base[0:i], base[i+1:]
		}
		vol, err := NewVolumeId(base)
		return vol, collection, err
	}

	return 0, "", fmt.Errorf("Path is not a volume: %s", name)
}

func (l *DiskLocation) loadExistingVolume(dir os.FileInfo, needleMapKind NeedleMapType, mutex *sync.RWMutex) {
	name := dir.Name()
	if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
		vid, collection, err := l.volumeIdFromPath(dir)
		if err == nil {
			mutex.RLock()
			_, found := l.volumes[vid]
			mutex.RUnlock()
			if !found {
				if v, e := NewVolume(l.Directory, collection, vid, needleMapKind, nil, nil, 0); e == nil {
					mutex.Lock()
					l.volumes[vid] = v
					mutex.Unlock()
					glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s",
						l.Directory+"/"+name, v.ReplicaPlacement, v.Version(), v.Size(), v.Ttl.String())
				} else {
					glog.V(0).Infof("new volume %s error %s", name, e)
				}
			}
		}
	}
}

func (l *DiskLocation) concurrentLoadingVolumes(needleMapKind NeedleMapType, concurrentFlag bool) {
	var concurrency int
	if concurrentFlag {
		//You could choose a better optimized concurency value after testing at your environment
		concurrency = 10
	} else {
		concurrency = 1
	}

	task_queue := make(chan os.FileInfo, 10*concurrency)
	go func() {
		if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
			for _, dir := range dirs {
				task_queue <- dir
			}
		}
		close(task_queue)
	}()

	var wg sync.WaitGroup
	var mutex sync.RWMutex
	for workerNum := 0; workerNum < concurrency; workerNum++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for dir := range task_queue {
				l.loadExistingVolume(dir, needleMapKind, &mutex)
			}
		}()
	}
	wg.Wait()

}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {
	l.Lock()
	defer l.Unlock()

	l.concurrentLoadingVolumes(needleMapKind, true)

	glog.V(0).Infoln("Store started on dir:", l.Directory, "with", len(l.volumes), "volumes", "max", l.MaxVolumeCount)
}

func (l *DiskLocation) DeleteCollectionFromDiskLocation(collection string) (e error) {
	l.Lock()
	defer l.Unlock()

	for k, v := range l.volumes {
		if v.Collection == collection {
			e = l.deleteVolumeById(k)
			if e != nil {
				return
			}
		}
	}
	return
}

func (l *DiskLocation) deleteVolumeById(vid VolumeId) (e error) {
	v, ok := l.volumes[vid]
	if !ok {
		return
	}
	e = v.Destroy()
	if e != nil {
		return
	}
	delete(l.volumes, vid)
	return
}

func (l *DiskLocation) LoadVolume(vid VolumeId, needleMapKind NeedleMapType) bool {
	if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, dir := range dirs {
			volId, _, err := l.volumeIdFromPath(dir)
			if vid == volId && err == nil {
				var mutex sync.RWMutex
				l.loadExistingVolume(dir, needleMapKind, &mutex)
				return true
			}
		}
	}

	return false
}

func (l *DiskLocation) DeleteVolume(vid VolumeId) error {
	l.Lock()
	defer l.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not found, VolumeId: %d", vid)
	}
	return l.deleteVolumeById(vid)
}

func (l *DiskLocation) UnloadVolume(vid VolumeId) error {
	l.Lock()
	defer l.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not loaded, VolumeId: %d", vid)
	}
	delete(l.volumes, vid)
	return nil
}

func (l *DiskLocation) SetVolume(vid VolumeId, volume *Volume) {
	l.Lock()
	defer l.Unlock()

	l.volumes[vid] = volume
}

func (l *DiskLocation) FindVolume(vid VolumeId) (*Volume, bool) {
	l.RLock()
	defer l.RUnlock()

	v, ok := l.volumes[vid]
	return v, ok
}

func (l *DiskLocation) VolumesLen() int {
	l.RLock()
	defer l.RUnlock()

	return len(l.volumes)
}

func (l *DiskLocation) Close() {
	l.Lock()
	defer l.Unlock()

	for _, v := range l.volumes {
		v.Close()
	}
	return
}
<commit_msg>fix leveldb mount bug<commit_after>package storage

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"

	"fmt"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

type DiskLocation struct {
	Directory      string
	MaxVolumeCount int
	volumes        map[VolumeId]*Volume
	sync.RWMutex
}

func NewDiskLocation(dir string, maxVolumeCount int) *DiskLocation {
	location := &DiskLocation{Directory: dir, MaxVolumeCount: maxVolumeCount}
	location.volumes = make(map[VolumeId]*Volume)
	return location
}

func (l *DiskLocation) volumeIdFromPath(dir os.FileInfo) (VolumeId, string, error) {
	name := dir.Name()
	if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
		collection := ""
		base := name[:len(name)-len(".dat")]
		i := strings.LastIndex(base, "_")
		if i > 0 {
			collection, base = base[0:i], base[i+1:]
		}
		vol, err := NewVolumeId(base)
		return vol, collection, err
	}

	return 0, "", fmt.Errorf("Path is not a volume: %s", name)
}

func (l *DiskLocation) loadExistingVolume(dir os.FileInfo, needleMapKind NeedleMapType, mutex *sync.RWMutex) {
	name := dir.Name()
	if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
		vid, collection, err := l.volumeIdFromPath(dir)
		if err == nil {
			mutex.RLock()
			_, found := l.volumes[vid]
			mutex.RUnlock()
			if !found {
				if v, e := NewVolume(l.Directory, collection, vid, needleMapKind, nil, nil, 0); e == nil {
					mutex.Lock()
					l.volumes[vid] = v
					mutex.Unlock()
					glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s",
						l.Directory+"/"+name, v.ReplicaPlacement, v.Version(), v.Size(), v.Ttl.String())
				} else {
					glog.V(0).Infof("new volume %s error %s", name, e)
				}
			}
		}
	}
}

func (l *DiskLocation) concurrentLoadingVolumes(needleMapKind NeedleMapType, concurrentFlag bool) {
	var concurrency int
	if concurrentFlag {
		//You could choose a better optimized concurency value after testing at your environment
		concurrency = 10
	} else {
		concurrency = 1
	}

	task_queue := make(chan os.FileInfo, 10*concurrency)
	go func() {
		if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
			for _, dir := range dirs {
				task_queue <- dir
			}
		}
		close(task_queue)
	}()

	var wg sync.WaitGroup
	var mutex sync.RWMutex
	for workerNum := 0; workerNum < concurrency; workerNum++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for dir := range task_queue {
				l.loadExistingVolume(dir, needleMapKind, &mutex)
			}
		}()
	}
	wg.Wait()

}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {
	l.Lock()
	defer l.Unlock()

	l.concurrentLoadingVolumes(needleMapKind, true)

	glog.V(0).Infoln("Store started on dir:", l.Directory, "with", len(l.volumes), "volumes", "max", l.MaxVolumeCount)
}

func (l *DiskLocation) DeleteCollectionFromDiskLocation(collection string) (e error) {
	l.Lock()
	defer l.Unlock()

	for k, v := range l.volumes {
		if v.Collection == collection {
			e = l.deleteVolumeById(k)
			if e != nil {
				return
			}
		}
	}
	return
}

func (l *DiskLocation) deleteVolumeById(vid VolumeId) (e error) {
	v, ok := l.volumes[vid]
	if !ok {
		return
	}
	e = v.Destroy()
	if e != nil {
		return
	}
	delete(l.volumes, vid)
	return
}

func (l *DiskLocation) LoadVolume(vid VolumeId, needleMapKind NeedleMapType) bool {
	if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, dir := range dirs {
			volId, _, err := l.volumeIdFromPath(dir)
			if vid == volId && err == nil {
				var mutex sync.RWMutex
				l.loadExistingVolume(dir, needleMapKind, &mutex)
				return true
			}
		}
	}

	return false
}

func (l *DiskLocation) DeleteVolume(vid VolumeId) error {
	l.Lock()
	defer l.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not found, VolumeId: %d", vid)
	}
	return l.deleteVolumeById(vid)
}

func (l *DiskLocation) UnloadVolume(vid VolumeId) error {
	l.Lock()
	defer l.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not loaded, VolumeId: %d", vid)
	}
	l.volumes[vid].Close()
	delete(l.volumes, vid)
	return nil
}

func (l *DiskLocation) SetVolume(vid VolumeId, volume *Volume) {
	l.Lock()
	defer l.Unlock()

	l.volumes[vid] = volume
}

func (l *DiskLocation) FindVolume(vid VolumeId) (*Volume, bool) {
	l.RLock()
	defer l.RUnlock()

	v, ok := l.volumes[vid]
	return v, ok
}

func (l *DiskLocation) VolumesLen() int {
	l.RLock()
	defer l.RUnlock()

	return len(l.volumes)
}

func (l *DiskLocation) Close() {
	l.Lock()
	defer l.Unlock()

	for _, v := range l.volumes {
		v.Close()
	}
	return
}
<|endoftext|>


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

package main

import (
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/codegangsta/cli"
	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
	"github.com/letsencrypt/boulder/policy"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/sa"
)

const (
	good = "valid"
	bad  = "invalid"

	filenameLayout = "20060102"

	checkPeriod = time.Hour * 24 * 90
)

type report struct {
	GoodCerts int64
	BadCerts  int64
	Entries   map[string]reportEntry
}

func (r *report) save(directory string) error {
	filename := path.Join(directory, fmt.Sprintf(
		"%s-%s-report.json",
		time.Now().Add(-time.Hour*24*90).Format(filenameLayout),
		time.Now().Format(filenameLayout),
	))
	content, err := json.Marshal(r)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, content, os.ModePerm)
}

type reportEntry struct {
	Valid    bool     `json:"validity"`
	Problems []string `json:"problem,omitempty"`
}

type certChecker struct {
	pa           core.PolicyAuthority
	dbMap        *gorp.DbMap
	certs        chan core.Certificate
	issuedReport report
}

func newChecker(dbMap *gorp.DbMap) certChecker {
	c := certChecker{
		pa:    policy.NewPolicyAuthorityImpl(),
		dbMap: dbMap,
	}
	c.issuedReport.Entries = make(map[string]reportEntry)
	return c
}

func (c *certChecker) getCerts() error {
	var certs []core.Certificate
	_, err := c.dbMap.Select(
		&certs,
		"SELECT * FROM certificates WHERE issued > :issued",
		map[string]interface{}{"issued": time.Now().Add(-checkPeriod)},
	)
	if err != nil {
		return err
	}
	c.certs = make(chan core.Certificate, len(certs))
	for _, cert := range certs {
		c.certs <- cert
	}
	// Close channel so range operations won't block when the channel empties out
	close(c.certs)
	return nil
}

func (c *certChecker) processCerts(wg *sync.WaitGroup) {
	for cert := range c.certs {
		problems := c.checkCert(cert)
		valid := len(problems) == 0
		c.issuedReport.Entries[cert.Serial] = reportEntry{
			Valid:    valid,
			Problems: problems,
		}
		if !valid {
			atomic.AddInt64(&c.issuedReport.BadCerts, 1)
		} else {
			atomic.AddInt64(&c.issuedReport.GoodCerts, 1)
		}
	}
	wg.Done()
}

func (c *certChecker) checkCert(cert core.Certificate) (problems []string) {
	// Check digests match
	if cert.Digest != core.Fingerprint256(cert.DER) {
		problems = append(problems, "Stored digest doesn't match certificate digest")
	}

	// Parse certificate
	parsedCert, err := x509.ParseCertificate(cert.DER)
	if err != nil {
		problems = append(problems, fmt.Sprintf("Couldn't parse stored certificate: %s", err))
	} else {
		// Check we have the right expiration time
		if parsedCert.NotAfter != cert.Expires {
			problems = append(problems, "Stored expiration doesn't match certificate NotAfter")
		}
		// Check basic constraints are set
		if !parsedCert.BasicConstraintsValid {
			problems = append(problems, "Certificate doesn't have basic constraints set")
		}
		// Check the cert isn't able to sign other certificates
		if parsedCert.IsCA {
			problems = append(problems, "Certificate can sign other certificates")
		}
		// Check the cert has the correct validity period
		if parsedCert.NotAfter.Sub(cert.Issued) > (checkPeriod) {
			problems = append(problems, "Certificate has a validity period longer than 90 days")
		}
		// Check that the PA is still willing to issue for each name in DNSNames + CommonName
		for _, name := range append(parsedCert.DNSNames, parsedCert.Subject.CommonName) {
			if err = c.pa.WillingToIssue(core.AcmeIdentifier{Type: core.IdentifierDNS, Value: name}); err != nil {
				problems = append(problems, fmt.Sprintf("Policy Authority isn't willing to issue for %s: %s", name, err))
			}
		}
		// Check the cert has the correct key usage extensions
		if !core.CmpExtKeyUsageSlice(parsedCert.ExtKeyUsage, []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}) {
			problems = append(problems, "Certificate has incorrect key usage extensions")
		}
	}
	return problems
}

func main() {
	app := cmd.NewAppShell("cert-checker", "Checks validity of certificates issued in the last 90 days")
	app.App.Flags = append(app.App.Flags, cli.IntFlag{
		Name:  "workers",
		Value: runtime.NumCPU(),
		Usage: "The number of concurrent workers used to process certificates",
	}, cli.StringFlag{
		Name:  "report-dir-path",
		Usage: "The path to write a JSON report on the certificates checks to (if no path is provided the report will not be written out)",
	}, cli.StringFlag{
		Name:  "db-connect",
		Usage: "SQL URI if not provided in the configuration file",
	})

	app.Config = func(c *cli.Context, config cmd.Config) cmd.Config {
		config.CertChecker.ReportDirectoryPath = c.GlobalString("report-dir-path")

		if connect := c.GlobalString("db-connect"); connect != "" {
			config.CertChecker.DBConnect = connect
		}

		return config
	}

	app.Action = func(c cmd.Config) {
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")

		blog.SetAuditLogger(auditlogger)
		auditlogger.Info(app.VersionString())

		dbMap, err := sa.NewDbMap(c.CertChecker.DBConnect)
		cmd.FailOnError(err, "Could not connect to database")

		checker := newChecker(dbMap)
		auditlogger.Info("# Getting certificates issued in the last 90 days")
		err = checker.getCerts()
		cmd.FailOnError(err, "Failed to get sample certificates")

		auditlogger.Info(fmt.Sprintf("# Processing sample, %d certificates using %d workers", len(checker.certs), c.CertChecker.Workers))
		wg := new(sync.WaitGroup)
		for i := 0; i < c.CertChecker.Workers; i++ {
			wg.Add(1)
			go checker.processCerts(wg)
		}
		wg.Wait()
		auditlogger.Info(fmt.Sprintf(
			"# Finished processing certificates, sample: %d, good: %d, bad: %d",
			len(checker.issuedReport.Entries),
			checker.issuedReport.GoodCerts,
			checker.issuedReport.BadCerts,
		))
		err = checker.issuedReport.save(c.CertChecker.ReportDirectoryPath)
		cmd.FailOnError(err, "Couldn't save issued certificate report")
	}

	app.Run()
}
<commit_msg>Standardized checkPeriod<commit_after>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cactus/go-statsd-client/statsd"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/codegangsta/cli"
	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
	"github.com/letsencrypt/boulder/policy"

	"github.com/letsencrypt/boulder/cmd"
	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/sa"
)

const (
	good = "valid"
	bad  = "invalid"

	filenameLayout = "20060102"

	checkPeriod = time.Hour * 24 * 90
)

type report struct {
	GoodCerts int64
	BadCerts  int64
	Entries   map[string]reportEntry
}

func (r *report) save(directory string) error {
	filename := path.Join(directory, fmt.Sprintf(
		"%s-%s-report.json",
		time.Now().Add(-checkPeriod).Format(filenameLayout),
		time.Now().Format(filenameLayout),
	))
	content, err := json.Marshal(r)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, content, os.ModePerm)
}

type reportEntry struct {
	Valid    bool     `json:"validity"`
	Problems []string `json:"problem,omitempty"`
}

type certChecker struct {
	pa           core.PolicyAuthority
	dbMap        *gorp.DbMap
	certs        chan core.Certificate
	issuedReport report
}

func newChecker(dbMap *gorp.DbMap) certChecker {
	c := certChecker{
		pa:    policy.NewPolicyAuthorityImpl(),
		dbMap: dbMap,
	}
	c.issuedReport.Entries = make(map[string]reportEntry)
	return c
}

func (c *certChecker) getCerts() error {
	var certs []core.Certificate
	_, err := c.dbMap.Select(
		&certs,
		"SELECT * FROM certificates WHERE issued > :issued",
		map[string]interface{}{"issued": time.Now().Add(-checkPeriod)},
	)
	if err != nil {
		return err
	}
	c.certs = make(chan core.Certificate, len(certs))
	for _, cert := range certs {
		c.certs <- cert
	}
	// Close channel so range operations won't block when the channel empties out
	close(c.certs)
	return nil
}

func (c *certChecker) processCerts(wg *sync.WaitGroup) {
	for cert := range c.certs {
		problems := c.checkCert(cert)
		valid := len(problems) == 0
		c.issuedReport.Entries[cert.Serial] = reportEntry{
			Valid:    valid,
			Problems: problems,
		}
		if !valid {
			atomic.AddInt64(&c.issuedReport.BadCerts, 1)
		} else {
			atomic.AddInt64(&c.issuedReport.GoodCerts, 1)
		}
	}
	wg.Done()
}

func (c *certChecker) checkCert(cert core.Certificate) (problems []string) {
	// Check digests match
	if cert.Digest != core.Fingerprint256(cert.DER) {
		problems = append(problems, "Stored digest doesn't match certificate digest")
	}

	// Parse certificate
	parsedCert, err := x509.ParseCertificate(cert.DER)
	if err != nil {
		problems = append(problems, fmt.Sprintf("Couldn't parse stored certificate: %s", err))
	} else {
		// Check we have the right expiration time
		if parsedCert.NotAfter != cert.Expires {
			problems = append(problems, "Stored expiration doesn't match certificate NotAfter")
		}
		// Check basic constraints are set
		if !parsedCert.BasicConstraintsValid {
			problems = append(problems, "Certificate doesn't have basic constraints set")
		}
		// Check the cert isn't able to sign other certificates
		if parsedCert.IsCA {
			problems = append(problems, "Certificate can sign other certificates")
		}
		// Check the cert has the correct validity period
		if parsedCert.NotAfter.Sub(cert.Issued) > (checkPeriod) {
			problems = append(problems, "Certificate has a validity period longer than 90 days")
		}
		// Check that the PA is still willing to issue for each name in DNSNames + CommonName
		for _, name := range append(parsedCert.DNSNames, parsedCert.Subject.CommonName) {
			if err = c.pa.WillingToIssue(core.AcmeIdentifier{Type: core.IdentifierDNS, Value: name}); err != nil {
				problems = append(problems, fmt.Sprintf("Policy Authority isn't willing to issue for %s: %s", name, err))
			}
		}
		// Check the cert has the correct key usage extensions
		if !core.CmpExtKeyUsageSlice(parsedCert.ExtKeyUsage, []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}) {
			problems = append(problems, "Certificate has incorrect key usage extensions")
		}
	}
	return problems
}

func main() {
	app := cmd.NewAppShell("cert-checker", "Checks validity of certificates issued in the last 90 days")
	app.App.Flags = append(app.App.Flags, cli.IntFlag{
		Name:  "workers",
		Value: runtime.NumCPU(),
		Usage: "The number of concurrent workers used to process certificates",
	}, cli.StringFlag{
		Name:  "report-dir-path",
		Usage: "The path to write a JSON report on the certificates checks to (if no path is provided the report will not be written out)",
	}, cli.StringFlag{
		Name:  "db-connect",
		Usage: "SQL URI if not provided in the configuration file",
	})

	app.Config = func(c *cli.Context, config cmd.Config) cmd.Config {
		config.CertChecker.ReportDirectoryPath = c.GlobalString("report-dir-path")

		if connect := c.GlobalString("db-connect"); connect != "" {
			config.CertChecker.DBConnect = connect
		}

		return config
	}

	app.Action = func(c cmd.Config) {
		stats, err := statsd.NewClient(c.Statsd.Server, c.Statsd.Prefix)
		cmd.FailOnError(err, "Couldn't connect to statsd")

		auditlogger, err := blog.Dial(c.Syslog.Network, c.Syslog.Server, c.Syslog.Tag, stats)
		cmd.FailOnError(err, "Could not connect to Syslog")

		blog.SetAuditLogger(auditlogger)
		auditlogger.Info(app.VersionString())

		dbMap, err := sa.NewDbMap(c.CertChecker.DBConnect)
		cmd.FailOnError(err, "Could not connect to database")

		checker := newChecker(dbMap)
		auditlogger.Info("# Getting certificates issued in the last 90 days")
		err = checker.getCerts()
		cmd.FailOnError(err, "Failed to get sample certificates")

		auditlogger.Info(fmt.Sprintf("# Processing sample, %d certificates using %d workers", len(checker.certs), c.CertChecker.Workers))
		wg := new(sync.WaitGroup)
		for i := 0; i < c.CertChecker.Workers; i++ {
			wg.Add(1)
			go checker.processCerts(wg)
		}
		wg.Wait()
		auditlogger.Info(fmt.Sprintf(
			"# Finished processing certificates, sample: %d, good: %d, bad: %d",
			len(checker.issuedReport.Entries),
			checker.issuedReport.GoodCerts,
			checker.issuedReport.BadCerts,
		))
		err = checker.issuedReport.save(c.CertChecker.ReportDirectoryPath)
		cmd.FailOnError(err, "Couldn't save issued certificate report")
	}

	app.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rsa"
	"fmt"
	"time"

	"github.com/pkg/browser"

	"github.com/datawire/ambassador/pkg/k8s"
	"github.com/dgrijalva/jwt-go"
	"github.com/gookit/color"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	k8sTypesMetaV1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8sClientCoreV1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
)

const SecretName = "ambassador-internal"

type LoginClaimsV1 struct {
	LoginTokenVersion string `json:"login_token_version"`
	jwt.StandardClaims
}

func aesLogin(cmd *cobra.Command, args []string) error {
	fmt.Println(color.Info.Sprintf("Connecting to the Ambassador Edge Policy Console in this cluster..."))

	// Grab options
	context, _ := cmd.Flags().GetString("context")
	namespace, _ := cmd.Flags().GetString("namespace")
	justShowURL, _ := cmd.Flags().GetBool("url")
	showToken, _ := cmd.Flags().GetBool("token")
	firstTime, _ := cmd.Flags().GetBool("firstTime")

	// Figure out the correct hostname
	hostname := args[0]

	// Prepare to talk to the cluster
	kubeinfo := k8s.NewKubeInfo("", context, namespace) // Default namespace is "ambassador"

	return do_login(kubeinfo, context, namespace, hostname, !justShowURL, justShowURL, showToken, firstTime)
}

func do_login(kubeinfo *k8s.KubeInfo, context, namespace, hostname string, openInBrowser, showURL, showToken, firstTime bool) error {
	restconfig, err := kubeinfo.GetRestConfig()
	if err != nil {
		return errors.Wrap(err, "Failed to connect to cluster (rest)")
	}

	// Obtain signing key
	// -> kubectl -n $namespace get secret $SecretName -o json
	privateKey, err := getSigningKey(restconfig, namespace)
	if err != nil {
		fmt.Println()
		fmt.Println("Failed to obtain expected information from the cluster.")
		fmt.Println("Is kubectl configured to connect to the correct cluster?")
		fmt.Printf("Is %s the namespace where Ambassador is installed?\n\n", namespace)
		fmt.Printf("    kubectl -n %s get svc,deploy\n\n", namespace)
		return err
	}

	// FIXME: validate that hostname by querying
	// https://{{hostname}}/edge_stack/admin/api/ambassador_cluster_id and
	// verifying that it returns the same UUID via direct access and via
	// port-forward/teleproxy. This avoids leaking login credentials to the
	// operator of a different website.

	// Construct claims
	now := time.Now()
	duration := 30 * time.Minute
	claims := &LoginClaimsV1{
		"v1",
		jwt.StandardClaims{
			IssuedAt:  now.Unix(),
			NotBefore: now.Unix(),
			ExpiresAt: (now.Add(duration)).Unix(),
		},
	}

	// Generate JWT
	token := jwt.NewWithClaims(jwt.GetSigningMethod("PS512"), claims)
	tokenString, err := token.SignedString(privateKey)
	if err != nil {
		return errors.Wrap(err, "Unexpected error generating JWT")
	}

	// Output
	url := fmt.Sprintf("https://%s/edge_stack/admin/first-install=%v#%s", hostname, firstTime, tokenString)

	// Remember if the browser successfully opened the URL
	browserOpened := false

	if openInBrowser {
		err = browser.OpenURL(url)
		if err == nil {
			browserOpened = true
		} else {
			fmt.Println("Unexpected error while trying to open your browser.")
			err = errors.Wrap(err, "browse")
		}
	}

	if showURL {
		if browserOpened {
			fmt.Println("We've opened the Ambassador Edge Policy Console for you in your browser:")
		} else {
			fmt.Println("Visit the following URL to access the Ambassador Edge Policy Console:")
		}

		fmt.Println(url)

		// Whitespace if we are also showing the token or opening in the browser
		if showToken {
			fmt.Println()
		}
	}

	if showToken {
		fmt.Println("The login token is")
		fmt.Println("    ", tokenString)
	}

	return err
}

// getSigningKey retrieves the designated secret from the cluster and returns
// the private key extracted from the secret data
func getSigningKey(restconfig *rest.Config, namespace string) (*rsa.PrivateKey, error) {
	coreClient, err := k8sClientCoreV1.NewForConfig(restconfig)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to connect to cluster (core)")
	}
	secretInterface := coreClient.Secrets(namespace)
	secret, err := secretInterface.Get(SecretName, k8sTypesMetaV1.GetOptions{})
	if err != nil {
		return nil, err
	}
	// Parse out the private key from the secret
	privatePEM, ok := secret.Data["rsa.key"]
	if !ok {
		return nil, errors.Errorf("secret name=%q namespace=%q exists but does not contain an %q %s field",
			SecretName, namespace, "rsa.key", "private-key")
	}
	return jwt.ParseRSAPrivateKeyFromPEM(privatePEM)
}
<commit_msg>checkpoint<commit_after>package main

import (
	"crypto/rsa"
	"fmt"
	"time"

	"github.com/pkg/browser"

	"github.com/datawire/ambassador/pkg/k8s"
	"github.com/dgrijalva/jwt-go"
	"github.com/gookit/color"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	k8sTypesMetaV1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8sClientCoreV1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
)

const SecretName = "ambassador-internal"

type LoginClaimsV1 struct {
	LoginTokenVersion string `json:"login_token_version"`
	jwt.StandardClaims
}

func aesLogin(cmd *cobra.Command, args []string) error {
	fmt.Println(color.Info.Sprintf("Connecting to the Ambassador Edge Policy Console in this cluster..."))

	// Grab options
	context, _ := cmd.Flags().GetString("context")
	namespace, _ := cmd.Flags().GetString("namespace")
	justShowURL, _ := cmd.Flags().GetBool("url")
	showToken, _ := cmd.Flags().GetBool("token")
	firstTime, _ := cmd.Flags().GetBool("firstTime")

	// Figure out the correct hostname
	hostname := args[0]

	// Prepare to talk to the cluster
	kubeinfo := k8s.NewKubeInfo("", context, namespace) // Default namespace is "ambassador"

	return do_login(kubeinfo, context, namespace, hostname, !justShowURL, justShowURL, showToken, firstTime)
}

func do_login(kubeinfo *k8s.KubeInfo, context, namespace, hostname string, openInBrowser, showURL, showToken, firstTime bool) error {
	restconfig, err := kubeinfo.GetRestConfig()
	if err != nil {
		return errors.Wrap(err, "Failed to connect to cluster (rest)")
	}

	// Obtain signing key
	// -> kubectl -n $namespace get secret $SecretName -o json
	privateKey, err := getSigningKey(restconfig, namespace)
	if err != nil {
		fmt.Println()
		fmt.Println("Failed to obtain expected information from the cluster.")
		fmt.Println("Is kubectl configured to connect to the correct cluster?")
		fmt.Printf("Is %s the namespace where Ambassador is installed?\n\n", namespace)
		fmt.Printf("    kubectl -n %s get svc,deploy\n\n", namespace)
		return err
	}

	// FIXME: validate that hostname by querying
	// https://{{hostname}}/edge_stack/admin/api/ambassador_cluster_id and
	// verifying that it returns the same UUID via direct access and via
	// port-forward/teleproxy. This avoids leaking login credentials to the
	// operator of a different website.

	// Construct claims
	now := time.Now()
	duration := 30 * time.Minute
	claims := &LoginClaimsV1{
		"v1",
		jwt.StandardClaims{
			IssuedAt:  now.Unix(),
			NotBefore: now.Unix(),
			ExpiresAt: (now.Add(duration)).Unix(),
		},
	}

	// Generate JWT
	token := jwt.NewWithClaims(jwt.GetSigningMethod("PS512"), claims)
	tokenString, err := token.SignedString(privateKey)
	if err != nil {
		return errors.Wrap(err, "Unexpected error generating JWT")
	}

	// Output
	url := fmt.Sprintf("https://%s/edge_stack/admin/?first-install=%v#%s", hostname, firstTime, tokenString)

	// Remember if the browser successfully opened the URL
	browserOpened := false

	if openInBrowser {
		err = browser.OpenURL(url)
		if err == nil {
			browserOpened = true
		} else {
			fmt.Println("Unexpected error while trying to open your browser.")
			err = errors.Wrap(err, "browse")
		}
	}

	if showURL {
		if browserOpened {
			fmt.Println("We've opened the Ambassador Edge Policy Console for you in your browser:")
		} else {
			fmt.Println("Visit the following URL to access the Ambassador Edge Policy Console:")
		}

		fmt.Println(url)

		// Whitespace if we are also showing the token or opening in the browser
		if showToken {
			fmt.Println()
		}
	}

	if showToken {
		fmt.Println("The login token is")
		fmt.Println("    ", tokenString)
	}

	return err
}

// getSigningKey retrieves the designated secret from the cluster and returns
// the private key extracted from the secret data
func getSigningKey(restconfig *rest.Config, namespace string) (*rsa.PrivateKey, error) {
	coreClient, err := k8sClientCoreV1.NewForConfig(restconfig)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to connect to cluster (core)")
	}
	secretInterface := coreClient.Secrets(namespace)
	secret, err := secretInterface.Get(SecretName, k8sTypesMetaV1.GetOptions{})
	if err != nil {
		return nil, err
	}
	// Parse out the private key from the secret
	privatePEM, ok := secret.Data["rsa.key"]
	if !ok {
		return nil, errors.Errorf("secret name=%q namespace=%q exists but does not contain an %q %s field",
			SecretName, namespace, "rsa.key", "private-key")
	}
	return jwt.ParseRSAPrivateKeyFromPEM(privatePEM)
}
<|endoftext|>


<commit_before>// Package main defines a service for handling various post-processing
// and house-keeping tasks associated with the pipelines.
// Most tasks will be run periodically, but some may be triggered
// by URL requests from authorized sources.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/pprof"
	"os"
	"runtime"
	"strconv"
	"time"

	"cloud.google.com/go/datastore"
	"github.com/m-lab/etl/batch"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// ###############################################################################
// All DataStore related code and variables
// ###############################################################################
const (
	dsNamespace   = "gardener"
	dsKind        = "gardener"
	batchStateKey = "batch-state"
)

var dsClient *datastore.Client

func getDSClient() (*datastore.Client, error) {
	var err error
	var client *datastore.Client
	if dsClient == nil {
		project, ok := os.LookupEnv("PROJECT")
		if !ok {
			return dsClient, errors.New("PROJECT env var not set")
		}
		client, err = datastore.NewClient(context.Background(), project)
		if err == nil {
			dsClient = client
		}
	}
	return dsClient, err
}

// Load retrieves an arbitrary record from datastore.
func Load(name string, obj interface{}) error {
	var client *datastore.Client
	client, err := getDSClient()
	if err != nil {
		return err
	}
	k := datastore.NameKey(dsKind, name, nil)
	k.Namespace = dsNamespace
	return client.Get(context.Background(), k, obj)
}

// Save stores an arbitrary object to kind/key in the default namespace.
// If a record already exists, then it is overwritten.
// TODO(gfr) Make an upsert version of this:
// https://cloud.google.com/datastore/docs/concepts/entities
func Save(key string, obj interface{}) error {
	client, err := getDSClient()
	if err != nil {
		return err
	}
	k := datastore.NameKey(dsKind, key, nil)
	k.Namespace = dsNamespace
	_, err = client.Put(context.Background(), k, obj)
	return err
}

// ###############################################################################
//  Batch processing task scheduling and support code
// ###############################################################################

// Persistent Queuer for use in handlers and gardener tasks.
var batchQueuer batch.Queuer

// QueueState holds the state information for each batch queue.
type QueueState struct {
	// Per queue, indicate which day is being processed in
	// that queue.  No need to process more than one day at a time.
	// The other queues will take up the slack while we add more
	// tasks when one queue is emptied.

	QueueName        string // Name of the batch queue.
	NextTask         string // Name of task file currently being enqueued.
	PendingPartition string // FQ Name of next table partition to be added.
}

// BatchState holds the entire batch processing state.
// It holds the state locally, and also is stored in DataStore for
// recovery when the instance is restarted, e.g. for weekly platform
// updates.
// At any given time, we restrict ourselves to a 14 day reprocessing window,
// and finish the beginning of that window before reprocessing any dates beyond it.
// When we determine that the first date in the window has been submitted for
// processing, we advance the window up to the next pending date.
type BatchState struct {
	Hostname       string       // Hostname of the gardener that saved this.
	InstanceID     string       // instance ID of the gardener that saved this.
	WindowStart    time.Time    // Start of two week window we are working on.
	QueueBase      string       // base name for queues.
	QStates        []QueueState // States for each queue.
	LastUpdateTime time.Time    // Time of last update.  (Is this in DS metadata?)
}

// MaybeScheduleMoreTasks will look for an empty task queue, and if it finds one, will look
// for corresponding days to add to the queue.
// Alternatively, it may look first for the N oldest days to be reprocessed, and will then
// check whether any of the task queues for those oldest days is empty, and conditionally add tasks.
func MaybeScheduleMoreTasks(queuer *batch.Queuer) {
	// GetTaskQueueDepth returns the number of pending items in a task queue.
	stats, err := queuer.GetTaskqueueStats()
	if err != nil {
		log.Println(err)
	} else {
		for k, v := range stats {
			if len(v) > 0 && v[0].Tasks == 0 && v[0].InFlight == 0 {
				log.Printf("Ready: %s: %v\n", k, v[0])
				// Should add more tasks now.
			}
		}
	}
}

// queuerFromEnv creates a Queuer struct initialized from environment variables.
// It uses TASKFILE_BUCKET, PROJECT, QUEUE_BASE, and NUM_QUEUES.
func queuerFromEnv() (batch.Queuer, error) {
	bucketName, ok := os.LookupEnv("TASKFILE_BUCKET")
	if !ok {
		return batch.Queuer{}, errors.New("TASKFILE_BUCKET not set")
	}
	project, ok := os.LookupEnv("PROJECT")
	if !ok {
		return batch.Queuer{}, errors.New("PROJECT not set")
	}
	queueBase, ok := os.LookupEnv("QUEUE_BASE")
	if !ok {
		return batch.Queuer{}, errors.New("QUEUE_BASE not set")
	}
	numQueues, err := strconv.Atoi(os.Getenv("NUM_QUEUES"))
	if err != nil {
		log.Println(err)
		return batch.Queuer{}, errors.New("Parse error on NUM_QUEUES")
	}

	return batch.CreateQueuer(http.DefaultClient, nil, queueBase, numQueues, project, bucketName, false)
}

// StartDateRFC3339 is the date at which reprocessing will start when it catches
// up to present.  For now, we are making this the beginning of the ETL timeframe,
// until we get annotation fixed to use the actual data date instead of NOW.
const StartDateRFC3339 = "2017-05-01T00:00:00Z00:00"

// startupBatch determines whether some other instance has control, and
// assumes control if not.
func startupBatch(base string, numQueues int) (BatchState, error) {
	hostname := os.Getenv("HOSTNAME")
	instance := os.Getenv("GAE_INSTANCE")
	queues := make([]QueueState, numQueues)
	var bs BatchState
	err := Load(batchStateKey, &bs)
	if err != nil {
		startDate, err := time.Parse(time.RFC3339, StartDateRFC3339)
		if err != nil {
			log.Println("Could not parse start time.  Not starting batch.")
			return bs, err
		}
		bs = BatchState{hostname, instance, startDate, base, queues, time.Now()}

	} else {
		// TODO - should check whether we should take over, or leave alone.
	}

	err = Save(batchStateKey, &bs)
	return bs, err
}

// ###############################################################################
//  Top level service control code.
// ###############################################################################

// periodic will run approximately every 5 minutes.
func periodic() {
	_, err := startupBatch(batchQueuer.QueueBase, batchQueuer.NumQueues)
	if err != nil {
		log.Fatal(err)
	}
	for {
		log.Println("Periodic is running")

		MaybeScheduleMoreTasks(&batchQueuer)

		// There is no need for randomness, since this is a singleton handler.
		time.Sleep(300 * time.Second)
	}
}

func setupPrometheus() {
	// Define a custom serve mux for prometheus to listen on a separate port.
	// We listen on a separate port so we can forward this port on the host VM.
	// We cannot forward port 8080 because it is used by AppEngine.
	mux := http.NewServeMux()
	// Assign the default prometheus handler to the standard exporter path.
	mux.Handle("/metrics", promhttp.Handler())
	// Assign the pprof handling paths to the external port to access individual
	// instances.
	mux.HandleFunc("/debug/pprof/", pprof.Index)
	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
	go http.ListenAndServe(":9090", mux)
}

// Status provides basic information about the service.  For now, it is just
// configuration and version info.  In future it will likely include more
// dynamic information.
// TODO(gfr) Add either a black list or a white list for the environment
// variables, so we can hide sensitive vars. https://github.com/m-lab/etl/issues/384
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "
NOTE: This is just one of potentially many instances.
\n")
	commit := os.Getenv("COMMIT_HASH")
	if len(commit) >= 8 {
		fmt.Fprintf(w, "Release: %s
  Commit: %s
\n",
			os.Getenv("RELEASE_TAG"), os.Getenv("COMMIT_HASH"), os.Getenv("COMMIT_HASH")[0:7])
	} else {
		fmt.Fprintf(w, "Release: %s   Commit: unknown\n", os.Getenv("RELEASE_TAG"))
	}

	env := os.Environ()
	for i := range env {
		fmt.Fprintf(w, "%s
\n", env[i])
	}
	fmt.Fprintf(w, "\n")
}

var healthy = false

// healthCheck, for now, used for both /ready and /alive.
func healthCheck(w http.ResponseWriter, r *http.Request) {
	if !healthy {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, `{"message": "Internal server error."}`)
	}
	fmt.Fprint(w, "ok")
}

// runService starts a service handler and runs forever.
// The configuration info comes from environment variables.
func runService() {
	// Enable block profiling
	runtime.SetBlockProfileRate(1000000) // One event per msec.

	setupPrometheus()
	// We also setup another prometheus handler on a non-standard path. This
	// path name will be accessible through the AppEngine service address,
	// however it will be served by a random instance.
	http.Handle("/random-metrics", promhttp.Handler())
	http.HandleFunc("/", Status)
	http.HandleFunc("/status", Status)

	http.HandleFunc("/alive", healthCheck)
	http.HandleFunc("/ready", healthCheck)

	var err error
	batchQueuer, err = queuerFromEnv()
	if err == nil {
		healthy = true
		log.Println("Running as a service.")

		// Run the background "periodic" function.
		go periodic()
	} else {
		// Leaving healthy == false
		// This will cause app-engine to roll back.
		log.Println(err)
		log.Println("Required environment variables are missing or invalid.")
	}

	// ListenAndServe, and terminate when it returns.
	log.Fatal(http.ListenAndServe(":8080", nil))
}

// ###############################################################################
//  Top level command line code.
// ###############################################################################

// These are used only for command line.  For service, environment variables are used
// for general parameters, and request parameter for month.
var (
	fProject = flag.String("project", "", "Project containing queues.")
	fQueue   = flag.String("queue", "etl-ndt-batch-", "Base of queue name.")
	// TODO implement listing queues to determine number of queue, and change this to 0
	fNumQueues = flag.Int("num_queues", 8, "Number of queues.  Normally determined by listing queues.")
	// Gardener will only read from this bucket, so its ok to use production bucket as default.
	fBucket = flag.String("bucket", "archive-mlab-oti", "Source bucket.")
	fExper  = flag.String("experiment", "ndt", "Experiment prefix, without trailing slash.")
	fMonth  = flag.String("month", "", "Single month spec, as YYYY/MM")
	fDay    = flag.String("day", "", "Single day spec, as YYYY/MM/DD")
	fDryRun = flag.Bool("dry_run", false, "Prevents all output to queue_pusher.")
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	// Check if invoked as a service.
	isService, _ := strconv.ParseBool(os.Getenv("GARDENER_SERVICE"))
	if isService {
		runService()
		return
	}

	// Otherwise this is a command line invocation...
	flag.Parse()
	// Check that either project or dry-run is set.
	// If dry-run, it is ok for the project to be unset, as the URLs
	// only are seen by a fake http client.
	if *fProject == "" && !*fDryRun {
		log.Println("Must specify project (or --dry_run)")
		flag.PrintDefaults()
		return
	}

	q, err := batch.CreateQueuer(http.DefaultClient, nil, *fQueue, *fNumQueues, *fProject, *fBucket, *fDryRun)
	if err != nil {
		// In command line mode, good to just fast fail.
		log.Fatal(err)
	}
	if *fMonth != "" {
		q.PostMonth(*fExper + "/" + *fMonth + "/")
	} else if *fDay != "" {
		q.PostDay(nil, *fExper+"/"+*fDay+"/")
	}
}
<commit_msg>fix gardener date bug<commit_after>// Package main defines a service for handling various post-processing
// and house-keeping tasks associated with the pipelines.
// Most tasks will be run periodically, but some may be triggered
// by URL requests from authorized sources.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/pprof"
	"os"
	"runtime"
	"strconv"
	"time"

	"cloud.google.com/go/datastore"
	"github.com/m-lab/etl/batch"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// ###############################################################################
// All DataStore related code and variables
// ###############################################################################
const (
	dsNamespace   = "gardener"
	dsKind        = "gardener"
	batchStateKey = "batch-state"
)

var dsClient *datastore.Client

func getDSClient() (*datastore.Client, error) {
	var err error
	var client *datastore.Client
	if dsClient == nil {
		project, ok := os.LookupEnv("PROJECT")
		if !ok {
			return dsClient, errors.New("PROJECT env var not set")
		}
		client, err = datastore.NewClient(context.Background(), project)
		if err == nil {
			dsClient = client
		}
	}
	return dsClient, err
}

// Load retrieves an arbitrary record from datastore.
func Load(name string, obj interface{}) error {
	var client *datastore.Client
	client, err := getDSClient()
	if err != nil {
		return err
	}
	k := datastore.NameKey(dsKind, name, nil)
	k.Namespace = dsNamespace
	return client.Get(context.Background(), k, obj)
}

// Save stores an arbitrary object to kind/key in the default namespace.
// If a record already exists, then it is overwritten.
// TODO(gfr) Make an upsert version of this:
// https://cloud.google.com/datastore/docs/concepts/entities
func Save(key string, obj interface{}) error {
	client, err := getDSClient()
	if err != nil {
		return err
	}
	k := datastore.NameKey(dsKind, key, nil)
	k.Namespace = dsNamespace
	_, err = client.Put(context.Background(), k, obj)
	return err
}

// ###############################################################################
//  Batch processing task scheduling and support code
// ###############################################################################

// Persistent Queuer for use in handlers and gardener tasks.
var batchQueuer batch.Queuer

// QueueState holds the state information for each batch queue.
type QueueState struct {
	// Per queue, indicate which day is being processed in
	// that queue.  No need to process more than one day at a time.
	// The other queues will take up the slack while we add more
	// tasks when one queue is emptied.

	QueueName        string // Name of the batch queue.
	NextTask         string // Name of task file currently being enqueued.
	PendingPartition string // FQ Name of next table partition to be added.
}

// BatchState holds the entire batch processing state.
// It holds the state locally, and also is stored in DataStore for
// recovery when the instance is restarted, e.g. for weekly platform
// updates.
// At any given time, we restrict ourselves to a 14 day reprocessing window,
// and finish the beginning of that window before reprocessing any dates beyond it.
// When we determine that the first date in the window has been submitted for
// processing, we advance the window up to the next pending date.
type BatchState struct {
	Hostname       string       // Hostname of the gardener that saved this.
	InstanceID     string       // instance ID of the gardener that saved this.
	WindowStart    time.Time    // Start of two week window we are working on.
	QueueBase      string       // base name for queues.
	QStates        []QueueState // States for each queue.
	LastUpdateTime time.Time    // Time of last update.  (Is this in DS metadata?)
}

// MaybeScheduleMoreTasks will look for an empty task queue, and if it finds one, will look
// for corresponding days to add to the queue.
// Alternatively, it may look first for the N oldest days to be reprocessed, and will then
// check whether any of the task queues for those oldest days is empty, and conditionally add tasks.
func MaybeScheduleMoreTasks(queuer *batch.Queuer) {
	// GetTaskQueueDepth returns the number of pending items in a task queue.
	stats, err := queuer.GetTaskqueueStats()
	if err != nil {
		log.Println(err)
	} else {
		for k, v := range stats {
			if len(v) > 0 && v[0].Tasks == 0 && v[0].InFlight == 0 {
				log.Printf("Ready: %s: %v\n", k, v[0])
				// Should add more tasks now.
			}
		}
	}
}

// queuerFromEnv creates a Queuer struct initialized from environment variables.
// It uses TASKFILE_BUCKET, PROJECT, QUEUE_BASE, and NUM_QUEUES.
func queuerFromEnv() (batch.Queuer, error) {
	bucketName, ok := os.LookupEnv("TASKFILE_BUCKET")
	if !ok {
		return batch.Queuer{}, errors.New("TASKFILE_BUCKET not set")
	}
	project, ok := os.LookupEnv("PROJECT")
	if !ok {
		return batch.Queuer{}, errors.New("PROJECT not set")
	}
	queueBase, ok := os.LookupEnv("QUEUE_BASE")
	if !ok {
		return batch.Queuer{}, errors.New("QUEUE_BASE not set")
	}
	numQueues, err := strconv.Atoi(os.Getenv("NUM_QUEUES"))
	if err != nil {
		log.Println(err)
		return batch.Queuer{}, errors.New("Parse error on NUM_QUEUES")
	}

	return batch.CreateQueuer(http.DefaultClient, nil, queueBase, numQueues, project, bucketName, false)
}

// StartDateRFC3339 is the date at which reprocessing will start when it catches
// up to present.  For now, we are making this the beginning of the ETL timeframe,
// until we get annotation fixed to use the actual data date instead of NOW.
const StartDateRFC3339 = "2017-05-01T00:00:00Z"

// startupBatch determines whether some other instance has control, and
// assumes control if not.
func startupBatch(base string, numQueues int) (BatchState, error) {
	hostname := os.Getenv("HOSTNAME")
	instance := os.Getenv("GAE_INSTANCE")
	queues := make([]QueueState, numQueues)
	var bs BatchState
	err := Load(batchStateKey, &bs)
	if err != nil {
		startDate, err := time.Parse(time.RFC3339, StartDateRFC3339)
		if err != nil {
			log.Println("Could not parse start time.  Not starting batch.")
			return bs, err
		}
		bs = BatchState{hostname, instance, startDate, base, queues, time.Now()}

	} else {
		// TODO - should check whether we should take over, or leave alone.
	}

	err = Save(batchStateKey, &bs)
	return bs, err
}

// ###############################################################################
//  Top level service control code.
// ###############################################################################

// periodic will run approximately every 5 minutes.
func periodic() {
	_, err := startupBatch(batchQueuer.QueueBase, batchQueuer.NumQueues)
	if err != nil {
		log.Fatal(err)
	}
	for {
		log.Println("Periodic is running")

		MaybeScheduleMoreTasks(&batchQueuer)

		// There is no need for randomness, since this is a singleton handler.
		time.Sleep(300 * time.Second)
	}
}

func setupPrometheus() {
	// Define a custom serve mux for prometheus to listen on a separate port.
	// We listen on a separate port so we can forward this port on the host VM.
	// We cannot forward port 8080 because it is used by AppEngine.
	mux := http.NewServeMux()
	// Assign the default prometheus handler to the standard exporter path.
	mux.Handle("/metrics", promhttp.Handler())
	// Assign the pprof handling paths to the external port to access individual
	// instances.
	mux.HandleFunc("/debug/pprof/", pprof.Index)
	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
	go http.ListenAndServe(":9090", mux)
}

// Status provides basic information about the service.  For now, it is just
// configuration and version info.  In future it will likely include more
// dynamic information.
// TODO(gfr) Add either a black list or a white list for the environment
// variables, so we can hide sensitive vars. https://github.com/m-lab/etl/issues/384
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "
NOTE: This is just one of potentially many instances.
\n")
	commit := os.Getenv("COMMIT_HASH")
	if len(commit) >= 8 {
		fmt.Fprintf(w, "Release: %s
  Commit: %s
\n",
			os.Getenv("RELEASE_TAG"), os.Getenv("COMMIT_HASH"), os.Getenv("COMMIT_HASH")[0:7])
	} else {
		fmt.Fprintf(w, "Release: %s   Commit: unknown\n", os.Getenv("RELEASE_TAG"))
	}

	env := os.Environ()
	for i := range env {
		fmt.Fprintf(w, "%s
\n", env[i])
	}
	fmt.Fprintf(w, "\n")
}

var healthy = false

// healthCheck, for now, used for both /ready and /alive.
func healthCheck(w http.ResponseWriter, r *http.Request) {
	if !healthy {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, `{"message": "Internal server error."}`)
	}
	fmt.Fprint(w, "ok")
}

// runService starts a service handler and runs forever.
// The configuration info comes from environment variables.
func runService() {
	// Enable block profiling
	runtime.SetBlockProfileRate(1000000) // One event per msec.

	setupPrometheus()
	// We also setup another prometheus handler on a non-standard path. This
	// path name will be accessible through the AppEngine service address,
	// however it will be served by a random instance.
	http.Handle("/random-metrics", promhttp.Handler())
	http.HandleFunc("/", Status)
	http.HandleFunc("/status", Status)

	http.HandleFunc("/alive", healthCheck)
	http.HandleFunc("/ready", healthCheck)

	var err error
	batchQueuer, err = queuerFromEnv()
	if err == nil {
		healthy = true
		log.Println("Running as a service.")

		// Run the background "periodic" function.
		go periodic()
	} else {
		// Leaving healthy == false
		// This will cause app-engine to roll back.
		log.Println(err)
		log.Println("Required environment variables are missing or invalid.")
	}

	// ListenAndServe, and terminate when it returns.
	log.Fatal(http.ListenAndServe(":8080", nil))
}

// ###############################################################################
//  Top level command line code.
// ###############################################################################

// These are used only for command line.  For service, environment variables are used
// for general parameters, and request parameter for month.
var (
	fProject = flag.String("project", "", "Project containing queues.")
	fQueue   = flag.String("queue", "etl-ndt-batch-", "Base of queue name.")
	// TODO implement listing queues to determine number of queue, and change this to 0
	fNumQueues = flag.Int("num_queues", 8, "Number of queues.  Normally determined by listing queues.")
	// Gardener will only read from this bucket, so its ok to use production bucket as default.
	fBucket = flag.String("bucket", "archive-mlab-oti", "Source bucket.")
	fExper  = flag.String("experiment", "ndt", "Experiment prefix, without trailing slash.")
	fMonth  = flag.String("month", "", "Single month spec, as YYYY/MM")
	fDay    = flag.String("day", "", "Single day spec, as YYYY/MM/DD")
	fDryRun = flag.Bool("dry_run", false, "Prevents all output to queue_pusher.")
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	// Check if invoked as a service.
	isService, _ := strconv.ParseBool(os.Getenv("GARDENER_SERVICE"))
	if isService {
		runService()
		return
	}

	// Otherwise this is a command line invocation...
	flag.Parse()
	// Check that either project or dry-run is set.
	// If dry-run, it is ok for the project to be unset, as the URLs
	// only are seen by a fake http client.
	if *fProject == "" && !*fDryRun {
		log.Println("Must specify project (or --dry_run)")
		flag.PrintDefaults()
		return
	}

	q, err := batch.CreateQueuer(http.DefaultClient, nil, *fQueue, *fNumQueues, *fProject, *fBucket, *fDryRun)
	if err != nil {
		// In command line mode, good to just fast fail.
		log.Fatal(err)
	}
	if *fMonth != "" {
		q.PostMonth(*fExper + "/" + *fMonth + "/")
	} else if *fDay != "" {
		q.PostDay(nil, *fExper+"/"+*fDay+"/")
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	"github.com/funkygao/gocli"
)

type Mirror struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Mirror) Run(args []string) (exitCode int) {
	return
}

func (*Mirror) Synopsis() string {
	return "Continuously copy data between two Kafka clusters"
}

func (this *Mirror) Help() string {
	help := fmt.Sprintf(`
Usage: %s mirror [options]

    Continuously copy data between two Kafka clusters

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>start on `gk mirror`: parse cli arguments<commit_after>package command

import (
	"flag"
	"fmt"
	"strings"

	"github.com/funkygao/gocli"
)

type Mirror struct {
	Ui  cli.Ui
	Cmd string

	zone1, zone2       string
	cluster1, cluster2 string
	excludes           string
	topicsExcluded     map[string]struct{}
}

func (this *Mirror) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("mirror", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone1, "z1", "", "")
	cmdFlags.StringVar(&this.zone2, "z2", "", "")
	cmdFlags.StringVar(&this.cluster1, "c1", "", "")
	cmdFlags.StringVar(&this.cluster2, "c2", "", "")
	cmdFlags.StringVar(&this.excludes, "excluded", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z1", "-z2", "-c1", "-c2").
		invalid(args) {
		return 2
	}

	this.topicsExcluded = make(map[string]struct{})
	for _, e := range strings.Split(this.excludes, ",") {
		this.topicsExcluded[e] = struct{}{}
	}

	return
}

func (*Mirror) Synopsis() string {
	return "Continuously copy data between two Kafka clusters"
}

func (this *Mirror) Help() string {
	help := fmt.Sprintf(`
Usage: %s mirror [options]

    Continuously copy data between two Kafka clusters

Options:

    -z1 from zone

    -z2 to zone

    -c1 from cluster

    -c2 to cluster

    -exclude comma seperated topic names

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>