<commit_before>package display

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/nanobox-io/nanobox/util/config"
)

var (
	// enable logging to a file
	Log = true

	// location of logfile
	LogFile = filepath.ToSlash(filepath.Join(config.GlobalDir(), "process.log"))

	// summarize the output and hide log details
	Summary = true

	// re-draw the summary when updates occur
	Interactive = terminal.IsTerminal(int(os.Stderr.Fd()))

	// info, warn, error, debug, trace
	Level = "info"

	// text, json
	Mode = "text"

	// writer to send output to
	Out = os.Stderr

	// internal
	logFile *os.File // open file descriptor of the log file
	// context
	context    int // track the context level
	topContext int // track the number to toplevel contexts
	// task
	taskStarted bool          // track if we're running a task
	taskLog     *bytes.Buffer // track the log of the current task, in case it fails
	prefixer    *Prefixer     // use a prefixer to prefix logs
	summarizer  *Summarizer   // summarizer to summarize the current task
)

// OpenContext opens a context level and prints the header
func OpenContext(format string, args ...interface{}) error {
	label := fmt.Sprintf(format, args...)

	// if the current context is 0, let's increment the topContext
	if context == 0 {
		topContext += 1
	}

	// increment the context level counter
	context += 1

	// if this is a subsequent top-level context, let's prefix with a newline
	if topContext > 1 && context == 1 {
		if err := printAll("\n"); err != nil {
			return err
		}
	}

	prefix := ""

	if context > 0 {
		prefix = strings.Repeat("  ", context-1)
	}

	header := fmt.Sprintf("%s+ %s :\n", prefix, label)

	if err := printAll(header); err != nil {
		return err
	}

	return nil
}

// CloseContext closes the context level and prints a newline
func CloseContext() error {

	// decrement the context level counter
	context -= 1

	// ensure the context doesn't drop below zero
	if context < 0 {
		context = 0
	}

	return nil
}

// StartTask starts a new task
func StartTask(format string, args ...interface{}) error {
	label := fmt.Sprintf(format, args...)

	// return an error if the current task has not ended
	if taskStarted {
		return errors.New("Current task has not been stopped")
	}

	// mark the task as started
	taskStarted = true

	// initialize the task log
	taskLog = bytes.NewBufferString("")

	// create a new prefixer
	prefixer = NewPrefixer(strings.Repeat("  ", context+1))

	// generate a header
	prefix := strings.Repeat("  ", context)
	header := fmt.Sprintf("%s+ %s :\n", prefix, label)

	// print the header to the logfile
	if err := printLogFile(header); err != nil {
		return err
	}

	if Summary {
		summarizer = NewSummarizer(label, prefix)
		summarizer.Start()
	} else {
		// print the header
		if err := printOut(header); err != nil {
			return err
		}
	}

	return nil
}

// StopTask stops the current task
func StopTask() error {

	// stop the task summarizer
	if Summary && summarizer != nil {
		summarizer.Stop()
		summarizer = nil
	}

	// mark task as stopped
	taskStarted = false

	// reset the task log
	taskLog = nil

	// reset the prefixer
	prefixer = nil

	return nil
}

// ErrorTask errors the current task
func ErrorTask() error {

	// stop the task summarizer
	if Summary {
		summarizer.Error()
		summarizer = nil

		// print the task log
		Out.Write(taskLog.Bytes())
	}

	// mark task as stopped
	taskStarted = false

	// reset the task log
	taskLog = nil

	// reset the prefixer
	prefixer = nil

	return nil
}

// Info sends an info level message to the current task
func Info(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 2 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Warn sends a warn level message to the current task
func Warn(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 3 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Error sends an error level message to the current task
func Error(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 4 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Debug sends a debug level message to the current task
func Debug(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 1 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Trace sends a trace level message to the current task
func Trace(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 0 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// log logs a message to the current task
func log(message string) error {

	// run the message through prefixer
	if prefixer != nil {
		message = prefixer.Parse(message)
	}

	// append to the taskLog
	if taskLog != nil {
		taskLog.WriteString(message)
	}

	// print message to logfile
	if err := printLogFile(message); err != nil {
		return err
	}

	if Summary && summarizer != nil {
		summarizer.Log(message)
	} else {
		// print the message
		if err := printOut(message); err != nil {
			return err
		}
	}

	return nil
}

// printAll prints a message to the Out channel and the logfile
func printAll(message string) error {

	// print to the Out writer
	if err := printOut(message); err != nil {
		return err
	}

	// print to the log file
	if err := printLogFile(message); err != nil {
		return err
	}

	return nil
}

// printOut will print a message to the out stream
func printOut(message string) error {
	_, err := io.WriteString(Out, message)
	return err
}

// printLogFile prints a message to the log file
func printLogFile(message string) error {
	// short-circuit if Log is set to false
	if !Log {
		return nil
	}

	// make sure the logfile is opened
	if err := openLogFile(); err != nil {
		return err
	}

	// print to the logfile
	logFile.WriteString(message)

	return nil
}

// openLogFile opens the logFile for writes
func openLogFile() error {

	// short-circuit if the logFile is already open
	if logFile != nil {
		return nil
	}

	truncate := os.O_RDWR | os.O_CREATE | os.O_TRUNC

	f, err := os.OpenFile(LogFile, truncate, 0644)
	if err != nil {
		return err
	}

	logFile = f

	return nil
}

// currentLogLevel returns the current log level as an int
func currentLogLevel() int {
	switch Level {
	case "error":
		return 4
	case "warn":
		return 3
	case "info":
		return 2
	case "debug":
		return 1
	case "trace":
		return 0
	}

	return 0
}
<commit_msg>update to remvoe summerizer<commit_after>package display

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/nanobox-io/nanobox/util/config"
)

var (
	// enable logging to a file
	Log = true

	// location of logfile
	LogFile = filepath.ToSlash(filepath.Join(config.GlobalDir(), "process.log"))

	// summarize the output and hide log details
	Summary = true

	// re-draw the summary when updates occur
	Interactive = terminal.IsTerminal(int(os.Stderr.Fd()))

	// info, warn, error, debug, trace
	Level = "info"

	// text, json
	Mode = "text"

	// writer to send output to
	Out = os.Stderr

	// internal
	logFile *os.File // open file descriptor of the log file
	// context
	context    int // track the context level
	topContext int // track the number to toplevel contexts
	// task
	taskStarted bool          // track if we're running a task
	taskLog     *bytes.Buffer // track the log of the current task, in case it fails
	prefixer    *Prefixer     // use a prefixer to prefix logs
	summarizer  *Summarizer   // summarizer to summarize the current task
)

// OpenContext opens a context level and prints the header
func OpenContext(format string, args ...interface{}) error {
	label := fmt.Sprintf(format, args...)

	// if the current context is 0, let's increment the topContext
	if context == 0 {
		topContext += 1
	}

	// increment the context level counter
	context += 1

	// if this is a subsequent top-level context, let's prefix with a newline
	if topContext > 1 && context == 1 {
		if err := printAll("\n"); err != nil {
			return err
		}
	}

	prefix := ""

	if context > 0 {
		prefix = strings.Repeat("  ", context-1)
	}

	header := fmt.Sprintf("%s+ %s :\n", prefix, label)

	if err := printAll(header); err != nil {
		return err
	}

	return nil
}

// CloseContext closes the context level and prints a newline
func CloseContext() error {

	// decrement the context level counter
	context -= 1

	// ensure the context doesn't drop below zero
	if context < 0 {
		context = 0
	}

	return nil
}

// StartTask starts a new task
func StartTask(format string, args ...interface{}) error {
	label := fmt.Sprintf(format, args...)

	// return an error if the current task has not ended
	if taskStarted {
		return errors.New("Current task has not been stopped")
	}

	// mark the task as started
	taskStarted = true

	// initialize the task log
	taskLog = bytes.NewBufferString("")

	// create a new prefixer
	prefixer = NewPrefixer(strings.Repeat("  ", context+1))

	// generate a header
	prefix := strings.Repeat("  ", context)
	header := fmt.Sprintf("%s+ %s :\n", prefix, label)

	// print the header to the logfile
	if err := printLogFile(header); err != nil {
		return err
	}

	if Summary {
		summarizer = NewSummarizer(label, prefix)
		summarizer.Start()
	} else {
		// print the header
		if err := printOut(header); err != nil {
			return err
		}
	}

	return nil
}

// StopTask stops the current task
func StopTask() error {

	// stop the task summarizer
	if Summary && summarizer != nil {
		summarizer.Stop()
		summarizer = nil
	}

	// mark task as stopped
	taskStarted = false

	// reset the task log
	taskLog = nil

	// reset the prefixer
	prefixer = nil

	return nil
}

// ErrorTask errors the current task
func ErrorTask() error {

	// stop the task summarizer
	if Summary && summarizer != nil {
		summarizer.Error()
		summarizer = nil

		// print the task log
		Out.Write(taskLog.Bytes())
	}

	// mark task as stopped
	taskStarted = false

	// reset the task log
	taskLog = nil

	// reset the prefixer
	prefixer = nil

	return nil
}

// Info sends an info level message to the current task
func Info(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 2 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Warn sends a warn level message to the current task
func Warn(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 3 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Error sends an error level message to the current task
func Error(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 4 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Debug sends a debug level message to the current task
func Debug(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 1 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// Trace sends a trace level message to the current task
func Trace(message string, args ...interface{}) error {
	if len(args) != 0 {
		message = fmt.Sprintf(message, args...)
	}

	// short-circuit if our log-level isn't high enough
	if currentLogLevel() > 0 {
		return nil
	}

	if err := log(message); err != nil {
		return err
	}

	return nil
}

// log logs a message to the current task
func log(message string) error {

	// run the message through prefixer
	if prefixer != nil {
		message = prefixer.Parse(message)
	}

	// append to the taskLog
	if taskLog != nil {
		taskLog.WriteString(message)
	}

	// print message to logfile
	if err := printLogFile(message); err != nil {
		return err
	}

	if Summary && summarizer != nil {
		summarizer.Log(message)
	} else {
		// print the message
		if err := printOut(message); err != nil {
			return err
		}
	}

	return nil
}

// printAll prints a message to the Out channel and the logfile
func printAll(message string) error {

	// print to the Out writer
	if err := printOut(message); err != nil {
		return err
	}

	// print to the log file
	if err := printLogFile(message); err != nil {
		return err
	}

	return nil
}

// printOut will print a message to the out stream
func printOut(message string) error {
	_, err := io.WriteString(Out, message)
	return err
}

// printLogFile prints a message to the log file
func printLogFile(message string) error {
	// short-circuit if Log is set to false
	if !Log {
		return nil
	}

	// make sure the logfile is opened
	if err := openLogFile(); err != nil {
		return err
	}

	// print to the logfile
	logFile.WriteString(message)

	return nil
}

// openLogFile opens the logFile for writes
func openLogFile() error {

	// short-circuit if the logFile is already open
	if logFile != nil {
		return nil
	}

	truncate := os.O_RDWR | os.O_CREATE | os.O_TRUNC

	f, err := os.OpenFile(LogFile, truncate, 0644)
	if err != nil {
		return err
	}

	logFile = f

	return nil
}

// currentLogLevel returns the current log level as an int
func currentLogLevel() int {
	switch Level {
	case "error":
		return 4
	case "warn":
		return 3
	case "info":
		return 2
	case "debug":
		return 1
	case "trace":
		return 0
	}

	return 0
}
<|endoftext|>


<commit_before>package stats

import (
	"fmt"
	"log"
	"net/http"
	"sync"
)

var _ = log.Println

type ServerStats struct {
	mutex       sync.Mutex
	phaseCounts map[string]int
	*CpuStatSampler
	PhaseOrder []string
}

func NewServerStats() *ServerStats {
	return &ServerStats{
		phaseCounts:    map[string]int{},
		CpuStatSampler: NewCpuStatSampler(),
	}
}

func (me *ServerStats) Enter(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]++
}

func (me *ServerStats) Exit(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]--
}

func CpuStatsWriteHttp(w http.ResponseWriter, stats []CpuStat) {
	if len(stats) == 0 {
		return
	}

	statm := CpuStat{}
	stat5 := CpuStat{}
	count5 := int64(0)
	for i, v := range stats {
		statm = statm.Add(v)
		if i >= len(stats)-5 {
			count5++
			stat5 = stat5.Add(v)
		}
	}

	printChild := statm.ChildCpu + statm.ChildSys > 0
	chHeader := ""
	if printChild {
		chHeader = "child cpu (ms)	child sys (ms)	"
	}
	fmt.Fprintf(w, "

self cpu (ms)	self sys (ms)	%stotal
",
		chHeader)
	for i, v := range stats {
		if i < len(stats)-5 {
			continue
		}

		chRow := ""
		if printChild {
			chRow = fmt.Sprintf("%d	%d	", v.ChildCpu/1e6, v.ChildSys/1e6)
		}

		fmt.Fprintf(w, "
%d	%d	%s%d
",
			v.SelfCpu/1e6, v.SelfSys/1e6, chRow,
			(v.Total())/1e6)
	}
	fmt.Fprintf(w, "
")

	fmt.Fprintf(w, "
CPU (last min): %d s self %d s sys", statm.SelfCpu/1e9, statm.SelfSys/1e9)
	if printChild {
		fmt.Fprintf(w, " %d s child %d s sys", statm.ChildCpu/1e9, statm.ChildSys/1e9)
	}

	fmt.Fprintf(w, " %.2f CPU", float64(statm.Total())/1e9 / float64(len(stats)))
	fmt.Fprintf(w, "
CPU (last %ds): %.2f self %.2f sys",
			count5, float64(stat5.SelfCpu)*1e-9, float64(stat5.SelfSys)*1e-9)
	if printChild {
		fmt.Fprintf(w, "%.2f s child %.2f s sys", float64(stat5.ChildCpu)*1e-9, float64(stat5.ChildSys)*1e-9)
	}
	fmt.Fprintf(w, " %.2f CPU", float64(statm.Total())/1e9 / float64(count5))
}

func CountStatsWriteHttp(w http.ResponseWriter, names []string, counts []int) {
	fmt.Fprintf(w, "
")
	for i, c := range counts {
		fmt.Fprintf(w, "
Jobs in phase %s: %d ", names[i], c)
	}
	fmt.Fprintf(w, "
")
}

func (me *ServerStats) WriteHttp(w http.ResponseWriter) {
	CpuStatsWriteHttp(w, me.CpuStats())
	CountStatsWriteHttp(w, me.PhaseOrder, me.PhaseCounts())
}

func (me *ServerStats) PhaseCounts() (r []int) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for _, n := range me.PhaseOrder {
		r = append(r, me.phaseCounts[n])
	}
	return r
}
<commit_msg>Fix thinko: use stat5 for 5s average of load.<commit_after>package stats

import (
	"fmt"
	"log"
	"net/http"
	"sync"
)

var _ = log.Println

type ServerStats struct {
	mutex       sync.Mutex
	phaseCounts map[string]int
	*CpuStatSampler
	PhaseOrder []string
}

func NewServerStats() *ServerStats {
	return &ServerStats{
		phaseCounts:    map[string]int{},
		CpuStatSampler: NewCpuStatSampler(),
	}
}

func (me *ServerStats) Enter(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]++
}

func (me *ServerStats) Exit(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]--
}

func CpuStatsWriteHttp(w http.ResponseWriter, stats []CpuStat) {
	if len(stats) == 0 {
		return
	}

	statm := CpuStat{}
	stat5 := CpuStat{}
	count5 := int64(0)
	for i, v := range stats {
		statm = statm.Add(v)
		if i >= len(stats)-5 {
			count5++
			stat5 = stat5.Add(v)
		}
	}

	printChild := statm.ChildCpu + statm.ChildSys > 0
	chHeader := ""
	if printChild {
		chHeader = "child cpu (ms)	child sys (ms)	"
	}
	fmt.Fprintf(w, "

self cpu (ms)	self sys (ms)	%stotal
",
		chHeader)
	for i, v := range stats {
		if i < len(stats)-5 {
			continue
		}

		chRow := ""
		if printChild {
			chRow = fmt.Sprintf("%d	%d	", v.ChildCpu/1e6, v.ChildSys/1e6)
		}

		fmt.Fprintf(w, "
%d	%d	%s%d
",
			v.SelfCpu/1e6, v.SelfSys/1e6, chRow,
			(v.Total())/1e6)
	}
	fmt.Fprintf(w, "
")

	fmt.Fprintf(w, "
CPU (last min): %d s self %d s sys", statm.SelfCpu/1e9, statm.SelfSys/1e9)
	if printChild {
		fmt.Fprintf(w, " %d s child %d s sys", statm.ChildCpu/1e9, statm.ChildSys/1e9)
	}

	fmt.Fprintf(w, " %.2f CPU", float64(statm.Total())/1e9 / float64(len(stats)))
	fmt.Fprintf(w, "
CPU (last %ds): %.2f self %.2f sys",
			count5, float64(stat5.SelfCpu)*1e-9, float64(stat5.SelfSys)*1e-9)
	if printChild {
		fmt.Fprintf(w, "%.2f s child %.2f s sys", float64(stat5.ChildCpu)*1e-9, float64(stat5.ChildSys)*1e-9)
	}
	fmt.Fprintf(w, " %.2f CPU", float64(stat5.Total())/1e9 / float64(count5))
}

func CountStatsWriteHttp(w http.ResponseWriter, names []string, counts []int) {
	fmt.Fprintf(w, "
")
	for i, c := range counts {
		fmt.Fprintf(w, "
Jobs in phase %s: %d ", names[i], c)
	}
	fmt.Fprintf(w, "
")
}

func (me *ServerStats) WriteHttp(w http.ResponseWriter) {
	CpuStatsWriteHttp(w, me.CpuStats())
	CountStatsWriteHttp(w, me.PhaseOrder, me.PhaseCounts())
}

func (me *ServerStats) PhaseCounts() (r []int) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for _, n := range me.PhaseOrder {
		r = append(r, me.phaseCounts[n])
	}
	return r
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cgotest

/*
#include <complex.h>

complex float complexFloatSquared(complex float a) { return a*a; }
complex double complexDoubleSquared(complex double a) { return a*a; }
*/
import "C"

import "testing"

func test8694(t *testing.T) {
	// Really just testing that this compiles, but check answer anyway.
	x := complex64(2 + 3i)
	x2 := x * x
	cx2 := C.complexFloatSquared(x)
	if cx2 != x2 {
		t.Errorf("C.complexFloatSquared(%v) = %v, want %v", x, cx2, x2)
	}

	y := complex128(2 + 3i)
	y2 := y * y
	cy2 := C.complexDoubleSquared(y)
	if cy2 != y2 {
		t.Errorf("C.complexDoubleSquared(%v) = %v, want %v", y, cy2, y2)
	}
}
<commit_msg>misc/cgo/test: skip test8694 on ARM.<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cgotest

/*
#include <complex.h>

complex float complexFloatSquared(complex float a) { return a*a; }
complex double complexDoubleSquared(complex double a) { return a*a; }
*/
import "C"

import (
	"runtime"
	"testing"
)

func test8694(t *testing.T) {
	if runtime.GOARCH == "arm" {
		t.Skip("test8694 is disabled on ARM because 5l cannot handle thumb library.")
	}
	// Really just testing that this compiles, but check answer anyway.
	x := complex64(2 + 3i)
	x2 := x * x
	cx2 := C.complexFloatSquared(x)
	if cx2 != x2 {
		t.Errorf("C.complexFloatSquared(%v) = %v, want %v", x, cx2, x2)
	}

	y := complex128(2 + 3i)
	y2 := y * y
	cy2 := C.complexDoubleSquared(y)
	if cy2 != y2 {
		t.Errorf("C.complexDoubleSquared(%v) = %v, want %v", y, cy2, y2)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"sync"

	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

type SecretRetriever interface {
	RetrieveSecret() (LKSecFullSecret, error)
}

type SecretStorer interface {
	StoreSecret(secret LKSecFullSecret) error
}

type SecretStore interface {
	SecretRetriever
	SecretStorer
}

type SecretStoreAll interface {
	RetrieveSecret(username NormalizedUsername) (LKSecFullSecret, error)
	StoreSecret(username NormalizedUsername, secret LKSecFullSecret) error
	ClearSecret(username NormalizedUsername) error
	GetUsersWithStoredSecrets() ([]string, error)
}

type SecretStoreContext interface {
	GetAllUserNames() (NormalizedUsername, []NormalizedUsername, error)
	GetStoredSecretServiceName() string
	GetStoredSecretAccessGroup() string
	GetLog() logger.Logger
}

type SecretStoreImp struct {
	username NormalizedUsername
	store    *SecretStoreLocked
	secret   LKSecFullSecret
	sync.Mutex
}

func (s *SecretStoreImp) RetrieveSecret() (LKSecFullSecret, error) {
	s.Lock()
	defer s.Unlock()

	if !s.secret.IsNil() {
		return s.secret, nil
	}
	sec, err := s.store.RetrieveSecret(s.username)
	if err != nil {
		return sec, err
	}
	s.secret = sec
	return sec, nil
}

func (s *SecretStoreImp) StoreSecret(secret LKSecFullSecret) error {
	s.Lock()
	defer s.Unlock()

	// clear out any in-memory secret in this instance
	s.secret = LKSecFullSecret{}
	return s.store.StoreSecret(s.username, secret)
}

// NewSecretStore returns a SecretStore interface that is only used for
// a short period of time (i.e. one function block).  Multiple calls to RetrieveSecret()
// will only call the underlying store.RetrieveSecret once.
func NewSecretStore(g *GlobalContext, username NormalizedUsername) SecretStore {
	store := g.SecretStore()
	if store != nil {
		return &SecretStoreImp{
			username: username,
			store:    store,
		}
	}
	return nil
}

func GetConfiguredAccounts(c SecretStoreContext, s SecretStoreAll) ([]keybase1.ConfiguredAccount, error) {
	currentUsername, otherUsernames, err := c.GetAllUserNames()
	if err != nil {
		return nil, err
	}

	allUsernames := append(otherUsernames, currentUsername)

	accounts := make(map[NormalizedUsername]keybase1.ConfiguredAccount)

	for _, username := range allUsernames {
		accounts[username] = keybase1.ConfiguredAccount{
			Username: username.String(),
		}
	}
	var storedSecretUsernames []string
	if s != nil {
		storedSecretUsernames, err = s.GetUsersWithStoredSecrets()
	}
	if err != nil {
		return nil, err
	}

	for _, username := range storedSecretUsernames {
		nu := NewNormalizedUsername(username)
		account, ok := accounts[nu]
		if ok {
			account.HasStoredSecret = true
			accounts[nu] = account
		}
	}

	configuredAccounts := make([]keybase1.ConfiguredAccount, 0, len(accounts))
	for _, account := range accounts {
		configuredAccounts = append(configuredAccounts, account)
	}

	return configuredAccounts, nil
}

func ClearStoredSecret(g *GlobalContext, username NormalizedUsername) error {
	ss := g.SecretStore()
	if ss == nil {
		return nil
	}
	return ss.ClearSecret(username)
}

// SecretStoreLocked protects a SecretStoreAll with a mutex.
type SecretStoreLocked struct {
	SecretStoreAll
	sync.Mutex
}

func NewSecretStoreLocked(g *GlobalContext) *SecretStoreLocked {
	var ss SecretStoreAll

	if g.Env.RememberPassphrase() {
		// use os-specific secret store
		g.Log.Debug("NewSecretStoreLocked: using os-specific SecretStore")
		ss = NewSecretStoreAll(g)
	} else {
		// config or command line flag said to use in-memory secret store
		g.Log.Debug("NewSecretStoreLocked: using memory-only SecretStore")
		ss = NewSecretStoreMem()
	}

	if ss == nil {
		// right now, some stuff depends on g.SecretStoreAll being nil or not
		return nil
	}

	return &SecretStoreLocked{
		SecretStoreAll: ss,
	}
}

func (s *SecretStoreLocked) RetrieveSecret(username NormalizedUsername) (LKSecFullSecret, error) {
	if s == nil || s.SecretStoreAll == nil {
		return LKSecFullSecret{}, nil
	}
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.RetrieveSecret(username)
}

func (s *SecretStoreLocked) StoreSecret(username NormalizedUsername, secret LKSecFullSecret) error {
	if s == nil || s.SecretStoreAll == nil {
		return nil
	}
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.StoreSecret(username, secret)
}

func (s *SecretStoreLocked) ClearSecret(username NormalizedUsername) error {
	if s == nil || s.SecretStoreAll == nil {
		return nil
	}
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.ClearSecret(username)
}

func (s *SecretStoreLocked) GetUsersWithStoredSecrets() ([]string, error) {
	if s == nil || s.SecretStoreAll == nil {
		return nil, nil
	}
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.GetUsersWithStoredSecrets()
}
<commit_msg>always use a memory SecretStore, sometimes a disk (#12456)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"sync"

	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

type SecretRetriever interface {
	RetrieveSecret() (LKSecFullSecret, error)
}

type SecretStorer interface {
	StoreSecret(secret LKSecFullSecret) error
}

type SecretStore interface {
	SecretRetriever
	SecretStorer
}

type SecretStoreAll interface {
	RetrieveSecret(username NormalizedUsername) (LKSecFullSecret, error)
	StoreSecret(username NormalizedUsername, secret LKSecFullSecret) error
	ClearSecret(username NormalizedUsername) error
	GetUsersWithStoredSecrets() ([]string, error)
}

type SecretStoreContext interface {
	GetAllUserNames() (NormalizedUsername, []NormalizedUsername, error)
	GetStoredSecretServiceName() string
	GetStoredSecretAccessGroup() string
	GetLog() logger.Logger
}

type SecretStoreImp struct {
	username NormalizedUsername
	store    *SecretStoreLocked
	secret   LKSecFullSecret
	sync.Mutex
}

func (s *SecretStoreImp) RetrieveSecret() (LKSecFullSecret, error) {
	s.Lock()
	defer s.Unlock()

	if !s.secret.IsNil() {
		return s.secret, nil
	}
	sec, err := s.store.RetrieveSecret(s.username)
	if err != nil {
		return sec, err
	}
	s.secret = sec
	return sec, nil
}

func (s *SecretStoreImp) StoreSecret(secret LKSecFullSecret) error {
	s.Lock()
	defer s.Unlock()

	// clear out any in-memory secret in this instance
	s.secret = LKSecFullSecret{}
	return s.store.StoreSecret(s.username, secret)
}

// NewSecretStore returns a SecretStore interface that is only used for
// a short period of time (i.e. one function block).  Multiple calls to RetrieveSecret()
// will only call the underlying store.RetrieveSecret once.
func NewSecretStore(g *GlobalContext, username NormalizedUsername) SecretStore {
	store := g.SecretStore()
	if store != nil {
		return &SecretStoreImp{
			username: username,
			store:    store,
		}
	}
	return nil
}

func GetConfiguredAccounts(c SecretStoreContext, s SecretStoreAll) ([]keybase1.ConfiguredAccount, error) {
	currentUsername, otherUsernames, err := c.GetAllUserNames()
	if err != nil {
		return nil, err
	}

	allUsernames := append(otherUsernames, currentUsername)

	accounts := make(map[NormalizedUsername]keybase1.ConfiguredAccount)

	for _, username := range allUsernames {
		accounts[username] = keybase1.ConfiguredAccount{
			Username: username.String(),
		}
	}
	var storedSecretUsernames []string
	if s != nil {
		storedSecretUsernames, err = s.GetUsersWithStoredSecrets()
	}
	if err != nil {
		return nil, err
	}

	for _, username := range storedSecretUsernames {
		nu := NewNormalizedUsername(username)
		account, ok := accounts[nu]
		if ok {
			account.HasStoredSecret = true
			accounts[nu] = account
		}
	}

	configuredAccounts := make([]keybase1.ConfiguredAccount, 0, len(accounts))
	for _, account := range accounts {
		configuredAccounts = append(configuredAccounts, account)
	}

	return configuredAccounts, nil
}

func ClearStoredSecret(g *GlobalContext, username NormalizedUsername) error {
	ss := g.SecretStore()
	if ss == nil {
		return nil
	}
	return ss.ClearSecret(username)
}

// SecretStoreLocked protects a SecretStoreAll with a mutex. It wraps two different
// SecretStoreAlls: one in memory and one in disk. In all cases, we always have a memory
// backing. If the OS and options provide one, we can additionally have a disk-backed
// secret store. It's a write-through cache, so on RetrieveSecret, the memory store
// will be checked first, and then the disk store.
type SecretStoreLocked struct {
	Contextified
	sync.Mutex
	mem  SecretStoreAll
	disk SecretStoreAll
}

func NewSecretStoreLocked(g *GlobalContext) *SecretStoreLocked {
	var disk SecretStoreAll

	mem := NewSecretStoreMem()

	if g.Env.RememberPassphrase() {
		// use os-specific secret store
		g.Log.Debug("NewSecretStoreLocked: using os-specific SecretStore")
		disk = NewSecretStoreAll(g)
	} else {
		// config or command line flag said to use in-memory secret store
		g.Log.Debug("NewSecretStoreLocked: using memory-only SecretStore")
	}

	return &SecretStoreLocked{
		Contextified: NewContextified(g),
		mem:          mem,
		disk:         disk,
	}
}

func (s *SecretStoreLocked) isNil() bool {
	return s.mem == nil && s.disk == nil
}

func (s *SecretStoreLocked) RetrieveSecret(username NormalizedUsername) (LKSecFullSecret, error) {
	if s == nil || s.isNil() {
		return LKSecFullSecret{}, nil
	}
	s.Lock()
	defer s.Unlock()

	res, err := s.mem.RetrieveSecret(username)
	if !res.IsNil() && err == nil {
		return res, nil
	}
	if err != nil {
		s.G().Log.Debug("SecretStoreLocked#RetrieveSecret: memory fetch error: %s", err.Error())
	}
	if s.disk == nil {
		return res, err
	}

	res, err = s.disk.RetrieveSecret(username)
	if err != nil {
		return res, err
	}
	tmp := s.mem.StoreSecret(username, res)
	if tmp != nil {
		s.G().Log.Debug("SecretStoreLocked#RetrieveSecret: failed to store secret in memory: %s", err.Error())
	}
	return res, err
}

func (s *SecretStoreLocked) StoreSecret(username NormalizedUsername, secret LKSecFullSecret) error {
	if s == nil || s.isNil() {
		return nil
	}
	s.Lock()
	defer s.Unlock()
	err := s.mem.StoreSecret(username, secret)
	if err != nil {
		s.G().Log.Debug("SecretStoreLocked#StoreSecret: failed to store secret in memory: %s", err.Error())
	}
	if s.disk == nil {
		return err
	}
	return s.disk.StoreSecret(username, secret)
}

func (s *SecretStoreLocked) ClearSecret(username NormalizedUsername) error {
	if s == nil || s.isNil() {
		return nil
	}
	s.Lock()
	defer s.Unlock()
	err := s.mem.ClearSecret(username)
	if err != nil {
		s.G().Log.Debug("SecretStoreLocked#ClearSecret: failed to clear memory: %s", err.Error())
	}
	if s.disk == nil {
		return err
	}
	return s.disk.ClearSecret(username)
}

func (s *SecretStoreLocked) GetUsersWithStoredSecrets() ([]string, error) {
	if s == nil || s.isNil() {
		return nil, nil
	}
	s.Lock()
	defer s.Unlock()
	if s.disk == nil {
		return s.mem.GetUsersWithStoredSecrets()
	}
	return s.disk.GetUsersWithStoredSecrets()
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// WorkerInstance encapsulate the execution state of vtworker.
type WorkerInstance struct {
	// global wrangler object we'll use
	Wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error

	TopoServer             topo.Server
	cell                   string
	LockTimeout            time.Duration
	commandDisplayInterval time.Duration
}

// NewWorkerInstance creates a new WorkerInstance.
func NewWorkerInstance(ts topo.Server, cell string, lockTimeout, commandDisplayInterval time.Duration) *WorkerInstance {
	return &WorkerInstance{TopoServer: ts, cell: cell, LockTimeout: lockTimeout, commandDisplayInterval: commandDisplayInterval}
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *WorkerInstance) setAndStartWorker(wrk Worker) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()
	if wi.currentWorker != nil {
		return nil, fmt.Errorf("A worker is already in progress: %v", wi.currentWorker)
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(context.Background())
	wi.lastRunError = nil
	done := make(chan struct{})
	wi.Wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, logutil.NewConsoleLogger()))

	// one go function runs the worker, changes state when done
	go func() {
		// run will take a long time
		log.Infof("Starting worker...")
		err := wrk.Run(wi.currentContext)

		// it's done, let's save our state
		wi.currentWorkerMutex.Lock()
		wi.currentContext = nil
		wi.currentCancelFunc = nil
		wi.lastRunError = err
		wi.currentWorkerMutex.Unlock()
		close(done)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *WorkerInstance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		<-sigChan
		// we got a signal, notify our modules
		wi.currentWorkerMutex.Lock()
		defer wi.currentWorkerMutex.Unlock()
		if wi.currentCancelFunc != nil {
			wi.currentCancelFunc()
		}
	}()
}
<commit_msg>vtworker: Shutdown idle worker if SIGTERM or SIGINT was received.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// WorkerInstance encapsulate the execution state of vtworker.
type WorkerInstance struct {
	// global wrangler object we'll use
	Wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error

	TopoServer             topo.Server
	cell                   string
	LockTimeout            time.Duration
	commandDisplayInterval time.Duration
}

// NewWorkerInstance creates a new WorkerInstance.
func NewWorkerInstance(ts topo.Server, cell string, lockTimeout, commandDisplayInterval time.Duration) *WorkerInstance {
	return &WorkerInstance{TopoServer: ts, cell: cell, LockTimeout: lockTimeout, commandDisplayInterval: commandDisplayInterval}
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *WorkerInstance) setAndStartWorker(wrk Worker) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()
	if wi.currentWorker != nil {
		return nil, fmt.Errorf("A worker is already in progress: %v", wi.currentWorker)
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(context.Background())
	wi.lastRunError = nil
	done := make(chan struct{})
	wi.Wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, logutil.NewConsoleLogger()))

	// one go function runs the worker, changes state when done
	go func() {
		// run will take a long time
		log.Infof("Starting worker...")
		err := wrk.Run(wi.currentContext)

		// it's done, let's save our state
		wi.currentWorkerMutex.Lock()
		wi.currentContext = nil
		wi.currentCancelFunc = nil
		wi.lastRunError = err
		wi.currentWorkerMutex.Unlock()
		close(done)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *WorkerInstance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		s := <-sigChan
		// we got a signal, notify our modules
		wi.currentWorkerMutex.Lock()
		defer wi.currentWorkerMutex.Unlock()
		if wi.currentCancelFunc != nil {
			wi.currentCancelFunc()
		} else {
			fmt.Printf("Shutting down idle worker after receiving signal: %v", s)
			os.Exit(0)
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"github.com/gorilla/feeds"
	"github.com/mvdan/xurls"
	"io"
	"net/http"
	"strings"
	"time"
)

func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

func isSpace(b byte) bool {
	return b == ' ' || b == '\t'
}

func scanContentLines(data []byte, atEOF bool) (int, []byte, error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	var (
		advance = 0
		buffer  bytes.Buffer
	)
	for {
		if i := bytes.IndexByte(data, '\n'); i >= 0 {
			advance += i + 1
			buffer.Write(dropCR(data[0:i]))

			data = data[i+1:]
			if len(data) > 0 && isSpace(data[0]) {
				j := 1
				for ; j < len(data); j++ {
					if !isSpace(data[j]) {
						break
					}
				}
				advance += j
				data = data[j:]
			} else {
				break
			}
		} else if atEOF {
			advance += len(data)
			buffer.Write(dropCR(data))
		} else {
			return 0, nil, nil
		}
	}
	if advance > 0 {
		return advance, buffer.Bytes(), nil
	}
	return 0, nil, nil
}

type Event struct {
	Summary      string
	Description  string
	Location     string
	Uid          string
	Start        time.Time
	End          time.Time
	Created      time.Time
	LastModified time.Time
}

func parseTime(dateString string, tzId string) (time.Time, error) {
	if tzId == "" {
		location, err := time.LoadLocation("UTC")
		if err != nil {
			return time.Time{}, err
		}
		result, err := time.ParseInLocation("20060102T150405Z", dateString, location)
		if err != nil {
			return time.ParseInLocation("20060102", dateString, location)
		}
		return result, nil
	} else {
		location, err := time.LoadLocation(tzId)
		if err != nil {
			return time.Time{}, err
		}
		result, err := time.ParseInLocation("20060102T150405", dateString, location)
		if err != nil {
			return time.ParseInLocation("20060102", dateString, location)
		}
		return result, nil
	}
}

func unescape(data string) (string, error) {
	tmp := data
	tmp = strings.Replace(tmp, "\\n", "\n", -1)
	tmp = strings.Replace(tmp, "\\", "", -1)
	return tmp, nil
}

func parse(scanner *bufio.Scanner) ([]Event, error) {
	var events []Event

	inEvent := false

	var (
		summary      string
		description  string
		location     string
		uid          string
		start        time.Time
		end          time.Time
		created      time.Time
		lastModified time.Time
	)
	for scanner.Scan() {
		var (
			key        string
			value      string
			err        error
			properties = map[string]string{}
		)

		array := strings.SplitN(scanner.Text(), ":", 2)
		if len(array) != 2 {
			return nil, errors.New("Error parsing a content line")
		}
		value = array[1]
		array = strings.SplitN(array[0], ";", 2)
		key = array[0]
		for i := 1; i < len(array); i++ {
			a := strings.SplitN(array[i], "=", 2)
			if len(array) != 2 {
				return nil, errors.New("Error parsing a property")
			}
			properties[a[0]] = a[1]
		}

		switch {
		case key == "BEGIN" && value == "VEVENT":
			if inEvent {
				return nil, errors.New("Error parsing BEGIN:VEVENT")
			}
			inEvent = true

			summary = ""
			description = ""
			location = ""
			uid = ""
			start = time.Time{}
			end = time.Time{}
			created = time.Time{}
			lastModified = time.Time{}

		case key == "END" && value == "VEVENT":
			if !inEvent {
				return nil, errors.New("Error parsing END:VEVENT")
			}
			inEvent = false

			events = append(events, Event{
				Summary:      summary,
				Description:  description,
				Location:     location,
				Uid:          uid,
				Start:        start,
				End:          end,
				Created:      created,
				LastModified: lastModified,
			})

		case key == "SUMMARY":
			if !inEvent {
				break
			}
			summary, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "DESCRIPTION":
			if !inEvent {
				break
			}
			description, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "LOCATION":
			if !inEvent {
				break
			}
			location, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "UID":
			if !inEvent {
				break
			}
			uid = value
		case key == "DTSTART":
			if !inEvent {
				break
			}
			start, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "DTEND":
			if !inEvent {
				break
			}
			end, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "CREATED":
			if !inEvent {
				break
			}
			created, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "LAST-MODIFIED":
			if !inEvent {
				break
			}
			lastModified, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		}

	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	if inEvent {
		return nil, errors.New("END:VEVENT is missing")
	}

	return events, nil
}

func GetSanrioEventsCalendar() (*feeds.Feed, error) {
	resp, err := http.Get("https://www.google.com/calendar/ical/qsqrk2emvnnvu45debac9dugr8%40group.calendar.google.com/public/basic.ics")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return GetSanrioEventsCalendarFromReader(resp.Body)
}

func GetSanrioEventsCalendarFromReader(reader io.Reader) (*feeds.Feed, error) {
	scanner := bufio.NewScanner(reader)
	scanner.Split(scanContentLines)

	events, err := parse(scanner)
	if err != nil {
		return nil, err
	}

	var items []*feeds.Item
	items = make([]*feeds.Item, len(events))
	for i, event := range events {
		items[i] = &feeds.Item{
			Title:       event.Summary,
			Description: fmt.Sprintf("%s\nStart: %s\nEnd: %s\n\n%s", event.Location, event.Start, event.End, event.Description),
			Link:        &feeds.Link{Href: xurls.Strict.FindString(event.Description)},
			Id:          event.Uid,
			Created:     event.Created,
			Updated:     event.LastModified,
		}
	}

	feed := &feeds.Feed{
		Title: "Sanrio Events",
		Link:  &feeds.Link{Href: "http://ameblo.jp/ohtaket/entry-12059393801.html"},
		Items: items,
	}
	return feed, nil
}
<commit_msg>Improve description<commit_after>package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"github.com/gorilla/feeds"
	"github.com/mvdan/xurls"
	"io"
	"net/http"
	"strings"
	"time"
)

func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

func isSpace(b byte) bool {
	return b == ' ' || b == '\t'
}

func scanContentLines(data []byte, atEOF bool) (int, []byte, error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	var (
		advance = 0
		buffer  bytes.Buffer
	)
	for {
		if i := bytes.IndexByte(data, '\n'); i >= 0 {
			advance += i + 1
			buffer.Write(dropCR(data[0:i]))

			data = data[i+1:]
			if len(data) > 0 && isSpace(data[0]) {
				j := 1
				for ; j < len(data); j++ {
					if !isSpace(data[j]) {
						break
					}
				}
				advance += j
				data = data[j:]
			} else {
				break
			}
		} else if atEOF {
			advance += len(data)
			buffer.Write(dropCR(data))
		} else {
			return 0, nil, nil
		}
	}
	if advance > 0 {
		return advance, buffer.Bytes(), nil
	}
	return 0, nil, nil
}

type Event struct {
	Summary      string
	Description  string
	Location     string
	Uid          string
	Start        time.Time
	End          time.Time
	Created      time.Time
	LastModified time.Time
}

func parseTime(dateString string, tzId string) (time.Time, error) {
	if tzId == "" {
		location, err := time.LoadLocation("UTC")
		if err != nil {
			return time.Time{}, err
		}
		result, err := time.ParseInLocation("20060102T150405Z", dateString, location)
		if err != nil {
			return time.ParseInLocation("20060102", dateString, location)
		}
		return result, nil
	} else {
		location, err := time.LoadLocation(tzId)
		if err != nil {
			return time.Time{}, err
		}
		result, err := time.ParseInLocation("20060102T150405", dateString, location)
		if err != nil {
			return time.ParseInLocation("20060102", dateString, location)
		}
		return result, nil
	}
}

func unescape(data string) (string, error) {
	tmp := data
	tmp = strings.Replace(tmp, "\\n", "\n", -1)
	tmp = strings.Replace(tmp, "\\", "", -1)
	return tmp, nil
}

func parse(scanner *bufio.Scanner) ([]Event, error) {
	var events []Event

	inEvent := false

	var (
		summary      string
		description  string
		location     string
		uid          string
		start        time.Time
		end          time.Time
		created      time.Time
		lastModified time.Time
	)
	for scanner.Scan() {
		var (
			key        string
			value      string
			err        error
			properties = map[string]string{}
		)

		array := strings.SplitN(scanner.Text(), ":", 2)
		if len(array) != 2 {
			return nil, errors.New("Error parsing a content line")
		}
		value = array[1]
		array = strings.SplitN(array[0], ";", 2)
		key = array[0]
		for i := 1; i < len(array); i++ {
			a := strings.SplitN(array[i], "=", 2)
			if len(array) != 2 {
				return nil, errors.New("Error parsing a property")
			}
			properties[a[0]] = a[1]
		}

		switch {
		case key == "BEGIN" && value == "VEVENT":
			if inEvent {
				return nil, errors.New("Error parsing BEGIN:VEVENT")
			}
			inEvent = true

			summary = ""
			description = ""
			location = ""
			uid = ""
			start = time.Time{}
			end = time.Time{}
			created = time.Time{}
			lastModified = time.Time{}

		case key == "END" && value == "VEVENT":
			if !inEvent {
				return nil, errors.New("Error parsing END:VEVENT")
			}
			inEvent = false

			events = append(events, Event{
				Summary:      summary,
				Description:  description,
				Location:     location,
				Uid:          uid,
				Start:        start,
				End:          end,
				Created:      created,
				LastModified: lastModified,
			})

		case key == "SUMMARY":
			if !inEvent {
				break
			}
			summary, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "DESCRIPTION":
			if !inEvent {
				break
			}
			description, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "LOCATION":
			if !inEvent {
				break
			}
			location, err = unescape(value)
			if err != nil {
				return nil, err
			}
		case key == "UID":
			if !inEvent {
				break
			}
			uid = value
		case key == "DTSTART":
			if !inEvent {
				break
			}
			start, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "DTEND":
			if !inEvent {
				break
			}
			end, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "CREATED":
			if !inEvent {
				break
			}
			created, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		case key == "LAST-MODIFIED":
			if !inEvent {
				break
			}
			lastModified, err = parseTime(value, properties["TZID"])
			if err != nil {
				return nil, err
			}
		}

	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	if inEvent {
		return nil, errors.New("END:VEVENT is missing")
	}

	return events, nil
}

func GetSanrioEventsCalendar() (*feeds.Feed, error) {
	resp, err := http.Get("https://www.google.com/calendar/ical/qsqrk2emvnnvu45debac9dugr8%40group.calendar.google.com/public/basic.ics")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return GetSanrioEventsCalendarFromReader(resp.Body)
}

func GetSanrioEventsCalendarFromReader(reader io.Reader) (*feeds.Feed, error) {
	scanner := bufio.NewScanner(reader)
	scanner.Split(scanContentLines)

	events, err := parse(scanner)
	if err != nil {
		return nil, err
	}

	var items []*feeds.Item
	items = make([]*feeds.Item, len(events))
	for i, event := range events {
		start := event.Start.Local()
		end := event.End.Local()
		var duration string
		if start.Format("20060102") == end.Format("20060102") {
			duration = start.Format("2006-01-02 (Mon) 15:04") + " - " + end.Format("15:04")
		} else {
			duration = start.Format("2006-01-02 (Mon) 15:04") + " - " + end.Format("2006-01-02 (Mon) 15:04")
		}
		items[i] = &feeds.Item{
			Title:       event.Summary,
			Description: fmt.Sprintf("Location: %s
Duration: %s

Description: %s", event.Location, duration, event.Description),
			Link:        &feeds.Link{Href: xurls.Strict.FindString(event.Description)},
			Id:          event.Uid,
			Created:     event.Created,
			Updated:     event.LastModified,
		}
	}

	feed := &feeds.Feed{
		Title: "Sanrio Events",
		Link:  &feeds.Link{Href: "http://ameblo.jp/ohtaket/entry-12059393801.html"},
		Items: items,
	}
	return feed, nil
}
<|endoftext|>


<commit_before>package gocbcore

import (
	"strings"
	"sync"
	"time"
)

type memdInitFunc func(*memdPipeline, time.Time) error

type CloseHandler func(*memdPipeline)
type BadRouteHandler func(*memdPipeline, *memdQRequest, *memdResponse)

type Callback func(*memdResponse, error)

type memdPipeline struct {
	lock sync.RWMutex

	queue *memdQueue

	address  string
	conn     memdReadWriteCloser
	isClosed bool
	ioDoneCh chan bool

	opList  memdOpMap

	handleBadRoute BadRouteHandler
	handleDeath    CloseHandler
}

func CreateMemdPipeline(address string) *memdPipeline {
	return &memdPipeline{
		address:  address,
		queue:    createMemdQueue(),
		ioDoneCh: make(chan bool, 1),
	}
}

func (s *memdPipeline) Address() string {
	return s.address
}

func (s *memdPipeline) Hostname() string {
	return strings.Split(s.address, ":")[0]
}

func (s *memdPipeline) IsClosed() bool {
	s.lock.Lock()
	rv := s.isClosed
	s.lock.Unlock()
	return rv
}

func (s *memdPipeline) SetHandlers(badRouteFn BadRouteHandler, deathFn CloseHandler) {
	s.lock.Lock()

	if s.isClosed {
		// We died between authentication and here, immediately notify the deathFn
		s.lock.Unlock()
		deathFn(s)
		return
	}

	s.handleBadRoute = badRouteFn
	s.handleDeath = deathFn
	s.lock.Unlock()
}

func (pipeline *memdPipeline) ExecuteRequest(req *memdQRequest, deadline time.Time) (respOut *memdResponse, errOut error) {
	if req.Callback != nil {
		panic("Tried to synchronously dispatch an operation with an async handler.")
	}

	signal := make(chan bool, 1)

	req.Callback = func(resp *memdResponse, err error) {
		respOut = resp
		errOut = err
		signal <- true
	}

	if !pipeline.queue.QueueRequest(req) {
		return nil, &generalError{"Failed to dispatch operation."}
	}

	timeoutTmr := AcquireTimer(deadline.Sub(time.Now()))
	select {
	case <-signal:
		ReleaseTimer(timeoutTmr, false)
		return
	case <-timeoutTmr.C:
		ReleaseTimer(timeoutTmr, true)
		req.Cancel()
		return nil, &timeoutError{}
	}
}

func (pipeline *memdPipeline) dispatchRequest(req *memdQRequest) error {
	// We do a cursory check of the server to avoid dispatching operations on the network
	//   that have already knowingly been cancelled.  This doesn't guarentee a cancelled
	//   operation from being sent, but it does reduce network IO when possible.
	if req.QueueOwner() != pipeline.queue {
		// Even though we failed to dispatch, this is not actually an error,
		//   we just consume the operation since its already been handled elsewhere
		return nil
	}

	pipeline.opList.Add(req)

	err := pipeline.conn.WritePacket(&req.memdRequest)
	if err != nil {
		logDebugf("Got write error")
		pipeline.opList.Remove(req)
		return err
	}

	return nil
}

func (s *memdPipeline) resolveRequest(resp *memdResponse) {
	opIndex := resp.Opaque

	// Find the request that goes with this response
	req := s.opList.FindAndMaybeRemove(opIndex)

	if req == nil {
		// There is no known request that goes with this response.  Ignore it.
		logDebugf("Received response with no corresponding request.")
		return
	}

	if !req.Persistent || (resp.Magic == ResMagic && resp.Status != StatusSuccess) {
		if !s.queue.UnqueueRequest(req) {
			// While we found a valid request, the request does not appear to be queued
			//   with this server anymore, this probably means that it has been cancelled.
			logDebugf("Received response for cancelled request.")
			return
		}
	}

	if resp.Magic == ResMagic && resp.Status == StatusNotMyVBucket {
		// If possible, lets backchannel our NMV back to the Agent of this memdQueueConn
		//   instance.  This is primarily meant to enhance performance, and allow the
		//   agent to be instantly notified upon a new configuration arriving.  If the
		//   backchannel isn't available, we just Callback with the NMV error.
		logDebugf("Received NMV response.")
		s.lock.RLock()
		badRouteFn := s.handleBadRoute
		s.lock.RUnlock()
		if badRouteFn != nil {
			badRouteFn(s, req, resp)
			return
		}
	}

	// Call the requests callback handler...  Ignore Status field for incoming requests.
	logDebugf("Dispatching response callback.")
	if resp.Magic == ReqMagic || resp.Status == StatusSuccess {
		req.Callback(resp, nil)
	} else {
		req.Callback(nil, &memdError{resp.Status})
	}
}

func (pipeline *memdPipeline) ioLoop() {
	killSig := make(chan bool)

	// Reading
	go func() {
		logDebugf("Reader loop starting...")
		for {
			resp := &memdResponse{}
			err := pipeline.conn.ReadPacket(resp)
			if err != nil {
				logDebugf("Server read error: %v", err)
				killSig <- true
				break
			}

			logDebugf("Got response to resolve.")
			pipeline.resolveRequest(resp)
		}
	}()

	// Writing
	logDebugf("Writer loop starting...")
	for {
		select {
		case req := <-pipeline.queue.reqsCh:
			logDebugf("Got a request to dispatch.")
			err := pipeline.dispatchRequest(req)
			if err != nil {
				// We can assume that the server is not fully drained yet, as the drainer blocks
				//   waiting for the IO goroutines to finish first.
				pipeline.queue.reqsCh <- req

				// We must wait for the receive goroutine to die as well before we can continue.
				<-killSig

				return
			}
		case <-killSig:
			return
		}
	}
}

func (pipeline *memdPipeline) Run() {
	logDebugf("Beginning pipeline runner")

	// Run the IO loop.  This will block until the connection has been closed.
	pipeline.ioLoop()

	// Now we must signal drainers that we are done!
	pipeline.ioDoneCh <- true

	// Signal the creator that we died :(
	pipeline.lock.Lock()
	pipeline.isClosed = true
	deathFn := pipeline.handleDeath
	pipeline.lock.Unlock()
	if deathFn != nil {
		deathFn(pipeline)
	} else {
		pipeline.Drain(nil)
	}
}

func (pipeline *memdPipeline) Close() {
	pipeline.Drain(nil)
}

func (pipeline *memdPipeline) Drain(reqCb drainedReqCallback) {
	// If the user does no pass a drain callback, we handle the requests
	//   by immediately failing them with a network error.
	if reqCb == nil {
		reqCb = func(req *memdQRequest) {
			req.Callback(nil, networkError{})
		}
	}

	// Make sure the connection is closed, which will signal the ioLoop
	//   to stop running and signal on ioDoneCh
	pipeline.conn.Close()

	// Drain the request queue, this will block until the io thread signals
	//   on ioDoneCh, and the queues have been completely emptied
	pipeline.queue.Drain(reqCb, pipeline.ioDoneCh)

	// As a last step, immediately notify all the requests that were
	//   on-the-wire that a network error has occurred.
	pipeline.opList.Drain(func(r *memdQRequest) {
		if pipeline.queue.UnqueueRequest(r) {
			r.Callback(nil, networkError{})
		}
	})
}
<commit_msg>Fix the race condition in ExecuteRequest differently<commit_after>package gocbcore

import (
	"strings"
	"sync"
	"time"
)

type memdInitFunc func(*memdPipeline, time.Time) error

type CloseHandler func(*memdPipeline)
type BadRouteHandler func(*memdPipeline, *memdQRequest, *memdResponse)

type Callback func(*memdResponse, error)

type memdPipeline struct {
	lock sync.RWMutex

	queue *memdQueue

	address  string
	conn     memdReadWriteCloser
	isClosed bool
	ioDoneCh chan bool

	opList  memdOpMap

	handleBadRoute BadRouteHandler
	handleDeath    CloseHandler
}

func CreateMemdPipeline(address string) *memdPipeline {
	return &memdPipeline{
		address:  address,
		queue:    createMemdQueue(),
		ioDoneCh: make(chan bool, 1),
	}
}

func (s *memdPipeline) Address() string {
	return s.address
}

func (s *memdPipeline) Hostname() string {
	return strings.Split(s.address, ":")[0]
}

func (s *memdPipeline) IsClosed() bool {
	s.lock.Lock()
	rv := s.isClosed
	s.lock.Unlock()
	return rv
}

func (s *memdPipeline) SetHandlers(badRouteFn BadRouteHandler, deathFn CloseHandler) {
	s.lock.Lock()

	if s.isClosed {
		// We died between authentication and here, immediately notify the deathFn
		s.lock.Unlock()
		deathFn(s)
		return
	}

	s.handleBadRoute = badRouteFn
	s.handleDeath = deathFn
	s.lock.Unlock()
}

func (pipeline *memdPipeline) ExecuteRequest(req *memdQRequest, deadline time.Time) (respOut *memdResponse, errOut error) {
	if req.Callback != nil {
		panic("Tried to synchronously dispatch an operation with an async handler.")
	}

	signal := make(chan bool)

	req.Callback = func(resp *memdResponse, err error) {
		respOut = resp
		errOut = err
		signal <- true
	}

	if !pipeline.queue.QueueRequest(req) {
		return nil, &generalError{"Failed to dispatch operation."}
	}

	timeoutTmr := AcquireTimer(deadline.Sub(time.Now()))
	select {
	case <-signal:
		ReleaseTimer(timeoutTmr, false)
		return
	case <-timeoutTmr.C:
		ReleaseTimer(timeoutTmr, true)
		if !req.Cancel() {
			<- signal
		}
		return nil, &timeoutError{}
	}
}

func (pipeline *memdPipeline) dispatchRequest(req *memdQRequest) error {
	// We do a cursory check of the server to avoid dispatching operations on the network
	//   that have already knowingly been cancelled.  This doesn't guarentee a cancelled
	//   operation from being sent, but it does reduce network IO when possible.
	if req.QueueOwner() != pipeline.queue {
		// Even though we failed to dispatch, this is not actually an error,
		//   we just consume the operation since its already been handled elsewhere
		return nil
	}

	pipeline.opList.Add(req)

	err := pipeline.conn.WritePacket(&req.memdRequest)
	if err != nil {
		logDebugf("Got write error")
		pipeline.opList.Remove(req)
		return err
	}

	return nil
}

func (s *memdPipeline) resolveRequest(resp *memdResponse) {
	opIndex := resp.Opaque

	// Find the request that goes with this response
	req := s.opList.FindAndMaybeRemove(opIndex)

	if req == nil {
		// There is no known request that goes with this response.  Ignore it.
		logDebugf("Received response with no corresponding request.")
		return
	}

	if !req.Persistent || (resp.Magic == ResMagic && resp.Status != StatusSuccess) {
		if !s.queue.UnqueueRequest(req) {
			// While we found a valid request, the request does not appear to be queued
			//   with this server anymore, this probably means that it has been cancelled.
			logDebugf("Received response for cancelled request.")
			return
		}
	}

	if resp.Magic == ResMagic && resp.Status == StatusNotMyVBucket {
		// If possible, lets backchannel our NMV back to the Agent of this memdQueueConn
		//   instance.  This is primarily meant to enhance performance, and allow the
		//   agent to be instantly notified upon a new configuration arriving.  If the
		//   backchannel isn't available, we just Callback with the NMV error.
		logDebugf("Received NMV response.")
		s.lock.RLock()
		badRouteFn := s.handleBadRoute
		s.lock.RUnlock()
		if badRouteFn != nil {
			badRouteFn(s, req, resp)
			return
		}
	}

	// Call the requests callback handler...  Ignore Status field for incoming requests.
	logDebugf("Dispatching response callback.")
	if resp.Magic == ReqMagic || resp.Status == StatusSuccess {
		req.Callback(resp, nil)
	} else {
		req.Callback(nil, &memdError{resp.Status})
	}
}

func (pipeline *memdPipeline) ioLoop() {
	killSig := make(chan bool)

	// Reading
	go func() {
		logDebugf("Reader loop starting...")
		for {
			resp := &memdResponse{}
			err := pipeline.conn.ReadPacket(resp)
			if err != nil {
				logDebugf("Server read error: %v", err)
				killSig <- true
				break
			}

			logDebugf("Got response to resolve.")
			pipeline.resolveRequest(resp)
		}
	}()

	// Writing
	logDebugf("Writer loop starting...")
	for {
		select {
		case req := <-pipeline.queue.reqsCh:
			logDebugf("Got a request to dispatch.")
			err := pipeline.dispatchRequest(req)
			if err != nil {
				// We can assume that the server is not fully drained yet, as the drainer blocks
				//   waiting for the IO goroutines to finish first.
				pipeline.queue.reqsCh <- req

				// We must wait for the receive goroutine to die as well before we can continue.
				<-killSig

				return
			}
		case <-killSig:
			return
		}
	}
}

func (pipeline *memdPipeline) Run() {
	logDebugf("Beginning pipeline runner")

	// Run the IO loop.  This will block until the connection has been closed.
	pipeline.ioLoop()

	// Now we must signal drainers that we are done!
	pipeline.ioDoneCh <- true

	// Signal the creator that we died :(
	pipeline.lock.Lock()
	pipeline.isClosed = true
	deathFn := pipeline.handleDeath
	pipeline.lock.Unlock()
	if deathFn != nil {
		deathFn(pipeline)
	} else {
		pipeline.Drain(nil)
	}
}

func (pipeline *memdPipeline) Close() {
	pipeline.Drain(nil)
}

func (pipeline *memdPipeline) Drain(reqCb drainedReqCallback) {
	// If the user does no pass a drain callback, we handle the requests
	//   by immediately failing them with a network error.
	if reqCb == nil {
		reqCb = func(req *memdQRequest) {
			req.Callback(nil, networkError{})
		}
	}

	// Make sure the connection is closed, which will signal the ioLoop
	//   to stop running and signal on ioDoneCh
	pipeline.conn.Close()

	// Drain the request queue, this will block until the io thread signals
	//   on ioDoneCh, and the queues have been completely emptied
	pipeline.queue.Drain(reqCb, pipeline.ioDoneCh)

	// As a last step, immediately notify all the requests that were
	//   on-the-wire that a network error has occurred.
	pipeline.opList.Drain(func(r *memdQRequest) {
		if pipeline.queue.UnqueueRequest(r) {
			r.Callback(nil, networkError{})
		}
	})
}
<|endoftext|>


<commit_before>package lmjrfll

import (
	"time"
)

// Expo is the day the lego expo will be held, where the teams present their project
type Expo struct {
	Year             int
	Name             string
	ExpoDate         time.Time
	ExpoEndTime      time.Time
	IsActive         bool
	IsCurrent        bool
	NumTeams         int
	NumRegistered    int
	RegistrationCost int // in dollars
	LocationName     string
}

// Team is the team name for a team of people
type Team struct {
	Name                 string
	CreateDate           time.Time
	ExpoKey              string
	JrFLLTeamNumber      string
	IsLookingForMemebers bool
	IsPaid               bool
	LeaderEmail          string
	NumMembers           int
	IsActive             bool
	IsWaitList           bool
	InviteCode           string
}

type TeamMember struct {
	TeamName  string
	TeamKey   string
	FirstName string
	LastName  string
	Email     string
}
<commit_msg>function to get the current expo<commit_after>package lmjrfll

import (
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	"time"
)

// Expo is the day the lego expo will be held, where the teams present their project
type Expo struct {
	Year             int
	Name             string
	ExpoDate         time.Time
	ExpoEndTime      time.Time
	IsActive         bool
	IsCurrent        bool
	NumTeams         int
	NumRegistered    int
	RegistrationCost int // in dollars
	LocationName     string
}

// Team is the team name for a team of people
type Team struct {
	Name                 string
	CreateDate           time.Time
	ExpoKey              string
	JrFLLTeamNumber      string
	IsLookingForMemebers bool
	IsPaid               bool
	LeaderEmail          string
	NumMembers           int
	IsActive             bool
	IsWaitList           bool
	InviteCode           string
}

type TeamMember struct {
	TeamName  string
	TeamKey   string
	FirstName string
	LastName  string
	Email     string
}

// Returns the current Expo, or an error that there is no current Expo
func GetCurrentExpo(c context.Context) (*Expo, error) {
	var expo Expo
	query := datastore.NewQuery("Expo").Filter("IsCurrent =", true)
	results := query.Run(c)
	_, err := results.Next(&expo)
	if err != nil {
		return nil, err
	}
	return &expo, nil
}
<|endoftext|>


<commit_before>package mrt

import (
	"github.com/osrg/gobgp/packet"
	"github.com/stretchr/testify/assert"
	"reflect"
	"testing"
	"time"
)

func TestMrtHdr(t *testing.T) {
	h1, err := NewHeader(10, TABLE_DUMPv2, RIB_IPV4_MULTICAST, 20)
	if err != nil {
		t.Fatal(err)
	}
	b1, err := h1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	h2 := &Header{}
	err = h2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(h1, h2), true)
}

func testPeer(t *testing.T, p1 *Peer) {
	b1, err := p1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	p2 := &Peer{}
	rest, err := p2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, len(rest), 0)
	assert.Equal(t, reflect.DeepEqual(p1, p2), true)
}

func TestMrtPeer(t *testing.T) {
	p := NewPeer("192.168.0.1", "10.0.0.1", 65000, false)
	testPeer(t, p)
}

func TestMrtPeerv6(t *testing.T) {
	p := NewPeer("192.168.0.1", "2001::1", 65000, false)
	testPeer(t, p)
}

func TestMrtPeerAS4(t *testing.T) {
	p := NewPeer("192.168.0.1", "2001::1", 135500, true)
	testPeer(t, p)
}

func TestMrtPeerIndexTable(t *testing.T) {
	p1 := NewPeer("192.168.0.1", "10.0.0.1", 65000, false)
	p2 := NewPeer("192.168.0.1", "2001::1", 65000, false)
	p3 := NewPeer("192.168.0.1", "2001::1", 135500, true)
	pt1 := NewPeerIndexTable("192.168.0.1", "test", []*Peer{p1, p2, p3})
	b1, err := pt1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	pt2 := &PeerIndexTable{}
	err = pt2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(pt1, pt2), true)
}

func TestMrtRibEntry(t *testing.T) {
	aspath1 := []bgp.AsPathParamInterface{
		bgp.NewAsPathParam(2, []uint16{1000}),
		bgp.NewAsPathParam(1, []uint16{1001, 1002}),
		bgp.NewAsPathParam(2, []uint16{1003, 1004}),
	}

	p := []bgp.PathAttributeInterface{
		bgp.NewPathAttributeOrigin(3),
		bgp.NewPathAttributeAsPath(aspath1),
		bgp.NewPathAttributeNextHop("129.1.1.2"),
		bgp.NewPathAttributeMultiExitDisc(1 << 20),
		bgp.NewPathAttributeLocalPref(1 << 22),
	}

	e1 := NewRibEntry(1, uint32(time.Now().Unix()), p)
	b1, err := e1.Serialize()
	if err != nil {
		t.Fatal(err)
	}

	e2 := &RibEntry{}
	rest, err := e2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, len(rest), 0)
	assert.Equal(t, reflect.DeepEqual(e1, e2), true)
}

func TestMrtRib(t *testing.T) {
	aspath1 := []bgp.AsPathParamInterface{
		bgp.NewAsPathParam(2, []uint16{1000}),
		bgp.NewAsPathParam(1, []uint16{1001, 1002}),
		bgp.NewAsPathParam(2, []uint16{1003, 1004}),
	}

	p := []bgp.PathAttributeInterface{
		bgp.NewPathAttributeOrigin(3),
		bgp.NewPathAttributeAsPath(aspath1),
		bgp.NewPathAttributeNextHop("129.1.1.2"),
		bgp.NewPathAttributeMultiExitDisc(1 << 20),
		bgp.NewPathAttributeLocalPref(1 << 22),
	}

	e1 := NewRibEntry(1, uint32(time.Now().Unix()), p)
	e2 := NewRibEntry(2, uint32(time.Now().Unix()), p)
	e3 := NewRibEntry(3, uint32(time.Now().Unix()), p)

	r1 := NewRib(1, bgp.NewIPAddrPrefix(24, "192.168.0.0"), []*RibEntry{e1, e2, e3})
	b1, err := r1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	r2 := &Rib{
		RouteFamily: bgp.RF_IPv4_UC,
	}
	err = r2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(r1, r2), true)
}

func TestMrtBgp4mpStateChange(t *testing.T) {
	c1 := NewBGP4MPStateChange(65000, 65001, 1, "192.168.0.1", "192.168.0.2", false, ACTIVE, ESTABLISHED)
	b1, err := c1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	c2 := &BGP4MPStateChange{BGP4MPHeader: &BGP4MPHeader{}}
	err = c2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	_, err = c2.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(c1, c2), true)
}

func TestMrtBgp4mpMessage(t *testing.T) {
	msg := bgp.NewBGPKeepAliveMessage()
	m1 := NewBGP4MPMessage(65000, 65001, 1, "192.168.0.1", "192.168.0.2", false, msg)
	b1, err := m1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	m2 := &BGP4MPMessage{BGP4MPHeader: &BGP4MPHeader{}}
	err = m2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(m1, m2), true)
}
<commit_msg>mrt: add Copyright<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 mrt

import (
	"github.com/osrg/gobgp/packet"
	"github.com/stretchr/testify/assert"
	"reflect"
	"testing"
	"time"
)

func TestMrtHdr(t *testing.T) {
	h1, err := NewHeader(10, TABLE_DUMPv2, RIB_IPV4_MULTICAST, 20)
	if err != nil {
		t.Fatal(err)
	}
	b1, err := h1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	h2 := &Header{}
	err = h2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(h1, h2), true)
}

func testPeer(t *testing.T, p1 *Peer) {
	b1, err := p1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	p2 := &Peer{}
	rest, err := p2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, len(rest), 0)
	assert.Equal(t, reflect.DeepEqual(p1, p2), true)
}

func TestMrtPeer(t *testing.T) {
	p := NewPeer("192.168.0.1", "10.0.0.1", 65000, false)
	testPeer(t, p)
}

func TestMrtPeerv6(t *testing.T) {
	p := NewPeer("192.168.0.1", "2001::1", 65000, false)
	testPeer(t, p)
}

func TestMrtPeerAS4(t *testing.T) {
	p := NewPeer("192.168.0.1", "2001::1", 135500, true)
	testPeer(t, p)
}

func TestMrtPeerIndexTable(t *testing.T) {
	p1 := NewPeer("192.168.0.1", "10.0.0.1", 65000, false)
	p2 := NewPeer("192.168.0.1", "2001::1", 65000, false)
	p3 := NewPeer("192.168.0.1", "2001::1", 135500, true)
	pt1 := NewPeerIndexTable("192.168.0.1", "test", []*Peer{p1, p2, p3})
	b1, err := pt1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	pt2 := &PeerIndexTable{}
	err = pt2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(pt1, pt2), true)
}

func TestMrtRibEntry(t *testing.T) {
	aspath1 := []bgp.AsPathParamInterface{
		bgp.NewAsPathParam(2, []uint16{1000}),
		bgp.NewAsPathParam(1, []uint16{1001, 1002}),
		bgp.NewAsPathParam(2, []uint16{1003, 1004}),
	}

	p := []bgp.PathAttributeInterface{
		bgp.NewPathAttributeOrigin(3),
		bgp.NewPathAttributeAsPath(aspath1),
		bgp.NewPathAttributeNextHop("129.1.1.2"),
		bgp.NewPathAttributeMultiExitDisc(1 << 20),
		bgp.NewPathAttributeLocalPref(1 << 22),
	}

	e1 := NewRibEntry(1, uint32(time.Now().Unix()), p)
	b1, err := e1.Serialize()
	if err != nil {
		t.Fatal(err)
	}

	e2 := &RibEntry{}
	rest, err := e2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, len(rest), 0)
	assert.Equal(t, reflect.DeepEqual(e1, e2), true)
}

func TestMrtRib(t *testing.T) {
	aspath1 := []bgp.AsPathParamInterface{
		bgp.NewAsPathParam(2, []uint16{1000}),
		bgp.NewAsPathParam(1, []uint16{1001, 1002}),
		bgp.NewAsPathParam(2, []uint16{1003, 1004}),
	}

	p := []bgp.PathAttributeInterface{
		bgp.NewPathAttributeOrigin(3),
		bgp.NewPathAttributeAsPath(aspath1),
		bgp.NewPathAttributeNextHop("129.1.1.2"),
		bgp.NewPathAttributeMultiExitDisc(1 << 20),
		bgp.NewPathAttributeLocalPref(1 << 22),
	}

	e1 := NewRibEntry(1, uint32(time.Now().Unix()), p)
	e2 := NewRibEntry(2, uint32(time.Now().Unix()), p)
	e3 := NewRibEntry(3, uint32(time.Now().Unix()), p)

	r1 := NewRib(1, bgp.NewIPAddrPrefix(24, "192.168.0.0"), []*RibEntry{e1, e2, e3})
	b1, err := r1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	r2 := &Rib{
		RouteFamily: bgp.RF_IPv4_UC,
	}
	err = r2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(r1, r2), true)
}

func TestMrtBgp4mpStateChange(t *testing.T) {
	c1 := NewBGP4MPStateChange(65000, 65001, 1, "192.168.0.1", "192.168.0.2", false, ACTIVE, ESTABLISHED)
	b1, err := c1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	c2 := &BGP4MPStateChange{BGP4MPHeader: &BGP4MPHeader{}}
	err = c2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	_, err = c2.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(c1, c2), true)
}

func TestMrtBgp4mpMessage(t *testing.T) {
	msg := bgp.NewBGPKeepAliveMessage()
	m1 := NewBGP4MPMessage(65000, 65001, 1, "192.168.0.1", "192.168.0.2", false, msg)
	b1, err := m1.Serialize()
	if err != nil {
		t.Fatal(err)
	}
	m2 := &BGP4MPMessage{BGP4MPHeader: &BGP4MPHeader{}}
	err = m2.DecodeFromBytes(b1)
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, reflect.DeepEqual(m1, m2), true)
}
<|endoftext|>


<commit_before>package gqt_test

import (
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Init process", func() {
	var (
		tmpDir        string
		parentCommand *exec.Cmd
	)

	BeforeEach(func() {
		tmpDir = tempDir("", "")

		cmd := exec.Command("gcc", "-static", "-o", "test_init", "test_init.c", "../../cmd/init/ignore_sigchild.c", "-I", "../../cmd/init")
		runCommandInDir(cmd, "cmd")
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	It("does not allow children to become zombies", func() {
		parentCommand = exec.Command("cmd/test_init")
		parentCommand.Start()

		Eventually(countPsOccurances).Should(Equal(1))

		psOut := string(runPs())
		fmt.Println(psOut)
		matchingPsLines := []string{}
		psLines := strings.Split(psOut, "\n")
		for _, psLine := range psLines {
			if !strings.Contains(psLine, "test_init") {
				continue
			}
			matchingPsLines = append(matchingPsLines, psLine)
		}

		Expect(strings.Join(matchingPsLines, "\n")).NotTo(ContainSubstring("defunct"))
	})
})

func countPsOccurances() int {
	psout := runPs()

	testInitRe := regexp.MustCompile("test_init")

	matches := testInitRe.FindAll(psout, -1)

	return len(matches)
}

func runPs() []byte {
	cmd := exec.Command("ps", "auxf")
	psout, err := cmd.Output()
	Expect(err).NotTo(HaveOccurred())

	return psout
}
<commit_msg>Do not print ps output unconditionally<commit_after>package gqt_test

import (
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Init process", func() {
	var (
		tmpDir        string
		parentCommand *exec.Cmd
	)

	BeforeEach(func() {
		tmpDir = tempDir("", "")

		cmd := exec.Command("gcc", "-static", "-o", "test_init", "test_init.c", "../../cmd/init/ignore_sigchild.c", "-I", "../../cmd/init")
		runCommandInDir(cmd, "cmd")
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	It("does not allow children to become zombies", func() {
		parentCommand = exec.Command("cmd/test_init")
		parentCommand.Start()

		Eventually(countPsOccurances).Should(Equal(1))

		psOut := string(runPs())
		matchingPsLines := []string{}
		psLines := strings.Split(psOut, "\n")
		for _, psLine := range psLines {
			if !strings.Contains(psLine, "test_init") {
				continue
			}
			matchingPsLines = append(matchingPsLines, psLine)
		}

		Expect(strings.Join(matchingPsLines, "\n")).NotTo(ContainSubstring("defunct"), fmt.Sprintf("\n\nps output:\n%s\n\n", psOut))
	})
})

func countPsOccurances() int {
	psout := runPs()

	testInitRe := regexp.MustCompile("test_init")

	matches := testInitRe.FindAll(psout, -1)

	return len(matches)
}

func runPs() []byte {
	cmd := exec.Command("ps", "auxf")
	psout, err := cmd.Output()
	Expect(err).NotTo(HaveOccurred())

	return psout
}
<|endoftext|>


<commit_before>// Package patricia implements a patricia tree, or a radix
// tree with a radix of 2 -- creating an uneven binary tree.
//
// Each entry is a key value pair. The key determines
// where the value is placed in the tree, with each bit
// of the key indicating a path. Values are arbitrary byte
// slices but only the SHA3-256 hash of the value is stored
// within the tree.
//
// The nodes in the tree form an immutable persistent data
// structure, therefore Copy is a O(1) operation.
package patricia

import (
	"bytes"

	"chain/crypto/sha3pool"
	"chain/errors"
	"chain/protocol/bc"
)

// ErrPrefix is returned from Insert if
// the key provided is a prefix to existing nodes.
var ErrPrefix = errors.New("key provided is a prefix to other keys")

var (
	leafPrefix     = []byte{0x00}
	interiorPrefix = []byte{0x01}
)

// Tree implements a patricia tree.
type Tree struct {
	root *node
}

// Copy returns a new tree with the same root as this tree. It
// is an O(1) operation.
func Copy(t *Tree) *Tree {
	newT := new(Tree)
	newT.root = t.root
	return newT
}

// WalkFunc is the type of the function called for each leaf
// visited by Walk. If an error is returned, processing stops.
type WalkFunc func(k []byte) error

// Walk walks the patricia tree calling walkFn for each leaf in
// the tree. If an error is returned by walkFn at any point,
// processing is stopped and the error is returned.
func Walk(t *Tree, walkFn WalkFunc) error {
	if t.root == nil {
		return nil
	}
	return walk(t.root, walkFn)
}

func walk(n *node, walkFn WalkFunc) error {
	if n.isLeaf {
		return walkFn(n.Key())
	}

	err := walk(n.children[0], walkFn)
	if err != nil {
		return err
	}

	err = walk(n.children[1], walkFn)
	return err
}

// ContainsKey returns true if the key contains the provided
// key, without checking its corresponding hash.
func (t *Tree) ContainsKey(bkey []byte) bool {
	if t.root == nil {
		return false
	}
	return lookup(t.root, bitKey(bkey)) != nil
}

// Contains returns true if the tree contains the provided
// key, value pair.
func (t *Tree) Contains(bkey, val []byte) bool {
	if t.root == nil {
		return false
	}

	key := bitKey(bkey)
	n := lookup(t.root, key)

	var hash bc.Hash
	h := sha3pool.Get256()
	h.Write(leafPrefix)
	h.Write(val[:])
	h.Read(hash[:])
	sha3pool.Put256(h)
	return n != nil && n.Hash() == hash
}

func lookup(n *node, key []uint8) *node {
	if bytes.Equal(n.key, key) {
		if !n.isLeaf {
			return nil
		}
		return n
	}
	if !bytes.HasPrefix(key, n.key) {
		return nil
	}

	bit := key[len(n.key)]
	return lookup(n.children[bit], key)
}

// Insert enters data into the tree.
// If the key is not already present in the tree,
// a new node will be created and inserted,
// rearranging the tree to the optimal structure.
// If the key is present, the existing node is found
// and its value is updated, leaving the structure of
// the tree alone.
func (t *Tree) Insert(bkey, val []byte) error {
	key := bitKey(bkey)

	var hash bc.Hash
	h := sha3pool.Get256()
	h.Write(leafPrefix)
	h.Write(val)
	h.Read(hash[:])
	sha3pool.Put256(h)

	if t.root == nil {
		t.root = &node{key: key, hash: &hash, isLeaf: true}
		return nil
	}

	var err error
	t.root, err = insert(t.root, key, &hash)
	return err
}

func insert(n *node, key []uint8, hash *bc.Hash) (*node, error) {
	if bytes.Equal(n.key, key) {
		if !n.isLeaf {
			return n, errors.Wrap(ErrPrefix)
		}

		n = &node{
			isLeaf: true,
			key:    n.key,
			hash:   hash,
		}
		return n, nil
	}

	if bytes.HasPrefix(key, n.key) {
		if n.isLeaf {
			return n, errors.Wrap(ErrPrefix)
		}
		bit := key[len(n.key)]

		child := n.children[bit]
		child, err := insert(child, key, hash)
		if err != nil {
			return n, err
		}
		newNode := new(node)
		*newNode = *n
		newNode.children[bit] = child // mutation is ok because newNode hasn't escaped yet
		newNode.hash = nil
		return newNode, nil
	}

	common := commonPrefixLen(n.key, key)
	newNode := &node{
		key: key[:common],
	}
	newNode.children[key[common]] = &node{
		key:    key,
		hash:   hash,
		isLeaf: true,
	}
	newNode.children[1-key[common]] = n
	return newNode, nil
}

// Delete removes up to one value with a matching key.
// After removing the node, it will rearrange the tree
// to the optimal structure.
func (t *Tree) Delete(bkey []byte) {
	key := bitKey(bkey)

	if t.root != nil {
		t.root = delete(t.root, key)
	}
}

func delete(n *node, key []uint8) *node {
	if bytes.Equal(key, n.key) {
		if !n.isLeaf {
			return n
		}
		return nil
	}

	if !bytes.HasPrefix(key, n.key) {
		return n
	}

	bit := key[len(n.key)]
	newChild := delete(n.children[bit], key)

	if newChild == nil {
		return n.children[1-bit]
	}

	newNode := new(node)
	*newNode = *n
	newNode.key = newChild.key[:len(n.key)] // only use slices of leaf node keys
	newNode.children[bit] = newChild
	newNode.hash = nil

	return newNode
}

// RootHash returns the merkle root of the tree.
func (t *Tree) RootHash() bc.Hash {
	root := t.root
	if root == nil {
		return bc.Hash{}
	}
	return root.Hash()
}

// bitKey takes a byte array and returns a key that can
// be used inside insert and delete operations.
func bitKey(byteKey []byte) []uint8 {
	key := make([]uint8, 0, len(byteKey)*8)
	for _, b := range byteKey {
		for i := uint(0); i < 8; i++ {
			key = append(key, (b>>(7-i))&1)
		}
	}
	return key
}

// byteKey is the inverse of bitKey.
func byteKey(bitKey []uint8) (key []byte) {
	key = make([]byte, len(bitKey)/8)
	for i := uint(0); i < uint(len(key)); i++ {
		var b byte
		for j := uint(0); j < 8; j++ {
			bit := bitKey[i*8+j]
			b |= bit << (7 - j)
		}
		key[i] = b
	}
	return key
}

func commonPrefixLen(a, b []uint8) int {
	var common int
	for i := 0; i < len(a) && i < len(b); i++ {
		if a[i] != b[i] {
			break
		}
		common++
	}
	return common
}

// node is a leaf or branch node in a tree
type node struct {
	key      []uint8
	hash     *bc.Hash
	isLeaf   bool
	children [2]*node
}

// Key returns the key for the current node as bytes, as it
// was provided to Insert.
func (n *node) Key() []byte { return byteKey(n.key) }

// Hash will return the hash for this node.
func (n *node) Hash() bc.Hash {
	n.calcHash()
	return *n.hash
}

func (n *node) calcHash() {
	if n.hash != nil {
		return
	}

	h := sha3pool.Get256()
	h.Write(interiorPrefix)
	for _, c := range n.children {
		c.calcHash()
		h.Write(c.hash[:])
	}

	var hash bc.Hash
	h.Read(hash[:])
	n.hash = &hash
	sha3pool.Put256(h)
}
<commit_msg>protocol/patricia: un-export ErrPrefix<commit_after>// Package patricia implements a patricia tree, or a radix
// tree with a radix of 2 -- creating an uneven binary tree.
//
// Each entry is a key value pair. The key determines
// where the value is placed in the tree, with each bit
// of the key indicating a path. Values are arbitrary byte
// slices but only the SHA3-256 hash of the value is stored
// within the tree.
//
// The nodes in the tree form an immutable persistent data
// structure, therefore Copy is a O(1) operation.
package patricia

import (
	"bytes"

	"chain/crypto/sha3pool"
	"chain/errors"
	"chain/protocol/bc"
)

var (
	leafPrefix     = []byte{0x00}
	interiorPrefix = []byte{0x01}
)

// Tree implements a patricia tree.
type Tree struct {
	root *node
}

// Copy returns a new tree with the same root as this tree. It
// is an O(1) operation.
func Copy(t *Tree) *Tree {
	newT := new(Tree)
	newT.root = t.root
	return newT
}

// WalkFunc is the type of the function called for each leaf
// visited by Walk. If an error is returned, processing stops.
type WalkFunc func(k []byte) error

// Walk walks the patricia tree calling walkFn for each leaf in
// the tree. If an error is returned by walkFn at any point,
// processing is stopped and the error is returned.
func Walk(t *Tree, walkFn WalkFunc) error {
	if t.root == nil {
		return nil
	}
	return walk(t.root, walkFn)
}

func walk(n *node, walkFn WalkFunc) error {
	if n.isLeaf {
		return walkFn(n.Key())
	}

	err := walk(n.children[0], walkFn)
	if err != nil {
		return err
	}

	err = walk(n.children[1], walkFn)
	return err
}

// ContainsKey returns true if the key contains the provided
// key, without checking its corresponding hash.
func (t *Tree) ContainsKey(bkey []byte) bool {
	if t.root == nil {
		return false
	}
	return lookup(t.root, bitKey(bkey)) != nil
}

// Contains returns true if the tree contains the provided
// key, value pair.
func (t *Tree) Contains(bkey, val []byte) bool {
	if t.root == nil {
		return false
	}

	key := bitKey(bkey)
	n := lookup(t.root, key)

	var hash bc.Hash
	h := sha3pool.Get256()
	h.Write(leafPrefix)
	h.Write(val[:])
	h.Read(hash[:])
	sha3pool.Put256(h)
	return n != nil && n.Hash() == hash
}

func lookup(n *node, key []uint8) *node {
	if bytes.Equal(n.key, key) {
		if !n.isLeaf {
			return nil
		}
		return n
	}
	if !bytes.HasPrefix(key, n.key) {
		return nil
	}

	bit := key[len(n.key)]
	return lookup(n.children[bit], key)
}

// Insert enters data into the tree.
// If the key is not already present in the tree,
// a new node will be created and inserted,
// rearranging the tree to the optimal structure.
// If the key is present, the existing node is found
// and its value is updated, leaving the structure of
// the tree alone.
// It is an error for bkey to be a prefix
// of a key already in t or to contain a key already
// in t as a prefix.
func (t *Tree) Insert(bkey, val []byte) error {
	key := bitKey(bkey)

	var hash bc.Hash
	h := sha3pool.Get256()
	h.Write(leafPrefix)
	h.Write(val)
	h.Read(hash[:])
	sha3pool.Put256(h)

	if t.root == nil {
		t.root = &node{key: key, hash: &hash, isLeaf: true}
		return nil
	}

	var err error
	t.root, err = insert(t.root, key, &hash)
	return err
}

func insert(n *node, key []uint8, hash *bc.Hash) (*node, error) {
	if bytes.Equal(n.key, key) {
		if !n.isLeaf {
			return n, errors.Wrap(errors.New("key provided is a prefix to other keys"))
		}

		n = &node{
			isLeaf: true,
			key:    n.key,
			hash:   hash,
		}
		return n, nil
	}

	if bytes.HasPrefix(key, n.key) {
		if n.isLeaf {
			return n, errors.Wrap(errors.New("key provided is a prefix to other keys"))
		}
		bit := key[len(n.key)]

		child := n.children[bit]
		child, err := insert(child, key, hash)
		if err != nil {
			return n, err
		}
		newNode := new(node)
		*newNode = *n
		newNode.children[bit] = child // mutation is ok because newNode hasn't escaped yet
		newNode.hash = nil
		return newNode, nil
	}

	common := commonPrefixLen(n.key, key)
	newNode := &node{
		key: key[:common],
	}
	newNode.children[key[common]] = &node{
		key:    key,
		hash:   hash,
		isLeaf: true,
	}
	newNode.children[1-key[common]] = n
	return newNode, nil
}

// Delete removes up to one value with a matching key.
// After removing the node, it will rearrange the tree
// to the optimal structure.
func (t *Tree) Delete(bkey []byte) {
	key := bitKey(bkey)

	if t.root != nil {
		t.root = delete(t.root, key)
	}
}

func delete(n *node, key []uint8) *node {
	if bytes.Equal(key, n.key) {
		if !n.isLeaf {
			return n
		}
		return nil
	}

	if !bytes.HasPrefix(key, n.key) {
		return n
	}

	bit := key[len(n.key)]
	newChild := delete(n.children[bit], key)

	if newChild == nil {
		return n.children[1-bit]
	}

	newNode := new(node)
	*newNode = *n
	newNode.key = newChild.key[:len(n.key)] // only use slices of leaf node keys
	newNode.children[bit] = newChild
	newNode.hash = nil

	return newNode
}

// RootHash returns the merkle root of the tree.
func (t *Tree) RootHash() bc.Hash {
	root := t.root
	if root == nil {
		return bc.Hash{}
	}
	return root.Hash()
}

// bitKey takes a byte array and returns a key that can
// be used inside insert and delete operations.
func bitKey(byteKey []byte) []uint8 {
	key := make([]uint8, 0, len(byteKey)*8)
	for _, b := range byteKey {
		for i := uint(0); i < 8; i++ {
			key = append(key, (b>>(7-i))&1)
		}
	}
	return key
}

// byteKey is the inverse of bitKey.
func byteKey(bitKey []uint8) (key []byte) {
	key = make([]byte, len(bitKey)/8)
	for i := uint(0); i < uint(len(key)); i++ {
		var b byte
		for j := uint(0); j < 8; j++ {
			bit := bitKey[i*8+j]
			b |= bit << (7 - j)
		}
		key[i] = b
	}
	return key
}

func commonPrefixLen(a, b []uint8) int {
	var common int
	for i := 0; i < len(a) && i < len(b); i++ {
		if a[i] != b[i] {
			break
		}
		common++
	}
	return common
}

// node is a leaf or branch node in a tree
type node struct {
	key      []uint8
	hash     *bc.Hash
	isLeaf   bool
	children [2]*node
}

// Key returns the key for the current node as bytes, as it
// was provided to Insert.
func (n *node) Key() []byte { return byteKey(n.key) }

// Hash will return the hash for this node.
func (n *node) Hash() bc.Hash {
	n.calcHash()
	return *n.hash
}

func (n *node) calcHash() {
	if n.hash != nil {
		return
	}

	h := sha3pool.Get256()
	h.Write(interiorPrefix)
	for _, c := range n.children {
		c.calcHash()
		h.Write(c.hash[:])
	}

	var hash bc.Hash
	h.Read(hash[:])
	n.hash = &hash
	sha3pool.Put256(h)
}
<|endoftext|>