<commit_before>package neptulon

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
)

// Conn is a full-duplex bidirectional client-server connection.
type Conn struct {
	ID                 string     // Randomly generated unique connection ID
	Data               *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session
	conn               *tls.Conn
	headerSize         int
	maxMsgSize         int
	readDeadline       time.Duration
	debug              bool
	err                error
	clientDisconnected bool // hack: Whether the client disconnected from server before server closed connection
}

// NewConn creates a new neptulon.Conn object which wraps a given tls.Conn object.
// Default values for headerSize, maxMsgSize, and readDeadline are 4 bytes, 4294967295 bytes (4GB), and 300 seconds, respectively.
// Debug mode logs all raw TCP communication.
func NewConn(conn *tls.Conn, headerSize, maxMsgSize, readDeadline int, debug bool) (*Conn, error) {
	if headerSize == 0 {
		headerSize = 4
	}
	if maxMsgSize == 0 {
		maxMsgSize = 4294967295
	}
	if readDeadline == 0 {
		readDeadline = 300
	}

	id, err := GenID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:           id,
		Data:         cmap.New(),
		conn:         conn,
		headerSize:   headerSize,
		maxMsgSize:   maxMsgSize,
		readDeadline: time.Second * time.Duration(readDeadline),
		debug:        debug,
	}, nil
}

// Dial creates a new client side connection to a server at the given network address,
// with optional CA and/or a client certificate (PEM encoded X.509 cert/key).
// Debug mode logs all raw TCP communication.
func Dial(addr string, ca []byte, clientCert []byte, clientCertKey []byte, debug bool) (*Conn, error) {
	var cas *x509.CertPool
	var certs []tls.Certificate
	if ca != nil {
		cas = x509.NewCertPool()
		ok := cas.AppendCertsFromPEM(ca)
		if !ok {
			return nil, errors.New("failed to parse the CA certificate")
		}
	}
	if clientCert != nil {
		tlsCert, err := tls.X509KeyPair(clientCert, clientCertKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		c, _ := pem.Decode(clientCert)
		if tlsCert.Leaf, err = x509.ParseCertificate(c.Bytes); err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		certs = []tls.Certificate{tlsCert}
	}

	// todo: dial timeout like that of net.Conn.DialTimeout
	c, err := tls.Dial("tcp", addr, &tls.Config{RootCAs: cas, Certificates: certs})
	if err != nil {
		return nil, err
	}

	return NewConn(c, 0, 0, 0, debug)
}

// SetReadDeadline set the read deadline for the connection in seconds.
func (c *Conn) SetReadDeadline(seconds int) {
	c.readDeadline = time.Second * time.Duration(seconds)
}

// Read waits for and reads the next incoming message from the TLS connection.
func (c *Conn) Read() (msg []byte, err error) {
	if err = c.conn.SetReadDeadline(time.Now().Add(c.readDeadline)); err != nil {
		return
	}

	// read the content length header
	h := make([]byte, c.headerSize)
	var n int
	n, err = c.conn.Read(h)
	if err != nil {
		return
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to read header size %v bytes but instead read %v bytes", c.headerSize, n)
		return
	}

	// calculate the content length
	n = readHeaderBytes(h)

	// read the message content
	msg = make([]byte, n)
	total := 0
	for total < n {
		// todo: log here in case it gets stuck, or there is a dos attack, pumping up cpu usage!
		i, err := c.conn.Read(msg[total:])
		if err != nil {
			err = fmt.Errorf("errored while reading incoming message: %v", err)
			break
		}
		total += i
	}
	if total != n {
		err = fmt.Errorf("expected to read %v bytes instead read %v bytes", n, total)
	}

	if c.debug {
		log.Println("Incoming message:", string(msg))
	}

	return
}

// Write writes given message to the connection.
func (c *Conn) Write(msg []byte) error {
	l := len(msg)
	h := makeHeaderBytes(l, c.headerSize)

	// write the header
	n, err := c.conn.Write(h)
	if err != nil {
		return err
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	// write the body
	// todo: do we need a write loop? bufio uses a loop but it might be due to buff length limitation
	// todo2: we might need a write loop according to this: http://linux.die.net/man/2/write
	n, err = c.conn.Write(msg)
	if err != nil {
		return err
	}
	if n != l {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// ConnectionState returns basic TLS details about the connection.
func (c *Conn) ConnectionState() tls.ConnectionState {
	return c.conn.ConnectionState()
}

// Close closes a connection.
// Note: TCP/IP stack does not guarantee delivery of messages before the connection is closed.
func (c *Conn) Close() error {
	return c.conn.Close() // todo: if conn.err is nil, send a close req and wait ack then close? (or even wait for everything else to finish?)
}

func makeHeaderBytes(h, size int) []byte {
	b := make([]byte, size)
	binary.LittleEndian.PutUint32(b, uint32(h))
	return b
}

func readHeaderBytes(h []byte) int {
	return int(binary.LittleEndian.Uint32(h))
}
<commit_msg>expand node<commit_after>package neptulon

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
)

// Conn is a full-duplex bidirectional client-server connection.
type Conn struct {
	ID                 string     // Randomly generated unique connection ID
	Data               *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session
	conn               *tls.Conn
	headerSize         int
	maxMsgSize         int
	readDeadline       time.Duration
	debug              bool
	err                error
	clientDisconnected bool // hack: Whether the client disconnected from server before server closed connection
}

// NewConn creates a new neptulon.Conn object which wraps a given tls.Conn object.
// Default values for headerSize, maxMsgSize, and readDeadline are 4 bytes, 4294967295 bytes (4GB), and 300 seconds, respectively.
// Debug mode logs all raw TCP communication.
func NewConn(conn *tls.Conn, headerSize, maxMsgSize, readDeadline int, debug bool) (*Conn, error) {
	if headerSize == 0 {
		headerSize = 4
	}
	if maxMsgSize == 0 {
		maxMsgSize = 4294967295
	}
	if readDeadline == 0 {
		readDeadline = 300
	}

	id, err := GenID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:           id,
		Data:         cmap.New(),
		conn:         conn,
		headerSize:   headerSize,
		maxMsgSize:   maxMsgSize,
		readDeadline: time.Second * time.Duration(readDeadline),
		debug:        debug,
	}, nil
}

// Dial creates a new client side connection to a server at the given network address,
// with optional CA and/or a client certificate (PEM encoded X.509 cert/key).
// Debug mode logs all raw TCP communication.
func Dial(addr string, ca []byte, clientCert []byte, clientCertKey []byte, debug bool) (*Conn, error) {
	var cas *x509.CertPool
	var certs []tls.Certificate
	if ca != nil {
		cas = x509.NewCertPool()
		ok := cas.AppendCertsFromPEM(ca)
		if !ok {
			return nil, errors.New("failed to parse the CA certificate")
		}
	}
	if clientCert != nil {
		tlsCert, err := tls.X509KeyPair(clientCert, clientCertKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		c, _ := pem.Decode(clientCert)
		if tlsCert.Leaf, err = x509.ParseCertificate(c.Bytes); err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		certs = []tls.Certificate{tlsCert}
	}

	// todo: dial timeout like that of net.Conn.DialTimeout
	c, err := tls.Dial("tcp", addr, &tls.Config{RootCAs: cas, Certificates: certs})
	if err != nil {
		return nil, err
	}

	return NewConn(c, 0, 0, 0, debug)
}

// SetReadDeadline set the read deadline for the connection in seconds.
func (c *Conn) SetReadDeadline(seconds int) {
	c.readDeadline = time.Second * time.Duration(seconds)
}

// Read waits for and reads the next incoming message from the TLS connection.
func (c *Conn) Read() (msg []byte, err error) {
	if err = c.conn.SetReadDeadline(time.Now().Add(c.readDeadline)); err != nil {
		return
	}

	// read the content length header
	h := make([]byte, c.headerSize)
	var n int
	n, err = c.conn.Read(h)
	if err != nil {
		return
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to read header size %v bytes but instead read %v bytes", c.headerSize, n)
		return
	}

	// calculate the content length
	n = readHeaderBytes(h)

	// read the message content
	msg = make([]byte, n)
	total := 0
	for total < n {
		// todo: log here in case it gets stuck, or there is a dos attack, pumping up cpu usage!
		i, err := c.conn.Read(msg[total:])
		if err != nil {
			err = fmt.Errorf("errored while reading incoming message: %v", err)
			break
		}
		total += i
	}
	if total != n {
		err = fmt.Errorf("expected to read %v bytes instead read %v bytes", n, total)
	}

	if c.debug {
		log.Println("Incoming message:", string(msg))
	}

	return
}

// Write writes given message to the connection.
func (c *Conn) Write(msg []byte) error {
	l := len(msg)
	h := makeHeaderBytes(l, c.headerSize)

	// write the header
	n, err := c.conn.Write(h)
	if err != nil {
		return err
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	// write the body
	// todo: do we need a write loop? bufio uses a loop but it might be due to buff length limitation
	// todo2: we might need a write loop according to this: http://linux.die.net/man/2/write but I don't know how Go uses it internally
	n, err = c.conn.Write(msg)
	if err != nil {
		return err
	}
	if n != l {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// ConnectionState returns basic TLS details about the connection.
func (c *Conn) ConnectionState() tls.ConnectionState {
	return c.conn.ConnectionState()
}

// Close closes a connection.
// Note: TCP/IP stack does not guarantee delivery of messages before the connection is closed.
func (c *Conn) Close() error {
	return c.conn.Close() // todo: if conn.err is nil, send a close req and wait ack then close? (or even wait for everything else to finish?)
}

func makeHeaderBytes(h, size int) []byte {
	b := make([]byte, size)
	binary.LittleEndian.PutUint32(b, uint32(h))
	return b
}

func readHeaderBytes(h []byte) int {
	return int(binary.LittleEndian.Uint32(h))
}
<|endoftext|>


<commit_before>// Supports Windows, Linux, Mac, and Raspberry Pi

package main

import (
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

type connection struct {
	// The websocket connection.
	ws *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte
}

func (c *connection) reader() {
	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			break
		}

		h.broadcast <- message
	}
	c.ws.Close()
}

func (c *connection) writer() {
	for message := range c.send {
		err := c.ws.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			break
		}
	}
	c.ws.Close()
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("Received a upload")
	port := r.FormValue("port")
	if port == "" {
		http.Error(w, "port is required", http.StatusBadRequest)
		return
	}
	board := r.FormValue("board")
	if board == "" {
		http.Error(w, "board is required", http.StatusBadRequest)
		return
	}
	sketch, header, err := r.FormFile("sketch_hex")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
	}

	path, err := saveFileonTempDir(header.Filename, sketch)

	go spProgram(port, board, path)
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("Started a new websocket handler")
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		return
	}
	//c := &connection{send: make(chan []byte, 256), ws: ws}
	c := &connection{send: make(chan []byte, 256*10), ws: ws}
	h.register <- c
	defer func() { h.unregister <- c }()
	go c.writer()
	c.reader()
}
<commit_msg>handle errors on POST request<commit_after>// Supports Windows, Linux, Mac, and Raspberry Pi

package main

import (
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

type connection struct {
	// The websocket connection.
	ws *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte
}

func (c *connection) reader() {
	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			break
		}

		h.broadcast <- message
	}
	c.ws.Close()
}

func (c *connection) writer() {
	for message := range c.send {
		err := c.ws.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			break
		}
	}
	c.ws.Close()
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("Received a upload")
	port := r.FormValue("port")
	if port == "" {
		http.Error(w, "port is required", http.StatusBadRequest)
		return
	}
	board := r.FormValue("board")
	if board == "" {
		http.Error(w, "board is required", http.StatusBadRequest)
		return
	}
	sketch, header, err := r.FormFile("sketch_hex")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
	}

	if header != nil {
		path, err := saveFileonTempDir(header.Filename, sketch)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadGateway)
		}

		go spProgram(port, board, path)
	}
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("Started a new websocket handler")
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		return
	}
	//c := &connection{send: make(chan []byte, 256), ws: ws}
	c := &connection{send: make(chan []byte, 256*10), ws: ws}
	h.register <- c
	defer func() { h.unregister <- c }()
	go c.writer()
	c.reader()
}
<|endoftext|>


<commit_before>package fs

import (
	"context"
	"fmt"
	"io"
)

const copyBufferSize = 1024 * 1024 * 4

// CopyFile copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
func CopyFile(src FileReader, dest File, perm ...Permissions) error {
	var buf []byte
	return CopyFileBuf(src, dest, &buf, perm...)
}

// CopyFileBuf copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
// An non nil pointer to a []byte variable must be passed for buf.
// If that variable holds a non zero length byte slice, then this slice will be used as buffer,
// else a byte slice will be allocated and assigned to the variable.
// Use this function to re-use buffers between CopyFileBuf calls.
func CopyFileBuf(src FileReader, dest File, buf *[]byte, perm ...Permissions) error {
	return CopyFileBufContext(context.Background(), src, dest, buf, perm...)
}

// CopyFileBufContext copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
// An non nil pointer to a []byte variable must be passed for buf.
// If that variable holds a non zero length byte slice, then this slice will be used as buffer,
// else a byte slice will be allocated and assigned to the variable.
// Use this function to re-use buffers between CopyFileBufContext calls.
func CopyFileBufContext(ctx context.Context, src FileReader, dest File, buf *[]byte, perm ...Permissions) error {
	if buf == nil {
		panic("CopyFileBuf: buf is nil") // not a file system error
	}

	if ctx.Err() != nil {
		return ctx.Err()
	}

	// Handle directories
	if dest.IsDir() {
		dest = dest.Join(src.Name())
	} else {
		err := dest.Dir().MakeDir()
		if err != nil {
			return fmt.Errorf("CopyFileBuf: can't make directory %q: %w", dest.Dir(), err)
		}
	}

	// Use same file system copy if possible
	srcFile, srcIsFile := src.(File)
	if srcIsFile {
		fs := srcFile.FileSystem()
		if fs == dest.FileSystem() {
			return fs.CopyFile(ctx, srcFile.Path(), dest.Path(), buf)
		}
	}

	r, err := src.OpenReader()
	if err != nil {
		return fmt.Errorf("CopyFileBuf: can't open src reader: %w", err)
	}
	defer r.Close()

	if len(perm) == 0 && srcIsFile {
		perm = []Permissions{srcFile.Permissions()}
	}
	w, err := dest.OpenWriter(perm...)
	if err != nil {
		return fmt.Errorf("CopyFileBuf: can't open dest writer: %w", err)
	}
	defer w.Close()

	if len(*buf) == 0 {
		*buf = make([]byte, copyBufferSize)
	}
	_, err = io.CopyBuffer(w, r, *buf)
	if err != nil {
		return fmt.Errorf("CopyFileBuf: error from io.CopyBuffer: %w", err)
	}
	return nil
}

// CopyRecursive can copy between files of different file systems.
// The filter patterns are applied on filename level, not the whole path.
func CopyRecursive(src, dest File, patterns ...string) error {
	var buf []byte
	return copyRecursive(context.Background(), src, dest, patterns, &buf)
}

// CopyRecursiveContext can copy between files of different file systems.
// The filter patterns are applied on filename level, not the whole path.
func CopyRecursiveContext(ctx context.Context, src, dest File, patterns ...string) error {
	var buf []byte
	return copyRecursive(ctx, src, dest, patterns, &buf)
}

func copyRecursive(ctx context.Context, src, dest File, patterns []string, buf *[]byte) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	if !src.IsDir() {
		// Just copy one file
		return CopyFileBufContext(ctx, src, dest, buf)
	}

	if dest.Exists() && !dest.IsDir() {
		return fmt.Errorf("Can't copy a directory (%s) over a file (%s)", src.URL(), dest.URL())
	}

	// TODO better check
	if !dest.Exists() {
		err := dest.MakeDir()
		if err != nil {
			return fmt.Errorf("copyRecursive: can't make dest dir %q: %w", dest, err)
		}
	}

	// Copy directories recursive
	return src.ListDirContext(ctx, func(file File) error {
		return copyRecursive(ctx, file, dest.Join(file.Name()), patterns, buf)
	}, patterns...)
}
<commit_msg>CopyFileBufContext optimization for src.(*MemFile)<commit_after>package fs

import (
	"context"
	"fmt"
	"io"
)

const copyBufferSize = 1024 * 1024 * 4

// CopyFile copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
func CopyFile(src FileReader, dest File, perm ...Permissions) error {
	var buf []byte
	return CopyFileBuf(src, dest, &buf, perm...)
}

// CopyFileBuf copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
// An non nil pointer to a []byte variable must be passed for buf.
// If that variable holds a non zero length byte slice, then this slice will be used as buffer,
// else a byte slice will be allocated and assigned to the variable.
// Use this function to re-use buffers between CopyFileBuf calls.
func CopyFileBuf(src FileReader, dest File, buf *[]byte, perm ...Permissions) error {
	return CopyFileBufContext(context.Background(), src, dest, buf, perm...)
}

// CopyFileBufContext copies a single file between different file systems.
// If dest has a path that does not exist, then the directories
// up to that path will be created.
// If dest is an existing directory, then a file with the base name
// of src will be created there.
// An non nil pointer to a []byte variable must be passed for buf.
// If that variable holds a non zero length byte slice, then this slice will be used as buffer,
// else a byte slice will be allocated and assigned to the variable.
// Use this function to re-use buffers between CopyFileBufContext calls.
func CopyFileBufContext(ctx context.Context, src FileReader, dest File, buf *[]byte, perm ...Permissions) error {
	if buf == nil {
		panic("CopyFileBuf: buf is nil") // not a file system error
	}

	if ctx.Err() != nil {
		return ctx.Err()
	}

	// Handle directories
	if dest.IsDir() {
		dest = dest.Join(src.Name())
	} else {
		err := dest.Dir().MakeDir()
		if err != nil {
			return fmt.Errorf("CopyFileBuf: can't make directory %q: %w", dest.Dir(), err)
		}
	}

	srcFile, srcIsFile := src.(File)
	if srcIsFile {
		// Use same file system copy if possible
		fs := srcFile.FileSystem()
		if fs == dest.FileSystem() {
			return fs.CopyFile(ctx, srcFile.Path(), dest.Path(), buf)
		}
	} else if srcMemFile, ok := src.(*MemFile); ok {
		// Don't use io.CopyBuffer in case of MemFile
		return dest.WriteAll(srcMemFile.FileData, perm...)
	}

	r, err := src.OpenReader()
	if err != nil {
		return fmt.Errorf("CopyFileBuf: can't open src reader: %w", err)
	}
	defer r.Close()

	if len(perm) == 0 && srcIsFile {
		perm = []Permissions{srcFile.Permissions()}
	}
	w, err := dest.OpenWriter(perm...)
	if err != nil {
		return fmt.Errorf("CopyFileBuf: can't open dest writer: %w", err)
	}
	defer w.Close()

	if len(*buf) == 0 {
		*buf = make([]byte, copyBufferSize)
	}
	_, err = io.CopyBuffer(w, r, *buf)
	if err != nil {
		return fmt.Errorf("CopyFileBuf: error from io.CopyBuffer: %w", err)
	}
	return nil
}

// CopyRecursive can copy between files of different file systems.
// The filter patterns are applied on filename level, not the whole path.
func CopyRecursive(src, dest File, patterns ...string) error {
	var buf []byte
	return copyRecursive(context.Background(), src, dest, patterns, &buf)
}

// CopyRecursiveContext can copy between files of different file systems.
// The filter patterns are applied on filename level, not the whole path.
func CopyRecursiveContext(ctx context.Context, src, dest File, patterns ...string) error {
	var buf []byte
	return copyRecursive(ctx, src, dest, patterns, &buf)
}

func copyRecursive(ctx context.Context, src, dest File, patterns []string, buf *[]byte) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	if !src.IsDir() {
		// Just copy one file
		return CopyFileBufContext(ctx, src, dest, buf)
	}

	if dest.Exists() && !dest.IsDir() {
		return fmt.Errorf("Can't copy a directory (%s) over a file (%s)", src.URL(), dest.URL())
	}

	// TODO better check
	if !dest.Exists() {
		err := dest.MakeDir()
		if err != nil {
			return fmt.Errorf("copyRecursive: can't make dest dir %q: %w", dest, err)
		}
	}

	// Copy directories recursive
	return src.ListDirContext(ctx, func(file File) error {
		return copyRecursive(ctx, file, dest.Join(file.Name()), patterns, buf)
	}, patterns...)
}
<|endoftext|>


<commit_before>/*
Package tasker is a light distribute producer&consumer task model based on beego.

Task

the main description of message or job. the state machine like this:

                                   +----> failed
                                   |
   pending ----+----> running -----+----> success
               ^                   |
               |                   v
               +---------------- retry

*/
package tasker

import (
	"errors"
	"math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
	"github.com/sony/sonyflake"
)

// Core is the task package config table.
//
// `MasterOutOfDate` means the time master state can be save, instance can race to be master when the duration of `Updated` to now bigger than this.
// `InstanceHeartbeat` is the max interval Instance should be check Master, should be less than `MasterOutOfDate`
type Core struct {
	Id                uint64
	MasterInstanceID  uint16    `orm:"column(master_instance_id)"`
	MasterFQDN        string    `orm:"column(master_fqdn)"`
	Updated           time.Time `orm:"auto_now"`
	MasterOutOfDate   int64     // ms
	InstanceHeartbeat int64     // ms
}

const (
	// defaultMasterOutOfDate is the default MasterOutOfDate value
	defaultMasterOutOfDate = 6000
	// defaultInstanceHeartbeat is the default InstanceHeartbeat value
	defaultInstanceHeartbeat = 3000
)

// UniqID use for generate worker id.
var UniqID *sonyflake.Sonyflake

// InstanceID is the tasker instance uniq key.
var InstanceID uint16

// IsMaster when true, the instance is master instance.
var IsMaster bool

/*
Init will initialize the tasker instance, include:

	1. generate the InstanceID use MachineID func, use instance private ip address when MachineID is nil
	2. start race master in goroutine
	3. initialize all task
*/
func Init(MachineID func() (uint16, error), CheckMachineID func(uint16) bool) (err error) {
	if MachineID == nil {
		MachineID = lower16BitPrivateIP
	}

	if InstanceID, err = MachineID(); err != nil {
		return
	}
	dora.Info().Msgf("tasker started with InstanceID %d FQDN %s", InstanceID, FQDN())

	if CheckMachineID != nil {
		if !CheckMachineID(InstanceID) {
			return errors.New("not valid MachineID via CheckMachineID method validate")
		}
	}

	if err = InitIDGEN(MachineID, CheckMachineID); err != nil {
		// InitIDGEN failed
		return
	}

	go keepMasterRace()

	InitAllTask()

	return nil
}

// InitIDGEN is initialize the ID generator, does not need invoke if had use `Init`
func InitIDGEN(MachineID func() (uint16, error), CheckMachineID func(uint16) bool) error {
	// init sonyflake.
	UniqID = sonyflake.NewSonyflake(sonyflake.Settings{
		StartTime:      time.Now(),
		MachineID:      MachineID,
		CheckMachineID: CheckMachineID,
	})
	if UniqID == nil {
		return errors.New("initialize unique id generate tool failed")
	}
	return nil
}

func RegisterModel() {
	orm.RegisterModel(new(Core), new(Task))
}

func getCore() (core *Core, err error) {
	o := orm.NewOrm()
	core = &Core{Id: 1} // alwary read id=1
	err = o.Read(core)
	return
}

func (self *Core) becomeMaster() {
	self.MasterInstanceID = InstanceID
	self.MasterFQDN = FQDN()

	o := orm.NewOrm()
	var err error
	if self.Id != 0 {
		// save
		_, err = o.Update(self, "MasterInstanceID", "MasterFQDN", "Updated")
	} else {
		// insert
		_, err = o.Insert(self)
	}

	if err != nil {
		dora.Error().Msgf("[tasker] try to becomeMaster failed: %s", err.Error())
	}
}

func (self *Core) heartbeatMaster() {
	o := orm.NewOrm()
	_, err := o.Update(self, "Updated")
	if err != nil {
		dora.Error().Msgf("[tasker] heartbeatMaster failed: %s", err.Error())
	}
}

// keepMasterRace always race to be master
func keepMasterRace() {
	rand.Seed(time.Now().Unix())

	for {
		core, err := getCore()
		if err != nil {
			if err == orm.ErrNoRows {
				// no core object, register self
				core.Id = 0
				core.MasterOutOfDate = defaultMasterOutOfDate
				core.InstanceHeartbeat = defaultInstanceHeartbeat
				core.becomeMaster()
				continue
			}

			// output the error log and sleep 1s
			dora.Error().Msgf("[tasker] get core config from db failed: %s", err.Error())
			time.Sleep(1 * time.Second)
			continue
		}

		if core.MasterInstanceID == InstanceID {
			// I'm master
			IsMaster = true
			core.heartbeatMaster()
		} else {
			// I'm not master
			IsMaster = false
			if time.Since(core.Updated).Nanoseconds()/int64(time.Millisecond) > core.MasterOutOfDate {
				// the master outofdate, I will be the master
				core.becomeMaster()
			}
		}

		time.Sleep(time.Duration(core.InstanceHeartbeat-rand.Int63n(2000)) * time.Millisecond)
	}
}

func lower16BitPrivateIP() (uint16, error) {
	ip, err := privateIPv4()
	if err != nil {
		return 0, err
	}

	return uint16(ip[2])<<8 + uint16(ip[3]), nil
}

func privateIPv4() (net.IP, error) {
	as, err := net.InterfaceAddrs()
	if err != nil {
		return nil, err
	}

	for _, a := range as {
		ipnet, ok := a.(*net.IPNet)
		if !ok || ipnet.IP.IsLoopback() {
			continue
		}

		ip := ipnet.IP.To4()
		if ip != nil {
			return ip, nil
		}
	}
	return nil, errors.New("no private ip address")
}

// FQDN Get Fully Qualified Domain Name
// returns "unknown" or hostanme in case of error
func FQDN() string {
	hostname, err := os.Hostname()
	if err != nil {
		return "unknown"

	}

	addrs, err := net.LookupIP(hostname)
	if err != nil {
		return hostname
	}

	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			ip, err := ipv4.MarshalText()
			if err != nil {
				return hostname
			}
			hosts, err := net.LookupAddr(string(ip))
			if err != nil {
				return hostname
			}
			fqdn := hosts[0]
			return strings.TrimSuffix(fqdn, ".") // return fqdn without trailing dot
		}

	}

	return hostname
}
<commit_msg>add cache for FQDN func<commit_after>/*
Package tasker is a light distribute producer&consumer task model based on beego.

Task

the main description of message or job. the state machine like this:

                                   +----> failed
                                   |
   pending ----+----> running -----+----> success
               ^                   |
               |                   v
               +---------------- retry

*/
package tasker

import (
	"errors"
	"math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
	"github.com/sony/sonyflake"
)

// Core is the task package config table.
//
// `MasterOutOfDate` means the time master state can be save, instance can race to be master when the duration of `Updated` to now bigger than this.
// `InstanceHeartbeat` is the max interval Instance should be check Master, should be less than `MasterOutOfDate`
type Core struct {
	Id                uint64
	MasterInstanceID  uint16    `orm:"column(master_instance_id)"`
	MasterFQDN        string    `orm:"column(master_fqdn)"`
	Updated           time.Time `orm:"auto_now"`
	MasterOutOfDate   int64 // ms
	InstanceHeartbeat int64 // ms
}

const (
	// defaultMasterOutOfDate is the default MasterOutOfDate value
	defaultMasterOutOfDate = 6000
	// defaultInstanceHeartbeat is the default InstanceHeartbeat value
	defaultInstanceHeartbeat = 3000
)

// UniqID use for generate worker id.
var UniqID *sonyflake.Sonyflake

// InstanceID is the tasker instance uniq key.
var InstanceID uint16

// IsMaster when true, the instance is master instance.
var IsMaster bool

/*
Init will initialize the tasker instance, include:

	1. generate the InstanceID use MachineID func, use instance private ip address when MachineID is nil
	2. start race master in goroutine
	3. initialize all task
*/
func Init(MachineID func() (uint16, error), CheckMachineID func(uint16) bool) (err error) {
	if MachineID == nil {
		MachineID = lower16BitPrivateIP
	}

	if InstanceID, err = MachineID(); err != nil {
		return
	}

	fqdn = getFQDN()
	dora.Info().Msgf("tasker started with InstanceID %d FQDN %s", InstanceID, FQDN())

	if CheckMachineID != nil {
		if !CheckMachineID(InstanceID) {
			return errors.New("not valid MachineID via CheckMachineID method validate")
		}
	}

	if err = InitIDGEN(MachineID, CheckMachineID); err != nil {
		// InitIDGEN failed
		return
	}

	go keepFQDNUpdating()
	go keepMasterRace()

	InitAllTask()

	return nil
}

// InitIDGEN is initialize the ID generator, does not need invoke if had use `Init`
func InitIDGEN(MachineID func() (uint16, error), CheckMachineID func(uint16) bool) error {
	// init sonyflake.
	UniqID = sonyflake.NewSonyflake(sonyflake.Settings{
		StartTime:      time.Now(),
		MachineID:      MachineID,
		CheckMachineID: CheckMachineID,
	})
	if UniqID == nil {
		return errors.New("initialize unique id generate tool failed")
	}
	return nil
}

func RegisterModel() {
	orm.RegisterModel(new(Core), new(Task))
}

func getCore() (core *Core, err error) {
	o := orm.NewOrm()
	core = &Core{Id: 1} // alwary read id=1
	err = o.Read(core)
	return
}

func (self *Core) becomeMaster() {
	self.MasterInstanceID = InstanceID
	self.MasterFQDN = FQDN()

	o := orm.NewOrm()
	var err error
	if self.Id != 0 {
		// save
		_, err = o.Update(self, "MasterInstanceID", "MasterFQDN", "Updated")
	} else {
		// insert
		_, err = o.Insert(self)
	}

	if err != nil {
		dora.Error().Msgf("[tasker] try to becomeMaster failed: %s", err.Error())
	}
}

func (self *Core) heartbeatMaster() {
	o := orm.NewOrm()
	_, err := o.Update(self, "Updated")
	if err != nil {
		dora.Error().Msgf("[tasker] heartbeatMaster failed: %s", err.Error())
	}
}

// keepMasterRace always race to be master
func keepMasterRace() {
	rand.Seed(time.Now().Unix())

	for {
		core, err := getCore()
		if err != nil {
			if err == orm.ErrNoRows {
				// no core object, register self
				core.Id = 0
				core.MasterOutOfDate = defaultMasterOutOfDate
				core.InstanceHeartbeat = defaultInstanceHeartbeat
				core.becomeMaster()
				continue
			}

			// output the error log and sleep 1s
			dora.Error().Msgf("[tasker] get core config from db failed: %s", err.Error())
			time.Sleep(1 * time.Second)
			continue
		}

		if core.MasterInstanceID == InstanceID {
			// I'm master
			IsMaster = true
			core.heartbeatMaster()
		} else {
			// I'm not master
			IsMaster = false
			if time.Since(core.Updated).Nanoseconds()/int64(time.Millisecond) > core.MasterOutOfDate {
				// the master outofdate, I will be the master
				core.becomeMaster()
			}
		}

		time.Sleep(time.Duration(core.InstanceHeartbeat-rand.Int63n(2000)) * time.Millisecond)
	}
}

func lower16BitPrivateIP() (uint16, error) {
	ip, err := privateIPv4()
	if err != nil {
		return 0, err
	}

	return uint16(ip[2])<<8 + uint16(ip[3]), nil
}

func privateIPv4() (net.IP, error) {
	as, err := net.InterfaceAddrs()
	if err != nil {
		return nil, err
	}

	for _, a := range as {
		ipnet, ok := a.(*net.IPNet)
		if !ok || ipnet.IP.IsLoopback() {
			continue
		}

		ip := ipnet.IP.To4()
		if ip != nil {
			return ip, nil
		}
	}
	return nil, errors.New("no private ip address")
}

var fqdn string

// FQDN Get Fully Qualified Domain Name
// returns "unknown" or hostanme in case of error
func FQDN() string {
	return fqdn
}

func keepFQDNUpdating() {
	for {
		fqdn = getFQDN()
		time.Sleep(180 * time.Second)
	}
}

func getFQDN() string {
	hostname, err := os.Hostname()
	if err != nil {
		return "unknown"
	}

	addrs, err := net.LookupIP(fqdn)
	if err != nil {
		return hostname
	}

	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			ip, err := ipv4.MarshalText()
			if err != nil {
				return hostname
			}
			hosts, err := net.LookupAddr(string(ip))
			if err != nil {
				return hostname
			}
			hostname = strings.TrimSuffix(hosts[0], ".") // return fqdn without trailing dot
			return hostname
		}
	}

	return hostname
}
<|endoftext|>


<commit_before>package cors

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	headerAllowOrigin      = "Access-Control-Allow-Origin"
	headerAllowCredentials = "Access-Control-Allow-Credentials"
	headerAllowHeaders     = "Access-Control-Allow-Headers"
	headerAllowMethods     = "Access-Control-Allow-Methods"
	headerMaxAge           = "Access-Control-Max-Age"

	headerOrigin         = "Origin"
	headerRequestMethod  = "Access-Control-Request-Method"
	headerRequestHeaders = "Access-Control-Request-Headers"
)

// Represents Access Control options.
type Opts struct {
	// If set, all origins are allowed.
	AllowAllOrigins bool
	// A list of allowed domain patterns.
	AllowOrigins []string
	// If set, allows to share auth credentials such as cookies.
	AllowCredentials bool
	// A list of allowed HTTP methods.
	AllowMethods []string
	// A list of allowed HTTP headers.
	AllowHeaders []string
	// Max age of the CORS headers.
	MaxAge time.Duration
}

// Converts options into a map of HTTP headers.
func (o *Opts) Header(origin string) (headers map[string]string) {
	headers = make(map[string]string)
	// if origin is not alowed, don't extend the headers
	// with CORS headers.
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}

	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allow credentials
	headers[headerAllowCredentials] = strconv.FormatBool(o.AllowCredentials)

	// add allow methods
	if len(o.AllowMethods) > 0 {
		headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
	}

	// add allow headers
	if len(o.AllowHeaders) > 0 {
		// TODO: Add default headers
		headers[headerAllowHeaders] = strings.Join(o.AllowHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

func (o *Opts) PreflightHeader(origin, rMethod, rHeaders string) (headers map[string]string) {
	headers = make(map[string]string)
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}
	// verify if requested method is allowed
	// TODO: Too many for loops
	for _, method := range o.AllowMethods {
		if method == rMethod {
			headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
			break
		}
	}

	// verify is requested headers are allowed
	var allowed []string
	for _, rHeader := range strings.Split(rHeaders, ",") {
	lookupLoop:
		for _, allowedHeader := range o.AllowHeaders {
			if rHeader == allowedHeader {
				allowed = append(allowed, rHeader)
				break lookupLoop
			}
		}
	}

	if len(allowed) > 0 {
		headers[headerAllowHeaders] = strings.Join(allowed, ",")
	}
	return
}

// Looks up if origin matches one of the patterns
// provided in Opts.AllowOrigins patterns.
func (o *Opts) IsOriginAllowed(origin string) (allowed bool) {
	for _, pattern := range o.AllowOrigins {
		allowed, _ = regexp.MatchString(pattern, origin)
		if allowed {
			return
		}
	}
	return
}

func Allow(opts *Opts) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		var (
			origin           = req.Header.Get(headerOrigin)
			requestedMethod  = req.Header.Get(headerRequestMethod)
			requestedHeaders = req.Header.Get(headerRequestHeaders)
			// additional headers to be added
			// to the response.
			headers map[string]string
		)

		if req.Method == "OPTIONS" &&
			(requestedMethod != "" || requestedHeaders != "") {
			// TODO: if preflight, respond with exact headers if allowed
			headers = opts.PreflightHeader(origin, requestedMethod, requestedHeaders)
		} else {
			headers = opts.Header(origin)
		}

		for key, value := range headers {
			res.Header().Set(key, value)
		}
	}
}
<commit_msg>-n [Migrated] Add max age to preflight responses.<commit_after>package cors

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	headerAllowOrigin      = "Access-Control-Allow-Origin"
	headerAllowCredentials = "Access-Control-Allow-Credentials"
	headerAllowHeaders     = "Access-Control-Allow-Headers"
	headerAllowMethods     = "Access-Control-Allow-Methods"
	headerMaxAge           = "Access-Control-Max-Age"

	headerOrigin         = "Origin"
	headerRequestMethod  = "Access-Control-Request-Method"
	headerRequestHeaders = "Access-Control-Request-Headers"
)

// Represents Access Control options.
type Opts struct {
	// If set, all origins are allowed.
	AllowAllOrigins bool
	// A list of allowed domain patterns.
	AllowOrigins []string
	// If set, allows to share auth credentials such as cookies.
	AllowCredentials bool
	// A list of allowed HTTP methods.
	AllowMethods []string
	// A list of allowed HTTP headers.
	AllowHeaders []string
	// Max age of the CORS headers.
	MaxAge time.Duration
}

// Converts options into a map of HTTP headers.
func (o *Opts) Header(origin string) (headers map[string]string) {
	headers = make(map[string]string)
	// if origin is not alowed, don't extend the headers
	// with CORS headers.
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}

	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allow credentials
	headers[headerAllowCredentials] = strconv.FormatBool(o.AllowCredentials)

	// add allow methods
	if len(o.AllowMethods) > 0 {
		headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
	}

	// add allow headers
	if len(o.AllowHeaders) > 0 {
		// TODO: Add default headers
		headers[headerAllowHeaders] = strings.Join(o.AllowHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

func (o *Opts) PreflightHeader(origin, rMethod, rHeaders string) (headers map[string]string) {
	headers = make(map[string]string)
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}
	// verify if requested method is allowed
	// TODO: Too many for loops
	for _, method := range o.AllowMethods {
		if method == rMethod {
			headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
			break
		}
	}

	// verify is requested headers are allowed
	var allowed []string
	for _, rHeader := range strings.Split(rHeaders, ",") {
	lookupLoop:
		for _, allowedHeader := range o.AllowHeaders {
			if rHeader == allowedHeader {
				allowed = append(allowed, rHeader)
				break lookupLoop
			}
		}
	}

	if len(allowed) > 0 {
		headers[headerAllowHeaders] = strings.Join(allowed, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

// Looks up if origin matches one of the patterns
// provided in Opts.AllowOrigins patterns.
func (o *Opts) IsOriginAllowed(origin string) (allowed bool) {
	for _, pattern := range o.AllowOrigins {
		allowed, _ = regexp.MatchString(pattern, origin)
		if allowed {
			return
		}
	}
	return
}

func Allow(opts *Opts) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		var (
			origin           = req.Header.Get(headerOrigin)
			requestedMethod  = req.Header.Get(headerRequestMethod)
			requestedHeaders = req.Header.Get(headerRequestHeaders)
			// additional headers to be added
			// to the response.
			headers map[string]string
		)

		if req.Method == "OPTIONS" &&
			(requestedMethod != "" || requestedHeaders != "") {
			// TODO: if preflight, respond with exact headers if allowed
			headers = opts.PreflightHeader(origin, requestedMethod, requestedHeaders)
		} else {
			headers = opts.Header(origin)
		}

		for key, value := range headers {
			res.Header().Set(key, value)
		}
	}
}
<|endoftext|>


<commit_before>// Package csrf generates and validates csrf tokens for martini.
// There are multiple methods of delivery including via a cookie or HTTP
// header.
// Validation occurs via a traditional hidden form key of "_csrf", or via
// a custom HTTP header "X-CSRFToken".
//
// package main
//
// import (
//     "github.com/codegangsta/martini"
//     "github.com/martini-contib/csrf"
//     "github.com/martini-contrib/render"
//     "github.com/martini-contib/sessions"
//     "net/http"
// )
//
// func main() {
//     m := martini.Classic()
//     store := sessions.NewCookieStore([]byte("secret123"))
//     m.Use(sessions.Sessions("my_session", store))
//     // Setup generation middleware.
//     m.Use(csrf.Generate(&csrf.Options{
//         Secret:     "token123",
//         SessionKey: "userID",
//     }))
//     m.Use(render.Renderer())
//
//     // Simulate the authentication of a session. If userID exists redirect
//     // to a form that requires csrf protection.
//     m.Get("/", func(s sessions.Session, r render.Render) {
//         if s.Get("userID") == nil {
//             r.Redirect("/login", 302)
//             return
//         }
//         r.Redirect("/protected", 302)
//     })
//
//     // Set userID for the session.
//     m.Get("/login", func(s sessions.Session, r render.Render) {
//         s.Set("userID", "123456")
//         r.Redirect("/", 302)
//     })
//
//     // Render a protected form. Passing a csrf token by calling x.GetToken()
//     m.Get("/protected", func(s sessions.Session, r render.Render, x csrf.CSRF) {
//         if s.Get("userID") == nil {
//             r.Redirect("/login", 401)
//             return
//         }
//         r.HTML(200, "protected", x.GetToken())
//     })
//
//     // Apply csrf validation to route.
//     m.Post("/protected", csrf.Validate, func(s sessions.Session, r render.Render) {
//         if s.Get("userID") != nil {
//             r.HTML(200, "result", "You submitted a valid token")
//             return
//         }
//         r.Redirect("/login", 401)
//     })
//
//     m.Run()
// }
package csrf

import (
	"code.google.com/p/xsrftoken"
	"fmt"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
	"net/http"
	"regexp"
	"strings"
	"time"
)

// CSRF is used to get the current token and validate a suspect token.
type CSRF interface {
	// Return HTTP header to search for token.
	GetHeaderName() string
	// Return form value to search for token.
	GetFormName() string
	// Return cookie name to search for token.
	GetCookieName() string
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
	// Error replies to the request with a custom function when ValidToken fails.
	Error(w http.ResponseWriter)
}

type csrf struct {
	// Header name value for setting and getting csrf token.
	Header string
	// Form name value for setting and getting csrf token.
	Form string
	// Cookie name value for setting and getting csrf token.
	Cookie string
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	ID string
	// Secret used along with the unique id above to generate the Token.
	Secret string
	// ErrorFunc is the custom function that replies to the request when ValidToken fails.
	ErrorFunc func(w http.ResponseWriter)
}

// Returns the name of the Http header for csrf token.
func (c *csrf) GetHeaderName() string {
	return c.Header
}

// Returns the name of the form value for csrf token.
func (c *csrf) GetFormName() string {
	return c.Form
}

// Returns the name of the cookie for csrf token.
func (c *csrf) GetCookieName() string {
	return c.Cookie
}

// Returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// Validates the passed token against the existing Secret and ID.
func (c *csrf) ValidToken(t string) bool {
	return xsrftoken.Valid(t, c.Secret, c.ID, "POST")
}

// Error replies to the request when ValidToken fails.
func (c *csrf) Error(w http.ResponseWriter) {
	c.ErrorFunc(w)
}

// Options maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// HTTP header used to set and get token.
	Header string
	// Form value used to set and get token.
	Form string
	// Cookie value used to set and get token.
	Cookie string
	// Key used for getting the unique ID per user.
	SessionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
	// The function called when Validate fails.
	ErrorFunc func(w http.ResponseWriter)
}

const domainReg = `/^\.?[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d]))(?:\.[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d])))*$/`

// Generate maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(opts *Options) martini.Handler {
	return func(s sessions.Session, c martini.Context, r *http.Request, w http.ResponseWriter) {
		if opts.Header == "" {
			opts.Header = "X-CSRFToken"
		}
		if opts.Form == "" {
			opts.Form = "_csrf"
		}
		if opts.Cookie == "" {
			opts.Cookie = "_csrf"
		}
		if opts.ErrorFunc == nil {
			opts.ErrorFunc = func(w http.ResponseWriter) {
				http.Error(w, "Invalid csrf token.", http.StatusBadRequest)
			}
		}

		x := &csrf{
			Secret:    opts.Secret,
			Header:    opts.Header,
			Form:      opts.Form,
			Cookie:    opts.Cookie,
			ErrorFunc: opts.ErrorFunc,
		}
		c.MapTo(x, (*CSRF)(nil))

		uid := s.Get(opts.SessionKey)
		if uid == nil {
			return
		}
		switch uid.(type) {
		case string:
			x.ID = uid.(string)
		default:
			return
		}

		if r.Method != "GET" {
			return
		}

		// If cookie present, map existing token, else generate a new one.
		if ex, err := r.Cookie(opts.Cookie); err == nil && ex.Value != "" {
			x.Token = ex.Value
		} else {
			x.Token = xsrftoken.Generate(x.Secret, x.ID, "POST")
			if opts.SetCookie {
				expire := time.Now().AddDate(0, 0, 1)
				// Verify the domain is valid. If it is not, set as empty.
				domain := strings.Split(r.Host, ":")[0]
				if ok, err := regexp.Match(domainReg, []byte(domain)); !ok || err != nil {
					domain = ""
				}

				cookie := &http.Cookie{
					Name:       opts.Cookie,
					Value:      x.Token,
					Path:       "/",
					Domain:     domain,
					Expires:    expire,
					RawExpires: expire.Format(time.UnixDate),
					MaxAge:     0,
					Secure:     opts.Secure,
					HttpOnly:   false,
					Raw:        fmt.Sprintf("%s=%s", opts.Cookie, x.Token),
					Unparsed:   []string{fmt.Sprintf("token=%s", x.Token)},
				}
				http.SetCookie(w, cookie)
			}
		}

		if opts.SetHeader {
			w.Header().Add(opts.Header, x.Token)
		}
	}

}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, custom Error is sent in the reply.
// If neither a header or form value is found, http.StatusBadRequest is sent.
func Validate(r *http.Request, w http.ResponseWriter, x CSRF) {
	if token := r.Header.Get(x.GetHeaderName()); token != "" {
		if !x.ValidToken(token) {
			x.Error(w)
		}
		return
	}
	if token := r.FormValue(x.GetFormName()); token != "" {
		if !x.ValidToken(token) {
			x.Error(w)
		}
		return
	}

	http.Error(w, "Bad Request", http.StatusBadRequest)
	return
}
<commit_msg>Fixing comment<commit_after>// Package csrf generates and validates csrf tokens for martini.
// There are multiple methods of delivery including via a cookie or HTTP
// header.
// Validation occurs via a traditional hidden form key of "_csrf", or via
// a custom HTTP header "X-CSRFToken".
//
// package main
//
// import (
//     "github.com/codegangsta/martini"
//     "github.com/martini-contib/csrf"
//     "github.com/martini-contrib/render"
//     "github.com/martini-contib/sessions"
//     "net/http"
// )
//
// func main() {
//     m := martini.Classic()
//     store := sessions.NewCookieStore([]byte("secret123"))
//     m.Use(sessions.Sessions("my_session", store))
//     // Setup generation middleware.
//     m.Use(csrf.Generate(&csrf.Options{
//         Secret:     "token123",
//         SessionKey: "userID",
//     }))
//     m.Use(render.Renderer())
//
//     // Simulate the authentication of a session. If userID exists redirect
//     // to a form that requires csrf protection.
//     m.Get("/", func(s sessions.Session, r render.Render) {
//         if s.Get("userID") == nil {
//             r.Redirect("/login", 302)
//             return
//         }
//         r.Redirect("/protected", 302)
//     })
//
//     // Set userID for the session.
//     m.Get("/login", func(s sessions.Session, r render.Render) {
//         s.Set("userID", "123456")
//         r.Redirect("/", 302)
//     })
//
//     // Render a protected form. Passing a csrf token by calling x.GetToken()
//     m.Get("/protected", func(s sessions.Session, r render.Render, x csrf.CSRF) {
//         if s.Get("userID") == nil {
//             r.Redirect("/login", 401)
//             return
//         }
//         r.HTML(200, "protected", x.GetToken())
//     })
//
//     // Apply csrf validation to route.
//     m.Post("/protected", csrf.Validate, func(s sessions.Session, r render.Render) {
//         if s.Get("userID") != nil {
//             r.HTML(200, "result", "You submitted a valid token")
//             return
//         }
//         r.Redirect("/login", 401)
//     })
//
//     m.Run()
// }
package csrf

import (
	"code.google.com/p/xsrftoken"
	"fmt"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
	"net/http"
	"regexp"
	"strings"
	"time"
)

// CSRF is used to get the current token and validate a suspect token.
type CSRF interface {
	// Return HTTP header to search for token.
	GetHeaderName() string
	// Return form value to search for token.
	GetFormName() string
	// Return cookie name to search for token.
	GetCookieName() string
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
	// Error replies to the request with a custom function when ValidToken fails.
	Error(w http.ResponseWriter)
}

type csrf struct {
	// Header name value for setting and getting csrf token.
	Header string
	// Form name value for setting and getting csrf token.
	Form string
	// Cookie name value for setting and getting csrf token.
	Cookie string
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	ID string
	// Secret used along with the unique id above to generate the Token.
	Secret string
	// ErrorFunc is the custom function that replies to the request when ValidToken fails.
	ErrorFunc func(w http.ResponseWriter)
}

// Returns the name of the HTTP header for csrf token.
func (c *csrf) GetHeaderName() string {
	return c.Header
}

// Returns the name of the form value for csrf token.
func (c *csrf) GetFormName() string {
	return c.Form
}

// Returns the name of the cookie for csrf token.
func (c *csrf) GetCookieName() string {
	return c.Cookie
}

// Returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// Validates the passed token against the existing Secret and ID.
func (c *csrf) ValidToken(t string) bool {
	return xsrftoken.Valid(t, c.Secret, c.ID, "POST")
}

// Error replies to the request when ValidToken fails.
func (c *csrf) Error(w http.ResponseWriter) {
	c.ErrorFunc(w)
}

// Options maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// HTTP header used to set and get token.
	Header string
	// Form value used to set and get token.
	Form string
	// Cookie value used to set and get token.
	Cookie string
	// Key used for getting the unique ID per user.
	SessionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
	// The function called when Validate fails.
	ErrorFunc func(w http.ResponseWriter)
}

const domainReg = `/^\.?[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d]))(?:\.[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d])))*$/`

// Generate maps CSRF to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(opts *Options) martini.Handler {
	return func(s sessions.Session, c martini.Context, r *http.Request, w http.ResponseWriter) {
		if opts.Header == "" {
			opts.Header = "X-CSRFToken"
		}
		if opts.Form == "" {
			opts.Form = "_csrf"
		}
		if opts.Cookie == "" {
			opts.Cookie = "_csrf"
		}
		if opts.ErrorFunc == nil {
			opts.ErrorFunc = func(w http.ResponseWriter) {
				http.Error(w, "Invalid csrf token.", http.StatusBadRequest)
			}
		}

		x := &csrf{
			Secret:    opts.Secret,
			Header:    opts.Header,
			Form:      opts.Form,
			Cookie:    opts.Cookie,
			ErrorFunc: opts.ErrorFunc,
		}
		c.MapTo(x, (*CSRF)(nil))

		uid := s.Get(opts.SessionKey)
		if uid == nil {
			return
		}
		switch uid.(type) {
		case string:
			x.ID = uid.(string)
		default:
			return
		}

		if r.Method != "GET" {
			return
		}

		// If cookie present, map existing token, else generate a new one.
		if ex, err := r.Cookie(opts.Cookie); err == nil && ex.Value != "" {
			x.Token = ex.Value
		} else {
			x.Token = xsrftoken.Generate(x.Secret, x.ID, "POST")
			if opts.SetCookie {
				expire := time.Now().AddDate(0, 0, 1)
				// Verify the domain is valid. If it is not, set as empty.
				domain := strings.Split(r.Host, ":")[0]
				if ok, err := regexp.Match(domainReg, []byte(domain)); !ok || err != nil {
					domain = ""
				}

				cookie := &http.Cookie{
					Name:       opts.Cookie,
					Value:      x.Token,
					Path:       "/",
					Domain:     domain,
					Expires:    expire,
					RawExpires: expire.Format(time.UnixDate),
					MaxAge:     0,
					Secure:     opts.Secure,
					HttpOnly:   false,
					Raw:        fmt.Sprintf("%s=%s", opts.Cookie, x.Token),
					Unparsed:   []string{fmt.Sprintf("token=%s", x.Token)},
				}
				http.SetCookie(w, cookie)
			}
		}

		if opts.SetHeader {
			w.Header().Add(opts.Header, x.Token)
		}
	}

}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, custom Error is sent in the reply.
// If neither a header or form value is found, http.StatusBadRequest is sent.
func Validate(r *http.Request, w http.ResponseWriter, x CSRF) {
	if token := r.Header.Get(x.GetHeaderName()); token != "" {
		if !x.ValidToken(token) {
			x.Error(w)
		}
		return
	}
	if token := r.FormValue(x.GetFormName()); token != "" {
		if !x.ValidToken(token) {
			x.Error(w)
		}
		return
	}

	http.Error(w, "Bad Request", http.StatusBadRequest)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/grafov/autograf/client"
)

/* This is a simple example of usage of Grafana client
for copying dashboards and saving them to a disk.
Usage:
  backup-dashboards http://grafana.host:3000 api-key-string-here
*/

func main() {
	var (
		boards        []client.FoundBoard
		boardWithMeta client.BoardWithMeta
		data          []byte
		err           error
	)
	if len(os.Args) != 2 {
		fmt.Fprint(os.Stderr, "Usage:  backup-dashboards http://grafana.host:3000 api-key-string-here\n")
		os.Exit(0)
	}
	c := client.New(os.Args[1], os.Args[2])
	if boards, err = c.SearchDashboards("", false); err != nil {
		fmt.Fprintf(os.Stderr, fmt.Sprintf("%s\n", err))
		os.Exit(1)
	}
	for _, link := range boards {
		if boardWithMeta, err = c.GetDashboard(link.URI); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, link.URI))
			continue
		}
		if data, err = json.MarshalIndent(boardWithMeta.Board, "", "  "); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, boardWithMeta.Board.Title))
			continue
		}
		if err = ioutil.WriteFile(fmt.Sprintf("%s.json", boardWithMeta.Meta.Slug), data, os.FileMode(int(0666))); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, boardWithMeta.Meta.Slug))
		}
	}
}
<commit_msg>Update utility code for reflect changed API.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/grafov/autograf/client"
	"github.com/grafov/autograf/grafana"
)

/* This is a simple example of usage of Grafana client
for copying dashboards and saving them to a disk.
Usage:
  backup-dashboards http://grafana.host:3000 api-key-string-here
*/

func main() {
	var (
		boards []client.FoundBoard
		board  grafana.Board
		meta   client.BoardProperties
		data   []byte
		err    error
	)
	if len(os.Args) != 2 {
		fmt.Fprint(os.Stderr, "Usage:  backup-dashboards http://grafana.host:3000 api-key-string-here\n")
		os.Exit(0)
	}
	c := client.New(os.Args[1], os.Args[2])
	if boards, err = c.SearchDashboards("", false); err != nil {
		fmt.Fprintf(os.Stderr, fmt.Sprintf("%s\n", err))
		os.Exit(1)
	}
	for _, link := range boards {
		if board, meta, err = c.GetDashboard(link.URI); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, link.URI))
			continue
		}
		if data, err = json.MarshalIndent(board, "", "  "); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, board.Title))
			continue
		}
		if err = ioutil.WriteFile(fmt.Sprintf("%s.json", meta.Slug), data, os.FileMode(int(0666))); err != nil {
			fmt.Fprintf(os.Stderr, fmt.Sprintf("%s for %s\n", err, meta.Slug))
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/pivotalservices/cfbackup/tileregistry"
	"github.com/xchapter7x/lo"
)

//CreateBURACliCommand - this will create a cli command object for backup / restore
func CreateBURACliCommand(name string, usage string, eh *ErrorHandler) (command cli.Command) {
	desc := fmt.Sprintf("%s --opsmanagerhost <host> --adminuser <usr> --adminpass <pass> --opsmanageruser <opsuser> --opsmanagerpass <opspass> -d <dir> --tile elastic-runtime", name)
	command = cli.Command{
		Name:        name,
		Usage:       usage,
		Description: desc,
		Flags:       buraFlags,
		Action:      buraAction(name, eh),
	}
	return
}

func buraAction(commandName string, eh *ErrorHandler) (action func(*cli.Context) error) {
	action = func(c *cli.Context) error {
		var (
			fs = &flagSet{
				host:                 c.String(flagList[opsManagerHost].Flag[0]),
				adminUser:            c.String(flagList[adminUser].Flag[0]),
				adminPass:            c.String(flagList[adminPass].Flag[0]),
				opsManagerUser:       c.String(flagList[opsManagerUser].Flag[0]),
				opsManagerPass:       c.String(flagList[opsManagerPass].Flag[0]),
				opsManagerPassphrase: c.String(flagList[opsManagerPassphrase].Flag[0]),
				dest:                 c.String(flagList[dest].Flag[0]),
				tile:                 c.String(flagList[tile].Flag[0]),
				encryptionKey:        c.String(flagList[encryptionKey].Flag[0]),
				clearBoshManifest:    c.Bool(flagList[clearBoshManifest].Flag[0]),
				pluginArgs:           c.String(flagList[pluginArgs].Flag[0]),
			}
		)

		if tileCloser, err := getTileFromRegistry(fs, commandName); err == nil {
			defer tileCloser.Close()
			if err = runTileAction(commandName, tileCloser); err != nil {
				lo.G.Errorf("there was an error: %s running %s on %s tile:%v", err.Error(), commandName, fs.Tile(), tile)
				exitOnError(eh, c, commandName, err)
				return err
			}
		} else {
			lo.G.Errorf("there was an error getting tile from registry: %s", err.Error())
			exitOnError(eh, c, commandName, err)
			return err
		}
		lo.G.Debug("Tile action completed successfully")
		return nil
	}
	return
}

func exitOnError(eh *ErrorHandler, c *cli.Context, commandName string, err error) {
	cli.ShowCommandHelp(c, commandName)
	eh.ExitCode = helpExitCode
	eh.Error = err
}

func runTileAction(commandName string, tile tileregistry.Tile) (err error) {
	lo.G.Debug("Running %s for tile: %+v", commandName, tile)
	switch commandName {
	case "backup":
		err = tile.Backup()
	case "restore":
		err = tile.Restore()
	}
	return
}

func getTileFromRegistry(fs *flagSet, commandName string) (tileCloser tileregistry.TileCloser, err error) {
	lo.G.Debug("checking registry for '%s' tile", fs.Tile())

	if tileBuilder, ok := tileregistry.GetRegistry()[fs.Tile()]; ok {
		lo.G.Debug("found tile in registry")

		if hasValidBackupRestoreFlags(fs) {
			lo.G.Debug("we have all required flags and a proper builder")
			tileCloser, err = tileBuilder.New(tileregistry.TileSpec{
				OpsManagerHost:       fs.Host(),
				AdminUser:            fs.AdminUser(),
				AdminPass:            fs.AdminPass(),
				OpsManagerUser:       fs.OpsManagerUser(),
				OpsManagerPass:       fs.OpsManagerPass(),
				OpsManagerPassphrase: fs.OpsManagerPassphrase(),
				ArchiveDirectory:     fs.Dest(),
				CryptKey:             fs.EncryptionKey(),
				ClearBoshManifest:    fs.ClearBoshManifest(),
				PluginArgs:           fs.PluginArgs(),
			})
			if err != nil {
				return nil, fmt.Errorf("failure to connect to ops manager host: %s", err.Error())
			}

		} else {
			err = ErrInvalidFlagArgs
		}

	} else {
		lo.G.Errorf("tile '%s' not found", fs.Tile())
		err = ErrInvalidTileSelection
	}
	return
}

var buraFlags = func() (flags []cli.Flag) {
	for _, v := range flagList {
		flags = append(flags, cli.StringFlag{
			Name:   strings.Join(v.Flag, ", "),
			Value:  "",
			Usage:  v.Desc,
			EnvVar: v.EnvVar,
		})
	}
	return
}()

const (
	errExitCode                 = 1
	helpExitCode                = 2
	cleanExitCode               = 0
	opsManagerHost       string = "opsmanagerHost"
	adminUser            string = "adminUser"
	adminPass            string = "adminPass"
	opsManagerUser       string = "opsManagerUser"
	opsManagerPass       string = "opsManagerPass"
	opsManagerPassphrase string = "opsManagerPassphrase"
	dest                 string = "destination"
	tile                 string = "tile"
	encryptionKey        string = "encryptionKey"
	clearBoshManifest    string = "clearboshmanifest"
	pluginArgs           string = "pluginArgs"
)

var (
	//ErrInvalidFlagArgs - error for invalid flags
	ErrInvalidFlagArgs = errors.New("invalid cli flag args")
	//ErrInvalidTileSelection - error for invalid tile
	ErrInvalidTileSelection = errors.New("invalid tile selected. try the 'list-tiles' option to see a list of available tiles")
	flagList                = map[string]flagBucket{
		opsManagerHost: flagBucket{
			Flag:   []string{"opsmanagerhost", "omh"},
			Desc:   "hostname for Ops Manager",
			EnvVar: "CFOPS_HOST",
		},
		adminUser: flagBucket{
			Flag:   []string{"adminuser", "du"},
			Desc:   "username for Ops Mgr admin (Ops Manager WebConsole Credentials)",
			EnvVar: "CFOPS_ADMIN_USER",
		},
		adminPass: flagBucket{
			Flag:   []string{"adminpass", "dp"},
			Desc:   "password for Ops Mgr admin (Ops Manager WebConsole Credentials)",
			EnvVar: "CFOPS_ADMIN_PASS",
		},
		opsManagerUser: flagBucket{
			Flag:   []string{"opsmanageruser", "omu"},
			Desc:   "username for Ops Manager VM Access (used for ssh connections)",
			EnvVar: "CFOPS_OM_USER",
		},
		opsManagerPass: flagBucket{
			Flag:   []string{"opsmanagerpass", "omp"},
			Desc:   "password for Ops Manager VM Access (used for ssh connections)",
			EnvVar: "CFOPS_OM_PASS",
		},
		opsManagerPassphrase: flagBucket{
			Flag:   []string{"opsmanagerpassphrase", "omr"},
			Desc:   "passphrase is used by Ops Manager 1.7 to decrypt the installation files during restore",
			EnvVar: "CFOPS_OM_PASSPHRASE",
		},
		dest: flagBucket{
			Flag:   []string{"destination", "d"},
			Desc:   "path of the Cloud Foundry archive",
			EnvVar: "CFOPS_DEST_PATH",
		},
		tile: flagBucket{
			Flag:   []string{"tile", "t"},
			Desc:   "a tile you would like to run the operation on",
			EnvVar: "CFOPS_TILE",
		},
		encryptionKey: flagBucket{
			Flag:   []string{"encryptionkey", "k"},
			Desc:   "encryption key to encrypt/decrypt your archive (key lengths supported are 16, 24, 32 for AES-128, AES-192, or AES-256)",
			EnvVar: "CFOPS_ENCRYPTION_KEY",
		},
		clearBoshManifest: flagBucket{
			Flag:   []string{"clear-bosh-manifest"},
			Desc:   "set this flag if you would like to clear the bosh-deployments.yml (this should only affect a restore of Ops-Manager)",
			EnvVar: "CFOPS_CLEAR_BOSH_MANIFEST",
		},
		pluginArgs: flagBucket{
			Flag:   []string{"pluginargs", "p"},
			Desc:   "Arguments for plugin to execute",
			EnvVar: "CFOPS_PLUGIN_ARGS",
		},
	}
)

type (
	flagSet struct {
		host                 string
		adminUser            string
		adminPass            string
		opsManagerUser       string
		opsManagerPass       string
		opsManagerPassphrase string
		dest                 string
		tile                 string
		encryptionKey        string
		clearBoshManifest    bool
		pluginArgs           string
	}

	flagBucket struct {
		Flag   []string
		Desc   string
		EnvVar string
	}
)

func (s *flagSet) Host() string {
	return s.host
}

func (s *flagSet) AdminUser() string {
	return s.adminUser
}

func (s *flagSet) AdminPass() string {
	return s.adminPass
}

func (s *flagSet) OpsManagerUser() string {
	return s.opsManagerUser
}

func (s *flagSet) OpsManagerPass() string {
	return s.opsManagerPass
}

func (s *flagSet) OpsManagerPassphrase() string {
	return s.opsManagerPassphrase
}

func (s *flagSet) Dest() string {
	return s.dest
}

func (s *flagSet) Tile() string {
	return s.tile
}

func (s *flagSet) EncryptionKey() string {
	return s.encryptionKey
}

func (s *flagSet) ClearBoshManifest() bool {
	return s.clearBoshManifest
}

func (s *flagSet) PluginArgs() string {
	return s.pluginArgs
}

func hasValidBackupRestoreFlags(fs *flagSet) bool {
	res := (fs.Host() != "" &&
		fs.AdminUser() != "" &&
		fs.AdminPass() != "" &&
		fs.OpsManagerUser() != "" &&
		fs.Dest() != "" &&
		fs.Tile() != "")

	if res == false {
		lo.G.Debug("OpsManagerHost: ", fs.Host())
		lo.G.Debug("AdminUser: ", fs.AdminUser())
		lo.G.Debug("AdminPass: ", fs.AdminPass())
		lo.G.Debug("OpsManagerUser: ", fs.OpsManagerUser())
		lo.G.Debug("OpsManagerPass: ", fs.OpsManagerPass())
		lo.G.Debug("Destination: ", fs.Dest())
	}
	return res
}
<commit_msg>Fix broken string interpolatoin in debug mode<commit_after>package main

import (
	"errors"
	"fmt"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/pivotalservices/cfbackup/tileregistry"
	"github.com/xchapter7x/lo"
)

//CreateBURACliCommand - this will create a cli command object for backup / restore
func CreateBURACliCommand(name string, usage string, eh *ErrorHandler) (command cli.Command) {
	desc := fmt.Sprintf("%s --opsmanagerhost <host> --adminuser <usr> --adminpass <pass> --opsmanageruser <opsuser> --opsmanagerpass <opspass> -d <dir> --tile elastic-runtime", name)
	command = cli.Command{
		Name:        name,
		Usage:       usage,
		Description: desc,
		Flags:       buraFlags,
		Action:      buraAction(name, eh),
	}
	return
}

func buraAction(commandName string, eh *ErrorHandler) (action func(*cli.Context) error) {
	action = func(c *cli.Context) error {
		var (
			fs = &flagSet{
				host:                 c.String(flagList[opsManagerHost].Flag[0]),
				adminUser:            c.String(flagList[adminUser].Flag[0]),
				adminPass:            c.String(flagList[adminPass].Flag[0]),
				opsManagerUser:       c.String(flagList[opsManagerUser].Flag[0]),
				opsManagerPass:       c.String(flagList[opsManagerPass].Flag[0]),
				opsManagerPassphrase: c.String(flagList[opsManagerPassphrase].Flag[0]),
				dest:                 c.String(flagList[dest].Flag[0]),
				tile:                 c.String(flagList[tile].Flag[0]),
				encryptionKey:        c.String(flagList[encryptionKey].Flag[0]),
				clearBoshManifest:    c.Bool(flagList[clearBoshManifest].Flag[0]),
				pluginArgs:           c.String(flagList[pluginArgs].Flag[0]),
			}
		)

		if tileCloser, err := getTileFromRegistry(fs, commandName); err == nil {
			defer tileCloser.Close()
			if err = runTileAction(commandName, tileCloser); err != nil {
				lo.G.Errorf("there was an error: %s running %s on %s tile:%v", err.Error(), commandName, fs.Tile(), tile)
				exitOnError(eh, c, commandName, err)
				return err
			}
		} else {
			lo.G.Errorf("there was an error getting tile from registry: %s", err.Error())
			exitOnError(eh, c, commandName, err)
			return err
		}
		lo.G.Debug("Tile action completed successfully")
		return nil
	}
	return
}

func exitOnError(eh *ErrorHandler, c *cli.Context, commandName string, err error) {
	cli.ShowCommandHelp(c, commandName)
	eh.ExitCode = helpExitCode
	eh.Error = err
}

func runTileAction(commandName string, tile tileregistry.Tile) (err error) {
	lo.G.Debugf("Running %s for tile: %+v", commandName, tile)
	switch commandName {
	case "backup":
		err = tile.Backup()
	case "restore":
		err = tile.Restore()
	}
	return
}

func getTileFromRegistry(fs *flagSet, commandName string) (tileCloser tileregistry.TileCloser, err error) {
	lo.G.Debugf("checking registry for '%s' tile", fs.Tile())

	if tileBuilder, ok := tileregistry.GetRegistry()[fs.Tile()]; ok {
		lo.G.Debug("found tile in registry")

		if hasValidBackupRestoreFlags(fs) {
			lo.G.Debug("we have all required flags and a proper builder")
			tileCloser, err = tileBuilder.New(tileregistry.TileSpec{
				OpsManagerHost:       fs.Host(),
				AdminUser:            fs.AdminUser(),
				AdminPass:            fs.AdminPass(),
				OpsManagerUser:       fs.OpsManagerUser(),
				OpsManagerPass:       fs.OpsManagerPass(),
				OpsManagerPassphrase: fs.OpsManagerPassphrase(),
				ArchiveDirectory:     fs.Dest(),
				CryptKey:             fs.EncryptionKey(),
				ClearBoshManifest:    fs.ClearBoshManifest(),
				PluginArgs:           fs.PluginArgs(),
			})
			if err != nil {
				return nil, fmt.Errorf("failure to connect to ops manager host: %s", err.Error())
			}

		} else {
			err = ErrInvalidFlagArgs
		}

	} else {
		lo.G.Errorf("tile '%s' not found", fs.Tile())
		err = ErrInvalidTileSelection
	}
	return
}

var buraFlags = func() (flags []cli.Flag) {
	for _, v := range flagList {
		flags = append(flags, cli.StringFlag{
			Name:   strings.Join(v.Flag, ", "),
			Value:  "",
			Usage:  v.Desc,
			EnvVar: v.EnvVar,
		})
	}
	return
}()

const (
	errExitCode                 = 1
	helpExitCode                = 2
	cleanExitCode               = 0
	opsManagerHost       string = "opsmanagerHost"
	adminUser            string = "adminUser"
	adminPass            string = "adminPass"
	opsManagerUser       string = "opsManagerUser"
	opsManagerPass       string = "opsManagerPass"
	opsManagerPassphrase string = "opsManagerPassphrase"
	dest                 string = "destination"
	tile                 string = "tile"
	encryptionKey        string = "encryptionKey"
	clearBoshManifest    string = "clearboshmanifest"
	pluginArgs           string = "pluginArgs"
)

var (
	//ErrInvalidFlagArgs - error for invalid flags
	ErrInvalidFlagArgs = errors.New("invalid cli flag args")
	//ErrInvalidTileSelection - error for invalid tile
	ErrInvalidTileSelection = errors.New("invalid tile selected. try the 'list-tiles' option to see a list of available tiles")
	flagList                = map[string]flagBucket{
		opsManagerHost: flagBucket{
			Flag:   []string{"opsmanagerhost", "omh"},
			Desc:   "hostname for Ops Manager",
			EnvVar: "CFOPS_HOST",
		},
		adminUser: flagBucket{
			Flag:   []string{"adminuser", "du"},
			Desc:   "username for Ops Mgr admin (Ops Manager WebConsole Credentials)",
			EnvVar: "CFOPS_ADMIN_USER",
		},
		adminPass: flagBucket{
			Flag:   []string{"adminpass", "dp"},
			Desc:   "password for Ops Mgr admin (Ops Manager WebConsole Credentials)",
			EnvVar: "CFOPS_ADMIN_PASS",
		},
		opsManagerUser: flagBucket{
			Flag:   []string{"opsmanageruser", "omu"},
			Desc:   "username for Ops Manager VM Access (used for ssh connections)",
			EnvVar: "CFOPS_OM_USER",
		},
		opsManagerPass: flagBucket{
			Flag:   []string{"opsmanagerpass", "omp"},
			Desc:   "password for Ops Manager VM Access (used for ssh connections)",
			EnvVar: "CFOPS_OM_PASS",
		},
		opsManagerPassphrase: flagBucket{
			Flag:   []string{"opsmanagerpassphrase", "omr"},
			Desc:   "passphrase is used by Ops Manager 1.7 to decrypt the installation files during restore",
			EnvVar: "CFOPS_OM_PASSPHRASE",
		},
		dest: flagBucket{
			Flag:   []string{"destination", "d"},
			Desc:   "path of the Cloud Foundry archive",
			EnvVar: "CFOPS_DEST_PATH",
		},
		tile: flagBucket{
			Flag:   []string{"tile", "t"},
			Desc:   "a tile you would like to run the operation on",
			EnvVar: "CFOPS_TILE",
		},
		encryptionKey: flagBucket{
			Flag:   []string{"encryptionkey", "k"},
			Desc:   "encryption key to encrypt/decrypt your archive (key lengths supported are 16, 24, 32 for AES-128, AES-192, or AES-256)",
			EnvVar: "CFOPS_ENCRYPTION_KEY",
		},
		clearBoshManifest: flagBucket{
			Flag:   []string{"clear-bosh-manifest"},
			Desc:   "set this flag if you would like to clear the bosh-deployments.yml (this should only affect a restore of Ops-Manager)",
			EnvVar: "CFOPS_CLEAR_BOSH_MANIFEST",
		},
		pluginArgs: flagBucket{
			Flag:   []string{"pluginargs", "p"},
			Desc:   "Arguments for plugin to execute",
			EnvVar: "CFOPS_PLUGIN_ARGS",
		},
	}
)

type (
	flagSet struct {
		host                 string
		adminUser            string
		adminPass            string
		opsManagerUser       string
		opsManagerPass       string
		opsManagerPassphrase string
		dest                 string
		tile                 string
		encryptionKey        string
		clearBoshManifest    bool
		pluginArgs           string
	}

	flagBucket struct {
		Flag   []string
		Desc   string
		EnvVar string
	}
)

func (s *flagSet) Host() string {
	return s.host
}

func (s *flagSet) AdminUser() string {
	return s.adminUser
}

func (s *flagSet) AdminPass() string {
	return s.adminPass
}

func (s *flagSet) OpsManagerUser() string {
	return s.opsManagerUser
}

func (s *flagSet) OpsManagerPass() string {
	return s.opsManagerPass
}

func (s *flagSet) OpsManagerPassphrase() string {
	return s.opsManagerPassphrase
}

func (s *flagSet) Dest() string {
	return s.dest
}

func (s *flagSet) Tile() string {
	return s.tile
}

func (s *flagSet) EncryptionKey() string {
	return s.encryptionKey
}

func (s *flagSet) ClearBoshManifest() bool {
	return s.clearBoshManifest
}

func (s *flagSet) PluginArgs() string {
	return s.pluginArgs
}

func hasValidBackupRestoreFlags(fs *flagSet) bool {
	res := (fs.Host() != "" &&
		fs.AdminUser() != "" &&
		fs.AdminPass() != "" &&
		fs.OpsManagerUser() != "" &&
		fs.Dest() != "" &&
		fs.Tile() != "")

	if res == false {
		lo.G.Debug("OpsManagerHost: ", fs.Host())
		lo.G.Debug("AdminUser: ", fs.AdminUser())
		lo.G.Debug("AdminPass: ", fs.AdminPass())
		lo.G.Debug("OpsManagerUser: ", fs.OpsManagerUser())
		lo.G.Debug("OpsManagerPass: ", fs.OpsManagerPass())
		lo.G.Debug("Destination: ", fs.Dest())
	}
	return res
}
<|endoftext|>


<commit_before>// (C) 2014 Mathias Dalheimer <md@gonium.net>.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"github.com/gonium/defluxio"
	"log"
	"sort"
	"strings"
	"time"
)

var configFile = flag.String("config", "defluxio-exporter.conf", "configuration file")
var meterID = flag.String("meter", "", "ID of the meter to query")
var startTimestamp = flag.Int("start", 0, "data export start: first unix timestamp to export")
var endTimestamp = flag.Int("end", 0, "data export end: last unix timestamp to export")
var cfg *defluxio.ExporterConfigurationData
var dbclient *defluxio.DBClient

func init() {
	flag.Parse()
	if strings.EqualFold(*meterID, "") {
		log.Fatal("You must specify the meter ID (i.e. -meter=foometer)")
	}
	if *startTimestamp == 0 {
		log.Fatal("You must specify the start timestamp( i.e. -start=1405607436)")
	}
	if *endTimestamp == 0 {
		log.Fatal("You must specify the end timestamp( i.e. -end=1405607465)")
	}
	if *startTimestamp >= *endTimestamp {
		log.Fatal("start timestamp cannot be after end timestamp.")
	}
	var err error
	cfg, err = defluxio.LoadExporterConfiguration(*configFile)
	if err != nil {
		log.Fatal("Error loading configuration: ", err.Error())
	}
	dbclient, err = defluxio.NewDBClient(&cfg.InfluxDB)
	if err != nil {
		log.Fatal("Cannot initialize database client:", err.Error())
	}
}

func main() {
	log.Printf("Attempting to export from meter %s\n", *meterID)
	//result, err := dbclient.GetLastFrequency(*meterID)
	//if err != nil {
	//	log.Fatal("Failed to query database: ", err.Error())
	//}
	//fmt.Printf("On %v, the frequency was recorded as %f\n",
	//	result.Reading.Timestamp, result.Reading.Value)
	//meterReadings, err := dbclient.GetLastFrequencies(*meterID, 10)
	//if err != nil {
	//	log.Fatal("Failed to query database: ", err.Error())
	//}
	//for _, element := range meterReadings {
	//	fmt.Printf("%v: %f\n", element.Reading.Timestamp,
	//		element.Reading.Value)
	//}

	// Hack for testing
	// TODO: Replace with real time.Unix foo from commandline
	timeReadings, terr := dbclient.GetFrequenciesBetween(*meterID,
		time.Unix(1405525188, 0), time.Unix(1405588163, 0))
	if terr != nil {
		log.Fatal("Failed to query database: ", terr.Error())
	}
	sort.Sort(defluxio.ByTimestamp(timeReadings))

	path := "fooexport.txt"
	tsve, eerr := defluxio.NewTsvExporter(path)
	if eerr != nil {
		log.Fatal("Cannot create exporter with file %s", path)
	}
	if eerr = tsve.ExportDataset(timeReadings); eerr != nil {
		log.Fatal("Failed to export dataset: %s", eerr.Error())
	} else {
		log.Printf("Export finished successfully.")
	}
}
<commit_msg>final fixes for the exporter - now ready for usage. closes #9.<commit_after>// (C) 2014 Mathias Dalheimer <md@gonium.net>.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"github.com/gonium/defluxio"
	"log"
	"os"
	"sort"
	"strings"
	"time"
)

var configFile = flag.String("config", "defluxio-exporter.conf", "configuration file")
var meterID = flag.String("meter", "", "ID of the meter to query")
var startTimestamp = flag.Int64("start", 0,
	"data export start: first unix timestamp to export")
var endTimestamp = flag.Int64("end", 0,
	"data export end: last unix timestamp to export")
var exportFilename = flag.String("file", "defluxio-export.txt",
	"path to file to use for export")
var force = flag.Bool("force", false,
	"force export, overwriting existing files")
var cfg *defluxio.ExporterConfigurationData
var dbclient *defluxio.DBClient

func init() {
	flag.Parse()
	if strings.EqualFold(*meterID, "") {
		log.Fatal("You must specify the meter ID (i.e. -meter=foometer)")
	}
	if *startTimestamp == 0 {
		log.Fatal("You must specify the start timestamp( i.e. -start=1405607436)")
	}
	if *endTimestamp == 0 {
		log.Fatal("You must specify the end timestamp( i.e. -end=1405607465)")
	}
	if *startTimestamp >= *endTimestamp {
		log.Fatal("start timestamp cannot be after end timestamp.")
	}
	if !*force {
		if _, err := os.Stat(*exportFilename); err == nil {
			log.Fatal("file ", *exportFilename, " exists - aborting.")
		}
	}
	var err error
	if cfg, err = defluxio.LoadExporterConfiguration(*configFile); err != nil {
		log.Fatal("Error loading configuration: ", err.Error())
	}
	if dbclient, err = defluxio.NewDBClient(&cfg.InfluxDB); err != nil {
		log.Fatal("Cannot initialize database client:", err.Error())
	}
}

func main() {
	log.Printf("Attempting to export from meter %s\n", *meterID)
	//result, err := dbclient.GetLastFrequency(*meterID)
	//if err != nil {
	//	log.Fatal("Failed to query database: ", err.Error())
	//}
	//fmt.Printf("On %v, the frequency was recorded as %f\n",
	//	result.Reading.Timestamp, result.Reading.Value)
	//meterReadings, err := dbclient.GetLastFrequencies(*meterID, 10)
	//if err != nil {
	//	log.Fatal("Failed to query database: ", err.Error())
	//}
	//for _, element := range meterReadings {
	//	fmt.Printf("%v: %f\n", element.Reading.Timestamp,
	//		element.Reading.Value)
	//}

	// Hack for testing
	// TODO: Replace with real time.Unix foo from commandline
	timeReadings, terr := dbclient.GetFrequenciesBetween(*meterID,
		time.Unix(*startTimestamp, 0), time.Unix(*endTimestamp, 0))
	if terr != nil {
		log.Fatal("Failed to query database: ", terr.Error())
	}
	sort.Sort(defluxio.ByTimestamp(timeReadings))

	tsve, eerr := defluxio.NewTsvExporter(*exportFilename)
	if eerr != nil {
		log.Fatal("Cannot create exporter with file %s", *exportFilename)
	}
	if eerr = tsve.ExportDataset(timeReadings); eerr != nil {
		log.Fatal("Failed to export dataset: %s", eerr.Error())
	} else {
		log.Printf("Export finished successfully.")
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"testing"
)

func TestEndToEnd(t *testing.T) {
	dir, err := ioutil.TempDir("", "fdroidcl")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	// Build fdroidcl in the temporary directory.
	fdroidcl := filepath.Join(dir, "fdroidcl")
	if out, err := exec.Command("go", "build",
		"-ldflags=-X main.testBasedir="+dir,
		"-o", fdroidcl).CombinedOutput(); err != nil {
		t.Fatalf("%s", out)
	}

	mustSucceed := func(t *testing.T, want string, args ...string) {
		mustRun(t, true, want, fdroidcl, args...)
	}
	mustFail := func(t *testing.T, want string, args ...string) {
		mustRun(t, false, want, fdroidcl, args...)
	}

	t.Run("Help", func(t *testing.T) {
		mustFail(t, `Usage: fdroidcl`, "-h")
	})
	t.Run("UnknownCommand", func(t *testing.T) {
		mustFail(t, `Unrecognised command`, "unknown")
	})
	t.Run("Version", func(t *testing.T) {
		mustSucceed(t, `^v`, "version")
	})

	t.Run("SearchBeforeUpdate", func(t *testing.T) {
		mustFail(t, `could not open index`, "search")
	})
	t.Run("UpdateFirst", func(t *testing.T) {
		mustSucceed(t, `done`, "update")
	})
	t.Run("UpdateCached", func(t *testing.T) {
		mustSucceed(t, `not modified`, "update")
	})
}

func mustRun(t *testing.T, success bool, wantRe, name string, args ...string) {
	cmd := exec.Command(name, args...)
	out, err := cmd.CombinedOutput()
	if success && err != nil {
		t.Fatalf("unexpected error: %v\n%s", err, out)
	} else if !success && err == nil {
		t.Fatalf("expected error, got none\n%s", out)
	}
	if !regexp.MustCompile(wantRe).Match(out) {
		t.Fatalf("output does not match %#q:\n%s", wantRe, out)
	}
}
<commit_msg>cmd/fdroidcl: add first search tests<commit_after>// Copyright (c) 2018, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"testing"
	"time"
)

func TestEndToEnd(t *testing.T) {
	url := config.Repos[0].URL
	client := http.Client{Timeout: 2 * time.Second}
	if _, err := client.Get(url); err != nil {
		t.Skipf("skipping since %s is unreachable: %v", url, err)
	}

	dir, err := ioutil.TempDir("", "fdroidcl")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	// Build fdroidcl in the temporary directory.
	fdroidcl := filepath.Join(dir, "fdroidcl")
	if out, err := exec.Command("go", "build",
		"-ldflags=-X main.testBasedir="+dir,
		"-o", fdroidcl).CombinedOutput(); err != nil {
		t.Fatalf("%s", out)
	}

	mustSucceed := func(t *testing.T, want string, args ...string) {
		mustRun(t, true, want, fdroidcl, args...)
	}
	mustFail := func(t *testing.T, want string, args ...string) {
		mustRun(t, false, want, fdroidcl, args...)
	}

	t.Run("Help", func(t *testing.T) {
		mustFail(t, `Usage: fdroidcl`, "-h")
	})
	t.Run("UnknownCommand", func(t *testing.T) {
		mustFail(t, `Unrecognised command`, "unknown")
	})
	t.Run("Version", func(t *testing.T) {
		mustSucceed(t, `^v`, "version")
	})

	t.Run("SearchBeforeUpdate", func(t *testing.T) {
		mustFail(t, `could not open index`, "search")
	})
	t.Run("UpdateFirst", func(t *testing.T) {
		mustSucceed(t, `done`, "update")
	})
	t.Run("UpdateCached", func(t *testing.T) {
		mustSucceed(t, `not modified`, "update")
	})

	t.Run("SearchNoArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, "search")
	})
	t.Run("SearchWithArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, "search", "fdroid.fdroid")
	})
	t.Run("SearchWithArgsNone", func(t *testing.T) {
		mustSucceed(t, `^$`, "search", "nomatches")
	})
	t.Run("SearchOnlyPackageNames", func(t *testing.T) {
		mustSucceed(t, `^[^ ]*$`, "search", "-q", "fdroid.fdroid")
	})
}

func mustRun(t *testing.T, success bool, wantRe, name string, args ...string) {
	cmd := exec.Command(name, args...)
	out, err := cmd.CombinedOutput()
	if success && err != nil {
		t.Fatalf("unexpected error: %v\n%s", err, out)
	} else if !success && err == nil {
		t.Fatalf("expected error, got none\n%s", out)
	}
	if !regexp.MustCompile(wantRe).Match(out) {
		t.Fatalf("output does not match %#q:\n%s", wantRe, out)
	}
}
<|endoftext|>


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

package main

import (
	"context"
	"log"
	"os"

	"google.golang.org/grpc"
)

var stdLog, errLog *log.Logger

func init() {
	stdLog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	errLog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
}

type loggerObserver struct{}

func (l *loggerObserver) GetName() string { return "loggerObserver" }

func (l *loggerObserver) ObserveUnary(
	ctx context.Context,
	req interface{},
	resp interface{},
	info *grpc.UnaryServerInfo,
	err error) {
	stdLog.Printf("Received Unary Request for Method: %s\n", info.FullMethod)
	stdLog.Printf("    Request:  %+v\n", req)
	if err == nil {
		stdLog.Printf("    Returning Response: %+v\n", resp)
	} else {
		stdLog.Printf("    Returning Error: %+v\n", err)
	}
	stdLog.Println("")
}

func (l *loggerObserver) ObserveStreamRequest(
	_ context.Context,
	req interface{},
	info *grpc.StreamServerInfo,
	_ error) {
	stdLog.Printf("%s Stream for Method: %s\n", streamType(info), info.FullMethod)
	stdLog.Printf("    Recieving Message:  %v\n", req)
	stdLog.Println("")
}

func (l *loggerObserver) ObserveStreamResponse(
	_ context.Context,
	resp interface{},
	info *grpc.StreamServerInfo,
	_ error) {
	stdLog.Printf("%s Stream for Method: %s\n", streamType(info), info.FullMethod)
	stdLog.Printf("    Sending Message:  %+v\n", resp)
	stdLog.Println("")
}

func streamType(info *grpc.StreamServerInfo) string {
	if info.IsClientStream && info.IsServerStream {
		return "Bi-directional"
	} else if info.IsClientStream {
		return "Client"
	}
	return "Server"
}
<commit_msg>feat: dump request headers in verbose mode (#340)<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"log"
	"os"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

var stdLog, errLog *log.Logger

func init() {
	stdLog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	errLog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
}

type loggerObserver struct{}

func (l *loggerObserver) GetName() string { return "loggerObserver" }

func dumpIncomingHeaders(ctx context.Context) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		stdLog.Printf("Cannot get metadata from the context.")
		return
	}

	stdLog.Printf("    Request headers:")
	for key, values := range md {
		for _, value := range values {
			stdLog.Printf("      %s: %s\n", key, value)
		}
	}
}

func (l *loggerObserver) ObserveUnary(
	ctx context.Context,
	req interface{},
	resp interface{},
	info *grpc.UnaryServerInfo,
	err error) {
	stdLog.Printf("Received Unary Request for Method: %s\n", info.FullMethod)
	if Verbose {
		dumpIncomingHeaders(ctx)
	}
	stdLog.Printf("    Request:  %+v\n", req)
	if err == nil {
		stdLog.Printf("    Returning Response: %+v\n", resp)
	} else {
		stdLog.Printf("    Returning Error: %+v\n", err)
	}
	stdLog.Println("")
}

func (l *loggerObserver) ObserveStreamRequest(
	ctx context.Context,
	req interface{},
	info *grpc.StreamServerInfo,
	_ error) {
	stdLog.Printf("%s Stream for Method: %s\n", streamType(info), info.FullMethod)
	if Verbose {
		dumpIncomingHeaders(ctx)
	}
	stdLog.Printf("    Receiving Message:  %v\n", req)
	stdLog.Println("")
}

func (l *loggerObserver) ObserveStreamResponse(
	_ context.Context,
	resp interface{},
	info *grpc.StreamServerInfo,
	_ error) {
	stdLog.Printf("%s Stream for Method: %s\n", streamType(info), info.FullMethod)
	stdLog.Printf("    Sending Message:  %+v\n", resp)
	stdLog.Println("")
}

func streamType(info *grpc.StreamServerInfo) string {
	if info.IsClientStream && info.IsServerStream {
		return "Bi-directional"
	} else if info.IsClientStream {
		return "Client"
	}
	return "Server"
}
<|endoftext|>