<commit_before>package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"url"
	"xml"
)

const (
	version   = "2.1"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 1, "pages to prime at once")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show information about primed pages")
)

type Sitemap struct {
	Loc string
	// Lastmod string
}

type Url struct {
	Loc string
	// Lastmod string
	// Changefreq string
	Priority float64
}

type Urlset struct {
	XMLName xml.Name "urlset"
	Sitemap []Sitemap
	Url     []Url
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (r *http.Response, err os.Error) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("User-Agent", useragent)
	if err != nil {
                return nil, err
	}
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, os.Error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    os.Error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		res, err = Get(path)
		if res.Status == "404 Not Found" {
			return nil, os.NewError("Web server returned 404 Not Found")
		} else if res.Status != "200 OK" {
			return nil, os.NewError("Web server did not serve the file")
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, os.NewError("Gzip decompression failed")
		}
	}
	err = xml.Unmarshal(f, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan []Url, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex. Fetching Urlsets from sitemaps...\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding Urlset from sitemap %s\n", v.Loc)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
				} else {
					ch <- ourlset.Url
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			urlset.Url = append(urlset.Url, childUrlset...)
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		log.Println("URLs in sitemap: ", urlset.Len())
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) os.Error {
	var (
		err   os.Error
		found bool = false
	)
	if *verbose {
		log.Printf("Get (weight %d) %s\n", int(u.Priority*100), u.Loc)
	}
	if *localDir != "" {
		parsed, err := url.ParseWithReference(u.Loc)
		joined := path.Join(*localDir, parsed.Path, *localSuffix)
		if _, err = os.Lstat(joined); err == nil {
			found = true
		}
	}
	if !found {
		_, err = Get(u.Loc)
	}
	wg.Done()
	<-sem
	return err
}

func main() {
	var (
		urlset *Urlset
		err    os.Error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		PrimeUrlset(urlset)
	}
}
<commit_msg>More aggressive default concurrency level (5 pages at once); toggleable reporting of dead/broken links and server errors; more clarity about child sitemaps internally<commit_after>package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"url"
	"xml"
)

const (
	version   = "2.1"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 5, "pages to prime at once")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show additional information about the priming process")
	nowarn      *bool   = flag.Bool("no-warn", false, "do not warn about pages that can't be loaded")
)

type Sitemap struct {
	Loc string
	// Lastmod string
}

type Url struct {
	Loc string
	// Lastmod string
	// Changefreq string
	Priority float64
}

type Urlset struct {
	XMLName xml.Name "urlset"
	Sitemap []Sitemap
	Url     []Url
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (r *http.Response, err os.Error) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("User-Agent", useragent)
	if err != nil {
		return nil, err
	}
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, os.Error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    os.Error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		if *verbose {
			log.Println("Downloading", path)
		}
		res, err = Get(path)
		if res.Status != "200 OK" {
			return nil, os.NewError("Could not fetch sitemap: " + res.Status)
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		if *verbose {
			log.Println("Extracting compressed data")
		}
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, os.NewError("Gzip decompression failed")
		}
	}
	err = xml.Unmarshal(f, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan *Urlset, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding URLs from child sitemap %s\n", v)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
				} else {
					ch <- ourlset
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			urlset.Url = append(urlset.Url, childUrlset.Url...)
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		log.Println("URLs to prime:", urlset.Len())
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) os.Error {
	var (
		err   os.Error
		found bool = false
	)
	if *verbose {
		log.Printf("Get (weight %d) %s\n", int(u.Priority*100), u.Loc)
	}
	if *localDir != "" {
		parsed, err := url.ParseWithReference(u.Loc)
		joined := path.Join(*localDir, parsed.Path, *localSuffix)
		if _, err = os.Lstat(joined); err == nil {
			found = true
		}
	}
	if !found {
		res, err := Get(u.Loc)
		if (err != nil || res.Status != "200 OK") && !*nowarn {
			var errmsg string
			if err != nil {
				errmsg = err.String()
			} else {
				errmsg = res.Status
			}
			log.Printf("Error priming %s: %s\n", u.Loc, errmsg)
		}
	}
	wg.Done()
	<-sem
	return err
}

func main() {
	var (
		urlset *Urlset
		err    os.Error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		PrimeUrlset(urlset)
	}
}
<|endoftext|>


<commit_before>package watchdog

import (
	"sync"
	"time"
)

// Basic scheduling unit
type Task struct {
	// How frequently the task should execute
	Schedule time.Duration
	// Function to invoke: each execution will be passed the time
	// it was originally scheduled for (which may be behind
	// wall-clock time in case of stalls).
	Command func(time.Time) error
	// How long to wait before considering an execution stalled
	Timeout time.Duration
}

// Information about each execution
type Execution struct {
	// Task being executed
	Task *Task
	// Time the Task was originally scheduled for
	StartedAt time.Time
	// Time the Task completed
	FinishedAt time.Time
	// Error returned by the Task Command
	Error error
}

// Information about each stall
type Stall struct {
	// Task which stalled
	Task *Task
	// Time the Task was originally scheduled for
	StartedAt time.Time
	// Time the task was considered stalled
	StalledAt time.Time
}

// Execution monitor
type Watchdog struct {
	tasks []*Task

	done chan bool
	sync sync.WaitGroup

	executions chan *Execution
	stalls     chan *Stall
}

// Create a new, running watchdog with the given task(s).
func Watch(tasks ...*Task) *Watchdog {
	w := &Watchdog{
		tasks:      tasks,
		done:       make(chan bool),
		executions: make(chan *Execution, 10),
		stalls:     make(chan *Stall, 10),
	}
	go w.run()
	return w
}

// Channel of executions for a given Watchdog. Note that the channel
// must be drained promptly while the Watchdog is running: the channel
// has a small buffer, but failure to keep up will apply backpressure
// to the system and delay scheduling.
func (w *Watchdog) Executions() <-chan *Execution {
	return w.executions
}

// Channel of stalls for a given Watchdog. As above, the channel must
// be drained while the Watchdog is running.
func (w *Watchdog) Stalls() <-chan *Stall {
	return w.stalls
}

func (w *Watchdog) run() {
	for _, task := range w.tasks {
		go w.runTask(task)
	}
	w.sync.Add(len(w.tasks))
}

func (w *Watchdog) runTask(task *Task) {
	ticker := time.NewTicker(task.Schedule)
	schedule := make(chan time.Time, 1)
	stallTimer := time.NewTimer(task.Schedule + 1*time.Millisecond)
	stallTimer.Stop()
	taskDone := make(chan bool, 1)
	go func() {
		for startedAt := range schedule {
			stallTimer.Reset(task.Timeout)
			err := task.Command(startedAt)
			stallTimer.Reset(task.Schedule)
			finishedAt := time.Now()
			select {
			case w.executions <- &Execution{task, startedAt, finishedAt, err}:
			default:
			}
		}
		taskDone <- true
	}()
	var startedAt time.Time
monitor:
	for {
		select {
		case <-w.done:
			ticker.Stop()
			close(schedule)
			break monitor
		case stalledAt := <-stallTimer.C:
			w.stalls <- &Stall{task, startedAt, stalledAt}
		case startedAt = <-ticker.C:
			select {
			case schedule <- startedAt:
			default:
			}
		}
	}
cleanup:
	for {
		select {
		case stalledAt := <-stallTimer.C:
			w.stalls <- &Stall{task, startedAt, stalledAt}
		case <-taskDone:
			stallTimer.Stop()
			break cleanup
		}
	}
	w.sync.Done()
}

// Stop a running Watchdog. Waits for any currently-executing tasks to
// complete, then closes the Executions and Stalls channels and
// returns.
func (w *Watchdog) Stop() {
	close(w.done)
	w.sync.Wait()
	close(w.executions)
	close(w.stalls)
}
<commit_msg>Use blocking send to executions, as documented<commit_after>package watchdog

import (
	"sync"
	"time"
)

// Basic scheduling unit
type Task struct {
	// How frequently the task should execute
	Schedule time.Duration
	// Function to invoke: each execution will be passed the time
	// it was originally scheduled for (which may be behind
	// wall-clock time in case of stalls).
	Command func(time.Time) error
	// How long to wait before considering an execution stalled
	Timeout time.Duration
}

// Information about each execution
type Execution struct {
	// Task being executed
	Task *Task
	// Time the Task was originally scheduled for
	StartedAt time.Time
	// Time the Task completed
	FinishedAt time.Time
	// Error returned by the Task Command
	Error error
}

// Information about each stall
type Stall struct {
	// Task which stalled
	Task *Task
	// Time the Task was originally scheduled for
	StartedAt time.Time
	// Time the task was considered stalled
	StalledAt time.Time
}

// Execution monitor
type Watchdog struct {
	tasks []*Task

	done chan bool
	sync sync.WaitGroup

	executions chan *Execution
	stalls     chan *Stall
}

// Create a new, running watchdog with the given task(s).
func Watch(tasks ...*Task) *Watchdog {
	w := &Watchdog{
		tasks:      tasks,
		done:       make(chan bool),
		executions: make(chan *Execution, 10),
		stalls:     make(chan *Stall, 10),
	}
	go w.run()
	return w
}

// Channel of executions for a given Watchdog. Note that the channel
// must be drained promptly while the Watchdog is running: the channel
// has a small buffer, but failure to keep up will apply backpressure
// to the system and delay scheduling.
func (w *Watchdog) Executions() <-chan *Execution {
	return w.executions
}

// Channel of stalls for a given Watchdog. As above, the channel must
// be drained while the Watchdog is running.
func (w *Watchdog) Stalls() <-chan *Stall {
	return w.stalls
}

func (w *Watchdog) run() {
	for _, task := range w.tasks {
		go w.runTask(task)
	}
	w.sync.Add(len(w.tasks))
}

func (w *Watchdog) runTask(task *Task) {
	ticker := time.NewTicker(task.Schedule)
	schedule := make(chan time.Time, 1)
	stallTimer := time.NewTimer(task.Schedule + 1*time.Millisecond)
	stallTimer.Stop()
	taskDone := make(chan bool, 1)
	go func() {
		for startedAt := range schedule {
			stallTimer.Reset(task.Timeout)
			err := task.Command(startedAt)
			stallTimer.Reset(task.Schedule)
			finishedAt := time.Now()
			w.executions <- &Execution{task, startedAt, finishedAt, err}
		}
		taskDone <- true
	}()
	var startedAt time.Time
monitor:
	for {
		select {
		case <-w.done:
			ticker.Stop()
			close(schedule)
			break monitor
		case stalledAt := <-stallTimer.C:
			w.stalls <- &Stall{task, startedAt, stalledAt}
		case startedAt = <-ticker.C:
			select {
			case schedule <- startedAt:
			default:
			}
		}
	}
cleanup:
	for {
		select {
		case stalledAt := <-stallTimer.C:
			w.stalls <- &Stall{task, startedAt, stalledAt}
		case <-taskDone:
			stallTimer.Stop()
			break cleanup
		}
	}
	w.sync.Done()
}

// Stop a running Watchdog. Waits for any currently-executing tasks to
// complete, then closes the Executions and Stalls channels and
// returns.
func (w *Watchdog) Stop() {
	close(w.done)
	w.sync.Wait()
	close(w.executions)
	close(w.stalls)
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
)

// EvalImportState is an EvalNode implementation that performs an
// ImportState operation on a provider. This will return the imported
// states but won't modify any actual state.
type EvalImportState struct {
	Provider *ResourceProvider
	Info     *InstanceInfo
	Output   *[]*InstanceState
}

// TODO: test
func (n *EvalImportState) Eval(ctx EvalContext) (interface{}, error) {
	provider := *n.Provider

	// Refresh!
	state, err := provider.ImportState(n.Info)
	if err != nil {
		return nil, fmt.Errorf(
			"import %s (id: %s): %s", n.Info.Type, n.Info.Id, err)
	}

	if n.Output != nil {
		*n.Output = state
	}

	return nil, nil
}
<commit_msg>terraform: call pre/post import hooks<commit_after>package terraform

import (
	"fmt"
)

// EvalImportState is an EvalNode implementation that performs an
// ImportState operation on a provider. This will return the imported
// states but won't modify any actual state.
type EvalImportState struct {
	Provider *ResourceProvider
	Info     *InstanceInfo
	Output   *[]*InstanceState
}

// TODO: test
func (n *EvalImportState) Eval(ctx EvalContext) (interface{}, error) {
	provider := *n.Provider

	{
		// Call pre-import hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PreImportState(n.Info)
		})
		if err != nil {
			return nil, err
		}
	}

	// Import!
	state, err := provider.ImportState(n.Info)
	if err != nil {
		return nil, fmt.Errorf(
			"import %s (id: %s): %s", n.Info.Type, n.Info.Id, err)
	}

	if n.Output != nil {
		*n.Output = state
	}

	{
		// Call post-import hook
		err := ctx.Hook(func(h Hook) (HookAction, error) {
			return h.PostImportState(n.Info, state)
		})
		if err != nil {
			return nil, err
		}
	}

	return nil, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/SlyMarbo/spdy"
	"io"
	"net/http"
	"strconv"
	"strings"
)

const HOST_PORT_API = "localhost:1443"
const HOST_PORT_SERVERS = "localhost:1444"

// responseCopier does the copying of the request
// from H to C and the response from C to H.
type responseCopier struct {
	w http.ResponseWriter
	s spdy.Stream
}

func (r *responseCopier) ReceiveData(_ *http.Request, data []byte, final bool) {
	if data == nil || len(data) == 0 {
		return
	}
	_, err := r.w.Write(data)
	if err != nil {
		r.s.Close()
	}
}

func (r *responseCopier) ReceiveHeader(_ *http.Request, header http.Header) {
	h := r.w.Header()
	for key, values := range header {
		for _, value := range values {
			h.Add(key, value)
			if key == ":status" {
				if i := strings.Index(value, " "); i > 0 {
					value = value[:i]
				}
				status, err := strconv.Atoi(value)
				if err != nil {
					fmt.Printf("Warning: Failed to parse status code %q.\n", value)
					continue
				}
				r.w.WriteHeader(status)
			}
		}
	}
}

func (r *responseCopier) ReceiveRequest(_ *http.Request) bool {
	return false
}

// Used in sending the response.
// Essentially, this is just adding
// the Close method so that it fulfils
// the io.ReadCloser interface.
type buffer struct {
	bytes.Buffer
}

func (b *buffer) Close() error {
	return nil
}

// placeholder for proper error handling.
func handle(err error) {
	if err != nil {
		panic(err)
	}
}

type Proxy struct {
	conn spdy.Conn
}

func (p *Proxy) RequestFromC(w http.ResponseWriter, r *http.Request) error {
	if p.conn == nil {
		fmt.Println("Warning: Could not serve request because C is not connected.")
		http.NotFound(w, r)
		return nil
	}

	copier := new(responseCopier)
	copier.w = w
	u := r.URL
	if u.Host == "" {
		u.Host = HOST_PORT_API
	}
	if u.Scheme == "" {
		u.Scheme = "https"
	}
	stream, err := p.conn.Request(r, copier, spdy.DefaultPriority(r.URL))
	if err != nil {
		return err
	}
	copier.s = stream
	return stream.Run()
}

func (p *Proxy) ServeC(w http.ResponseWriter, r *http.Request) {
	// clean up the old connection
	if p.conn != nil {
		handle(p.conn.Close())
	}

	// Read in the request body.
	buf := new(buffer)
	_, err := io.Copy(buf, r.Body)
	handle(err)
	handle(r.Body.Close())
	fmt.Printf("%q from C: %q.\n", r.Method, buf.String())

	// re-purpose the connection.
	conn, _, err := w.(http.Hijacker).Hijack()
	handle(err)

	// send the 200 to C.
	res := new(http.Response)
	res.Status = "200 Connection Established"
	res.StatusCode = 200
	res.Proto = "HTTP/1.1"
	res.ProtoMajor = 1
	res.ProtoMinor = 1
	buf.Reset()
	message := "Hello from P"
	buf.WriteString(message)
	res.Body = buf
	res.ContentLength = int64(len(message))
	handle(res.Write(conn))

	// prepare for serving requests from H.
	client, err := spdy.NewClientConn(conn, nil, 3)
	handle(err)
	p.conn = client
	fmt.Println("Ready")
	client.Run()
}

func (p *Proxy) ServeH(w http.ResponseWriter, r *http.Request) {
	handle(p.RequestFromC(w, r))
}

func main() {
	certFile := "cert.pem"
	keyFile := "cert.key"

	proxy := new(Proxy)
	http.HandleFunc("/", proxy.ServeC)
	go http.ListenAndServeTLS(HOST_PORT_SERVERS, certFile, keyFile, nil) // Serve C

	hServe := new(http.Server)
	mux := http.NewServeMux()
	mux.HandleFunc("/", proxy.ServeH)
	hServe.Handler = mux
	hServe.Addr = HOST_PORT_API
	spdy.AddSPDY(hServe)
	handle(hServe.ListenAndServeTLS(certFile, keyFile)) // Serve H
}
<commit_msg>Update p.go<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/SlyMarbo/spdy"
	"io"
	"net/http"
	"strconv"
	"strings"
)

const HOST_PORT_API = "localhost:1443"
const HOST_PORT_SERVERS = "localhost:1444"

// responseCopier does the copying of the request
// from H to C and the response from C to H.
type responseCopier struct {
	w http.ResponseWriter
	s spdy.Stream
}

func (r *responseCopier) ReceiveData(_ *http.Request, data []byte, final bool) {
	if data == nil || len(data) == 0 {
		return
	}
	_, err := r.w.Write(data)
	if err != nil {
		r.s.Close()
	}
}

func (r *responseCopier) ReceiveHeader(_ *http.Request, header http.Header) {
	h := r.w.Header()
	for key, values := range header {
		for _, value := range values {
			h.Add(key, value)
			if key == ":status" {
				if i := strings.Index(value, " "); i > 0 {
					value = value[:i]
				}
				status, err := strconv.Atoi(value)
				if err != nil {
					fmt.Printf("Warning: Failed to parse status code %q.\n", value)
					continue
				}
				r.w.WriteHeader(status)
			}
		}
	}
}

func (r *responseCopier) ReceiveRequest(_ *http.Request) bool {
	return false
}

// Used in sending the response.
// Essentially, this is just adding
// the Close method so that it fulfils
// the io.ReadCloser interface.
type buffer struct {
	bytes.Buffer
}

func (b *buffer) Close() error {
	return nil
}

// placeholder for proper error handling.
func handle(err error) {
	if err != nil {
		panic(err)
	}
}

type Proxy struct {
	conn spdy.Conn
}

func (p *Proxy) RequestFromC(w http.ResponseWriter, r *http.Request) error {
	if p.conn == nil {
		fmt.Println("Warning: Could not serve request because C is not connected.")
		http.NotFound(w, r)
		return nil
	}

	copier := new(responseCopier)
	copier.w = w
	u := r.URL
	if u.Host == "" {
		u.Host = HOST_PORT_API
	}
	if u.Scheme == "" {
		u.Scheme = "https"
	}
	stream, err := p.conn.Request(r, copier, spdy.DefaultPriority(r.URL))
	if err != nil {
		return err
	}
	copier.s = stream
	return stream.Run()
}

func (p *Proxy) ServeC(w http.ResponseWriter, r *http.Request) {
	// clean up the old connection
	if p.conn != nil {
		handle(p.conn.Close())
	}

	// Read in the request body.
	buf := new(buffer)
	_, err := io.Copy(buf, r.Body)
	handle(err)
	handle(r.Body.Close())
	fmt.Printf("%q from C: %q.\n", r.Method, buf.String())

	// re-purpose the connection.
	conn, _, err := w.(http.Hijacker).Hijack()
	handle(err)

	// send the 200 to C.
	res := new(http.Response)
	res.Status = "200 Connection Established"
	res.StatusCode = 200
	res.Proto = "HTTP/1.1"
	res.ProtoMajor = 1
	res.ProtoMinor = 1
	buf.Reset()
	message := "Hello from P"
	buf.WriteString(message)
	res.Body = buf
	res.ContentLength = int64(len(message))
	handle(res.Write(conn))

	// prepare for serving requests from H.
	client, err := spdy.NewClientConn(conn, nil, 3)
	handle(err)
	p.conn = client
	fmt.Println("Ready")
	client.Run()
}

func (p *Proxy) ServeH(w http.ResponseWriter, r *http.Request) {
	err := p.RequestFromC(w, r)
	if err != nil {
		fmt.Println("Encountered an error serving API request:", err)
	}
}

func main() {
	certFile := "cert.pem"
	keyFile := "cert.key"

	proxy := new(Proxy)
	http.HandleFunc("/", proxy.ServeC)
	go http.ListenAndServeTLS(HOST_PORT_SERVERS, certFile, keyFile, nil) // Serve C

	hServe := new(http.Server)
	mux := http.NewServeMux()
	mux.HandleFunc("/", proxy.ServeH)
	hServe.Handler = mux
	hServe.Addr = HOST_PORT_API
	spdy.AddSPDY(hServe)
	handle(hServe.ListenAndServeTLS(certFile, keyFile)) // Serve H
}
<|endoftext|>


<commit_before>package vpki

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"time"

	"github.com/hashicorp/vault/api"
)

// RawPair is a simple explicitly-named pair of byte slices returned by
// the RawPair function.
type RawPair struct {
	Private, Public []byte
}

// RawSignCSR takes a certificate request template, private keye, and ttl, and
// returns the private/public keypair, unparsed, for any applications which may
// need to consume the certificates directly in their PEM form. The RawPair
// struct is used to help prevent transposition errors by explicitly naming the
// public/private pairs rather than returning two byte slices.
func (c *Client) RawSignCSR(csr *x509.CertificateRequest, k *rsa.PrivateKey, ttl time.Duration) (*RawPair, error) {
	csrBs, err := x509.CreateCertificateRequest(rand.Reader, csr, k)
	if err != nil {
		return nil, err
	}

	pubBs, err := c.rawCSR(csrBs, csr.Subject.CommonName, ttl)
	if err != nil {
		return nil, err
	}

	return &RawPair{Public: pubBs, Private: pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(k),
	})}, nil
}

// RawCert is a very high-level method used to obtain the raw certificate data.
func (c *Client) RawCert(cn string) (*RawPair, error) {
	csr, k, err := c.getCSR(cn)
	if err != nil {
		return nil, err
	}
	return c.RawSignCSR(csr, k, c.TTL)
}

func (c *Client) rawCSR(csr []byte, cn string, ttl time.Duration) ([]byte, error) {
	pemB := &pem.Block{
		Bytes: csr,
		Type:  csrName,
	}

	return c.RawSignCSRBytes(pem.EncodeToMemory(pemB), cn, ttl)
}

func (c *Client) write(path string, data map[string]interface{}) (*api.Secret, error) {
	return c.sw.Write(c.Mount+"/"+path, data)
}

//RawSignCSRBytes takes the bytes of a Certificate Signing Request, the CN and
//the ttl, and returns raw bytes of the signed certificate bundle.
func (c *Client) RawSignCSRBytes(csr []byte, cn string, ttl time.Duration) ([]byte, error) {

	data := map[string]interface{}{
		"csr":         string(csr),
		"common_name": cn,
		"format":      "pem_bundle",
		"ttl":         ttl.String(),
	}

	if c.sw == nil {
		c.init()
	}

	secret, err := c.write("sign", data)
	if err != nil {
		return nil, err
	}

	return []byte(secret.Data["certificate"].(string)), nil
}

//RawSignIntermediateCSRBytes takes the bytes of a Certificate Signing Request, the CN and
//the ttl, and returns raw bytes of the signed certificate bundle.
func (c *Client) RawSignIntermediateCSRBytes(csr []byte, cn string, ttl time.Duration) ([]byte, error) {

	data := map[string]interface{}{
		"csr":            string(csr),
		"use_csr_values": true,
		"common_name":    cn,
		"format":         "pem_bundle",
		"ttl":            ttl.String(),
	}

	if c.sw == nil {
		c.init()
	}

	secret, err := c.write("sign-intermediate", data)
	if err != nil {
		return nil, err
	}

	return []byte(secret.Data["certificate"].(string)), nil
}
<commit_msg>Write to role endpoint<commit_after>package vpki

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"time"

	"github.com/hashicorp/vault/api"
)

// RawPair is a simple explicitly-named pair of byte slices returned by
// the RawPair function.
type RawPair struct {
	Private, Public []byte
}

// RawSignCSR takes a certificate request template, private keye, and ttl, and
// returns the private/public keypair, unparsed, for any applications which may
// need to consume the certificates directly in their PEM form. The RawPair
// struct is used to help prevent transposition errors by explicitly naming the
// public/private pairs rather than returning two byte slices.
func (c *Client) RawSignCSR(csr *x509.CertificateRequest, k *rsa.PrivateKey, ttl time.Duration) (*RawPair, error) {
	csrBs, err := x509.CreateCertificateRequest(rand.Reader, csr, k)
	if err != nil {
		return nil, err
	}

	pubBs, err := c.rawCSR(csrBs, csr.Subject.CommonName, ttl)
	if err != nil {
		return nil, err
	}

	return &RawPair{Public: pubBs, Private: pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(k),
	})}, nil
}

// RawCert is a very high-level method used to obtain the raw certificate data.
func (c *Client) RawCert(cn string) (*RawPair, error) {
	csr, k, err := c.getCSR(cn)
	if err != nil {
		return nil, err
	}
	return c.RawSignCSR(csr, k, c.TTL)
}

func (c *Client) rawCSR(csr []byte, cn string, ttl time.Duration) ([]byte, error) {
	pemB := &pem.Block{
		Bytes: csr,
		Type:  csrName,
	}

	return c.RawSignCSRBytes(pem.EncodeToMemory(pemB), cn, ttl)
}

func (c *Client) write(path string, data map[string]interface{}) (*api.Secret, error) {
	return c.sw.Write(c.Mount+"/"+path, data)
}

func (c *Client) writeRole(path string, data map[string]interface{}) (*api.Secret, error) {
	return c.sw.Write(c.Mount+"/"+c.Role+"/"+path, data)
}

//RawSignCSRBytes takes the bytes of a Certificate Signing Request, the CN and
//the ttl, and returns raw bytes of the signed certificate bundle.
func (c *Client) RawSignCSRBytes(csr []byte, cn string, ttl time.Duration) ([]byte, error) {

	data := map[string]interface{}{
		"csr":         string(csr),
		"common_name": cn,
		"format":      "pem_bundle",
		"ttl":         ttl.String(),
	}

	if c.sw == nil {
		c.init()
	}

	secret, err := c.writeRole("sign", data)
	if err != nil {
		return nil, err
	}

	return []byte(secret.Data["certificate"].(string)), nil
}

//RawSignIntermediateCSRBytes takes the bytes of a Certificate Signing Request, the CN and
//the ttl, and returns raw bytes of the signed certificate bundle.
func (c *Client) RawSignIntermediateCSRBytes(csr []byte, cn string, ttl time.Duration) ([]byte, error) {

	data := map[string]interface{}{
		"csr":            string(csr),
		"use_csr_values": true,
		"common_name":    cn,
		"format":         "pem_bundle",
		"ttl":            ttl.String(),
	}

	if c.sw == nil {
		c.init()
	}

	secret, err := c.writeRole("sign-intermediate", data)
	if err != nil {
		return nil, err
	}

	return []byte(secret.Data["certificate"].(string)), nil
}
<|endoftext|>


<commit_before>package tests

import (
	randc "crypto/rand"
	"math"
	randm "math/rand"
	"sync"
	"time"
)

// Алфавит
type Abc struct {
	alphabet []rune
	bits     float64
}

// Идентификатор
type Shortid struct {
	abc   Abc
	rand  uint
	epoch time.Time  // Начало времен
	ms    uint       // Время генерации последнего значения
	count uint       // Кол-во идентификаторов в рамках одой и той же мс.
	mx    sync.Mutex // Блокировка для конкурентного доступа
}

// Создание нового генератора
func NewShortid(alphabet string) *Shortid {
	return &Shortid{
		abc:   NewAbc(alphabet),
		epoch: time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC),
		ms:    0,
		count: 0,
	}
}

func (sid *Shortid) Generate() string {
	for true {
		idrunes := sid.GenerateRandom()
		valid := false
		for c := range idrunes {
			if (c >= 'a') || (c <= 'z') {
				valid = true
				break
			}
		}
		if valid {
			for i, c := range idrunes {
				if (c >= 'a') && (c <= 'z') {
					if randm.Intn(2) == 0 {
						idrunes[i] = c + 'A' - 'a'
					}
				}
			}
			return string(idrunes)
		}
	}
	panic("Unreacheable code")
}

// Generate generates a new short Id.
func (sid *Shortid) GenerateRandom() []rune {
	ms, count := sid.getMsAndCounter(sid.epoch)
	idrunes := sid.abc.Encode(ms, 40)
	if count > 0 {
		idrunes = append(idrunes, sid.abc.Encode(count, 0)...)
	}
	return idrunes
}

func (sid *Shortid) getMsAndCounter(epoch time.Time) (uint, uint) {
	sid.mx.Lock()
	defer sid.mx.Unlock()
	ms := uint(time.Now().Sub(epoch).Nanoseconds() / 1000000)
	if ms <= sid.ms {
		sid.count++
	} else {
		sid.count = 0
		sid.ms = ms
	}
	return sid.ms, sid.count
}

// Abc returns the instance of alphabet used for representing the Ids.
func (sid *Shortid) Abc() Abc {
	return sid.abc
}

// Epoch returns the value of epoch used as the beginning of millisecond counting (normally
// 2016-01-01 00:00:00 local time)
func (sid *Shortid) Epoch() time.Time {
	return sid.epoch
}

// NewAbc constructs a new instance of shuffled alphabet to be used for Id representation.
func NewAbc(alphabet string) Abc {
	abc := Abc{alphabet: nil, bits: math.Log2(float64(len(alphabet)))}
	abc.shuffle(alphabet, 0)
	return abc
}

func (abc *Abc) shuffle(alphabet string, seed uint64) {
	source := []rune(alphabet)
	for len(source) > 1 {
		seed = (seed*9301 + 49297) % 233280
		i := int(seed * uint64(len(source)) / 233280)

		abc.alphabet = append(abc.alphabet, source[i])
		source = append(source[:i], source[i+1:]...)
	}
	abc.alphabet = append(abc.alphabet, source[0])
}

// Encode encodes a given value into a slice of runes of length nsymbols. In case nsymbols==0, the
// length of the result is automatically computed from data. Even if fewer symbols is required to
// encode the data than nsymbols, all positions are used encoding 0 where required to guarantee
// uniqueness in case further data is added to the sequence. The value of digits [4,6] represents
// represents n in 2^n, which defines how much randomness flows into the algorithm: 4 -- every value
// can be represented by 4 symbols in the alphabet (permitting at most 16 values), 5 -- every value
// can be represented by 2 symbols in the alphabet (permitting at most 32 values), 6 -- every value
// is represented by exactly 1 symbol with no randomness (permitting 64 values).
func (abc *Abc) Encode(val, bits uint) []rune {
	nsymbols := uint(0)
	randBits := uint(2)
	if bits > 0 {
		nsymbols = uint(math.Ceil(float64(bits) / (abc.bits - float64(randBits))))
	} else if val > 0 {
		nsymbols = uint(math.Ceil(math.Log2(float64(val+1)) / (abc.bits - float64(randBits))))
	}
	if nsymbols == 0 {
		return []rune{}
	}
	// no random component if digits == 6
	res := make([]rune, int(nsymbols))
	data := val
	for i := range res {
		index := data % uint(len(abc.alphabet)>>randBits)
		index = (index << randBits) | uint(randm.Int31n(1<<randBits))
		res[i] = abc.alphabet[index]
		data /= uint(len(abc.alphabet) >> randBits)
	}
	return res
}

func maskedRandomInts(size, mask int) []int {
	ints := make([]int, size)
	bytes := make([]byte, size)
	if _, err := randc.Read(bytes); err == nil {
		for i, b := range bytes {
			ints[i] = int(b) & mask
		}
	} else {
		for i := range ints {
			ints[i] = randm.Intn(0xff) & mask
		}
	}
	return ints
}
<commit_msg>Небольшое улучшение производительности<commit_after>package tests

import (
	"math"
	"math/rand"
	"sync"
	"time"
)

// Алфавит
type Abc struct {
	alphabet []rune
	bits     float64
}

// Идентификатор
type Shortid struct {
	abc   Abc
	rnd   *rand.Rand
	epoch time.Time  // Начало времен
	ms    uint       // Время генерации последнего значения
	count uint       // Кол-во идентификаторов в рамках одой и той же мс.
	mx    sync.Mutex // Блокировка для конкурентного доступа
}

// Создание нового генератора
func NewShortid(alphabet string) *Shortid {
	return &Shortid{
		rnd:   rand.New(rand.NewSource(time.Now().UnixNano())),
		abc:   NewAbc(alphabet),
		epoch: time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC),
		ms:    0,
		count: 0,
	}
}

func (sid *Shortid) Generate() string {
	for true {
		idrunes := sid.GenerateRandom()
		valid := false
		for c := range idrunes {
			if (c >= 'a') || (c <= 'z') {
				valid = true
				break
			}
		}
		if valid {
			for i, c := range idrunes {
				if (c >= 'a') && (c <= 'z') {
					if sid.rnd.Intn(2) == 0 {
						idrunes[i] = c + 'A' - 'a'
					}
				}
			}
			return string(idrunes)
		}
	}
	panic("Unreacheable code")
}

// Generate generates a new short Id.
func (sid *Shortid) GenerateRandom() []rune {
	ms, count := sid.getMsAndCounter(sid.epoch)
	idrunes := sid.abc.Encode(sid.rnd, ms, 40)
	if count > 0 {
		idrunes = append(idrunes, sid.abc.Encode(sid.rnd, count, 0)...)
	}
	return idrunes
}

func (sid *Shortid) getMsAndCounter(epoch time.Time) (uint, uint) {
	sid.mx.Lock()
	defer sid.mx.Unlock()
	ms := uint(time.Now().Sub(epoch).Nanoseconds() / 1000000)
	if ms <= sid.ms {
		sid.count++
	} else {
		sid.count = 0
		sid.ms = ms
	}
	return sid.ms, sid.count
}

// Abc returns the instance of alphabet used for representing the Ids.
func (sid *Shortid) Abc() Abc {
	return sid.abc
}

// Epoch returns the value of epoch used as the beginning of millisecond counting (normally
// 2016-01-01 00:00:00 local time)
func (sid *Shortid) Epoch() time.Time {
	return sid.epoch
}

// NewAbc constructs a new instance of shuffled alphabet to be used for Id representation.
func NewAbc(alphabet string) Abc {
	abc := Abc{alphabet: nil, bits: math.Log2(float64(len(alphabet)))}
	abc.shuffle(alphabet, 0)
	return abc
}

func (abc *Abc) shuffle(alphabet string, seed uint64) {
	source := []rune(alphabet)
	for len(source) > 1 {
		seed = (seed*9301 + 49297) % 233280
		i := int(seed * uint64(len(source)) / 233280)

		abc.alphabet = append(abc.alphabet, source[i])
		source = append(source[:i], source[i+1:]...)
	}
	abc.alphabet = append(abc.alphabet, source[0])
}

// Encode encodes a given value into a slice of runes of length nsymbols. In case nsymbols==0, the
// length of the result is automatically computed from data. Even if fewer symbols is required to
// encode the data than nsymbols, all positions are used encoding 0 where required to guarantee
// uniqueness in case further data is added to the sequence. The value of digits [4,6] represents
// represents n in 2^n, which defines how much randomness flows into the algorithm: 4 -- every value
// can be represented by 4 symbols in the alphabet (permitting at most 16 values), 5 -- every value
// can be represented by 2 symbols in the alphabet (permitting at most 32 values), 6 -- every value
// is represented by exactly 1 symbol with no randomness (permitting 64 values).
func (abc *Abc) Encode(rnd *rand.Rand, val, bits uint) []rune {
	nsymbols := uint(0)
	randBits := uint(2)
	if bits > 0 {
		nsymbols = uint(math.Ceil(float64(bits) / (abc.bits - float64(randBits))))
	} else if val > 0 {
		nsymbols = uint(math.Ceil(math.Log2(float64(val+1)) / (abc.bits - float64(randBits))))
	}
	if nsymbols == 0 {
		return []rune{}
	}
	// no random component if digits == 6
	res := make([]rune, int(nsymbols))
	data := val
	for i := range res {
		index := data % uint(len(abc.alphabet)>>randBits)
		index = (index << randBits) | uint(rnd.Int31n(1<<randBits))
		res[i] = abc.alphabet[index]
		data /= uint(len(abc.alphabet) >> randBits)
	}
	return res
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"

	"github.com/robustirc/robustirc/ircserver"
	"github.com/robustirc/robustirc/types"

	"github.com/hashicorp/raft"
)

//go:generate go run gentmpl.go status irclog

func handleStatus(res http.ResponseWriter, req *http.Request) {
	p, _ := peerStore.Peers()

	lo, err := ircStore.FirstIndex()
	if err != nil {
		log.Printf("Could not get first index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}
	hi, err := ircStore.LastIndex()
	if err != nil {
		log.Printf("Could not get last index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}

	var entries []*raft.Log
	if lo != 0 && hi != 0 {
		for i := lo; i <= hi; i++ {
			l := new(raft.Log)

			if err := ircStore.GetLog(i, l); err != nil {
				log.Printf("Could not get entry %d: %v", i, err)
				http.Error(res, "internal error", 500)
				return
			}
			entries = append(entries, l)
		}
	}

	args := struct {
		Addr               string
		State              raft.RaftState
		Leader             net.Addr
		Peers              []net.Addr
		First              uint64
		Last               uint64
		Entries            []*raft.Log
		Stats              map[string]string
		Sessions           map[types.RobustId]*ircserver.Session
		GetMessageRequests map[string]GetMessageStats
	}{
		*peerAddr,
		node.State(),
		node.Leader(),
		p,
		lo,
		hi,
		entries,
		node.Stats(),
		ircserver.Sessions,
		GetMessageRequests,
	}

	statusTpl.Execute(res, args)
}

func handleIrclog(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(r.FormValue("sessionid"), 0, 64)
	if err != nil || id == 0 {
		http.Error(w, "Invalid session", http.StatusBadRequest)
		return
	}

	session := types.RobustId{Id: id}

	s, err := ircserver.GetSession(session)
	if err != nil {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	// TODO(secure): pagination

	lastSeen := s.StartId
	var messages []*types.RobustMessage
	// TODO(secure): input messages (i.e. raft log entries) which don’t result
	// in an output message in the current session (e.g. PRIVMSGs) don’t show
	// up in here at all.
	// XXX: The following code is pretty horrible. It iterates through _all_
	// log messages to create a map from id to decoded message, in order to add
	// them to the messages slice in the second loop. We should come up with a
	// better way that is lighter on resources. Perhaps store the processed
	// indexes in the session?
	inputs := make(map[types.RobustId]*types.RobustMessage)
	first, _ := ircStore.FirstIndex()
	last, _ := ircStore.LastIndex()
	for idx := first; idx <= last; idx++ {
		var elog raft.Log

		if err := ircStore.GetLog(idx, &elog); err != nil {
			http.Error(w, fmt.Sprintf("Cannot read log: %v", err), http.StatusInternalServerError)
			return
		}
		if elog.Type != raft.LogCommand {
			continue
		}
		msg := types.NewRobustMessageFromBytes(elog.Data)
		if msg.Session.Id != session.Id {
			continue
		}
		inputs[msg.Id] = &msg
	}

	for {
		if msg := ircserver.GetMessageNonBlocking(lastSeen); msg != nil {
			if msg.Type == types.RobustIRCToClient && s.InterestedIn(msg) {
				if msg.Id.Reply == 1 {
					if inputmsg, ok := inputs[types.RobustId{Id: msg.Id.Id}]; ok {
						messages = append(messages, inputmsg)
					}
				}
				messages = append(messages, msg)
			}
			lastSeen = msg.Id
		} else {
			break
		}
	}

	args := struct {
		Session  types.RobustId
		Messages []*types.RobustMessage
	}{
		session,
		messages,
	}
	if err := irclogTpl.Execute(w, args); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<commit_msg>status: just skip missing ircstore entries<commit_after>package main

import (
	"log"
	"net"
	"net/http"
	"strconv"

	"github.com/robustirc/robustirc/ircserver"
	"github.com/robustirc/robustirc/types"

	"github.com/hashicorp/raft"
)

//go:generate go run gentmpl.go status irclog

func handleStatus(res http.ResponseWriter, req *http.Request) {
	p, _ := peerStore.Peers()

	lo, err := ircStore.FirstIndex()
	if err != nil {
		log.Printf("Could not get first index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}
	hi, err := ircStore.LastIndex()
	if err != nil {
		log.Printf("Could not get last index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}

	var entries []*raft.Log
	if lo != 0 && hi != 0 {
		for i := lo; i <= hi; i++ {
			l := new(raft.Log)

			if err := ircStore.GetLog(i, l); err != nil {
				// Not every message goes into the ircStore (e.g. raft peer change
				// messages do not).
				continue
			}
			entries = append(entries, l)
		}
	}

	args := struct {
		Addr               string
		State              raft.RaftState
		Leader             net.Addr
		Peers              []net.Addr
		First              uint64
		Last               uint64
		Entries            []*raft.Log
		Stats              map[string]string
		Sessions           map[types.RobustId]*ircserver.Session
		GetMessageRequests map[string]GetMessageStats
	}{
		*peerAddr,
		node.State(),
		node.Leader(),
		p,
		lo,
		hi,
		entries,
		node.Stats(),
		ircserver.Sessions,
		GetMessageRequests,
	}

	statusTpl.Execute(res, args)
}

func handleIrclog(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(r.FormValue("sessionid"), 0, 64)
	if err != nil || id == 0 {
		http.Error(w, "Invalid session", http.StatusBadRequest)
		return
	}

	session := types.RobustId{Id: id}

	s, err := ircserver.GetSession(session)
	if err != nil {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	// TODO(secure): pagination

	lastSeen := s.StartId
	var messages []*types.RobustMessage
	// TODO(secure): input messages (i.e. raft log entries) which don’t result
	// in an output message in the current session (e.g. PRIVMSGs) don’t show
	// up in here at all.
	// XXX: The following code is pretty horrible. It iterates through _all_
	// log messages to create a map from id to decoded message, in order to add
	// them to the messages slice in the second loop. We should come up with a
	// better way that is lighter on resources. Perhaps store the processed
	// indexes in the session?
	inputs := make(map[types.RobustId]*types.RobustMessage)
	first, _ := ircStore.FirstIndex()
	last, _ := ircStore.LastIndex()
	for idx := first; idx <= last; idx++ {
		var elog raft.Log

		if err := ircStore.GetLog(idx, &elog); err != nil {
			// Not every message goes into the ircStore (e.g. raft peer change
			// messages do not).
			continue
		}
		if elog.Type != raft.LogCommand {
			continue
		}
		msg := types.NewRobustMessageFromBytes(elog.Data)
		if msg.Session.Id != session.Id {
			continue
		}
		inputs[msg.Id] = &msg
	}

	for {
		if msg := ircserver.GetMessageNonBlocking(lastSeen); msg != nil {
			if msg.Type == types.RobustIRCToClient && s.InterestedIn(msg) {
				if msg.Id.Reply == 1 {
					if inputmsg, ok := inputs[types.RobustId{Id: msg.Id.Id}]; ok {
						messages = append(messages, inputmsg)
					}
				}
				messages = append(messages, msg)
			}
			lastSeen = msg.Id
		} else {
			break
		}
	}

	args := struct {
		Session  types.RobustId
		Messages []*types.RobustMessage
	}{
		session,
		messages,
	}
	if err := irclogTpl.Execute(w, args); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

type statusHandler struct {
	// Mutex for locking
	mu *sync.Mutex
	// Backend dialer to check if target is up and running
	dial func() (net.Conn, error)
	// Current status
	listening bool
	reloading bool
}

type statusResponse struct {
	Ok            bool      `json:"ok"`
	Status        string    `json:"status"`
	BackendOk     bool      `json:"backend_ok"`
	BackendStatus string    `json:"backend_status"`
	BackendError  string    `json:"backend_error,omitempty"`
	Time          time.Time `json:"time"`
	Hostname      string    `json:"hostname,omitempty"`
	Message       string    `json:"message"`
	Revision      string    `json:"revision"`
	Compiler      string    `json:"compiler"`
}

func newStatusHandler(dial func() (net.Conn, error)) *statusHandler {
	return &statusHandler{&sync.Mutex{}, dial, false, false}
}

func (s *statusHandler) Listening() {
	s.mu.Lock()
	s.listening = true
	s.reloading = false
	s.mu.Unlock()
}

func (s *statusHandler) Reloading() {
	s.mu.Lock()
	s.reloading = true
	s.mu.Unlock()
}

func (s *statusHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	resp := statusResponse{
		Time: time.Now(),
	}

	conn, err := s.dial()
	resp.BackendOk = err == nil
	resp.Revision = buildRevision
	resp.Compiler = buildCompiler

	if resp.BackendOk {
		defer conn.Close()
		resp.BackendStatus = "ok"
	} else {
		resp.BackendError = err.Error()
		resp.BackendStatus = "critical"
	}

	s.mu.Lock()
	resp.Ok = s.listening && resp.BackendOk
	if !s.listening {
		resp.Message = "initializing"
	} else if s.reloading {
		resp.Message = "reloading"
	} else {
		resp.Message = "listening"
	}
	s.mu.Unlock()

	if resp.Ok && resp.BackendOk {
		resp.Status = "ok"
	} else {
		resp.Status = "critical"
	}

	hostname, err := os.Hostname()
	if err == nil {
		resp.Hostname = hostname
	}

	out, err := json.Marshal(resp)
	if err != nil {
		panic(err)
	}

	w.Write(out)
}
<commit_msg>Return 503 on status if tunnel down<commit_after>/*-
 * Copyright 2015 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

type statusHandler struct {
	// Mutex for locking
	mu *sync.Mutex
	// Backend dialer to check if target is up and running
	dial func() (net.Conn, error)
	// Current status
	listening bool
	reloading bool
}

type statusResponse struct {
	Ok            bool      `json:"ok"`
	Status        string    `json:"status"`
	BackendOk     bool      `json:"backend_ok"`
	BackendStatus string    `json:"backend_status"`
	BackendError  string    `json:"backend_error,omitempty"`
	Time          time.Time `json:"time"`
	Hostname      string    `json:"hostname,omitempty"`
	Message       string    `json:"message"`
	Revision      string    `json:"revision"`
	Compiler      string    `json:"compiler"`
}

func newStatusHandler(dial func() (net.Conn, error)) *statusHandler {
	return &statusHandler{&sync.Mutex{}, dial, false, false}
}

func (s *statusHandler) Listening() {
	s.mu.Lock()
	s.listening = true
	s.reloading = false
	s.mu.Unlock()
}

func (s *statusHandler) Reloading() {
	s.mu.Lock()
	s.reloading = true
	s.mu.Unlock()
}

func (s *statusHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	resp := statusResponse{
		Time: time.Now(),
	}

	conn, err := s.dial()
	resp.BackendOk = err == nil
	resp.Revision = buildRevision
	resp.Compiler = buildCompiler

	if resp.BackendOk {
		defer conn.Close()
		resp.BackendStatus = "ok"
	} else {
		resp.BackendError = err.Error()
		resp.BackendStatus = "critical"
	}

	s.mu.Lock()
	resp.Ok = s.listening && resp.BackendOk
	if !s.listening {
		resp.Message = "initializing"
	} else if s.reloading {
		resp.Message = "reloading"
	} else {
		resp.Message = "listening"
	}
	s.mu.Unlock()

	if resp.Ok && resp.BackendOk {
		resp.Status = "ok"
	} else {
		resp.Status = "critical"
	}

	hostname, err := os.Hostname()
	if err == nil {
		resp.Hostname = hostname
	}

	out, err := json.Marshal(resp)
	if err != nil {
		panic(err)
	}

	if !resp.Ok {
		w.WriteHeader(http.StatusServiceUnavailable)
	}

	w.Write(out)
}
<|endoftext|>


<commit_before>package bidirpc

import (
	"bytes"
	"errors"
)

type stream struct {
	session *Session
	id      byte
	inC     chan *bytes.Buffer
	reader  *bytes.Buffer
	writer  bytes.Buffer
}

func newStream(session *Session, id byte) *stream {
	s := &stream{
		session: session,
		id:      id,
		inC:     make(chan *bytes.Buffer, 1),
	}
	return s
}

func (s *stream) Read(p []byte) (n int, err error) {
	if s.reader != nil && s.reader.Len() == 0 {
		s.session.bp.Put(s.reader) // returns the buffer to the bufferPool
		s.reader = nil
	}

	if s.reader == nil {
		select {
		case <-s.session.closedC:
			return 0, errors.New("stream read from a closed session")
		case s.reader = <-s.inC: // only the message body
		}
	}

	return s.reader.Read(p)
}

func (s *stream) Write(p []byte) (n int, err error) {
	if s.writer.Len() == 0 {
		var dummyHeader [4]byte
		if n, err := s.writer.Write(dummyHeader[:]); err != nil {
			return n, err
		}
	}
	return s.writer.Write(p)
}

func (s *stream) flush() (err error) {
	buffer := s.writer.Bytes()
	bodyLen := s.writer.Len() - 4
	encodeHeader(buffer, s.id, bodyLen)

	err = s.session.write(buffer)

	s.writer.Reset()
	return
}
<commit_msg>We can safely assume Buffer.Write will not returns a  error<commit_after>package bidirpc

import (
	"bytes"
	"errors"
)

type stream struct {
	session *Session
	id      byte
	inC     chan *bytes.Buffer
	reader  *bytes.Buffer
	writer  bytes.Buffer
}

func newStream(session *Session, id byte) *stream {
	s := &stream{
		session: session,
		id:      id,
		inC:     make(chan *bytes.Buffer, 1),
	}
	return s
}

func (s *stream) Read(p []byte) (n int, err error) {
	if s.reader != nil && s.reader.Len() == 0 {
		s.session.bp.Put(s.reader) // returns the buffer to the bufferPool
		s.reader = nil
	}

	if s.reader == nil {
		select {
		case <-s.session.closedC:
			return 0, errors.New("stream read from a closed session")
		case s.reader = <-s.inC: // only the message body
		}
	}

	return s.reader.Read(p)
}

func (s *stream) Write(p []byte) (n int, err error) {
	if s.writer.Len() == 0 {
		var dummyHeader [4]byte
		s.writer.Write(dummyHeader[:])
	}
	return s.writer.Write(p)
}

func (s *stream) flush() (err error) {
	buffer := s.writer.Bytes()
	bodyLen := s.writer.Len() - 4
	encodeHeader(buffer, s.id, bodyLen)

	err = s.session.write(buffer)

	s.writer.Reset()
	return
}
<|endoftext|>


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

package badger

import (
	"bytes"
	"context"
	"sync"
	"time"

	"github.com/dgraph-io/badger/pb"
	"github.com/dgraph-io/badger/y"
	humanize "github.com/dustin/go-humanize"
)

const pageSize = 4 << 20 // 4MB

// Stream provides a framework to concurrently iterate over a snapshot of Badger, pick up
// key-values, batch them up and call Send. Stream does concurrent iteration over many smaller key
// ranges. It does NOT send keys in lexicographical sorted order. To get keys in sorted
// order, use Iterator.
type Stream struct {
	// Prefix to only iterate over certain range of keys. If set to nil (default), Stream would
	// iterate over the entire DB.
	Prefix []byte

	// Number of goroutines to use for iterating over key ranges. Defaults to 16.
	NumGo int

	// Badger would produce log entries in Infof to indicate the progress of Stream. LogPrefix can
	// be used to help differentiate them from other activities. Default is "Badger.Stream".
	LogPrefix string

	// ChooseKey is invoked each time a new key is encountered. Note that this is not called
	// on every version of the value, only the first encountered version (i.e. the highest version
	// of the value a key has). ChooseKey can be left nil to select all keys.
	//
	// Note: Calls to ChooseKey are concurrent.
	ChooseKey func(item *Item) bool

	// KeyToList, similar to ChooseKey, is only invoked on the highest version of the value. It
	// is upto the caller to iterate over the versions and generate zero, one or more KVs. It
	// is expected that the user would advance the iterator to go through the versions of the
	// values. However, the user MUST immediately return from this function on the first encounter
	// with a mismatching key. See example usage in ToList function. Can be left nil to use ToList
	// function by default.
	//
	// Note: Calls to KeyToList are concurrent.
	KeyToList func(key []byte, itr *Iterator) (*pb.KVList, error)

	// This is the method where Stream sends the final output. All calls to Send are done by a
	// single goroutine, i.e. logic within Send method can expect single threaded execution.
	Send func(*pb.KVList) error

	readTs  uint64
	db      *DB
	rangeCh chan keyRange
	kvChan  chan *pb.KVList
}

// ToList is a default implementation of KeyToList. It picks up all valid versions of the key,
// skipping over deleted or expired keys.
func (st *Stream) ToList(key []byte, itr *Iterator) (*pb.KVList, error) {
	list := &pb.KVList{}
	for ; itr.Valid(); itr.Next() {
		item := itr.Item()
		if item.IsDeletedOrExpired() {
			break
		}
		if !bytes.Equal(key, item.Key()) {
			// Break out on the first encounter with another key.
			break
		}

		valCopy, err := item.ValueCopy(nil)
		if err != nil {
			return nil, err
		}
		kv := &pb.KV{
			Key:       item.KeyCopy(nil),
			Value:     valCopy,
			UserMeta:  []byte{item.UserMeta()},
			Version:   item.Version(),
			ExpiresAt: item.ExpiresAt(),
		}
		list.Kv = append(list.Kv, kv)
		if st.db.opt.NumVersionsToKeep == 1 {
			break
		}

		if item.DiscardEarlierVersions() {
			break
		}
	}
	return list, nil
}

// keyRange is [start, end), including start, excluding end. Do ensure that the start,
// end byte slices are owned by keyRange struct.
func (st *Stream) produceRanges(ctx context.Context) {
	splits := st.db.KeySplits(st.Prefix)
	start := y.SafeCopy(nil, st.Prefix)
	for _, key := range splits {
		st.rangeCh <- keyRange{left: start, right: y.SafeCopy(nil, []byte(key))}
		start = y.SafeCopy(nil, []byte(key))
	}
	// Edge case: prefix is empty and no splits exist. In that case, we should have at least one
	// keyRange output.
	st.rangeCh <- keyRange{left: start}
	close(st.rangeCh)
}

// produceKVs picks up ranges from rangeCh, generates KV lists and sends them to kvChan.
func (st *Stream) produceKVs(ctx context.Context) error {
	var size int
	var txn *Txn
	if st.readTs > 0 {
		txn = st.db.NewTransactionAt(st.readTs, false)
	} else {
		txn = st.db.NewTransaction(false)
	}
	defer txn.Discard()

	iterate := func(kr keyRange) error {
		iterOpts := DefaultIteratorOptions
		iterOpts.AllVersions = true
		iterOpts.Prefix = st.Prefix
		iterOpts.PrefetchValues = false
		itr := txn.NewIterator(iterOpts)
		defer itr.Close()

		outList := new(pb.KVList)
		var prevKey []byte
		for itr.Seek(kr.left); itr.Valid(); {
			// it.Valid would only return true for keys with the provided Prefix in iterOpts.
			item := itr.Item()
			if bytes.Equal(item.Key(), prevKey) {
				itr.Next()
				continue
			}
			prevKey = append(prevKey[:0], item.Key()...)

			// Check if we reached the end of the key range.
			if len(kr.right) > 0 && bytes.Compare(item.Key(), kr.right) >= 0 {
				break
			}
			// Check if we should pick this key.
			if st.ChooseKey != nil && !st.ChooseKey(item) {
				continue
			}

			// Now convert to key value.
			list, err := st.KeyToList(item.KeyCopy(nil), itr)
			if err != nil {
				return err
			}
			if list == nil || len(list.Kv) == 0 {
				continue
			}
			outList.Kv = append(outList.Kv, list.Kv...)
			size += list.Size()
			if size >= pageSize {
				st.kvChan <- outList
				outList = new(pb.KVList)
				size = 0
			}
		}
		if len(outList.Kv) > 0 {
			st.kvChan <- outList
		}
		return nil
	}

	for {
		select {
		case kr, ok := <-st.rangeCh:
			if !ok {
				// Done with the keys.
				return nil
			}
			if err := iterate(kr); err != nil {
				return err
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (st *Stream) streamKVs(ctx context.Context) error {
	var count int
	var bytesSent uint64
	t := time.NewTicker(time.Second)
	defer t.Stop()
	now := time.Now()

	slurp := func(batch *pb.KVList) error {
	loop:
		for {
			select {
			case kvs, ok := <-st.kvChan:
				if !ok {
					break loop
				}
				y.AssertTrue(kvs != nil)
				batch.Kv = append(batch.Kv, kvs.Kv...)
			default:
				break loop
			}
		}
		sz := uint64(batch.Size())
		bytesSent += sz
		count += len(batch.Kv)
		t := time.Now()
		if err := st.Send(batch); err != nil {
			return err
		}
		st.db.opt.Infof("%s Created batch of size: %s in %s.\n",
			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
		return nil
	}

outer:
	for {
		var batch *pb.KVList
		select {
		case <-ctx.Done():
			return ctx.Err()

		case <-t.C:
			dur := time.Since(now)
			durSec := uint64(dur.Seconds())
			if durSec == 0 {
				continue
			}
			speed := bytesSent / durSec
			st.db.opt.Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n", st.LogPrefix,
				y.FixedDuration(dur), humanize.Bytes(bytesSent), humanize.Bytes(speed))

		case kvs, ok := <-st.kvChan:
			if !ok {
				break outer
			}
			y.AssertTrue(kvs != nil)
			batch = kvs
			if err := slurp(batch); err != nil {
				return err
			}
		}
	}

	st.db.opt.Infof("%s Sent %d keys\n", st.LogPrefix, count)
	return nil
}

// Orchestrate runs Stream. It picks up ranges from the SSTables, then runs NumGo number of
// goroutines to iterate over these ranges and batch up KVs in lists. It concurrently runs a single
// goroutine to pick these lists, batch them up further and send to Output.Send. Orchestrate also
// spits logs out to Infof, using provided LogPrefix. Note that all calls to Output.Send
// are serial. In case any of these steps encounter an error, Orchestrate would stop execution and
// return that error. Orchestrate can be called multiple times, but in serial order.
func (st *Stream) Orchestrate(ctx context.Context) error {
	st.rangeCh = make(chan keyRange, 3) // Contains keys for posting lists.

	// kvChan should only have a small capacity to ensure that we don't buffer up too much data if
	// sending is slow. Page size is set to 4MB, which is used to lazily cap the size of each
	// KVList. To get around 64MB buffer, we can set the channel size to 16.
	st.kvChan = make(chan *pb.KVList, 16)

	if st.KeyToList == nil {
		st.KeyToList = st.ToList
	}

	// Picks up ranges from Badger, and sends them to rangeCh.
	go st.produceRanges(ctx)

	errCh := make(chan error, 1) // Stores error by consumeKeys.
	var wg sync.WaitGroup
	for i := 0; i < st.NumGo; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// Picks up ranges from rangeCh, generates KV lists, and sends them to kvChan.
			if err := st.produceKVs(ctx); err != nil {
				select {
				case errCh <- err:
				default:
				}
			}
		}()
	}

	// Pick up key-values from kvChan and send to stream.
	kvErr := make(chan error, 1)
	go func() {
		// Picks up KV lists from kvChan, and sends them to Output.
		kvErr <- st.streamKVs(ctx)
	}()
	wg.Wait()        // Wait for produceKVs to be over.
	close(st.kvChan) // Now we can close kvChan.

	select {
	case err := <-errCh: // Check error from produceKVs.
		return err
	default:
	}

	// Wait for key streaming to be over.
	err := <-kvErr
	return err
}

func (db *DB) newStream() *Stream {
	return &Stream{db: db, NumGo: 16, LogPrefix: "Badger.Stream"}
}

// NewStream creates a new Stream.
func (db *DB) NewStream() *Stream {
	if db.opt.managedTxns {
		panic("This API can not be called in managed mode.")
	}
	return db.newStream()
}

// NewStreamAt creates a new Stream at a particular timestamp. Should only be used with managed DB.
func (db *DB) NewStreamAt(readTs uint64) *Stream {
	if !db.opt.managedTxns {
		panic("This API can only be called in managed mode.")
	}
	stream := db.newStream()
	stream.readTs = readTs
	return stream
}
<commit_msg>Bug fix in Stream Framework:<commit_after>/*
 * Copyright 2018 Dgraph Labs, Inc. and Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package badger

import (
	"bytes"
	"context"
	"sync"
	"time"

	"github.com/dgraph-io/badger/pb"
	"github.com/dgraph-io/badger/y"
	humanize "github.com/dustin/go-humanize"
)

const pageSize = 4 << 20 // 4MB

// Stream provides a framework to concurrently iterate over a snapshot of Badger, pick up
// key-values, batch them up and call Send. Stream does concurrent iteration over many smaller key
// ranges. It does NOT send keys in lexicographical sorted order. To get keys in sorted
// order, use Iterator.
type Stream struct {
	// Prefix to only iterate over certain range of keys. If set to nil (default), Stream would
	// iterate over the entire DB.
	Prefix []byte

	// Number of goroutines to use for iterating over key ranges. Defaults to 16.
	NumGo int

	// Badger would produce log entries in Infof to indicate the progress of Stream. LogPrefix can
	// be used to help differentiate them from other activities. Default is "Badger.Stream".
	LogPrefix string

	// ChooseKey is invoked each time a new key is encountered. Note that this is not called
	// on every version of the value, only the first encountered version (i.e. the highest version
	// of the value a key has). ChooseKey can be left nil to select all keys.
	//
	// Note: Calls to ChooseKey are concurrent.
	ChooseKey func(item *Item) bool

	// KeyToList, similar to ChooseKey, is only invoked on the highest version of the value. It
	// is upto the caller to iterate over the versions and generate zero, one or more KVs. It
	// is expected that the user would advance the iterator to go through the versions of the
	// values. However, the user MUST immediately return from this function on the first encounter
	// with a mismatching key. See example usage in ToList function. Can be left nil to use ToList
	// function by default.
	//
	// Note: Calls to KeyToList are concurrent.
	KeyToList func(key []byte, itr *Iterator) (*pb.KVList, error)

	// This is the method where Stream sends the final output. All calls to Send are done by a
	// single goroutine, i.e. logic within Send method can expect single threaded execution.
	Send func(*pb.KVList) error

	readTs  uint64
	db      *DB
	rangeCh chan keyRange
	kvChan  chan *pb.KVList
}

// ToList is a default implementation of KeyToList. It picks up all valid versions of the key,
// skipping over deleted or expired keys.
func (st *Stream) ToList(key []byte, itr *Iterator) (*pb.KVList, error) {
	list := &pb.KVList{}
	for ; itr.Valid(); itr.Next() {
		item := itr.Item()
		if item.IsDeletedOrExpired() {
			break
		}
		if !bytes.Equal(key, item.Key()) {
			// Break out on the first encounter with another key.
			break
		}

		valCopy, err := item.ValueCopy(nil)
		if err != nil {
			return nil, err
		}
		kv := &pb.KV{
			Key:       item.KeyCopy(nil),
			Value:     valCopy,
			UserMeta:  []byte{item.UserMeta()},
			Version:   item.Version(),
			ExpiresAt: item.ExpiresAt(),
		}
		list.Kv = append(list.Kv, kv)
		if st.db.opt.NumVersionsToKeep == 1 {
			break
		}

		if item.DiscardEarlierVersions() {
			break
		}
	}
	return list, nil
}

// keyRange is [start, end), including start, excluding end. Do ensure that the start,
// end byte slices are owned by keyRange struct.
func (st *Stream) produceRanges(ctx context.Context) {
	splits := st.db.KeySplits(st.Prefix)
	start := y.SafeCopy(nil, st.Prefix)
	for _, key := range splits {
		st.rangeCh <- keyRange{left: start, right: y.SafeCopy(nil, []byte(key))}
		start = y.SafeCopy(nil, []byte(key))
	}
	// Edge case: prefix is empty and no splits exist. In that case, we should have at least one
	// keyRange output.
	st.rangeCh <- keyRange{left: start}
	close(st.rangeCh)
}

// produceKVs picks up ranges from rangeCh, generates KV lists and sends them to kvChan.
func (st *Stream) produceKVs(ctx context.Context) error {
	var size int
	var txn *Txn
	if st.readTs > 0 {
		txn = st.db.NewTransactionAt(st.readTs, false)
	} else {
		txn = st.db.NewTransaction(false)
	}
	defer txn.Discard()

	iterate := func(kr keyRange) error {
		iterOpts := DefaultIteratorOptions
		iterOpts.AllVersions = true
		iterOpts.Prefix = st.Prefix
		iterOpts.PrefetchValues = false
		itr := txn.NewIterator(iterOpts)
		defer itr.Close()

		outList := new(pb.KVList)
		var prevKey []byte
		for itr.Seek(kr.left); itr.Valid(); {
			// it.Valid would only return true for keys with the provided Prefix in iterOpts.
			item := itr.Item()
			if bytes.Equal(item.Key(), prevKey) {
				itr.Next()
				continue
			}
			prevKey = append(prevKey[:0], item.Key()...)

			// Check if we reached the end of the key range.
			if len(kr.right) > 0 && bytes.Compare(item.Key(), kr.right) >= 0 {
				break
			}
			// Check if we should pick this key.
			if st.ChooseKey != nil && !st.ChooseKey(item) {
				continue
			}

			// Now convert to key value.
			list, err := st.KeyToList(item.KeyCopy(nil), itr)
			if err != nil {
				return err
			}
			if list == nil || len(list.Kv) == 0 {
				continue
			}
			outList.Kv = append(outList.Kv, list.Kv...)
			size += list.Size()
			if size >= pageSize {
				select {
				case st.kvChan <- outList:
				case <-ctx.Done():
					return ctx.Err()
				}
				outList = new(pb.KVList)
				size = 0
			}
		}
		if len(outList.Kv) > 0 {
			select {
			case st.kvChan <- outList:
			case <-ctx.Done():
				return ctx.Err()
			}
		}
		return nil
	}

	for {
		select {
		case kr, ok := <-st.rangeCh:
			if !ok {
				// Done with the keys.
				return nil
			}
			if err := iterate(kr); err != nil {
				return err
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (st *Stream) streamKVs(ctx context.Context) error {
	var count int
	var bytesSent uint64
	t := time.NewTicker(time.Second)
	defer t.Stop()
	now := time.Now()

	slurp := func(batch *pb.KVList) error {
	loop:
		for {
			select {
			case kvs, ok := <-st.kvChan:
				if !ok {
					break loop
				}
				y.AssertTrue(kvs != nil)
				batch.Kv = append(batch.Kv, kvs.Kv...)
			default:
				break loop
			}
		}
		sz := uint64(batch.Size())
		bytesSent += sz
		count += len(batch.Kv)
		t := time.Now()
		if err := st.Send(batch); err != nil {
			return err
		}
		st.db.opt.Infof("%s Created batch of size: %s in %s.\n",
			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
		return nil
	}

outer:
	for {
		var batch *pb.KVList
		select {
		case <-ctx.Done():
			return ctx.Err()

		case <-t.C:
			dur := time.Since(now)
			durSec := uint64(dur.Seconds())
			if durSec == 0 {
				continue
			}
			speed := bytesSent / durSec
			st.db.opt.Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n", st.LogPrefix,
				y.FixedDuration(dur), humanize.Bytes(bytesSent), humanize.Bytes(speed))

		case kvs, ok := <-st.kvChan:
			if !ok {
				break outer
			}
			y.AssertTrue(kvs != nil)
			batch = kvs
			if err := slurp(batch); err != nil {
				return err
			}
		}
	}

	st.db.opt.Infof("%s Sent %d keys\n", st.LogPrefix, count)
	return nil
}

// Orchestrate runs Stream. It picks up ranges from the SSTables, then runs NumGo number of
// goroutines to iterate over these ranges and batch up KVs in lists. It concurrently runs a single
// goroutine to pick these lists, batch them up further and send to Output.Send. Orchestrate also
// spits logs out to Infof, using provided LogPrefix. Note that all calls to Output.Send
// are serial. In case any of these steps encounter an error, Orchestrate would stop execution and
// return that error. Orchestrate can be called multiple times, but in serial order.
func (st *Stream) Orchestrate(ctx context.Context) error {
	st.rangeCh = make(chan keyRange, 3) // Contains keys for posting lists.

	// kvChan should only have a small capacity to ensure that we don't buffer up too much data if
	// sending is slow. Page size is set to 4MB, which is used to lazily cap the size of each
	// KVList. To get around 64MB buffer, we can set the channel size to 16.
	st.kvChan = make(chan *pb.KVList, 16)

	if st.KeyToList == nil {
		st.KeyToList = st.ToList
	}

	// Picks up ranges from Badger, and sends them to rangeCh.
	go st.produceRanges(ctx)

	errCh := make(chan error, 1) // Stores error by consumeKeys.
	var wg sync.WaitGroup
	for i := 0; i < st.NumGo; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// Picks up ranges from rangeCh, generates KV lists, and sends them to kvChan.
			if err := st.produceKVs(ctx); err != nil {
				select {
				case errCh <- err:
				default:
				}
			}
		}()
	}

	// Pick up key-values from kvChan and send to stream.
	kvErr := make(chan error, 1)
	go func() {
		// Picks up KV lists from kvChan, and sends them to Output.
		kvErr <- st.streamKVs(ctx)
	}()
	wg.Wait()        // Wait for produceKVs to be over.
	close(st.kvChan) // Now we can close kvChan.

	select {
	case err := <-errCh: // Check error from produceKVs.
		return err
	default:
	}

	// Wait for key streaming to be over.
	err := <-kvErr
	return err
}

func (db *DB) newStream() *Stream {
	return &Stream{db: db, NumGo: 16, LogPrefix: "Badger.Stream"}
}

// NewStream creates a new Stream.
func (db *DB) NewStream() *Stream {
	if db.opt.managedTxns {
		panic("This API can not be called in managed mode.")
	}
	return db.newStream()
}

// NewStreamAt creates a new Stream at a particular timestamp. Should only be used with managed DB.
func (db *DB) NewStreamAt(readTs uint64) *Stream {
	if !db.opt.managedTxns {
		panic("This API can only be called in managed mode.")
	}
	stream := db.newStream()
	stream.readTs = readTs
	return stream
}
<|endoftext|>


<commit_before>package xlsx

import (
	"archive/zip"
	"encoding/xml"
	"fmt"
	"io"
)

// xlsxWorkbook directly maps the workbook element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxWorkbook struct {
	FileVersion  xlsxFileVersion  `xml:"fileVersion"`
	WorkbookPr   xlsxWorkbookPr   `xml:"workbookPr"`
	BookViews    xlsxBookViews    `xml:"bookViews"`
	Sheets       xlsxSheets       `xml:"sheets"`
	DefinedNames xlsxDefinedNames `xml:"definedNames"`
	CalcPr       xlsxCalcPr       `xml:"calcPr"`
}

// xlsxFileVersion directly maps the fileVersion element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxFileVersion struct {
	AppName      string `xml:"appName,attr"`
	LastEdited   string `xml:"lastEdited,attr"`
	LowestEdited string `xml:"lowestEdited,attr"`
	RupBuild     string `xml:"rupBuild,attr"`
}

// xlsxWorkbookPr directly maps the workbookPr element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxWorkbookPr struct {
	DefaultThemeVersion string `xml:"defaultThemeVersion,attr"`
}

// xlsxBookViews directly maps the bookViews element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxBookViews struct {
	WorkBookView []xlsxWorkBookView `xml:"workbookView"`
}

// xlsxWorkBookView directly maps the workbookView element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxWorkBookView struct {
	XWindow      string `xml:"xWindow,attr"`
	YWindow      string `xml:"yWindow,attr"`
	WindowWidth  string `xml:"windowWidth,attr"`
	WindowHeight string `xml:"windowHeight,attr"`
}

// xlsxSheets directly maps the sheets element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxSheets struct {
	Sheet []xlsxSheet `xml:"sheet"`
}

// xlsxSheet directly maps the sheet element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxSheet struct {
	Name    string `xml:"name,attr"`
	SheetId string `xml:"sheetId,attr"`
	Id      string `xml:"id,attr"`
}

// xlsxDefinedNames directly maps the definedNames element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxDefinedNames struct {
	DefinedName []xlsxDefinedName `xml:"definedName"`
}

// xlsxDefinedName directly maps the definedName element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxDefinedName struct {
	Data         string `xml:",chardata"`
	Name         string `xml:"name,attr"`
	LocalSheetID string `xml:"localSheetId,attr"`
}

// xlsxCalcPr directly maps the calcPr element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxCalcPr struct {
	CalcId string `xml:"calcId,attr"`
}

// getWorksheetFromSheet() is an internal helper function to open a sheetN.xml file, refered to by an xlsx.xlsxSheet struct, from the XLSX file and unmarshal it an xlsx.xlsxWorksheet struct
func getWorksheetFromSheet(sheet xlsxSheet, worksheets map[string]*zip.File) (*xlsxWorksheet, error) {
	var rc io.ReadCloser
	var decoder *xml.Decoder
	var worksheet *xlsxWorksheet
	var error error
	worksheet = new(xlsxWorksheet)
	sheetName := fmt.Sprintf("sheet%s", sheet.SheetId)
	f := worksheets[sheetName]
	rc, error = f.Open()
	if error != nil {
		return nil, error
	}
	decoder = xml.NewDecoder(rc)
	error = decoder.Decode(worksheet)
	if error != nil {
		return nil, error
	}
	return worksheet, nil
}
<commit_msg>bug fix<commit_after>package xlsx

import (
	"archive/zip"
	"encoding/xml"
	"fmt"
	"io"
)

// xlsxWorkbook directly maps the workbook element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxWorkbook struct {
	FileVersion  xlsxFileVersion  `xml:"fileVersion"`
	WorkbookPr   xlsxWorkbookPr   `xml:"workbookPr"`
	BookViews    xlsxBookViews    `xml:"bookViews"`
	Sheets       xlsxSheets       `xml:"sheets"`
	DefinedNames xlsxDefinedNames `xml:"definedNames"`
	CalcPr       xlsxCalcPr       `xml:"calcPr"`
}

// xlsxFileVersion directly maps the fileVersion element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxFileVersion struct {
	AppName      string `xml:"appName,attr"`
	LastEdited   string `xml:"lastEdited,attr"`
	LowestEdited string `xml:"lowestEdited,attr"`
	RupBuild     string `xml:"rupBuild,attr"`
}

// xlsxWorkbookPr directly maps the workbookPr element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxWorkbookPr struct {
	DefaultThemeVersion string `xml:"defaultThemeVersion,attr"`
}

// xlsxBookViews directly maps the bookViews element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxBookViews struct {
	WorkBookView []xlsxWorkBookView `xml:"workbookView"`
}

// xlsxWorkBookView directly maps the workbookView element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxWorkBookView struct {
	XWindow      string `xml:"xWindow,attr"`
	YWindow      string `xml:"yWindow,attr"`
	WindowWidth  string `xml:"windowWidth,attr"`
	WindowHeight string `xml:"windowHeight,attr"`
}

// xlsxSheets directly maps the sheets element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxSheets struct {
	Sheet []xlsxSheet `xml:"sheet"`
}

// xlsxSheet directly maps the sheet element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxSheet struct {
	Name    string `xml:"name,attr"`
	SheetId string `xml:"sheetId,attr"`
	Id      string `xml:"id,attr"`
}

// xlsxDefinedNames directly maps the definedNames element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxDefinedNames struct {
	DefinedName []xlsxDefinedName `xml:"definedName"`
}

// xlsxDefinedName directly maps the definedName element from the
// namespace http://schemas.openxmlformats.org/spreadsheetml/2006/main
// - currently I have not checked it for completeness - it does as
// much as I need.
type xlsxDefinedName struct {
	Data         string `xml:",chardata"`
	Name         string `xml:"name,attr"`
	LocalSheetID string `xml:"localSheetId,attr"`
}

// xlsxCalcPr directly maps the calcPr element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main -
// currently I have not checked it for completeness - it does as much
// as I need.
type xlsxCalcPr struct {
	CalcId string `xml:"calcId,attr"`
}

// getWorksheetFromSheet() is an internal helper function to open a sheetN.xml file, refered to by an xlsx.xlsxSheet struct, from the XLSX file and unmarshal it an xlsx.xlsxWorksheet struct
func getWorksheetFromSheet(sheet xlsxSheet, worksheets map[string]*zip.File) (*xlsxWorksheet, error) {
	var rc io.ReadCloser
	var decoder *xml.Decoder
	var worksheet *xlsxWorksheet
	var error error
	worksheet = new(xlsxWorksheet)
	sheetName := fmt.Sprintf("sheet%s", sheet.Id[3:])
	f := worksheets[sheetName]
	rc, error = f.Open()
	if error != nil {
		return nil, error
	}
	decoder = xml.NewDecoder(rc)
	error = decoder.Decode(worksheet)
	if error != nil {
		return nil, error
	}
	return worksheet, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Davin Hills. All rights reserved.
// MIT license. License details can be found in the LICENSE file.

package goauto

import (
	"fmt"
	"regexp"
	"time"
)

// A Workflow represents a set of tasks for files matching one or more regex patterns
type Workflow struct {
	Name       string
	Concurrent bool
	Op         Op
	Regexs     []*regexp.Regexp
	Tasks      []Tasker
}

// NewWorkflow returns a Workflow with tasks
func NewWorkflow(tasks ...Tasker) *Workflow {
	wf := new(Workflow)
	wf.WatchOp(Create | Write | Remove | Rename)
	wf.Add(tasks...)
	return wf
}

// WatchPattern adds one or more regex for matching files for this workflow
// An invalid regexp pattern will return an error
// Sets file operations to Create|Write|Remove|Rename if not set
func (wf *Workflow) WatchPattern(patterns ...string) error {
	if wf.Op == 0 {
		wf.WatchOp(Create | Write | Remove | Rename)
	}
	for _, p := range patterns {
		r, err := regexp.Compile(p)
		if err != nil {
			return err
		}
		wf.Regexs = append(wf.Regexs, r)
	}
	return nil
}

// WatchOp sets the file operations to match
// The default is Create | Write | Remove | Rename
func (wf *Workflow) WatchOp(op Op) {
	wf.Op = op
}

// Match checks a file name against the regexp of the Workflow and the file operation
func (wf *Workflow) Match(fpath string, op uint32) bool {
	if uint32(wf.Op)&op == op {
		for _, r := range wf.Regexs {
			if r.MatchString(fpath) {
				return true
			}
		}
	}
	return false
}

// Add adds a task to the workflow
func (wf *Workflow) Add(tasks ...Tasker) {
	for _, t := range tasks {
		wf.Tasks = append(wf.Tasks, t)
	}
}

func (wf *Workflow) runner(info *TaskInfo) {
	if info.Verbose {
		fmt.Fprintf(info.Tout, ">> %v %v for %v\n\n", time.Now().Format("2006/01/02 3:04pm"), wf.Name, info.Src)
	}
	fname := info.Src
	info.Collect = []string{fname}
	var err error
	for _, t := range wf.Tasks {
		info.Target = "" // reset the Target
		if err = t.Run(info); err != nil {
			fmt.Fprintln(info.Terr, err)
			fmt.Fprintf(info.Terr, "Fail! Workflow %v did not complete for %v\n\n\n", wf.Name, fname)
			return
		}
		if info.Target != "" {
			// if the task set a target use it for the Src in the next task
			info.Src = info.Target
			info.Collect = append(info.Collect, info.Target)
		}
	}
	fmt.Fprintln(info.Tout, "")
}

// Run will start the execution of tasks
func (wf *Workflow) Run(info *TaskInfo) {
	if wf.Concurrent {
		go wf.runner(info)
		return
	}
	wf.runner(info)
}
<commit_msg>Remove the new line from the the workflow completion. If a user wants one they can add it to a task.<commit_after>// Copyright 2015 Davin Hills. All rights reserved.
// MIT license. License details can be found in the LICENSE file.

package goauto

import (
	"fmt"
	"regexp"
	"time"
)

// A Workflow represents a set of tasks for files matching one or more regex patterns
type Workflow struct {
	Name       string
	Concurrent bool
	Op         Op
	Regexs     []*regexp.Regexp
	Tasks      []Tasker
}

// NewWorkflow returns a Workflow with tasks
func NewWorkflow(tasks ...Tasker) *Workflow {
	wf := new(Workflow)
	wf.WatchOp(Create | Write | Remove | Rename)
	wf.Add(tasks...)
	return wf
}

// WatchPattern adds one or more regex for matching files for this workflow
// An invalid regexp pattern will return an error
// Sets file operations to Create|Write|Remove|Rename if not set
func (wf *Workflow) WatchPattern(patterns ...string) error {
	if wf.Op == 0 {
		wf.WatchOp(Create | Write | Remove | Rename)
	}
	for _, p := range patterns {
		r, err := regexp.Compile(p)
		if err != nil {
			return err
		}
		wf.Regexs = append(wf.Regexs, r)
	}
	return nil
}

// WatchOp sets the file operations to match
// The default is Create | Write | Remove | Rename
func (wf *Workflow) WatchOp(op Op) {
	wf.Op = op
}

// Match checks a file name against the regexp of the Workflow and the file operation
func (wf *Workflow) Match(fpath string, op uint32) bool {
	if uint32(wf.Op)&op == op {
		for _, r := range wf.Regexs {
			if r.MatchString(fpath) {
				return true
			}
		}
	}
	return false
}

// Add adds a task to the workflow
func (wf *Workflow) Add(tasks ...Tasker) {
	for _, t := range tasks {
		wf.Tasks = append(wf.Tasks, t)
	}
}

func (wf *Workflow) runner(info *TaskInfo) {
	if info.Verbose {
		fmt.Fprintf(info.Tout, ">> %v %v for %v\n\n", time.Now().Format("2006/01/02 3:04pm"), wf.Name, info.Src)
	}
	fname := info.Src
	info.Collect = []string{fname}
	var err error
	for _, t := range wf.Tasks {
		info.Target = "" // reset the Target
		if err = t.Run(info); err != nil {
			fmt.Fprintln(info.Terr, err)
			fmt.Fprintf(info.Terr, "Fail! Workflow %v did not complete for %v\n\n\n", wf.Name, fname)
			return
		}
		if info.Target != "" {
			// if the task set a target use it for the Src in the next task
			info.Src = info.Target
			info.Collect = append(info.Collect, info.Target)
		}
	}
}

// Run will start the execution of tasks
func (wf *Workflow) Run(info *TaskInfo) {
	if wf.Concurrent {
		go wf.runner(info)
		return
	}
	wf.runner(info)
}
<|endoftext|>


<commit_before>package dry

import (
	"bytes"
	"crypto/md5"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func StringPrettifyJSON(compactJSON string) string {
	var buf bytes.Buffer
	err := json.Indent(&buf, []byte(compactJSON), "", "\t")
	if err != nil {
		return err.Error()
	}
	return buf.String()
}

func StringEscapeJSON(jsonString string) string {
	jsonString = strings.Replace(jsonString, `\`, `\\`, -1)
	jsonString = strings.Replace(jsonString, `"`, `\"`, -1)
	return jsonString
}

func StringStripHTMLTags(text string) (plainText string) {
	chars := []byte(text)
	tagStart := -1
	for i := 0; i < len(chars); i++ {
		if chars[i] == '<' {
			tagStart = i
		} else if chars[i] == '>' && tagStart != -1 {
			chars = append(chars[:tagStart], chars[i+1:]...)
			i, tagStart = tagStart-1, -1
		}
	}
	return string(chars)
}

func StringMD5(data string) string {
	hash := md5.New()
	hash.Write([]byte(data))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func StringAddURLParam(url, name, value string) string {
	var separator string
	if strings.IndexRune(url, '?') == -1 {
		separator = "?"
	} else {
		separator = "&"
	}
	return url + separator + name + "=" + value
}

func StringConvertTime(timeString, formatIn, formatOut string) (resultTime string, err error) {
	if timeString == "" {
		return "", nil
	}
	t, err := time.Parse(formatIn, timeString)
	if err != nil {
		return "", err
	}
	return t.Format(formatOut), nil
}

func StringCSV(records [][]string) string {
	var buf bytes.Buffer
	writer := csv.NewWriter(&buf)
	err := writer.WriteAll(records)
	if err != nil {
		return ""
	}
	return buf.String()
}

func StringToInt(s string) int {
	i, _ := strconv.ParseInt(s, 10, 64)
	return int(i)
}

func StringToFloat(s string) float64 {
	f, _ := strconv.ParseFloat(s, 64)
	return f
}

func StringToBool(s string) bool {
	b, _ := strconv.ParseBool(s)
	return b
}

func StringInSlice(s string, slice []string) bool {
	for i := range slice {
		if slice[i] == s {
			return true
		}
	}
	return false
}

// StringJoinFormat formats every value in values with format
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoinFormat(format string, values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprintf(format, v.Index(i).Interface()))
	}
	return buffer.String()
}

// StringJoin formats every value in values according to its default formatting
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoin(values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprint(v.Index(i).Interface()))
	}
	return buffer.String()
}

func StringFormatBigInt(mem uint64) string {
	switch {
	case mem >= 10e12:
		return fmt.Sprintf("%dT", mem/1e12)
	case mem >= 1e12:
		return strings.TrimSuffix(fmt.Sprintf("%.1fT", float64(mem)/1e12), ".0")

	case mem >= 10e9:
		return fmt.Sprintf("%dG", mem/1e9)
	case mem >= 1e9:
		return strings.TrimSuffix(fmt.Sprintf("%.1fG", float64(mem)/1e9), ".0")

	case mem >= 10e6:
		return fmt.Sprintf("%dM", mem/1e6)
	case mem >= 1e6:
		return strings.TrimSuffix(fmt.Sprintf("%.1fM", float64(mem)/1e6), ".0")

	case mem >= 10e3:
		return fmt.Sprintf("%dk", mem/1e3)
	case mem >= 1e3:
		return strings.TrimSuffix(fmt.Sprintf("%.1fk", float64(mem)/1e3), ".0")
	}
	return fmt.Sprintf("%d", mem)
}

func StringFormatMemory(mem uint64) string {
	return StringFormatBigInt(mem) + "B"
}
<commit_msg>StringMD5Hex and StringSHA1Base64<commit_after>package dry

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"encoding/base64"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func StringPrettifyJSON(compactJSON string) string {
	var buf bytes.Buffer
	err := json.Indent(&buf, []byte(compactJSON), "", "\t")
	if err != nil {
		return err.Error()
	}
	return buf.String()
}

func StringEscapeJSON(jsonString string) string {
	jsonString = strings.Replace(jsonString, `\`, `\\`, -1)
	jsonString = strings.Replace(jsonString, `"`, `\"`, -1)
	return jsonString
}

func StringStripHTMLTags(text string) (plainText string) {
	chars := []byte(text)
	tagStart := -1
	for i := 0; i < len(chars); i++ {
		if chars[i] == '<' {
			tagStart = i
		} else if chars[i] == '>' && tagStart != -1 {
			chars = append(chars[:tagStart], chars[i+1:]...)
			i, tagStart = tagStart-1, -1
		}
	}
	return string(chars)
}

// StringMD5Hex returns the hex encoded MD5 hash of data
func StringMD5Hex(data string) string {
	hash := md5.New()
	hash.Write([]byte(data))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

// StringSHA1Base64 returns the base64 encoded SHA1 hash of data
func StringSHA1Base64(data string) string {
	hash := sha1.Sum([]byte(data))
	return base64.StdEncoding.EncodeToString(hash[:])
}

func StringAddURLParam(url, name, value string) string {
	var separator string
	if strings.IndexRune(url, '?') == -1 {
		separator = "?"
	} else {
		separator = "&"
	}
	return url + separator + name + "=" + value
}

func StringConvertTime(timeString, formatIn, formatOut string) (resultTime string, err error) {
	if timeString == "" {
		return "", nil
	}
	t, err := time.Parse(formatIn, timeString)
	if err != nil {
		return "", err
	}
	return t.Format(formatOut), nil
}

func StringCSV(records [][]string) string {
	var buf bytes.Buffer
	writer := csv.NewWriter(&buf)
	err := writer.WriteAll(records)
	if err != nil {
		return ""
	}
	return buf.String()
}

func StringToInt(s string) int {
	i, _ := strconv.ParseInt(s, 10, 64)
	return int(i)
}

func StringToFloat(s string) float64 {
	f, _ := strconv.ParseFloat(s, 64)
	return f
}

func StringToBool(s string) bool {
	b, _ := strconv.ParseBool(s)
	return b
}

func StringInSlice(s string, slice []string) bool {
	for i := range slice {
		if slice[i] == s {
			return true
		}
	}
	return false
}

// StringJoinFormat formats every value in values with format
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoinFormat(format string, values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprintf(format, v.Index(i).Interface()))
	}
	return buffer.String()
}

// StringJoin formats every value in values according to its default formatting
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoin(values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprint(v.Index(i).Interface()))
	}
	return buffer.String()
}

func StringFormatBigInt(mem uint64) string {
	switch {
	case mem >= 10e12:
		return fmt.Sprintf("%dT", mem/1e12)
	case mem >= 1e12:
		return strings.TrimSuffix(fmt.Sprintf("%.1fT", float64(mem)/1e12), ".0")

	case mem >= 10e9:
		return fmt.Sprintf("%dG", mem/1e9)
	case mem >= 1e9:
		return strings.TrimSuffix(fmt.Sprintf("%.1fG", float64(mem)/1e9), ".0")

	case mem >= 10e6:
		return fmt.Sprintf("%dM", mem/1e6)
	case mem >= 1e6:
		return strings.TrimSuffix(fmt.Sprintf("%.1fM", float64(mem)/1e6), ".0")

	case mem >= 10e3:
		return fmt.Sprintf("%dk", mem/1e3)
	case mem >= 1e3:
		return strings.TrimSuffix(fmt.Sprintf("%.1fk", float64(mem)/1e3), ".0")
	}
	return fmt.Sprintf("%d", mem)
}

func StringFormatMemory(mem uint64) string {
	return StringFormatBigInt(mem) + "B"
}
<|endoftext|>