<commit_before>package stamp

import (
	"net"
	"strconv"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/dedis/prifi/coco"
	"github.com/dedis/prifi/coco/coconet"
	"github.com/dedis/prifi/coco/hashid"
	"github.com/dedis/prifi/coco/proof"
	"github.com/dedis/prifi/coco/sign"
	"github.com/dedis/prifi/coco/test/logutils"
)

type Server struct {
	coco.Signer
	name    string
	Clients map[string]coconet.Conn

	// for aggregating messages from clients
	mux        sync.Mutex
	Queue      [][]MustReplyMessage
	READING    int
	PROCESSING int

	// Leaves, Root and Proof for a round
	Leaves []hashid.HashId // can be removed after we verify protocol
	Root   hashid.HashId
	Proofs []proof.Proof

	rLock     sync.Mutex
	maxRounds int
	closeChan chan bool

	Logger   string
	Hostname string
	App      string
}

func NewServer(signer coco.Signer) *Server {
	s := &Server{}

	s.Clients = make(map[string]coconet.Conn)
	s.Queue = make([][]MustReplyMessage, 2)
	s.READING = 0
	s.PROCESSING = 1

	s.Signer = signer
	s.Signer.RegisterAnnounceFunc(s.OnAnnounce())
	s.Signer.RegisterDoneFunc(s.OnDone())
	s.rLock = sync.Mutex{}

	// listen for client requests at one port higher
	// than the signing node
	h, p, err := net.SplitHostPort(s.Signer.Name())
	if err == nil {
		i, err := strconv.Atoi(p)
		if err != nil {
			log.Fatal(err)
		}
		s.name = net.JoinHostPort(h, strconv.Itoa(i+1))
	}
	s.Queue[s.READING] = make([]MustReplyMessage, 0)
	s.Queue[s.PROCESSING] = make([]MustReplyMessage, 0)
	s.closeChan = make(chan bool, 5)
	return s
}

var clientNumber int = 0

func (s *Server) Close() {
	log.Printf("closing stampserver: %p", s)
	s.closeChan <- true
	s.Signer.Close()
}

// listen for clients connections
// this server needs to be running on a different port
// than the Signer that is beneath it
func (s *Server) Listen() error {
	// log.Println("Listening @ ", s.name)
	ln, err := net.Listen("tcp4", s.name)
	if err != nil {
		panic(err)
	}

	go func() {
		for {
			// log.Printf("LISTENING TO CLIENTS: %p", s)
			conn, err := ln.Accept()
			if err != nil {
				// handle error
				log.Errorln("failed to accept connection")
				continue
			}

			c := coconet.NewTCPConnFromNet(conn)
			// log.Println("CLIENT TCP CONNECTION SUCCESSFULLY ESTABLISHED:", c)

			if _, ok := s.Clients[c.Name()]; !ok {
				s.Clients[c.Name()] = c

				go func(c coconet.Conn) {
					for {
						tsm := TimeStampMessage{}
						err := c.Get(&tsm)
						if err != nil {
							log.Errorf("%p Failed to get from child:", s, err)
							s.Close()
							return
						}
						switch tsm.Type {
						default:
							log.Errorf("Message of unknown type: %v\n", tsm.Type)
						case StampRequestType:
							// log.Println("RECEIVED STAMP REQUEST")
							s.mux.Lock()
							READING := s.READING
							s.Queue[READING] = append(s.Queue[READING],
								MustReplyMessage{Tsm: tsm, To: c.Name()})
							s.mux.Unlock()
						}
					}
				}(c)
			}

		}
	}()

	return nil
}

// Used for goconns
// should only be used if clients are created in batch
func (s *Server) ListenToClients() {
	// log.Printf("LISTENING TO CLIENTS: %p", s, s.Clients)
	for _, c := range s.Clients {
		go func(c coconet.Conn) {
			for {
				tsm := TimeStampMessage{}
				err := c.Get(&tsm)
				if err == coconet.ErrClosed {
					log.Errorf("%p Failed to get from client:", s, err)
					s.Close()
					return
				}
				if err != nil {
					log.WithFields(log.Fields{
						"file": logutils.File(),
					}).Errorf("%p failed To get message:", s, err)
				}
				switch tsm.Type {
				default:
					log.Errorln("Message of unknown type")
				case StampRequestType:
					// log.Println("STAMP REQUEST")
					s.mux.Lock()
					READING := s.READING
					s.Queue[READING] = append(s.Queue[READING],
						MustReplyMessage{Tsm: tsm, To: c.Name()})
					s.mux.Unlock()
				}
			}
		}(c)
	}
}

func (s *Server) ConnectToLogger() {
	return
	if s.Logger == "" || s.Hostname == "" || s.App == "" {
		log.Println("skipping connect to logger")
		return
	}
	log.Println("Connecting to Logger")
	lh, _ := logutils.NewLoggerHook(s.Logger, s.Hostname, s.App)
	log.Println("Connected to Logger")
	log.AddHook(lh)
}

func (s *Server) LogReRun(nextRole string, curRole string) {
	if nextRole == "root" {
		var messg = s.Name() + " became root"
		if curRole == "root" {
			messg = s.Name() + " remained root"
		}

		go s.ConnectToLogger()

		log.WithFields(log.Fields{
			"file": logutils.File(),
			"type": "role_change",
		}).Infoln(messg)
		// log.Printf("role change: %p", s)

	} else {
		var messg = s.Name() + " remained regular"
		if curRole == "root" {
			messg = s.Name() + " became regular"
		}

		if curRole == "root" {
			log.WithFields(log.Fields{
				"file": logutils.File(),
				"type": "role_change",
			}).Infoln(messg)
			log.Printf("role change: %p", s)
		}

	}

}

func (s *Server) runAsRoot(nRounds int) string {
	// every 5 seconds start a new round
	ticker := time.Tick(ROUND_TIME)
	log.Infoln(s.Name(), "running as root", s.LastRound(), int64(nRounds))

	for {
		select {
		case nextRole := <-s.ViewChangeCh():
			return nextRole
			// s.reRunWith(nextRole, nRounds, true)
		case <-ticker:

			start := time.Now()
			log.Println(s.Name(), "is STAMP SERVER STARTING SIGNING ROUND FOR:", s.LastRound()+1, "of", nRounds)

			err := s.StartSigningRound()
			if err == sign.ChangingViewError {
				// report change in view, and continue with the select
				log.WithFields(log.Fields{
					"file": logutils.File(),
					"type": "view_change",
				}).Info("Tried to stary signing round on " + s.Name() + " but it reports view change in progress")
				// skip # of failed round
				// s.SetLastSeenRound(s.LastRound() + 1)
				break
			}

			if err != nil {
				log.Errorln(err)
				break
			}

			elapsed := time.Since(start)
			log.WithFields(log.Fields{
				"file":  logutils.File(),
				"type":  "root_round",
				"round": s.LastRound(),
				"time":  elapsed,
			}).Info("root round")

			if s.LastRound() >= int64(nRounds) {
				log.Errorln(s.Name(), "reports exceeded the max round: terminating", s.LastRound(), ">=", nRounds)
				return "close"
			}
		}
	}
}

func (s *Server) runAsRegular() string {
	select {
	case <-s.closeChan:
		log.WithFields(log.Fields{
			"file": logutils.File(),
			"type": "close",
		}).Infoln("server" + s.Name() + "has closed")
		return ""

	case nextRole := <-s.ViewChangeCh():
		return nextRole
	}
}

// Listen on client connections. If role is root also send annoucement
// for all of the nRounds
func (s *Server) Run(role string, nRounds int) {
	// defer func() {
	// 	log.Infoln(s.Name(), "CLOSE AFTER RUN")
	// 	s.Close()
	// }()
	go func() { err := s.Signer.Listen(); s.Close(); log.Error(err) }()

	s.rLock.Lock()
	s.maxRounds = nRounds
	s.rLock.Unlock()

	var nextRole string // next role when view changes
	for {
		switch role {

		case "root":
			nextRole = s.runAsRoot(nRounds)
		case "regular":
			nextRole = s.runAsRegular()
		case "test":
			ticker := time.Tick(2000 * time.Millisecond)
			for _ = range ticker {
				s.AggregateCommits(0)
			}
		}

		log.Println(s.Name(), "nextRole: ", nextRole)
		if nextRole == "close" {
			s.Close()
			return
		}
		if nextRole == "" {
			return
		}
		s.LogReRun(nextRole, role)
		role = nextRole
	}

}

func (s *Server) OnAnnounce() coco.CommitFunc {
	return func(view int) []byte {
		//log.Println("Aggregating Commits")
		return s.AggregateCommits(view)
	}
}

func (s *Server) OnDone() coco.DoneFunc {
	return func(view int, SNRoot hashid.HashId, LogHash hashid.HashId, p proof.Proof) {
		s.mux.Lock()
		for i, msg := range s.Queue[s.PROCESSING] {
			// proof to get from s.Root to big root
			combProof := make(proof.Proof, len(p))
			copy(combProof, p)

			// add my proof to get from a leaf message to my root s.Root
			combProof = append(combProof, s.Proofs[i]...)

			// proof that i can get from a leaf message to the big root
			if coco.DEBUG == true {
				proof.CheckProof(s.Signer.(*sign.Node).Suite().Hash, SNRoot, s.Leaves[i], combProof)
			}

			respMessg := TimeStampMessage{
				Type:  StampReplyType,
				ReqNo: msg.Tsm.ReqNo,
				Srep:  &StampReply{Sig: SNRoot, Prf: combProof}}

			s.PutToClient(msg.To, respMessg)
		}
		s.mux.Unlock()
	}

}

func (s *Server) AggregateCommits(view int) []byte {
	//log.Println(s.Name(), "calling AggregateCommits")
	s.mux.Lock()
	// get data from s once to avoid refetching from structure
	Queue := s.Queue
	READING := s.READING
	PROCESSING := s.PROCESSING
	// messages read will now be processed
	READING, PROCESSING = PROCESSING, READING
	s.READING, s.PROCESSING = s.PROCESSING, s.READING
	s.Queue[READING] = s.Queue[READING][:0]

	// give up if nothing to process
	if len(Queue[PROCESSING]) == 0 {
		s.mux.Unlock()
		s.Root = make([]byte, hashid.Size)
		s.Proofs = make([]proof.Proof, 1)
		return s.Root
	}

	// pull out to be Merkle Tree leaves
	s.Leaves = make([]hashid.HashId, 0)
	for _, msg := range Queue[PROCESSING] {
		s.Leaves = append(s.Leaves, hashid.HashId(msg.Tsm.Sreq.Val))
	}
	s.mux.Unlock()

	// non root servers keep track of rounds here
	if !s.IsRoot(view) {
		s.rLock.Lock()
		lsr := s.LastRound()
		mr := s.maxRounds
		s.rLock.Unlock()
		// if this is our last round then close the connections
		if lsr >= int64(mr) && mr >= 0 {
			s.closeChan <- true
		}
	}

	// create Merkle tree for this round's messages and check corectness
	s.Root, s.Proofs = proof.ProofTree(s.Suite().Hash, s.Leaves)
	if coco.DEBUG == true {
		if proof.CheckLocalProofs(s.Suite().Hash, s.Root, s.Leaves, s.Proofs) == true {
			log.Println("Local Proofs of", s.Name(), "successful for round "+strconv.Itoa(int(s.LastRound())))
		} else {
			panic("Local Proofs" + s.Name() + " unsuccessful for round " + strconv.Itoa(int(s.LastRound())))
		}
	}

	return s.Root
}

// Send message to client given by name
func (s *Server) PutToClient(name string, data coconet.BinaryMarshaler) {
	err := s.Clients[name].Put(data)
	if err == coconet.ErrClosed {
		s.Close()
		return
	}
	if err != nil && err != coconet.ErrNotEstablished {
		log.Warnf("%p error putting to client: %v", s, err)
	}
}
<commit_msg>removed extra logging<commit_after>package stamp

import (
	"net"
	"strconv"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/dedis/prifi/coco"
	"github.com/dedis/prifi/coco/coconet"
	"github.com/dedis/prifi/coco/hashid"
	"github.com/dedis/prifi/coco/proof"
	"github.com/dedis/prifi/coco/sign"
	"github.com/dedis/prifi/coco/test/logutils"
)

type Server struct {
	coco.Signer
	name    string
	Clients map[string]coconet.Conn

	// for aggregating messages from clients
	mux        sync.Mutex
	Queue      [][]MustReplyMessage
	READING    int
	PROCESSING int

	// Leaves, Root and Proof for a round
	Leaves []hashid.HashId // can be removed after we verify protocol
	Root   hashid.HashId
	Proofs []proof.Proof

	rLock     sync.Mutex
	maxRounds int
	closeChan chan bool

	Logger   string
	Hostname string
	App      string
}

func NewServer(signer coco.Signer) *Server {
	s := &Server{}

	s.Clients = make(map[string]coconet.Conn)
	s.Queue = make([][]MustReplyMessage, 2)
	s.READING = 0
	s.PROCESSING = 1

	s.Signer = signer
	s.Signer.RegisterAnnounceFunc(s.OnAnnounce())
	s.Signer.RegisterDoneFunc(s.OnDone())
	s.rLock = sync.Mutex{}

	// listen for client requests at one port higher
	// than the signing node
	h, p, err := net.SplitHostPort(s.Signer.Name())
	if err == nil {
		i, err := strconv.Atoi(p)
		if err != nil {
			log.Fatal(err)
		}
		s.name = net.JoinHostPort(h, strconv.Itoa(i+1))
	}
	s.Queue[s.READING] = make([]MustReplyMessage, 0)
	s.Queue[s.PROCESSING] = make([]MustReplyMessage, 0)
	s.closeChan = make(chan bool, 5)
	return s
}

var clientNumber int = 0

func (s *Server) Close() {
	log.Printf("closing stampserver: %p", s)
	s.closeChan <- true
	s.Signer.Close()
}

// listen for clients connections
// this server needs to be running on a different port
// than the Signer that is beneath it
func (s *Server) Listen() error {
	// log.Println("Listening @ ", s.name)
	ln, err := net.Listen("tcp4", s.name)
	if err != nil {
		panic(err)
	}

	go func() {
		for {
			// log.Printf("LISTENING TO CLIENTS: %p", s)
			conn, err := ln.Accept()
			if err != nil {
				// handle error
				log.Errorln("failed to accept connection")
				continue
			}

			c := coconet.NewTCPConnFromNet(conn)
			// log.Println("CLIENT TCP CONNECTION SUCCESSFULLY ESTABLISHED:", c)

			if _, ok := s.Clients[c.Name()]; !ok {
				s.Clients[c.Name()] = c

				go func(c coconet.Conn) {
					for {
						tsm := TimeStampMessage{}
						err := c.Get(&tsm)
						if err != nil {
							log.Errorf("%p Failed to get from child:", s, err)
							s.Close()
							return
						}
						switch tsm.Type {
						default:
							log.Errorf("Message of unknown type: %v\n", tsm.Type)
						case StampRequestType:
							// log.Println("RECEIVED STAMP REQUEST")
							s.mux.Lock()
							READING := s.READING
							s.Queue[READING] = append(s.Queue[READING],
								MustReplyMessage{Tsm: tsm, To: c.Name()})
							s.mux.Unlock()
						}
					}
				}(c)
			}

		}
	}()

	return nil
}

// Used for goconns
// should only be used if clients are created in batch
func (s *Server) ListenToClients() {
	// log.Printf("LISTENING TO CLIENTS: %p", s, s.Clients)
	for _, c := range s.Clients {
		go func(c coconet.Conn) {
			for {
				tsm := TimeStampMessage{}
				err := c.Get(&tsm)
				if err == coconet.ErrClosed {
					log.Errorf("%p Failed to get from client:", s, err)
					s.Close()
					return
				}
				if err != nil {
					log.WithFields(log.Fields{
						"file": logutils.File(),
					}).Errorf("%p failed To get message:", s, err)
				}
				switch tsm.Type {
				default:
					log.Errorln("Message of unknown type")
				case StampRequestType:
					// log.Println("STAMP REQUEST")
					s.mux.Lock()
					READING := s.READING
					s.Queue[READING] = append(s.Queue[READING],
						MustReplyMessage{Tsm: tsm, To: c.Name()})
					s.mux.Unlock()
				}
			}
		}(c)
	}
}

func (s *Server) ConnectToLogger() {
	return
	if s.Logger == "" || s.Hostname == "" || s.App == "" {
		log.Println("skipping connect to logger")
		return
	}
	log.Println("Connecting to Logger")
	lh, _ := logutils.NewLoggerHook(s.Logger, s.Hostname, s.App)
	log.Println("Connected to Logger")
	log.AddHook(lh)
}

func (s *Server) LogReRun(nextRole string, curRole string) {
	if nextRole == "root" {
		var messg = s.Name() + " became root"
		if curRole == "root" {
			messg = s.Name() + " remained root"
		}

		go s.ConnectToLogger()

		log.WithFields(log.Fields{
			"file": logutils.File(),
			"type": "role_change",
		}).Infoln(messg)
		// log.Printf("role change: %p", s)

	} else {
		var messg = s.Name() + " remained regular"
		if curRole == "root" {
			messg = s.Name() + " became regular"
		}

		if curRole == "root" {
			log.WithFields(log.Fields{
				"file": logutils.File(),
				"type": "role_change",
			}).Infoln(messg)
			log.Printf("role change: %p", s)
		}

	}

}

func (s *Server) runAsRoot(nRounds int) string {
	// every 5 seconds start a new round
	ticker := time.Tick(ROUND_TIME)
	log.Infoln(s.Name(), "running as root", s.LastRound(), int64(nRounds))

	for {
		select {
		case nextRole := <-s.ViewChangeCh():
			return nextRole
			// s.reRunWith(nextRole, nRounds, true)
		case <-ticker:

			start := time.Now()
			log.Println(s.Name(), "is STAMP SERVER STARTING SIGNING ROUND FOR:", s.LastRound()+1, "of", nRounds)

			err := s.StartSigningRound()
			if err == sign.ChangingViewError {
				// report change in view, and continue with the select
				log.WithFields(log.Fields{
					"file": logutils.File(),
					"type": "view_change",
				}).Info("Tried to stary signing round on " + s.Name() + " but it reports view change in progress")
				// skip # of failed round
				// s.SetLastSeenRound(s.LastRound() + 1)
				break
			}

			if err != nil {
				log.Errorln(err)
				break
			}

			elapsed := time.Since(start)
			log.WithFields(log.Fields{
				"file":  logutils.File(),
				"type":  "root_round",
				"round": s.LastRound(),
				"time":  elapsed,
			}).Info("root round")

			if s.LastRound() >= int64(nRounds) {
				log.Errorln(s.Name(), "reports exceeded the max round: terminating", s.LastRound(), ">=", nRounds)
				return "close"
			}
		}
	}
}

func (s *Server) runAsRegular() string {
	select {
	case <-s.closeChan:
		log.WithFields(log.Fields{
			"file": logutils.File(),
			"type": "close",
		}).Infoln("server" + s.Name() + "has closed")
		return ""

	case nextRole := <-s.ViewChangeCh():
		return nextRole
	}
}

// Listen on client connections. If role is root also send annoucement
// for all of the nRounds
func (s *Server) Run(role string, nRounds int) {
	// defer func() {
	// 	log.Infoln(s.Name(), "CLOSE AFTER RUN")
	// 	s.Close()
	// }()
	go func() { err := s.Signer.Listen(); s.Close(); log.Error(err) }()

	s.rLock.Lock()
	s.maxRounds = nRounds
	s.rLock.Unlock()

	var nextRole string // next role when view changes
	for {
		switch role {

		case "root":
			nextRole = s.runAsRoot(nRounds)
		case "regular":
			nextRole = s.runAsRegular()
		case "test":
			ticker := time.Tick(2000 * time.Millisecond)
			for _ = range ticker {
				s.AggregateCommits(0)
			}
		}

		// log.Println(s.Name(), "nextRole: ", nextRole)
		if nextRole == "close" {
			s.Close()
			return
		}
		if nextRole == "" {
			return
		}
		s.LogReRun(nextRole, role)
		role = nextRole
	}

}

func (s *Server) OnAnnounce() coco.CommitFunc {
	return func(view int) []byte {
		//log.Println("Aggregating Commits")
		return s.AggregateCommits(view)
	}
}

func (s *Server) OnDone() coco.DoneFunc {
	return func(view int, SNRoot hashid.HashId, LogHash hashid.HashId, p proof.Proof) {
		s.mux.Lock()
		for i, msg := range s.Queue[s.PROCESSING] {
			// proof to get from s.Root to big root
			combProof := make(proof.Proof, len(p))
			copy(combProof, p)

			// add my proof to get from a leaf message to my root s.Root
			combProof = append(combProof, s.Proofs[i]...)

			// proof that i can get from a leaf message to the big root
			if coco.DEBUG == true {
				proof.CheckProof(s.Signer.(*sign.Node).Suite().Hash, SNRoot, s.Leaves[i], combProof)
			}

			respMessg := TimeStampMessage{
				Type:  StampReplyType,
				ReqNo: msg.Tsm.ReqNo,
				Srep:  &StampReply{Sig: SNRoot, Prf: combProof}}

			s.PutToClient(msg.To, respMessg)
		}
		s.mux.Unlock()
	}

}

func (s *Server) AggregateCommits(view int) []byte {
	//log.Println(s.Name(), "calling AggregateCommits")
	s.mux.Lock()
	// get data from s once to avoid refetching from structure
	Queue := s.Queue
	READING := s.READING
	PROCESSING := s.PROCESSING
	// messages read will now be processed
	READING, PROCESSING = PROCESSING, READING
	s.READING, s.PROCESSING = s.PROCESSING, s.READING
	s.Queue[READING] = s.Queue[READING][:0]

	// give up if nothing to process
	if len(Queue[PROCESSING]) == 0 {
		s.mux.Unlock()
		s.Root = make([]byte, hashid.Size)
		s.Proofs = make([]proof.Proof, 1)
		return s.Root
	}

	// pull out to be Merkle Tree leaves
	s.Leaves = make([]hashid.HashId, 0)
	for _, msg := range Queue[PROCESSING] {
		s.Leaves = append(s.Leaves, hashid.HashId(msg.Tsm.Sreq.Val))
	}
	s.mux.Unlock()

	// non root servers keep track of rounds here
	if !s.IsRoot(view) {
		s.rLock.Lock()
		lsr := s.LastRound()
		mr := s.maxRounds
		s.rLock.Unlock()
		// if this is our last round then close the connections
		if lsr >= int64(mr) && mr >= 0 {
			s.closeChan <- true
		}
	}

	// create Merkle tree for this round's messages and check corectness
	s.Root, s.Proofs = proof.ProofTree(s.Suite().Hash, s.Leaves)
	if coco.DEBUG == true {
		if proof.CheckLocalProofs(s.Suite().Hash, s.Root, s.Leaves, s.Proofs) == true {
			log.Println("Local Proofs of", s.Name(), "successful for round "+strconv.Itoa(int(s.LastRound())))
		} else {
			panic("Local Proofs" + s.Name() + " unsuccessful for round " + strconv.Itoa(int(s.LastRound())))
		}
	}

	return s.Root
}

// Send message to client given by name
func (s *Server) PutToClient(name string, data coconet.BinaryMarshaler) {
	err := s.Clients[name].Put(data)
	if err == coconet.ErrClosed {
		s.Close()
		return
	}
	if err != nil && err != coconet.ErrNotEstablished {
		log.Warnf("%p error putting to client: %v", s, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"text/template"
	"time"

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

	"github.com/codedellemc/gocsi"
	"github.com/codedellemc/gocsi/csi"
)

const (
	// defaultVersion is the default CSI_VERSION string if none
	// is provided via a CLI argument or environment variable
	defaultVersion = "0.1.0"

	// maxUint32 is the maximum value for a uint32. this is
	// defined as math.MaxUint32, but it's redefined here
	// in order to avoid importing the math package for just
	// a constant value
	maxUint32 = 4294967295

	// maxInt32 is the maximum value for an int32. this is
	// defined as math.MaxInt32, but it's redefined here
	// in order to avoid importing the math package for just
	// a constant value
	maxInt32 = 2147483647
)

var appName = path.Base(os.Args[0])

func main() {

	// the program should have at least two args:
	//
	//     args[0]  path of executable
	//     args[1]  csi rpc
	if len(os.Args) < 2 {
		usage(os.Stderr)
		os.Exit(1)
	}

	// match the name of the rpc or one of its aliases
	rpc := os.Args[1]
	c := func(ccc ...[]*cmd) *cmd {
		for _, cc := range ccc {
			for _, c := range cc {
				if strings.EqualFold(rpc, c.Name) {
					rpc = c.Name
					return c
				}
				for _, a := range c.Aliases {
					if strings.EqualFold(rpc, a) {
						rpc = a
						return c
					}
				}
			}
		}
		return nil
	}(controllerCmds, identityCmds, nodeCmds)

	// assert that a command for the requested rpc was found
	if c == nil {
		fmt.Fprintf(os.Stderr, "error: invalid rpc: %s\n", rpc)
		usage(os.Stderr)
		os.Exit(1)
	}

	if c.Action == nil {
		panic("nil rpc action")
	}
	if c.Flags == nil {
		panic("nil rpc flags")
	}

	ctx := context.Background()

	// parse the command line with the command's flag set
	cflags := c.Flags(ctx, rpc)
	if err := cflags.Parse(os.Args[2:]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// assert that the endpoint value is required
	if args.endpoint == "" {
		fmt.Fprintln(os.Stderr, "error: endpoint is required")
		cflags.Usage()
		os.Exit(1)
	}

	// assert that the version is required and valid
	versionRX := regexp.MustCompile(`(\d+)\.(\d+)\.(\d+)`)
	versionMatch := versionRX.FindStringSubmatch(args.szVersion)
	if len(versionMatch) == 0 {
		fmt.Fprintf(
			os.Stderr,
			"error: invalid version: %s\n",
			args.szVersion)
		os.Exit(1)
	}
	versionMajor, _ := strconv.Atoi(versionMatch[1])
	if versionMajor > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: MAJOR > uint32: %v\n", versionMajor)
		os.Exit(1)
	}
	versionMinor, _ := strconv.Atoi(versionMatch[2])
	if versionMinor > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: MINOR > uint32: %v\n", versionMinor)
		os.Exit(1)
	}
	versionPatch, _ := strconv.Atoi(versionMatch[3])
	if versionPatch > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: PATCH > uint32: %v\n", versionPatch)
		os.Exit(1)
	}
	args.version = &csi.Version{
		Major: uint32(versionMajor),
		Minor: uint32(versionMinor),
		Patch: uint32(versionPatch),
	}

	// initialize a grpc client
	gclient, err := newGrpcClient(ctx, args.endpoint, args.insecure)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// if a service is specified then add it to the context
	// as gRPC metadata
	if args.service != "" {
		ctx = metadata.NewContext(
			ctx, metadata.Pairs("csi.service", args.service))
	}

	// execute the command
	if err := c.Action(ctx, cflags, gclient); err != nil {
		fmt.Fprintln(os.Stderr, err)
		if _, ok := err.(*errUsage); ok {
			cflags.Usage()
		}
		os.Exit(1)
	}
}

func newGrpcClient(
	ctx context.Context,
	endpoint string,
	insecure bool) (*grpc.ClientConn, error) {

	dialOpts := []grpc.DialOption{
		grpc.WithUnaryInterceptor(gocsi.ChainUnaryClient(
			gocsi.ClientCheckReponseError,
			gocsi.ClientResponseValidator)),
		grpc.WithDialer(
			func(target string, timeout time.Duration) (net.Conn, error) {
				proto, addr, err := gocsi.ParseProtoAddr(target)
				if err != nil {
					return nil, err
				}
				return net.DialTimeout(proto, addr, timeout)
			}),
	}

	if insecure {
		dialOpts = append(dialOpts, grpc.WithInsecure())
	}

	return grpc.DialContext(ctx, endpoint, dialOpts...)
}

///////////////////////////////////////////////////////////////////////////////
//                            Default Formats                                //
///////////////////////////////////////////////////////////////////////////////

// mapSzOfSzFormat is the default Go template format for
// emitting a map[string]string
const mapSzOfSzFormat = `{{range $k, $v := .}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{"\n"}}`

// volumeInfoFormat is the default Go template format for
// emitting a *csi.VolumeInfo
const volumeInfoFormat = `{{with .GetId}}{{range $k, $v := .GetValues}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{end}}{{"\n"}}`

// versionFormat is the default Go template format for emitting a *csi.Version
const versionFormat = `{{.GetMajor}}.{{.GetMinor}}.{{.GetPatch}}{{"\n"}}`

// pluginInfoFormat is the default Go template format for
// emitting a *csi.GetPluginInfoResponse_Result
const pluginInfoFormat = `{{.Name}}{{print "\t"}}{{.VendorVersion}}{{print "\t"}}` +
	`{{with .GetManifest}}{{range $k, $v := .}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{end}}{{"\n"}}`

// capFormat is the default Go template for emitting a
// *csi.{Controller,Node}ServiceCapability
const capFormat = `{{with .GetRpc}}{{.Type}}{{end}}{{"\n"}}`

// valCapFormat is the default Go tempate for emitting a
// *csi.ValidateVolumeCapabilitiesResponse_Result
const valCapFormat = `{{with .GetSupported}}{{print "supported: "}}{{.}}` +
	`{{print "\n"}}{{end}}{{with .GetMessage}}{{print "\tmessage: "}}` +
	`{{.}}{{end}}{{"\n"}}`

///////////////////////////////////////////////////////////////////////////////
//                                Commands                                   //
///////////////////////////////////////////////////////////////////////////////
type errUsage struct {
	msg string
}

func (e *errUsage) Error() string {
	return e.msg
}

type cmd struct {
	Name    string
	Aliases []string
	Action  func(context.Context, *flag.FlagSet, *grpc.ClientConn) error
	Flags   func(context.Context, string) *flag.FlagSet
}

///////////////////////////////////////////////////////////////////////////////
//                                Usage                                      //
///////////////////////////////////////////////////////////////////////////////
func usage(w io.Writer) {
	const h = `usage: {{.Name}} RPC [ARGS...]{{range $Name, $Cmds := .Categories}}

       {{$Name}} RPCs{{range $Cmds}}
         {{.Name}}{{if .Aliases}} ({{join .Aliases ", "}}){{end}}{{end}}{{end}}

Use the -? flag with an RPC for additional help.
`
	f := template.FuncMap{"join": strings.Join}
	t := template.Must(template.New(appName).Funcs(f).Parse(h))
	d := struct {
		Name       string
		Categories map[string][]*cmd
	}{
		appName,
		map[string][]*cmd{
			"CONTROLLER": controllerCmds,
			"IDENTITY":   identityCmds,
			"NODE":       nodeCmds,
		},
	}
	t.Execute(w, d)
}

///////////////////////////////////////////////////////////////////////////////
//                               Global Flags                                //
///////////////////////////////////////////////////////////////////////////////
var args struct {
	service   string
	endpoint  string
	format    string
	help      bool
	insecure  bool
	szVersion string
	version   *csi.Version
}

func flagsGlobal(
	fs *flag.FlagSet,
	formatDefault, formatObjectType string) {

	fs.StringVar(
		&args.endpoint,
		"endpoint",
		os.Getenv("CSI_ENDPOINT"),
		"The endpoint address")

	fs.StringVar(
		&args.service,
		"service",
		"",
		"The name of the CSD service to use.")

	version := defaultVersion
	if v := os.Getenv("CSI_VERSION"); v != "" {
		version = v
	}
	fs.StringVar(
		&args.szVersion,
		"version",
		version,
		"The API version string")

	insecure := true
	if v := os.Getenv("CSI_INSECURE"); v != "" {
		insecure, _ = strconv.ParseBool(v)
	}
	fs.BoolVar(
		&args.insecure,
		"insecure",
		insecure,
		"Disables transport security")

	fmtMsg := &bytes.Buffer{}
	fmt.Fprint(fmtMsg, "The Go template used to print an object.")
	if formatObjectType != "" {
		fmt.Fprintf(fmtMsg, " This command emits a %s.", formatObjectType)
	}
	fs.StringVar(
		&args.format,
		"format",
		formatDefault,
		fmtMsg.String())
}

// stringSliceArg is used for parsing a csv arg into a string slice
type stringSliceArg struct {
	szVal string
	vals  []string
}

func (s *stringSliceArg) String() string {
	return s.szVal
}

func (s *stringSliceArg) Set(val string) error {
	s.vals = append(s.vals, strings.Split(val, ",")...)
	return nil
}

// mapOfStringArg is used for parsing a csv, key=value arg into
// a map[string]string
type mapOfStringArg struct {
	szVal string
	vals  map[string]string
}

func (s *mapOfStringArg) String() string {
	return s.szVal
}

func (s *mapOfStringArg) Set(val string) error {
	if s.vals == nil {
		s.vals = map[string]string{}
	}
	vals := strings.Split(val, ",")
	for _, v := range vals {
		vp := strings.SplitN(v, "=", 2)
		switch len(vp) {
		case 1:
			s.vals[vp[0]] = ""
		case 2:
			s.vals[vp[0]] = vp[1]
		}
	}
	return nil
}
<commit_msg>Print VolumeInfo.Metdata for csc ls<commit_after>package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"text/template"
	"time"

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

	"github.com/codedellemc/gocsi"
	"github.com/codedellemc/gocsi/csi"
)

const (
	// defaultVersion is the default CSI_VERSION string if none
	// is provided via a CLI argument or environment variable
	defaultVersion = "0.1.0"

	// maxUint32 is the maximum value for a uint32. this is
	// defined as math.MaxUint32, but it's redefined here
	// in order to avoid importing the math package for just
	// a constant value
	maxUint32 = 4294967295

	// maxInt32 is the maximum value for an int32. this is
	// defined as math.MaxInt32, but it's redefined here
	// in order to avoid importing the math package for just
	// a constant value
	maxInt32 = 2147483647
)

var appName = path.Base(os.Args[0])

func main() {

	// the program should have at least two args:
	//
	//     args[0]  path of executable
	//     args[1]  csi rpc
	if len(os.Args) < 2 {
		usage(os.Stderr)
		os.Exit(1)
	}

	// match the name of the rpc or one of its aliases
	rpc := os.Args[1]
	c := func(ccc ...[]*cmd) *cmd {
		for _, cc := range ccc {
			for _, c := range cc {
				if strings.EqualFold(rpc, c.Name) {
					rpc = c.Name
					return c
				}
				for _, a := range c.Aliases {
					if strings.EqualFold(rpc, a) {
						rpc = a
						return c
					}
				}
			}
		}
		return nil
	}(controllerCmds, identityCmds, nodeCmds)

	// assert that a command for the requested rpc was found
	if c == nil {
		fmt.Fprintf(os.Stderr, "error: invalid rpc: %s\n", rpc)
		usage(os.Stderr)
		os.Exit(1)
	}

	if c.Action == nil {
		panic("nil rpc action")
	}
	if c.Flags == nil {
		panic("nil rpc flags")
	}

	ctx := context.Background()

	// parse the command line with the command's flag set
	cflags := c.Flags(ctx, rpc)
	if err := cflags.Parse(os.Args[2:]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// assert that the endpoint value is required
	if args.endpoint == "" {
		fmt.Fprintln(os.Stderr, "error: endpoint is required")
		cflags.Usage()
		os.Exit(1)
	}

	// assert that the version is required and valid
	versionRX := regexp.MustCompile(`(\d+)\.(\d+)\.(\d+)`)
	versionMatch := versionRX.FindStringSubmatch(args.szVersion)
	if len(versionMatch) == 0 {
		fmt.Fprintf(
			os.Stderr,
			"error: invalid version: %s\n",
			args.szVersion)
		os.Exit(1)
	}
	versionMajor, _ := strconv.Atoi(versionMatch[1])
	if versionMajor > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: MAJOR > uint32: %v\n", versionMajor)
		os.Exit(1)
	}
	versionMinor, _ := strconv.Atoi(versionMatch[2])
	if versionMinor > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: MINOR > uint32: %v\n", versionMinor)
		os.Exit(1)
	}
	versionPatch, _ := strconv.Atoi(versionMatch[3])
	if versionPatch > maxUint32 {
		fmt.Fprintf(
			os.Stderr, "error: PATCH > uint32: %v\n", versionPatch)
		os.Exit(1)
	}
	args.version = &csi.Version{
		Major: uint32(versionMajor),
		Minor: uint32(versionMinor),
		Patch: uint32(versionPatch),
	}

	// initialize a grpc client
	gclient, err := newGrpcClient(ctx, args.endpoint, args.insecure)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// if a service is specified then add it to the context
	// as gRPC metadata
	if args.service != "" {
		ctx = metadata.NewContext(
			ctx, metadata.Pairs("csi.service", args.service))
	}

	// execute the command
	if err := c.Action(ctx, cflags, gclient); err != nil {
		fmt.Fprintln(os.Stderr, err)
		if _, ok := err.(*errUsage); ok {
			cflags.Usage()
		}
		os.Exit(1)
	}
}

func newGrpcClient(
	ctx context.Context,
	endpoint string,
	insecure bool) (*grpc.ClientConn, error) {

	dialOpts := []grpc.DialOption{
		grpc.WithUnaryInterceptor(gocsi.ChainUnaryClient(
			gocsi.ClientCheckReponseError,
			gocsi.ClientResponseValidator)),
		grpc.WithDialer(
			func(target string, timeout time.Duration) (net.Conn, error) {
				proto, addr, err := gocsi.ParseProtoAddr(target)
				if err != nil {
					return nil, err
				}
				return net.DialTimeout(proto, addr, timeout)
			}),
	}

	if insecure {
		dialOpts = append(dialOpts, grpc.WithInsecure())
	}

	return grpc.DialContext(ctx, endpoint, dialOpts...)
}

///////////////////////////////////////////////////////////////////////////////
//                            Default Formats                                //
///////////////////////////////////////////////////////////////////////////////

// mapSzOfSzFormat is the default Go template format for
// emitting a map[string]string
const mapSzOfSzFormat = `{{range $k, $v := .}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{"\n"}}`

// volumeInfoFormat is the default Go template format for
// emitting a *csi.VolumeInfo
const volumeInfoFormat = `{{with .Id}}{{range $k, $v := .Values}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{end}}` +
	`{{if .Metadata}}{{with .Metadata}}{{range $k, $v := .Values}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{end}}{{end}}` +
	`{{"\n"}}`

// versionFormat is the default Go template format for emitting a *csi.Version
const versionFormat = `{{.GetMajor}}.{{.GetMinor}}.{{.GetPatch}}{{"\n"}}`

// pluginInfoFormat is the default Go template format for
// emitting a *csi.GetPluginInfoResponse_Result
const pluginInfoFormat = `{{.Name}}{{print "\t"}}{{.VendorVersion}}{{print "\t"}}` +
	`{{with .GetManifest}}{{range $k, $v := .}}` +
	`{{printf "%s=%s\t" $k $v}}{{end}}{{end}}{{"\n"}}`

// capFormat is the default Go template for emitting a
// *csi.{Controller,Node}ServiceCapability
const capFormat = `{{with .GetRpc}}{{.Type}}{{end}}{{"\n"}}`

// valCapFormat is the default Go tempate for emitting a
// *csi.ValidateVolumeCapabilitiesResponse_Result
const valCapFormat = `{{with .GetSupported}}{{print "supported: "}}{{.}}` +
	`{{print "\n"}}{{end}}{{with .GetMessage}}{{print "\tmessage: "}}` +
	`{{.}}{{end}}{{"\n"}}`

///////////////////////////////////////////////////////////////////////////////
//                                Commands                                   //
///////////////////////////////////////////////////////////////////////////////
type errUsage struct {
	msg string
}

func (e *errUsage) Error() string {
	return e.msg
}

type cmd struct {
	Name    string
	Aliases []string
	Action  func(context.Context, *flag.FlagSet, *grpc.ClientConn) error
	Flags   func(context.Context, string) *flag.FlagSet
}

///////////////////////////////////////////////////////////////////////////////
//                                Usage                                      //
///////////////////////////////////////////////////////////////////////////////
func usage(w io.Writer) {
	const h = `usage: {{.Name}} RPC [ARGS...]{{range $Name, $Cmds := .Categories}}

       {{$Name}} RPCs{{range $Cmds}}
         {{.Name}}{{if .Aliases}} ({{join .Aliases ", "}}){{end}}{{end}}{{end}}

Use the -? flag with an RPC for additional help.
`
	f := template.FuncMap{"join": strings.Join}
	t := template.Must(template.New(appName).Funcs(f).Parse(h))
	d := struct {
		Name       string
		Categories map[string][]*cmd
	}{
		appName,
		map[string][]*cmd{
			"CONTROLLER": controllerCmds,
			"IDENTITY":   identityCmds,
			"NODE":       nodeCmds,
		},
	}
	t.Execute(w, d)
}

///////////////////////////////////////////////////////////////////////////////
//                               Global Flags                                //
///////////////////////////////////////////////////////////////////////////////
var args struct {
	service   string
	endpoint  string
	format    string
	help      bool
	insecure  bool
	szVersion string
	version   *csi.Version
}

func flagsGlobal(
	fs *flag.FlagSet,
	formatDefault, formatObjectType string) {

	fs.StringVar(
		&args.endpoint,
		"endpoint",
		os.Getenv("CSI_ENDPOINT"),
		"The endpoint address")

	fs.StringVar(
		&args.service,
		"service",
		"",
		"The name of the CSD service to use.")

	version := defaultVersion
	if v := os.Getenv("CSI_VERSION"); v != "" {
		version = v
	}
	fs.StringVar(
		&args.szVersion,
		"version",
		version,
		"The API version string")

	insecure := true
	if v := os.Getenv("CSI_INSECURE"); v != "" {
		insecure, _ = strconv.ParseBool(v)
	}
	fs.BoolVar(
		&args.insecure,
		"insecure",
		insecure,
		"Disables transport security")

	fmtMsg := &bytes.Buffer{}
	fmt.Fprint(fmtMsg, "The Go template used to print an object.")
	if formatObjectType != "" {
		fmt.Fprintf(fmtMsg, " This command emits a %s.", formatObjectType)
	}
	fs.StringVar(
		&args.format,
		"format",
		formatDefault,
		fmtMsg.String())
}

// stringSliceArg is used for parsing a csv arg into a string slice
type stringSliceArg struct {
	szVal string
	vals  []string
}

func (s *stringSliceArg) String() string {
	return s.szVal
}

func (s *stringSliceArg) Set(val string) error {
	s.vals = append(s.vals, strings.Split(val, ",")...)
	return nil
}

// mapOfStringArg is used for parsing a csv, key=value arg into
// a map[string]string
type mapOfStringArg struct {
	szVal string
	vals  map[string]string
}

func (s *mapOfStringArg) String() string {
	return s.szVal
}

func (s *mapOfStringArg) Set(val string) error {
	if s.vals == nil {
		s.vals = map[string]string{}
	}
	vals := strings.Split(val, ",")
	for _, v := range vals {
		vp := strings.SplitN(v, "=", 2)
		switch len(vp) {
		case 1:
			s.vals[vp[0]] = ""
		case 2:
			s.vals[vp[0]] = vp[1]
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package httpstream

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"

	"code.google.com/p/go.net/websocket"
	"github.com/gorilla/mux"

	"github.com/gliderlabs/logspout/router"
)

func init() {
	router.HttpHandlers.Register(LogStreamer, "logs")
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func LogStreamer() http.Handler {
	logs := mux.NewRouter()
	logsHandler := func(w http.ResponseWriter, req *http.Request) {
		params := mux.Vars(req)
		route := new(router.Route)

		if params["value"] != "" {
			switch params["predicate"] {
			case "id":
				route.FilterID = params["value"]
				if len(route.ID) > 12 {
					route.FilterID = route.FilterID[:12]
				}
			case "name":
				route.FilterName = params["value"]
			}
		}

		if route.FilterID != "" && !router.Routes.RoutingFrom(route.FilterID) {
			http.NotFound(w, req)
			return
		}

		defer debug("http: logs streamer disconnected")
		logstream := make(chan *router.Message)
		defer close(logstream)

		var closer <-chan bool
		if req.Header.Get("Upgrade") == "websocket" {
			debug("http: logs streamer connected [websocket]")
			closerBi := make(chan bool)
			defer websocketStreamer(w, req, logstream, closerBi)
			closer = closerBi
		} else {
			debug("http: logs streamer connected [http]")
			defer httpStreamer(w, req, logstream, route.MultiContainer())
			closer = w.(http.CloseNotifier).CloseNotify()
		}
		route.OverrideCloser(closer)

		router.Routes.Route(route, logstream)
	}
	logs.HandleFunc("/logs/{predicate:[a-zA-Z]+}:{value}", logsHandler).Methods("GET")
	logs.HandleFunc("/logs", logsHandler).Methods("GET")
	return logs
}

type Colorizer map[string]int

// returns up to 14 color escape codes (then repeats) for each unique key
func (c Colorizer) Get(key string) string {
	i, exists := c[key]
	if !exists {
		c[key] = len(c)
		i = c[key]
	}
	bright := "1;"
	if i%14 > 6 {
		bright = ""
	}
	return "\x1b[" + bright + "3" + strconv.Itoa(7-(i%7)) + "m"
}

func marshal(obj interface{}) []byte {
	bytes, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		log.Println("marshal:", err)
	}
	return bytes
}

func normalName(name string) string {
	return name[1:]
}

func websocketStreamer(w http.ResponseWriter, req *http.Request, logstream chan *router.Message, closer chan bool) {
	websocket.Handler(func(conn *websocket.Conn) {
		for logline := range logstream {
			if req.URL.Query().Get("source") != "" && logline.Source != req.URL.Query().Get("source") {
				continue
			}
			_, err := conn.Write(append(marshal(logline), '\n'))
			if err != nil {
				closer <- true
				return
			}
		}
	}).ServeHTTP(w, req)
}

func httpStreamer(w http.ResponseWriter, req *http.Request, logstream chan *router.Message, multi bool) {
	var colors Colorizer
	var usecolor, usejson bool
	nameWidth := 16
	if req.URL.Query().Get("colors") != "off" {
		colors = make(Colorizer)
		usecolor = true
	}
	if req.Header.Get("Accept") == "application/json" {
		w.Header().Add("Content-Type", "application/json")
		usejson = true
	} else {
		w.Header().Add("Content-Type", "text/plain")
	}
	for logline := range logstream {
		if req.URL.Query().Get("sources") != "" && logline.Source != req.URL.Query().Get("sources") {
			continue
		}
		if usejson {
			w.Write(append(marshal(logline), '\n'))
		} else {
			if multi {
				name := normalName(logline.Container.Name)
				if len(name) > nameWidth {
					nameWidth = len(name)
				}
				if usecolor {
					w.Write([]byte(fmt.Sprintf(
						"%s%"+strconv.Itoa(nameWidth)+"s|%s\x1b[0m\n",
						colors.Get(name), name, logline.Data,
					)))
				} else {
					w.Write([]byte(fmt.Sprintf(
						"%"+strconv.Itoa(nameWidth)+"s|%s\n", name, logline.Data,
					)))
				}
			} else {
				w.Write(append([]byte(logline.Data), '\n'))
			}
		}
		w.(http.Flusher).Flush()
	}
}
<commit_msg>Removed deprecated library hosted in google code in favor of its new home<commit_after>package httpstream

import (
	"encoding/json"
	"fmt"
	"github.com/gliderlabs/logspout/router"
	"github.com/gorilla/mux"
	"golang.org/x/net/websocket"
	"log"
	"net/http"
	"os"
	"strconv"
)

func init() {
	router.HttpHandlers.Register(LogStreamer, "logs")
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func LogStreamer() http.Handler {
	logs := mux.NewRouter()
	logsHandler := func(w http.ResponseWriter, req *http.Request) {
		params := mux.Vars(req)
		route := new(router.Route)

		if params["value"] != "" {
			switch params["predicate"] {
			case "id":
				route.FilterID = params["value"]
				if len(route.ID) > 12 {
					route.FilterID = route.FilterID[:12]
				}
			case "name":
				route.FilterName = params["value"]
			}
		}

		if route.FilterID != "" && !router.Routes.RoutingFrom(route.FilterID) {
			http.NotFound(w, req)
			return
		}

		defer debug("http: logs streamer disconnected")
		logstream := make(chan *router.Message)
		defer close(logstream)

		var closer <-chan bool
		if req.Header.Get("Upgrade") == "websocket" {
			debug("http: logs streamer connected [websocket]")
			closerBi := make(chan bool)
			defer websocketStreamer(w, req, logstream, closerBi)
			closer = closerBi
		} else {
			debug("http: logs streamer connected [http]")
			defer httpStreamer(w, req, logstream, route.MultiContainer())
			closer = w.(http.CloseNotifier).CloseNotify()
		}
		route.OverrideCloser(closer)

		router.Routes.Route(route, logstream)
	}
	logs.HandleFunc("/logs/{predicate:[a-zA-Z]+}:{value}", logsHandler).Methods("GET")
	logs.HandleFunc("/logs", logsHandler).Methods("GET")
	return logs
}

type Colorizer map[string]int

// returns up to 14 color escape codes (then repeats) for each unique key
func (c Colorizer) Get(key string) string {
	i, exists := c[key]
	if !exists {
		c[key] = len(c)
		i = c[key]
	}
	bright := "1;"
	if i%14 > 6 {
		bright = ""
	}
	return "\x1b[" + bright + "3" + strconv.Itoa(7-(i%7)) + "m"
}

func marshal(obj interface{}) []byte {
	bytes, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		log.Println("marshal:", err)
	}
	return bytes
}

func normalName(name string) string {
	return name[1:]
}

func websocketStreamer(w http.ResponseWriter, req *http.Request, logstream chan *router.Message, closer chan bool) {
	websocket.Handler(func(conn *websocket.Conn) {
		for logline := range logstream {
			if req.URL.Query().Get("source") != "" && logline.Source != req.URL.Query().Get("source") {
				continue
			}
			_, err := conn.Write(append(marshal(logline), '\n'))
			if err != nil {
				closer <- true
				return
			}
		}
	}).ServeHTTP(w, req)
}

func httpStreamer(w http.ResponseWriter, req *http.Request, logstream chan *router.Message, multi bool) {
	var colors Colorizer
	var usecolor, usejson bool
	nameWidth := 16
	if req.URL.Query().Get("colors") != "off" {
		colors = make(Colorizer)
		usecolor = true
	}
	if req.Header.Get("Accept") == "application/json" {
		w.Header().Add("Content-Type", "application/json")
		usejson = true
	} else {
		w.Header().Add("Content-Type", "text/plain")
	}
	for logline := range logstream {
		if req.URL.Query().Get("sources") != "" && logline.Source != req.URL.Query().Get("sources") {
			continue
		}
		if usejson {
			w.Write(append(marshal(logline), '\n'))
		} else {
			if multi {
				name := normalName(logline.Container.Name)
				if len(name) > nameWidth {
					nameWidth = len(name)
				}
				if usecolor {
					w.Write([]byte(fmt.Sprintf(
						"%s%"+strconv.Itoa(nameWidth)+"s|%s\x1b[0m\n",
						colors.Get(name), name, logline.Data,
					)))
				} else {
					w.Write([]byte(fmt.Sprintf(
						"%"+strconv.Itoa(nameWidth)+"s|%s\n", name, logline.Data,
					)))
				}
			} else {
				w.Write(append([]byte(logline.Data), '\n'))
			}
		}
		w.(http.Flusher).Flush()
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// +build integration

package integration

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"

	"github.com/syncthing/protocol"
	"github.com/syncthing/syncthing/internal/rc"
)

var jsonEndpoints = []string{
	"/rest/db/completion?device=I6KAH76-66SLLLB-5PFXSOA-UFJCDZC-YAOMLEK-CP2GB32-BV5RQST-3PSROAU&folder=default",
	"/rest/db/ignores?folder=default",
	"/rest/db/need?folder=default",
	"/rest/db/status?folder=default",
	"/rest/db/browse?folder=default",
	"/rest/events?since=-1&limit=5",
	"/rest/stats/device",
	"/rest/stats/folder",
	"/rest/svc/deviceid?id=I6KAH76-66SLLLB-5PFXSOA-UFJCDZC-YAOMLEK-CP2GB32-BV5RQST-3PSROAU",
	"/rest/svc/lang",
	"/rest/svc/report",
	"/rest/system/browse?current=.",
	"/rest/system/config",
	"/rest/system/config/insync",
	"/rest/system/connections",
	"/rest/system/discovery",
	"/rest/system/error",
	"/rest/system/ping",
	"/rest/system/status",
	"/rest/system/upgrade",
	"/rest/system/version",
}

func TestGetIndex(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	// Check for explicint index.html

	res, err := http.Get("http://localhost:8082/index.html")
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Errorf("Status %d != 200", res.StatusCode)
	}
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	if len(bs) < 1024 {
		t.Errorf("Length %d < 1024", len(bs))
	}
	if !bytes.Contains(bs, []byte("")) {
		t.Error("Incorrect response")
	}
	if res.Header.Get("Set-Cookie") == "" {
		t.Error("No set-cookie header")
	}
	res.Body.Close()

	// Check for implicit index.html

	res, err = http.Get("http://localhost:8082/")
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Errorf("Status %d != 200", res.StatusCode)
	}
	bs, err = ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	if len(bs) < 1024 {
		t.Errorf("Length %d < 1024", len(bs))
	}
	if !bytes.Contains(bs, []byte("")) {
		t.Error("Incorrect response")
	}
	if res.Header.Get("Set-Cookie") == "" {
		t.Error("No set-cookie header")
	}
	res.Body.Close()
}

func TestGetIndexAuth(t *testing.T) {
	p := startInstance(t, 1)
	defer checkedStop(t, p)

	// Without auth should give 401

	res, err := http.Get("http://127.0.0.1:8081/")
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 401 {
		t.Errorf("Status %d != 401", res.StatusCode)
	}

	// With wrong username/password should give 401

	req, err := http.NewRequest("GET", "http://127.0.0.1:8081/", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.SetBasicAuth("testuser", "wrongpass")

	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 401 {
		t.Fatalf("Status %d != 401", res.StatusCode)
	}

	// With correct username/password should succeed

	req, err = http.NewRequest("GET", "http://127.0.0.1:8081/", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.SetBasicAuth("testuser", "testpass")

	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 200 {
		t.Fatalf("Status %d != 200", res.StatusCode)
	}
}

func TestGetJSON(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	for _, path := range jsonEndpoints {
		res, err := http.Get("http://127.0.0.1:8082" + path)
		if err != nil {
			t.Error(path, err)
			continue
		}

		if ct := res.Header.Get("Content-Type"); ct != "application/json; charset=utf-8" {
			t.Errorf("Incorrect Content-Type %q for %q", ct, path)
			continue
		}

		var intf interface{}
		err = json.NewDecoder(res.Body).Decode(&intf)
		res.Body.Close()

		if err != nil {
			t.Error(path, err)
		}
	}
}

func TestPOSTWithoutCSRF(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	// Should fail without CSRF

	req, err := http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 403 {
		t.Fatalf("Status %d != 403 for POST", res.StatusCode)
	}

	// Get CSRF

	req, err = http.NewRequest("GET", "http://127.0.0.1:8082/", nil)
	if err != nil {
		t.Fatal(err)
	}
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	hdr := res.Header.Get("Set-Cookie")
	if !strings.Contains(hdr, "CSRF-Token") {
		t.Error("Missing CSRF-Token in", hdr)
	}

	// Should succeed with CSRF

	req, err = http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("X-CSRF-Token", hdr[len("CSRF-Token="):])
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 200 {
		t.Fatalf("Status %d != 200 for POST", res.StatusCode)
	}

	// Should fail with incorrect CSRF

	req, err = http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("X-CSRF-Token", hdr[len("CSRF-Token="):]+"X")
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 403 {
		t.Fatalf("Status %d != 403 for POST", res.StatusCode)
	}
}

func setupAPIBench() *rc.Process {
	err := removeAll("s1", "s2", "h1/index*", "h2/index*")
	if err != nil {
		panic(err)
	}

	err = generateFiles("s1", 25000, 20, "../LICENSE")
	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile("s1/knownfile", []byte("somedatahere"), 0644)
	if err != nil {
		panic(err)
	}

	// This will panic if there is an actual failure to start, when we try to
	// call nil.Fatal(...)
	return startInstance(nil, 1)
}

func benchmarkURL(b *testing.B, url string) {
	p := setupAPIBench()
	defer p.Stop()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := p.Get(url)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkAPI_db_completion(b *testing.B) {
	benchmarkURL(b, "/rest/db/completion?folder=default&device="+protocol.LocalDeviceID.String())
}

func BenchmarkAPI_db_file(b *testing.B) {
	benchmarkURL(b, "/rest/db/file?folder=default&file=knownfile")
}

func BenchmarkAPI_db_ignores(b *testing.B) {
	benchmarkURL(b, "/rest/db/ignores?folder=default")
}

func BenchmarkAPI_db_need(b *testing.B) {
	benchmarkURL(b, "/rest/db/need?folder=default")
}

func BenchmarkAPI_db_status(b *testing.B) {
	benchmarkURL(b, "/rest/db/status?folder=default")
}

func BenchmarkAPI_db_browse_dirsonly(b *testing.B) {
	benchmarkURL(b, "/rest/db/browse?folder=default&dirsonly=true")
}
<commit_msg>Fix CSRF tests (fixes #2009)<commit_after>// Copyright (C) 2014 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// +build integration

package integration

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"

	"github.com/syncthing/protocol"
	"github.com/syncthing/syncthing/internal/rc"
)

var jsonEndpoints = []string{
	"/rest/db/completion?device=I6KAH76-66SLLLB-5PFXSOA-UFJCDZC-YAOMLEK-CP2GB32-BV5RQST-3PSROAU&folder=default",
	"/rest/db/ignores?folder=default",
	"/rest/db/need?folder=default",
	"/rest/db/status?folder=default",
	"/rest/db/browse?folder=default",
	"/rest/events?since=-1&limit=5",
	"/rest/stats/device",
	"/rest/stats/folder",
	"/rest/svc/deviceid?id=I6KAH76-66SLLLB-5PFXSOA-UFJCDZC-YAOMLEK-CP2GB32-BV5RQST-3PSROAU",
	"/rest/svc/lang",
	"/rest/svc/report",
	"/rest/system/browse?current=.",
	"/rest/system/config",
	"/rest/system/config/insync",
	"/rest/system/connections",
	"/rest/system/discovery",
	"/rest/system/error",
	"/rest/system/ping",
	"/rest/system/status",
	"/rest/system/upgrade",
	"/rest/system/version",
}

func TestGetIndex(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	// Check for explicint index.html

	res, err := http.Get("http://localhost:8082/index.html")
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Errorf("Status %d != 200", res.StatusCode)
	}
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	if len(bs) < 1024 {
		t.Errorf("Length %d < 1024", len(bs))
	}
	if !bytes.Contains(bs, []byte("")) {
		t.Error("Incorrect response")
	}
	if res.Header.Get("Set-Cookie") == "" {
		t.Error("No set-cookie header")
	}
	res.Body.Close()

	// Check for implicit index.html

	res, err = http.Get("http://localhost:8082/")
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != 200 {
		t.Errorf("Status %d != 200", res.StatusCode)
	}
	bs, err = ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	if len(bs) < 1024 {
		t.Errorf("Length %d < 1024", len(bs))
	}
	if !bytes.Contains(bs, []byte("")) {
		t.Error("Incorrect response")
	}
	if res.Header.Get("Set-Cookie") == "" {
		t.Error("No set-cookie header")
	}
	res.Body.Close()
}

func TestGetIndexAuth(t *testing.T) {
	p := startInstance(t, 1)
	defer checkedStop(t, p)

	// Without auth should give 401

	res, err := http.Get("http://127.0.0.1:8081/")
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 401 {
		t.Errorf("Status %d != 401", res.StatusCode)
	}

	// With wrong username/password should give 401

	req, err := http.NewRequest("GET", "http://127.0.0.1:8081/", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.SetBasicAuth("testuser", "wrongpass")

	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 401 {
		t.Fatalf("Status %d != 401", res.StatusCode)
	}

	// With correct username/password should succeed

	req, err = http.NewRequest("GET", "http://127.0.0.1:8081/", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.SetBasicAuth("testuser", "testpass")

	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 200 {
		t.Fatalf("Status %d != 200", res.StatusCode)
	}
}

func TestGetJSON(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	for _, path := range jsonEndpoints {
		res, err := http.Get("http://127.0.0.1:8082" + path)
		if err != nil {
			t.Error(path, err)
			continue
		}

		if ct := res.Header.Get("Content-Type"); ct != "application/json; charset=utf-8" {
			t.Errorf("Incorrect Content-Type %q for %q", ct, path)
			continue
		}

		var intf interface{}
		err = json.NewDecoder(res.Body).Decode(&intf)
		res.Body.Close()

		if err != nil {
			t.Error(path, err)
		}
	}
}

func TestPOSTWithoutCSRF(t *testing.T) {
	p := startInstance(t, 2)
	defer checkedStop(t, p)

	// Should fail without CSRF

	req, err := http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 403 {
		t.Fatalf("Status %d != 403 for POST", res.StatusCode)
	}

	// Get CSRF

	req, err = http.NewRequest("GET", "http://127.0.0.1:8082/", nil)
	if err != nil {
		t.Fatal(err)
	}
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	hdr := res.Header.Get("Set-Cookie")
	id := res.Header.Get("X-Syncthing-ID")[:5]
	if !strings.Contains(hdr, "CSRF-Token") {
		t.Error("Missing CSRF-Token in", hdr)
	}

	// Should succeed with CSRF

	req, err = http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}

	req.Header.Set("X-CSRF-Token-"+id, hdr[len("CSRF-Token-"+id+"="):])
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 200 {
		t.Fatalf("Status %d != 200 for POST", res.StatusCode)
	}

	// Should fail with incorrect CSRF

	req, err = http.NewRequest("POST", "http://127.0.0.1:8082/rest/system/error/clear", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("X-CSRF-Token-"+id, hdr[len("CSRF-Token-"+id+"="):]+"X")
	res, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if res.StatusCode != 403 {
		t.Fatalf("Status %d != 403 for POST", res.StatusCode)
	}
}

func setupAPIBench() *rc.Process {
	err := removeAll("s1", "s2", "h1/index*", "h2/index*")
	if err != nil {
		panic(err)
	}

	err = generateFiles("s1", 25000, 20, "../LICENSE")
	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile("s1/knownfile", []byte("somedatahere"), 0644)
	if err != nil {
		panic(err)
	}

	// This will panic if there is an actual failure to start, when we try to
	// call nil.Fatal(...)
	return startInstance(nil, 1)
}

func benchmarkURL(b *testing.B, url string) {
	p := setupAPIBench()
	defer p.Stop()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := p.Get(url)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkAPI_db_completion(b *testing.B) {
	benchmarkURL(b, "/rest/db/completion?folder=default&device="+protocol.LocalDeviceID.String())
}

func BenchmarkAPI_db_file(b *testing.B) {
	benchmarkURL(b, "/rest/db/file?folder=default&file=knownfile")
}

func BenchmarkAPI_db_ignores(b *testing.B) {
	benchmarkURL(b, "/rest/db/ignores?folder=default")
}

func BenchmarkAPI_db_need(b *testing.B) {
	benchmarkURL(b, "/rest/db/need?folder=default")
}

func BenchmarkAPI_db_status(b *testing.B) {
	benchmarkURL(b, "/rest/db/status?folder=default")
}

func BenchmarkAPI_db_browse_dirsonly(b *testing.B) {
	benchmarkURL(b, "/rest/db/browse?folder=default&dirsonly=true")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
)

func main() {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Println(r)
		w.Write([]byte("fart"))
	})

	http.ListenAndServe(":8080", nil)
}
<commit_msg>Switched to type handler instead of func<commit_after>package main

import (
	"fmt"
	"net/http"
)

type karmas map[string]int

func main() {
	k := karmas(make(map[string]int))
	http.Handle("/", k)

	http.ListenAndServe(":8080", nil)
}

func (k karmas) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Println(r)
	w.Write([]byte("fart"))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"io"
	"log"
	"regexp"
	"strconv"
)

func colPos(slice []string, value string) int {
	for p, v := range slice {
		if v == value {
			return p
		}
	}

	return -1
}

func getColByName(name string, cols []string, vals []interface{}) *string {
	if cmdi := colPos(cols, name); cmdi != -1 {
		if bytes, ok := vals[cmdi].(*sql.RawBytes); ok {
			str := string(*bytes)
			return &str
		} else {
			panic("not raw bytes")
		}
	}

	return nil
}

type selectQuery string

type procEntry struct {
	Id    int64
	Time  int
	Query selectQuery
}

func (qry *selectQuery) c14n() string {
	out := string(*qry)

	out = regexp.MustCompile(`"(?:\\"|""|[^"])+"|'(?:\\'|''|[^'])+'`).ReplaceAllString(out, "[[string]]")

	// @todo negative numbers present interesting problems

	lastOut := out
	for { //solves a problem with sets like 10,20,30 when there are no lookaround options as in go
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\d+\.\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[float]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	lastOut = out
	for {
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[int]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	out = regexp.MustCompile(`\((?:\s*\[\[([a-z]+)\]\]\s*,?\s*)+\)`).ReplaceAllString(out, `[[$1-list]]`)

	return out
}

func (qry *selectQuery) csha1() string {
	h := sha1.New()
	io.WriteString(h, qry.c14n())
	return fmt.Sprintf("%x", h.Sum(nil))
}

type explainEntry struct {
	Table string
	Rows  int
}

func (qry *selectQuery) explain(db *sql.DB) ([]explainEntry, error) {
	output := make([]explainEntry, 0)

	rows, err := db.Query("EXPLAIN " + string(*qry))
	if err != nil {
		return output, fmt.Errorf("Explain Error, %s", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i, _ := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		tbl := getColByName("table", cols, vals)
		rows := getColByName("rows", cols, vals)

		rowInt, err := strconv.Atoi(*rows)
		if err != nil {
			rowInt = 0
		}

		output = append(output, explainEntry{
			Table: *tbl,
			Rows:  rowInt,
		})
	}

	return output, nil
}

func getActiveQueries(db *sql.DB) []procEntry {
	rows, err := db.Query("SHOW FULL PROCESSLIST")
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	output := make([]procEntry, 0)

	isSelect := regexp.MustCompile("(?i)^\\s*select\\s")

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i, _ := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		id := getColByName("Id", cols, vals)
		timez := getColByName("Time", cols, vals)
		cmd := getColByName("Command", cols, vals)
		info := getColByName("Info", cols, vals)

		if *cmd != "Query" || !isSelect.MatchString(*info) {
			continue
		}

		idInt, err := strconv.ParseInt(*id, 10, 64)
		if err != nil {
			idInt = int64(0)
		}

		timeInt, err := strconv.Atoi(*timez)
		if err != nil {
			panic(err)
		}

		output = append(output, procEntry{
			Id:    idInt,
			Time:  timeInt,
			Query: selectQuery(*info),
		})
	}

	return output
}
<commit_msg>Fixes negatives<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"io"
	"log"
	"regexp"
	"strconv"
)

func colPos(slice []string, value string) int {
	for p, v := range slice {
		if v == value {
			return p
		}
	}

	return -1
}

func getColByName(name string, cols []string, vals []interface{}) *string {
	if cmdi := colPos(cols, name); cmdi != -1 {
		if bytes, ok := vals[cmdi].(*sql.RawBytes); ok {
			str := string(*bytes)
			return &str
		} else {
			panic("not raw bytes")
		}
	}

	return nil
}

type selectQuery string

type procEntry struct {
	Id    int64
	Time  int
	Query selectQuery
}

func (qry *selectQuery) c14n() string {
	out := string(*qry)

	out = regexp.MustCompile(`"(?:\\"|""|[^"])+"|'(?:\\'|''|[^'])+'`).ReplaceAllString(out, "[[string]]")

	// @todo negative numbers present interesting problems

	lastOut := out
	for { //solves a problem with sets like 10,20,30 when there are no lookaround options as in go
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\d+\.\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[float]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	lastOut = out
	for {
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\-?\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[int]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	out = regexp.MustCompile(`\((?:\s*\[\[([a-z]+)\]\]\s*,?\s*)+\)`).ReplaceAllString(out, `[[$1-list]]`)

	return out
}

func (qry *selectQuery) csha1() string {
	h := sha1.New()
	io.WriteString(h, qry.c14n())
	return fmt.Sprintf("%x", h.Sum(nil))
}

type explainEntry struct {
	Table string
	Rows  int
}

func (qry *selectQuery) explain(db *sql.DB) ([]explainEntry, error) {
	output := make([]explainEntry, 0)

	rows, err := db.Query("EXPLAIN " + string(*qry))
	if err != nil {
		return output, fmt.Errorf("Explain Error, %s", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i, _ := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		tbl := getColByName("table", cols, vals)
		rows := getColByName("rows", cols, vals)

		rowInt, err := strconv.Atoi(*rows)
		if err != nil {
			rowInt = 0
		}

		output = append(output, explainEntry{
			Table: *tbl,
			Rows:  rowInt,
		})
	}

	return output, nil
}

func getActiveQueries(db *sql.DB) []procEntry {
	rows, err := db.Query("SHOW FULL PROCESSLIST")
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	output := make([]procEntry, 0)

	isSelect := regexp.MustCompile("(?i)^\\s*select\\s")

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i, _ := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		id := getColByName("Id", cols, vals)
		timez := getColByName("Time", cols, vals)
		cmd := getColByName("Command", cols, vals)
		info := getColByName("Info", cols, vals)

		if *cmd != "Query" || !isSelect.MatchString(*info) {
			continue
		}

		idInt, err := strconv.ParseInt(*id, 10, 64)
		if err != nil {
			idInt = int64(0)
		}

		timeInt, err := strconv.Atoi(*timez)
		if err != nil {
			panic(err)
		}

		output = append(output, procEntry{
			Id:    idInt,
			Time:  timeInt,
			Query: selectQuery(*info),
		})
	}

	return output
}
<|endoftext|>


<commit_before>// Package GobDB implements a persistant key-value store of
// gob-compatible types. This is accomplished with a light
// wrapper around leveldb and Go's gob encoding library.
package GobDB

import (
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/util"
	"strconv"
)

// DB is a LevelDB wrapper that stores key-value pairs of
// gob-compatible types.
type DB struct {
	internal *leveldb.DB
	location string
	encoder  FilteredEncoder
	decoder  Decoder
	prepared bool
}

// At returns an unopened database at with given datafile.
func At(path string) *DB {
	return &DB{location: path}
}

// Open sets up the internal leveldb if not done already.
func (db *DB) Open() error {
	if db.IsOpen() {
		return nil
	}

	ret, err := leveldb.OpenFile(db.location, nil)
	if err == nil {
		db.internal = ret
		db.prepare()
	}
	return err
}

// IsOpen checks whether or not the database is open.
func (db DB) IsOpen() bool {
	return db.internal != nil
}

// Close tears down the internal leveldb, writing all contents
// to file.
func (db *DB) Close() {
	if db.IsOpen() {
		db.internal.Close()
		db.internal = nil
	}
}

// Put encodes given key and value through gob, inserting resulting
// byte slices into the database's internal leveldb.
func (db *DB) Put(key, value interface{}) error {
	// Encode key via gob, registering types if necessary.
	o1, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form prefixed key.
	pkey := []byte("GobDB:key:")
	pkey = append(pkey, o1...)

	// Encode value via gob, registering types if necessary.
	val, err := db.encode(value)
	if err != nil {
		return err
	}

	// Insert gobbed values into leveldb.
	return db.internal.Put(pkey, val, nil)
}

// Get ncodes given key via gob, fetches the corresponding
// value from within leveldb, and decodes that value into
// parameter two.
func (db *DB) Get(key, value interface{}) error {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form prefixed key.
	pkey := []byte("GobDB:key:")
	pkey = append(pkey, obj...)

	// Fetch gob-encoded value.
	val, err := db.internal.Get(pkey, nil)
	if err != nil {
		return err
	}

	// Decode value into second paramater (which should be a
	// pointer)
	return db.decoder.Decode(val, value)
}

// Has encodes given key via gob and checks if the resulting
// byte slice exists in the database's internal leveldb.
func (db DB) Has(key interface{}) bool {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return false
	}

	// Note: this is niave - should check error type.
	_, err = db.internal.Get(obj, nil)
	return err == nil
}

// Delete encodes given key via gob, deleting the resulting
// byte slice from the database's internal leveldb.
func (db *DB) Delete(key interface{}) error {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form key bytes.
	kbytes := []byte("GobDB:key:")
	kbytes = append(kbytes, obj...)

	// Delete!
	return db.internal.Delete(kbytes, nil)
}

// Entries counts key-value pairs in the database. This
// includes only pairs written through GobDB.Put.
func (db *DB) Entries() int {
	i := 0
	iter := db.internal.NewIterator(util.BytesPrefix([]byte("GobDB:key:")), nil)
	for iter.Next() {
		i++
	}
	iter.Release()
	return i
}

// Reset erases caches and closes leveldb. This way, the db
// is forced to reload gobbed values as though it had just
// been opened for the first time.
func (db *DB) Reset() {
	db.Close()
	db.prepared = false
}

// Internal opens and fetches the underlying leveldb. Clients
// may use this to perform direct writing of byte slices, or
// to access leveldb APIs left out of this wrapper.
//
// Note: GobDB stores its mappings in the prefix "GobDB:", so
// that prefix should be avoided.
//
// Note: closing the parent GobDB will invalidate the returned
// value of this function.
func (db *DB) Internal() *leveldb.DB {
	err := db.Open()
	if err != nil {
		return nil
	}

	return db.internal
}

// Encodes given key via gob, registers its type if necessary,
// and routes any errors outward.
func (db *DB) encode(key interface{}) ([]byte, error) {
	err := db.Open()
	if err != nil {
		return []byte{}, err
	}

	// Gob encode key.
	def, obj, err := db.encoder.Encode(key)
	if err != nil {
		return []byte{}, err
	}

	// Register key with decoder.
	db.decoder.Register(append(def, obj...))

	// Register key type.
	err = db.registerType(def, obj)
	if err != nil {
		return []byte{}, err
	}

	// err = db.decoder.Decode(append(def, obj...), nil)
	// if err != nil {
	// 	return []byte{}, err
	// }

	return obj, nil
}

// When the database is opened for the first time, scrolls through
// all entries to form the same encoder that was used before the
// previous db.Close() call.
//
// Note that this is an expensive operation, scaling linearly with
// dataset size. Accordingly, you should utilize the open and close
// methods instead of initializing new DB objects all the time. For
// each initialization, the decoder is literally thrown in the
// garbage.
func (db *DB) prepare() error {
	if db.prepared == true {
		return nil
	}

	iter := db.internal.NewIterator(util.BytesPrefix([]byte("GobDB:prep#")), nil)
	for iter.Next() {
		value := iter.Value()
		db.encoder.Encode(value)
		db.decoder.Register(value)
	}
	iter.Release()
	err := iter.Error()
	if err == nil {
		db.prepared = true
	}
	return err
}

func (db *DB) setPrepCount(value int) error {
	key := []byte("GobDB:prep-count")
	data := []byte(strconv.Itoa(value))
	return db.internal.Put(key, data, nil)
}

func (db *DB) incPrepCount(i int) error {
	size := db.prepCount()
	if size == -1 {
		return db.setPrepCount(1)
	}
	return db.setPrepCount(size + i)
}

func (db *DB) prepCount() int {
	err := db.Open()
	if err != nil {
		return -1
	}

	val, err := db.internal.Get([]byte("GobDB:prep-count"), nil)
	if err != nil {
		return 0
	}
	n, _ := strconv.Atoi(string(val))
	return n
}

// Checks if definition is present in current db.
func (db *DB) isPresent(def []byte) bool {
	err := db.Open()
	if err != nil {
		return false
	}

	key := []byte("GobDB:prep:")
	key = append(key, def...)
	_, err = db.internal.Get(key, nil)
	return err == nil
}

// If not done already, registers type and example object in both
// the encoder and decoder.
func (db *DB) registerType(def, obj []byte) error {
	// Ensure that database is open.
	err := db.Open()
	if err != nil {
		return err
	}

	// Stop if type is already registered in database.
	if db.isPresent(def) {
		return nil
	}

	// Map prep#<n> to type definition bytes.
	k1 := []byte("GobDB:prep#" + strconv.Itoa(db.prepCount()))
	err = db.incPrepCount(1)
	if err != nil {
		return err
	}
	v1 := []byte{}
	v1 = append(v1, def...)
	v1 = append(v1, obj...)
	err = db.internal.Put(k1, v1, nil)
	if err != nil {
		return err
	}

	// Map prep:<def> to empty string (for checking duplicate keys).
	k2 := []byte("GobDB:prep:")
	k2 = append(k2, def...)
	v2 := []byte("")
	err = db.internal.Put(k2, v2, nil)
	if err != nil {
		db.internal.Delete(k1, nil)
		db.incPrepCount(-1)
		return err
	}
	return nil
}
<commit_msg>+ CompactRange wrapper<commit_after>// Package GobDB implements a persistant key-value store of
// gob-compatible types. This is accomplished with a light
// wrapper around leveldb and Go's gob encoding library.
package GobDB

import (
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/util"
	"strconv"
)

// DB is a LevelDB wrapper that stores key-value pairs of
// gob-compatible types.
type DB struct {
	internal *leveldb.DB
	location string
	encoder  FilteredEncoder
	decoder  Decoder
	prepared bool
}

// At returns an unopened database at with given datafile.
func At(path string) *DB {
	return &DB{location: path}
}

// Open sets up the internal leveldb if not done already.
func (db *DB) Open() error {
	if db.IsOpen() {
		return nil
	}

	ret, err := leveldb.OpenFile(db.location, nil)
	if err == nil {
		db.internal = ret
		db.prepare()
	}
	return err
}

// IsOpen checks whether or not the database is open.
func (db DB) IsOpen() bool {
	return db.internal != nil
}

// Close tears down the internal leveldb, writing all contents
// to file.
func (db *DB) Close() {
	if db.IsOpen() {
		db.internal.Close()
		db.internal = nil
	}
}

// Put encodes given key and value through gob, inserting resulting
// byte slices into the database's internal leveldb.
func (db *DB) Put(key, value interface{}) error {
	// Encode key via gob, registering types if necessary.
	o1, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form prefixed key.
	pkey := []byte("GobDB:key:")
	pkey = append(pkey, o1...)

	// Encode value via gob, registering types if necessary.
	val, err := db.encode(value)
	if err != nil {
		return err
	}

	// Insert gobbed values into leveldb.
	return db.internal.Put(pkey, val, nil)
}

// Get ncodes given key via gob, fetches the corresponding
// value from within leveldb, and decodes that value into
// parameter two.
func (db *DB) Get(key, value interface{}) error {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form prefixed key.
	pkey := []byte("GobDB:key:")
	pkey = append(pkey, obj...)

	// Fetch gob-encoded value.
	val, err := db.internal.Get(pkey, nil)
	if err != nil {
		return err
	}

	// Decode value into second paramater (which should be a
	// pointer)
	return db.decoder.Decode(val, value)
}

// Has encodes given key via gob and checks if the resulting
// byte slice exists in the database's internal leveldb.
func (db DB) Has(key interface{}) bool {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return false
	}

	// Note: this is niave - should check error type.
	_, err = db.internal.Get(obj, nil)
	return err == nil
}

// Delete encodes given key via gob, deleting the resulting
// byte slice from the database's internal leveldb.
func (db *DB) Delete(key interface{}) error {
	// Encode key via gob, registering its type if necessary.
	obj, err := db.encode(key)
	if err != nil {
		return err
	}

	// Form key bytes.
	kbytes := []byte("GobDB:key:")
	kbytes = append(kbytes, obj...)

	// Delete!
	return db.internal.Delete(kbytes, nil)
}

// Entries counts key-value pairs in the database. This
// includes only pairs written through GobDB.Put.
func (db *DB) Entries() int {
	i := 0
	iter := db.internal.NewIterator(util.BytesPrefix([]byte("GobDB:key:")), nil)
	for iter.Next() {
		i++
	}
	iter.Release()
	return i
}

// Reset erases caches and closes leveldb. This way, the db
// is forced to reload gobbed values as though it had just
// been opened for the first time.
func (db *DB) Reset() {
	db.Close()
	db.prepared = false
}

// Internal opens and fetches the underlying leveldb. Clients
// may use this to perform direct writing of byte slices, or
// to access leveldb APIs left out of this wrapper.
//
// Note: GobDB stores its mappings in the prefix "GobDB:", so
// that prefix should be avoided.
//
// Note: closing the parent GobDB will invalidate the returned
// value of this function.
func (db *DB) Internal() *leveldb.DB {
	err := db.Open()
	if err != nil {
		return nil
	}

	return db.internal
}

// Compact performs leveldb.CompactRange on the range of
// key-value mappings maintained by GobDB. Pairs outside
// the 'GobDB:' prefix aren't affected.
func (db *DB) Compact() error {
	err := db.Open()
	if err != nil {
		return err
	}

	prefix := util.BytesPrefix([]byte("GobDB:"))
	return db.internal.CompactRange(*prefix)
}

// Encodes given key via gob, registers its type if necessary,
// and routes any errors outward.
func (db *DB) encode(key interface{}) ([]byte, error) {
	err := db.Open()
	if err != nil {
		return []byte{}, err
	}

	// Gob encode key.
	def, obj, err := db.encoder.Encode(key)
	if err != nil {
		return []byte{}, err
	}

	// Register key with decoder.
	db.decoder.Register(append(def, obj...))

	// Register key type.
	err = db.registerType(def, obj)
	if err != nil {
		return []byte{}, err
	}

	// err = db.decoder.Decode(append(def, obj...), nil)
	// if err != nil {
	// 	return []byte{}, err
	// }

	return obj, nil
}

// When the database is opened for the first time, scrolls through
// all entries to form the same encoder that was used before the
// previous db.Close() call.
//
// Note that this is an expensive operation, scaling linearly with
// dataset size. Accordingly, you should utilize the open and close
// methods instead of initializing new DB objects all the time. For
// each initialization, the decoder is literally thrown in the
// garbage.
func (db *DB) prepare() error {
	if db.prepared == true {
		return nil
	}

	iter := db.internal.NewIterator(util.BytesPrefix([]byte("GobDB:prep#")), nil)
	for iter.Next() {
		value := iter.Value()
		db.encoder.Encode(value)
		db.decoder.Register(value)
	}
	iter.Release()
	err := iter.Error()
	if err == nil {
		db.prepared = true
	}
	return err
}

func (db *DB) setPrepCount(value int) error {
	key := []byte("GobDB:prep-count")
	data := []byte(strconv.Itoa(value))
	return db.internal.Put(key, data, nil)
}

func (db *DB) incPrepCount(i int) error {
	size := db.prepCount()
	if size == -1 {
		return db.setPrepCount(1)
	}
	return db.setPrepCount(size + i)
}

func (db *DB) prepCount() int {
	err := db.Open()
	if err != nil {
		return -1
	}

	val, err := db.internal.Get([]byte("GobDB:prep-count"), nil)
	if err != nil {
		return 0
	}
	n, _ := strconv.Atoi(string(val))
	return n
}

// Checks if definition is present in current db.
func (db *DB) isPresent(def []byte) bool {
	err := db.Open()
	if err != nil {
		return false
	}

	key := []byte("GobDB:prep:")
	key = append(key, def...)
	_, err = db.internal.Get(key, nil)
	return err == nil
}

// If not done already, registers type and example object in both
// the encoder and decoder.
func (db *DB) registerType(def, obj []byte) error {
	// Ensure that database is open.
	err := db.Open()
	if err != nil {
		return err
	}

	// Stop if type is already registered in database.
	if db.isPresent(def) {
		return nil
	}

	// Map prep#<n> to type definition bytes.
	k1 := []byte("GobDB:prep#" + strconv.Itoa(db.prepCount()))
	err = db.incPrepCount(1)
	if err != nil {
		return err
	}
	v1 := []byte{}
	v1 = append(v1, def...)
	v1 = append(v1, obj...)
	err = db.internal.Put(k1, v1, nil)
	if err != nil {
		return err
	}

	// Map prep:<def> to empty string (for checking duplicate keys).
	k2 := []byte("GobDB:prep:")
	k2 = append(k2, def...)
	v2 := []byte("")
	err = db.internal.Put(k2, v2, nil)
	if err != nil {
		db.internal.Delete(k1, nil)
		db.incPrepCount(-1)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package sqlgen2

import (
	"github.com/rickb777/sqlgen2/schema"
	"github.com/rickb777/sqlgen2/util"
	"log"
	"context"
)

type Database struct {
	db      Execer
	dialect schema.Dialect
	ctx     context.Context
	logger  *log.Logger
	wrapper interface{}
}

// NewDatabase createa a new database handler, which wraps the core *sql.DB.
func NewDatabase(db Execer, dialect schema.Dialect) *Database {
	return &Database{db, dialect, context.Background(), nil, nil}
}

func (database *Database) DB() Execer {
	return database.db
}

// Wrapper gets whatever structure is present, as needed.
func (database *Database) Dialect() schema.Dialect {
	return database.dialect
}

// SetContext sets the context for subsequent queries.
func (database *Database) SetContext(ctx context.Context) *Database {
	database.ctx = ctx
	return database
}

// Logger gets the trace logger.
func (database *Database) Logger() *log.Logger {
	return database.logger
}

// SetLogger sets the logger for subsequent queries, returning the interface.
func (database *Database) SetLogger(logger *log.Logger) *Database {
	database.logger = logger
	return database
}

// Wrapper gets whatever structure is present, as needed.
func (database *Database) Wrapper() interface{} {
	return database.wrapper
}

// SetWrapper sets a user-defined wrapper or container.
func (database *Database) SetWrapper(wrapper interface{}) *Database {
	database.wrapper = wrapper
	return database
}

// DoesTableExist gets all the table names in the database/schema.
func (database *Database) TableExists(name TableName) (yes bool, err error) {
	wanted := name.String()
	rows, err := database.db.QueryContext(database.ctx, showTables(database.dialect))
	if err != nil {
		return false, err
	}
	defer rows.Close()

	for rows.Next() {
		var s string
		rows.Scan(&s)
		if s == wanted {
			return true, rows.Err()
		}
	}
	return false, rows.Err()
}

// ListTables gets all the table names in the database/schema.
func (database *Database) ListTables(dialect schema.Dialect) (util.StringList, error) {
	ss := make(util.StringList, 0)
	rows, err := database.db.QueryContext(database.ctx, showTables(dialect))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var s string
		rows.Scan(&s)
		ss = append(ss, s)
	}
	return ss, rows.Err()
}

func showTables(dialect schema.Dialect) string {
	switch dialect.Index() {
	case schema.SqliteIndex:
		return `SELECT name FROM sqlite_master WHERE type = "table"`
	case schema.MysqlIndex:
		return `SHOW TABLES`
	case schema.PostgresIndex:
		return `SELECT tablename FROM pg_catalog.pg_tables`
	}
	panic(dialect.String())
}
<commit_msg>Database is now an Execer<commit_after>package sqlgen2

import (
	"github.com/rickb777/sqlgen2/schema"
	"github.com/rickb777/sqlgen2/util"
	"log"
	"context"
	"database/sql"
)

type Database struct {
	db      Execer
	dialect schema.Dialect
	ctx     context.Context
	logger  *log.Logger
	wrapper interface{}
}

// NewDatabase createa a new database handler, which wraps the core *sql.DB.
func NewDatabase(db Execer, dialect schema.Dialect) *Database {
	return &Database{db, dialect, context.Background(), nil, nil}
}

func (database *Database) DB() Execer {
	return database.db
}

// Wrapper gets whatever structure is present, as needed.
func (database *Database) Dialect() schema.Dialect {
	return database.dialect
}

// SetContext sets the context for subsequent queries.
func (database *Database) SetContext(ctx context.Context) *Database {
	database.ctx = ctx
	return database
}

// Logger gets the trace logger.
func (database *Database) Logger() *log.Logger {
	return database.logger
}

// SetLogger sets the logger for subsequent queries, returning the interface.
func (database *Database) SetLogger(logger *log.Logger) *Database {
	database.logger = logger
	return database
}

// Wrapper gets whatever structure is present, as needed.
func (database *Database) Wrapper() interface{} {
	return database.wrapper
}

// SetWrapper sets a user-defined wrapper or container.
func (database *Database) SetWrapper(wrapper interface{}) *Database {
	database.wrapper = wrapper
	return database
}

//-------------------------------------------------------------------------------------------------

func (database *Database) Exec(query string, args ...interface{}) (sql.Result, error) {
	return database.db.ExecContext(database.ctx, query, args...)
}

func (database *Database) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	return database.db.ExecContext(ctx, query, args...)
}

func (database *Database) Prepare(query string) (*sql.Stmt, error) {
	return database.db.PrepareContext(database.ctx, query)
}

func (database *Database) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
	return database.db.PrepareContext(ctx, query)
}

func (database *Database) Query(query string, args ...interface{}) (*sql.Rows, error) {
	return database.db.QueryContext(database.ctx, query, args...)
}

func (database *Database) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	return database.db.QueryContext(ctx, query, args...)
}

func (database *Database) QueryRow(query string, args ...interface{}) *sql.Row {
	return database.db.QueryRowContext(database.ctx, query, args...)
}

func (database *Database) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
	return database.db.QueryRowContext(ctx, query, args...)
}

//-------------------------------------------------------------------------------------------------

// DoesTableExist gets all the table names in the database/schema.
func (database *Database) TableExists(name TableName) (yes bool, err error) {
	wanted := name.String()
	rows, err := database.db.QueryContext(database.ctx, showTables(database.dialect))
	if err != nil {
		return false, err
	}
	defer rows.Close()

	for rows.Next() {
		var s string
		rows.Scan(&s)
		if s == wanted {
			return true, rows.Err()
		}
	}
	return false, rows.Err()
}

// ListTables gets all the table names in the database/schema.
func (database *Database) ListTables(dialect schema.Dialect) (util.StringList, error) {
	ss := make(util.StringList, 0)
	rows, err := database.db.QueryContext(database.ctx, showTables(dialect))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var s string
		rows.Scan(&s)
		ss = append(ss, s)
	}
	return ss, rows.Err()
}

func showTables(dialect schema.Dialect) string {
	switch dialect.Index() {
	case schema.SqliteIndex:
		return `SELECT name FROM sqlite_master WHERE type = "table"`
	case schema.MysqlIndex:
		return `SHOW TABLES`
	case schema.PostgresIndex:
		return `SELECT tablename FROM pg_catalog.pg_tables`
	}
	panic(dialect.String())
}
<|endoftext|>