<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	b.Run("read-random-badger", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count int
			for pb.Next() {
				key := newKey()
				var val badger.KVItem
				if bdb.Get(key, &val); val.Value() != nil {
					count++
				}
			}
			if count > 100000 {
				b.Logf("badger %d keys had valid values.", count)
			}
		})
	})
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count int
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			if count > 100000 {
				b.Logf("rocks %d keys had valid values.", count)
			}
		})
	})
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	b.Run("read-random-lmdb", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count int
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					count++
					return nil
				})
			}
			if count > 100000 {
				b.Logf("lmdb %d keys had valid values.", count)
			}
		})
	})
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>Tweaks to benchmarks.<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var totalCount uint64
	b.Run("read-random-badger", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			var val badger.KVItem
			for pb.Next() {
				key := newKey()
				if bdb.Get(key, &val); val.Value() != nil {
					count++
				}
			}

		})
	})
	b.Logf("badger %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var totalCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("rocks %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var totalCount uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					count++
					return nil
				})
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("lmdb %d keys had valid values.", totalCount)
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"os"
	"testing"
)

func BenchmarkByteWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		byteWrite(f, str)
	}
}

func BenchmarkStringWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		stringWrite(f, str)
	}
}

func BenchmarkBufioWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	w := bufio.NewWriterSize(f, 8192)
	for i := 0; i < b.N; i++ {
		bufioWrite(w, f, str)
	}
}

func BenchmarkDirectWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		directWrite(f, str)
	}
}
<commit_msg>add close<commit_after>package main

import (
	"bufio"
	"os"
	"testing"
)

func BenchmarkCloseEvery(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		byteWrite(f, str)
		f.Close()
		f, ferr = os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
		if ferr != nil {
			panic(ferr.Error())
		}
	}
}

func BenchmarkByteWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		byteWrite(f, str)
	}
}

func BenchmarkStringWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		stringWrite(f, str)
	}
}

func BenchmarkBufioWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	w := bufio.NewWriterSize(f, 8192)
	for i := 0; i < b.N; i++ {
		bufioWrite(w, f, str)
	}
}

func BenchmarkDirectWrite(b *testing.B) {
	str := []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")

	fileName := "tmp.log"
	ferr := os.Remove(fileName)
	if ferr != nil {
		if !os.IsNotExist(ferr) {
			panic(ferr.Error())
		}
	}
	f, ferr := os.OpenFile(fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
	defer f.Close()
	if ferr != nil {
		panic(ferr.Error())
	}
	for i := 0; i < b.N; i++ {
		directWrite(f, str)
	}
}
<|endoftext|>


<commit_before>package xxhash

import (
	"strings"
	"testing"
)

var benchmarks = []struct {
	name string
	n    int64
}{
	{"4B", 4},
	{"100B", 100},
	{"4KB", 4e3},
	{"10MB", 10e6},
}

func BenchmarkSum64(b *testing.B) {
	for _, bb := range benchmarks {
		in := make([]byte, bb.n)
		for i := range in {
			in[i] = byte(i)
		}
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				_ = Sum64(in)
			}
		})
	}
}

func BenchmarkSum64String(b *testing.B) {
	for _, bb := range benchmarks {
		s := strings.Repeat("a", int(bb.n))
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				_ = Sum64String(s)
			}
		})
	}
}

func BenchmarkDigestBytes(b *testing.B) {
	for _, bb := range benchmarks {
		in := make([]byte, bb.n)
		for i := range in {
			in[i] = byte(i)
		}
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				h := New()
				h.Write(in)
				_ = h.Sum64()
			}
		})
	}
}

func BenchmarkDigestString(b *testing.B) {
	for _, bb := range benchmarks {
		s := strings.Repeat("a", int(bb.n))
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				h := New()
				h.WriteString(s)
				_ = h.Sum64()
			}
		})
	}
}
<commit_msg>Add 16-byte hash benchmarks<commit_after>package xxhash

import (
	"strings"
	"testing"
)

var benchmarks = []struct {
	name string
	n    int64
}{
	{"4B", 4},
	{"16B", 16},
	{"100B", 100},
	{"4KB", 4e3},
	{"10MB", 10e6},
}

func BenchmarkSum64(b *testing.B) {
	for _, bb := range benchmarks {
		in := make([]byte, bb.n)
		for i := range in {
			in[i] = byte(i)
		}
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				_ = Sum64(in)
			}
		})
	}
}

func BenchmarkSum64String(b *testing.B) {
	for _, bb := range benchmarks {
		s := strings.Repeat("a", int(bb.n))
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				_ = Sum64String(s)
			}
		})
	}
}

func BenchmarkDigestBytes(b *testing.B) {
	for _, bb := range benchmarks {
		in := make([]byte, bb.n)
		for i := range in {
			in[i] = byte(i)
		}
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				h := New()
				h.Write(in)
				_ = h.Sum64()
			}
		})
	}
}

func BenchmarkDigestString(b *testing.B) {
	for _, bb := range benchmarks {
		s := strings.Repeat("a", int(bb.n))
		b.Run(bb.name, func(b *testing.B) {
			b.SetBytes(bb.n)
			for i := 0; i < b.N; i++ {
				h := New()
				h.WriteString(s)
				_ = h.Sum64()
			}
		})
	}
}
<|endoftext|>


<commit_before>package gohm

import (
	"compress/flate"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/golang/snappy"
)

type compressionResponseWriter struct {
	http.ResponseWriter
	compressionWriter io.Writer
}

func (g compressionResponseWriter) Write(b []byte) (int, error) {
	return g.compressionWriter.Write(b)
}

// WithGzip returns a new http.Handler that optionally compresses the response
// text using the gzip compression algorithm when the HTTP request's
// `Accept-Encoding` header includes the string `gzip`.
//
// NOTE: The specified next http.Handler ought not set `Content-Length` header,
// or the value reported will be wrong.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithGzip(someHandler))
func WithGzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}
		gz := gzip.NewWriter(w)
		defer func() {
			if err := gz.Close(); err != nil {
				Error(w, fmt.Sprintf("cannot compress stream: %s", err), http.StatusInternalServerError)
			}
		}()
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Vary", "Accept-Encoding")
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: gz}, r)
	})
}

// WithCompression returns a new http.Handler that optionally compresses the
// response text using either the gzip or deflate compression algorithm when the
// HTTP request's `Accept-Encoding` header includes the string `gzip` or
// `deflate`. To prevent the specified next http.Handler from also seeing the
// `Accept-Encoding` request header, and possibly also compressing the data a
// second time, this function removes that header from the request.
//
// NOTE: The specified next http.Handler ought not set `Content-Length` header,
// or the value reported will be wrong.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithCompression(someHandler))
func WithCompression(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		var newWriteCloser io.WriteCloser
		var encodingName, responseHeaderName string

		requestHeaderName := "TE"
		acceptableEncodings := r.Header.Get(requestHeaderName)
		if acceptableEncodings != "" {
			// If Transfer-Encoding specified, then completely ignore
			// Accept-Encoding, because the upstream node has specifically
			// requested a node-to-node transfer compression algorithm.
			responseHeaderName = "Transfer-Encoding"
		} else {
			responseHeaderName = "Content-Encoding"
			requestHeaderName = "Accept-Encoding"
			acceptableEncodings = r.Header.Get(requestHeaderName)
		}

		// Because many browsers include a buggy deflate compression algorithm,
		// prefer `gzip` over `deflate` if both are acceptable.
		if encodingName = "snappy"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = snappy.NewBufferedWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using snappy: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "gzip"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = gzip.NewWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using gzip: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "deflate"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser, err = flate.NewWriter(w, -1)
			if err != nil {
				next.ServeHTTP(w, r)
				return
			}
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using deflate: %s", err), http.StatusInternalServerError)
				}
			}()
		} else {
			next.ServeHTTP(w, r)
			return
		}
		r.Header.Del(requestHeaderName)
		w.Header().Set(responseHeaderName, encodingName)
		if responseHeaderName == "Content-Encoding" {
			w.Header().Set("Vary", responseHeaderName)
		}
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: newWriteCloser}, r)
	})
}
<commit_msg>WithCompression no longer triggers on TE request header<commit_after>package gohm

import (
	"compress/flate"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/golang/snappy"
)

type compressionResponseWriter struct {
	http.ResponseWriter
	compressionWriter io.Writer
}

func (g compressionResponseWriter) Write(b []byte) (int, error) {
	return g.compressionWriter.Write(b)
}

// WithGzip returns a new http.Handler that optionally compresses the response
// text using the gzip compression algorithm when the HTTP request's
// `Accept-Encoding` header includes the string `gzip`.
//
// NOTE: The specified next http.Handler ought not set `Content-Length` header,
// or the reported length value will be wrong. As a matter of fact, all HTTP
// response handlers ought to allow net/http library to set `Content-Length`
// response header or not based on a handful of RFCs.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithGzip(someHandler))
func WithGzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}
		gz := gzip.NewWriter(w)
		defer func() {
			if err := gz.Close(); err != nil {
				Error(w, fmt.Sprintf("cannot compress stream: %s", err), http.StatusInternalServerError)
			}
		}()
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Vary", "Accept-Encoding")
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: gz}, r)
	})
}

// WithCompression returns a new http.Handler that optionally compresses the
// response text using either the snapper, gzip, or deflate compression
// algorithm when the HTTP request's `Accept-Encoding` header includes the
// string `snappy, `gzip`, or `deflate`. To prevent the downstream next
// http.Handler from also seeing the `Accept-Encoding` request header, and
// possibly also compressing the data a second time, this function removes that
// header from the request.
//
// NOTE: The specified next http.Handler ought not set `Content-Length` header,
// or the reported length value will be wrong. As a matter of fact, all HTTP
// response handlers ought to allow net/http library to set `Content-Length`
// response header or not based on a handful of RFCs.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithCompression(someHandler))
func WithCompression(next http.Handler) http.Handler {
	const requestHeader = "Accept-Encoding"
	const responseHeader = "Content-Encoding"

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var newWriteCloser io.WriteCloser
		var encodingAlgorithm string

		acceptableEncodings := r.Header.Get(requestHeader)

		// Shortcut if no Accept-Encoding header
		if acceptableEncodings == "" {
			next.ServeHTTP(w, r)
			return
		}

		// Offer snappy, gzip, and deflate compression. Because many browsers
		// include a buggy deflate compression algorithm, prefer gzip over
		// deflate if both are acceptable. TODO: include support for brotli
		// algorithm: br.
		if encodingAlgorithm = "snappy"; strings.Contains(acceptableEncodings, encodingAlgorithm) {
			newWriteCloser = snappy.NewBufferedWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using snappy: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingAlgorithm = "gzip"; strings.Contains(acceptableEncodings, encodingAlgorithm) {
			newWriteCloser = gzip.NewWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using gzip: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingAlgorithm = "deflate"; strings.Contains(acceptableEncodings, encodingAlgorithm) {
			var err error
			newWriteCloser, err = flate.NewWriter(w, flate.DefaultCompression)
			if err != nil {
				// This should never happen, but if cannot create a new deflate
				// writer, then ignore the Accept-Encoding header and send the
				// unchanged request to the downstream handler.
				next.ServeHTTP(w, r)
				return
			}
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using deflate: %s", err), http.StatusInternalServerError)
				}
			}()
		} else {
			// Upstream requests a compression algorithms that is not
			// supported. Ignore the Accept-Encoding header and send the
			// unchanged request to the downstream handler.
			next.ServeHTTP(w, r)
			return
		}

		// Delete the Accept-Encoding header from the request to prevent
		// downstream handler from seeing it and possibly also compressing data,
		// resulting in a payload that needs to be decompressed twice.
		r.Header.Del(requestHeader)

		// Set the response headers accordingly.
		w.Header().Set(responseHeader, encodingAlgorithm)
		w.Header().Set("Vary", responseHeader)

		// Have the downstream handler service this request, writing the
		// response to our compression writer.
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: newWriteCloser}, r)
	})
}
<|endoftext|>


<commit_before>package authentication

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"

	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure/cli"
	"github.com/hashicorp/go-multierror"
)

type azureCLIProfile struct {
	subscription *cli.Subscription

	clientId       string
	environment    string
	subscriptionId string
	tenantId       string
}

type azureCliTokenAuth struct {
	profile                      *azureCLIProfile
	servicePrincipalAuthDocsLink string
}

func (a azureCliTokenAuth) build(b Builder) (authMethod, error) {
	auth := azureCliTokenAuth{
		profile: &azureCLIProfile{
			subscriptionId: b.SubscriptionID,
			tenantId:       b.TenantID,
			clientId:       "04b07795-8ddb-461a-bbee-02f9e1bf7b46", // fixed first party client id for Az CLI
		},
		servicePrincipalAuthDocsLink: b.ClientSecretDocsLink,
	}

	sub, err := obtainSubscription(b.SubscriptionID)
	if err != nil {
		return nil, fmt.Errorf("obtain subscription(%s) from Azure CLI: %+v", b.SubscriptionID, err)
	}
	auth.profile.subscription = sub

	// Authenticating as a Service Principal doesn't return all of the information we need for authentication purposes
	// as such Service Principal authentication is supported using the specific auth method
	if sub.User == nil || !strings.EqualFold(sub.User.Type, "user") {
		return nil, fmt.Errorf(`Authenticating using the Azure CLI is only supported as a User (not a Service Principal).

To authenticate to Azure using a Service Principal, you can use the separate 'Authenticate using a Service Principal'
auth method - instructions for which can be found here: %s

Alternatively you can authenticate using the Azure CLI by using a User Account.`, auth.servicePrincipalAuthDocsLink)
	}

	// Populate fields
	if auth.profile.subscriptionId == "" {
		auth.profile.subscriptionId = sub.ID
	}
	if auth.profile.tenantId == "" {
		auth.profile.tenantId = sub.TenantID
	}
	// always pull the environment from the Azure CLI, since the Access Token's associated with it
	auth.profile.environment = normalizeEnvironmentName(sub.EnvironmentName)

	return auth, nil
}

func (a azureCliTokenAuth) isApplicable(b Builder) bool {
	return b.SupportsAzureCliToken
}

func (a azureCliTokenAuth) getAuthorizationToken(sender autorest.Sender, oauth *OAuthConfig, endpoint string) (autorest.Authorizer, error) {
	if oauth.OAuth == nil {
		return nil, fmt.Errorf("Error getting Authorization Token for cli auth: an OAuth token wasn't configured correctly; please file a bug with more details")
	}

	// the Azure CLI appears to cache these, so to maintain compatibility with the interface this method is intentionally not on the pointer
	token, err := obtainAuthorizationToken(endpoint, a.profile.subscriptionId)
	if err != nil {
		return nil, fmt.Errorf("Error obtaining Authorization Token from the Azure CLI: %s", err)
	}

	adalToken, err := token.ToADALToken()
	if err != nil {
		return nil, fmt.Errorf("Error converting Authorization Token to an ADAL Token: %s", err)
	}

	spt, err := adal.NewServicePrincipalTokenFromManualToken(*oauth.OAuth, a.profile.clientId, endpoint, adalToken)
	if err != nil {
		return nil, err
	}

	var refreshFunc adal.TokenRefresh = func(ctx context.Context, resource string) (*adal.Token, error) {
		token, err := obtainAuthorizationToken(resource, a.profile.subscriptionId)
		if err != nil {
			return nil, err
		}

		adalToken, err := token.ToADALToken()
		if err != nil {
			return nil, err
		}

		return &adalToken, nil
	}
	spt.SetCustomRefreshFunc(refreshFunc)

	auth := autorest.NewBearerAuthorizer(spt)
	return auth, nil
}

func (a azureCliTokenAuth) name() string {
	return "Obtaining a token from the Azure CLI"
}

func (a azureCliTokenAuth) populateConfig(c *Config) error {
	c.ClientID = a.profile.clientId
	c.TenantID = a.profile.tenantId
	c.Environment = a.profile.environment
	c.SubscriptionID = a.profile.subscriptionId

	c.GetAuthenticatedObjectID = func(ctx context.Context) (string, error) {
		objectId, err := obtainAuthenticatedObjectID()
		if err != nil {
			return "", err
		}

		return objectId, nil
	}

	return nil
}

func (a azureCliTokenAuth) validate() error {
	var err *multierror.Error

	errorMessageFmt := "A %s was not found in your Azure CLI Credentials.\n\nPlease login to the Azure CLI again via `az login`"

	if a.profile == nil {
		return fmt.Errorf("Azure CLI Profile is nil - this is an internal error and should be reported.")
	}

	if a.profile.clientId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Client ID"))
	}

	if a.profile.subscriptionId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Subscription ID"))
	}

	if a.profile.tenantId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Tenant ID"))
	}

	return err.ErrorOrNil()
}

func obtainAuthenticatedObjectID() (string, error) {

	var json struct {
		ObjectId string `json:"objectId"`
	}

	err := jsonUnmarshalAzCmd(&json, "ad", "signed-in-user", "show", "-o=json")
	if err != nil {
		return "", fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return json.ObjectId, nil
}

func obtainAuthorizationToken(endpoint string, subscriptionId string) (*cli.Token, error) {
	var token cli.Token
	err := jsonUnmarshalAzCmd(&token, "account", "get-access-token", "--resource", endpoint, "--subscription", subscriptionId, "-o=json")
	if err != nil {
		return nil, fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return &token, nil
}

// obtainSubscription return a subscription object of the specified subscriptionId.
// If the subscriptionId is empty, it returns the default subscription.
func obtainSubscription(subscriptionId string) (*cli.Subscription, error) {
	var sub cli.Subscription
	cmd := make([]string, 2, 4)
	cmd = []string{"account", "show"}
	if subscriptionId != "" {
		cmd = append(cmd, "-s", subscriptionId)
	}
	err := jsonUnmarshalAzCmd(&sub, cmd...)
	if err != nil {
		return nil, fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return &sub, nil
}

func jsonUnmarshalAzCmd(i interface{}, arg ...string) error {
	var stderr bytes.Buffer
	var stdout bytes.Buffer

	cmd := exec.Command("az", arg...)

	cmd.Stderr = &stderr
	cmd.Stdout = &stdout

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("Error launching Azure CLI: %+v", err)
	}

	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("Error waiting for the Azure CLI: %+v", err)
	}

	stdOutStr := stdout.String()
	stdErrStr := stderr.String()
	if stdErrStr != "" {
		return fmt.Errorf("Error retrieving running Azure CLI: %s", strings.TrimSpace(stdErrStr))
	}

	if err := json.Unmarshal([]byte(stdOutStr), &i); err != nil {
		return fmt.Errorf("Error unmarshaling the result of Azure CLI: %v", err)
	}

	return nil
}
<commit_msg>explicitly indicate the `az acccount show` format as json<commit_after>package authentication

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"

	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure/cli"
	"github.com/hashicorp/go-multierror"
)

type azureCLIProfile struct {
	subscription *cli.Subscription

	clientId       string
	environment    string
	subscriptionId string
	tenantId       string
}

type azureCliTokenAuth struct {
	profile                      *azureCLIProfile
	servicePrincipalAuthDocsLink string
}

func (a azureCliTokenAuth) build(b Builder) (authMethod, error) {
	auth := azureCliTokenAuth{
		profile: &azureCLIProfile{
			subscriptionId: b.SubscriptionID,
			tenantId:       b.TenantID,
			clientId:       "04b07795-8ddb-461a-bbee-02f9e1bf7b46", // fixed first party client id for Az CLI
		},
		servicePrincipalAuthDocsLink: b.ClientSecretDocsLink,
	}

	sub, err := obtainSubscription(b.SubscriptionID)
	if err != nil {
		return nil, fmt.Errorf("obtain subscription(%s) from Azure CLI: %+v", b.SubscriptionID, err)
	}
	auth.profile.subscription = sub

	// Authenticating as a Service Principal doesn't return all of the information we need for authentication purposes
	// as such Service Principal authentication is supported using the specific auth method
	if sub.User == nil || !strings.EqualFold(sub.User.Type, "user") {
		return nil, fmt.Errorf(`Authenticating using the Azure CLI is only supported as a User (not a Service Principal).

To authenticate to Azure using a Service Principal, you can use the separate 'Authenticate using a Service Principal'
auth method - instructions for which can be found here: %s

Alternatively you can authenticate using the Azure CLI by using a User Account.`, auth.servicePrincipalAuthDocsLink)
	}

	// Populate fields
	if auth.profile.subscriptionId == "" {
		auth.profile.subscriptionId = sub.ID
	}
	if auth.profile.tenantId == "" {
		auth.profile.tenantId = sub.TenantID
	}
	// always pull the environment from the Azure CLI, since the Access Token's associated with it
	auth.profile.environment = normalizeEnvironmentName(sub.EnvironmentName)

	return auth, nil
}

func (a azureCliTokenAuth) isApplicable(b Builder) bool {
	return b.SupportsAzureCliToken
}

func (a azureCliTokenAuth) getAuthorizationToken(sender autorest.Sender, oauth *OAuthConfig, endpoint string) (autorest.Authorizer, error) {
	if oauth.OAuth == nil {
		return nil, fmt.Errorf("Error getting Authorization Token for cli auth: an OAuth token wasn't configured correctly; please file a bug with more details")
	}

	// the Azure CLI appears to cache these, so to maintain compatibility with the interface this method is intentionally not on the pointer
	token, err := obtainAuthorizationToken(endpoint, a.profile.subscriptionId)
	if err != nil {
		return nil, fmt.Errorf("Error obtaining Authorization Token from the Azure CLI: %s", err)
	}

	adalToken, err := token.ToADALToken()
	if err != nil {
		return nil, fmt.Errorf("Error converting Authorization Token to an ADAL Token: %s", err)
	}

	spt, err := adal.NewServicePrincipalTokenFromManualToken(*oauth.OAuth, a.profile.clientId, endpoint, adalToken)
	if err != nil {
		return nil, err
	}

	var refreshFunc adal.TokenRefresh = func(ctx context.Context, resource string) (*adal.Token, error) {
		token, err := obtainAuthorizationToken(resource, a.profile.subscriptionId)
		if err != nil {
			return nil, err
		}

		adalToken, err := token.ToADALToken()
		if err != nil {
			return nil, err
		}

		return &adalToken, nil
	}
	spt.SetCustomRefreshFunc(refreshFunc)

	auth := autorest.NewBearerAuthorizer(spt)
	return auth, nil
}

func (a azureCliTokenAuth) name() string {
	return "Obtaining a token from the Azure CLI"
}

func (a azureCliTokenAuth) populateConfig(c *Config) error {
	c.ClientID = a.profile.clientId
	c.TenantID = a.profile.tenantId
	c.Environment = a.profile.environment
	c.SubscriptionID = a.profile.subscriptionId

	c.GetAuthenticatedObjectID = func(ctx context.Context) (string, error) {
		objectId, err := obtainAuthenticatedObjectID()
		if err != nil {
			return "", err
		}

		return objectId, nil
	}

	return nil
}

func (a azureCliTokenAuth) validate() error {
	var err *multierror.Error

	errorMessageFmt := "A %s was not found in your Azure CLI Credentials.\n\nPlease login to the Azure CLI again via `az login`"

	if a.profile == nil {
		return fmt.Errorf("Azure CLI Profile is nil - this is an internal error and should be reported.")
	}

	if a.profile.clientId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Client ID"))
	}

	if a.profile.subscriptionId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Subscription ID"))
	}

	if a.profile.tenantId == "" {
		err = multierror.Append(err, fmt.Errorf(errorMessageFmt, "Tenant ID"))
	}

	return err.ErrorOrNil()
}

func obtainAuthenticatedObjectID() (string, error) {

	var json struct {
		ObjectId string `json:"objectId"`
	}

	err := jsonUnmarshalAzCmd(&json, "ad", "signed-in-user", "show", "-o=json")
	if err != nil {
		return "", fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return json.ObjectId, nil
}

func obtainAuthorizationToken(endpoint string, subscriptionId string) (*cli.Token, error) {
	var token cli.Token
	err := jsonUnmarshalAzCmd(&token, "account", "get-access-token", "--resource", endpoint, "--subscription", subscriptionId, "-o=json")
	if err != nil {
		return nil, fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return &token, nil
}

// obtainSubscription return a subscription object of the specified subscriptionId.
// If the subscriptionId is empty, it returns the default subscription.
func obtainSubscription(subscriptionId string) (*cli.Subscription, error) {
	var sub cli.Subscription
	cmd := make([]string, 2, 4)
	cmd = []string{"account", "show", "-o=json"}
	if subscriptionId != "" {
		cmd = append(cmd, "-s", subscriptionId)
	}
	err := jsonUnmarshalAzCmd(&sub, cmd...)
	if err != nil {
		return nil, fmt.Errorf("Error parsing json result from the Azure CLI: %v", err)
	}

	return &sub, nil
}

func jsonUnmarshalAzCmd(i interface{}, arg ...string) error {
	var stderr bytes.Buffer
	var stdout bytes.Buffer

	cmd := exec.Command("az", arg...)

	cmd.Stderr = &stderr
	cmd.Stdout = &stdout

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("Error launching Azure CLI: %+v", err)
	}

	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("Error waiting for the Azure CLI: %+v", err)
	}

	stdOutStr := stdout.String()
	stdErrStr := stderr.String()
	if stdErrStr != "" {
		return fmt.Errorf("Error retrieving running Azure CLI: %s", strings.TrimSpace(stdErrStr))
	}

	if err := json.Unmarshal([]byte(stdOutStr), &i); err != nil {
		return fmt.Errorf("Error unmarshaling the result of Azure CLI: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package engine

import (
	"context"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"

	"github.com/webx-top/echo/logger"
)

type (
	// Engine defines an interface for HTTP server.
	Engine interface {
		SetHandler(Handler)
		SetLogger(logger.Logger)
		Start() error
		Stop() error
		Shutdown(ctx context.Context) error
	}

	// Request defines an interface for HTTP request.
	Request interface {
		// Scheme returns the HTTP protocol scheme, `http` or `https`.
		Scheme() string

		// Host returns HTTP request host. Per RFC 2616, this is either the value of
		// the `Host` header or the host name given in the URL itself.
		Host() string

		// SetHost sets the host of the request.
		SetHost(string)

		// URI returns the unmodified `Request-URI` sent by the client.
		URI() string

		// SetURI sets the URI of the request.
		SetURI(string)

		// URL returns `engine.URL`.
		URL() URL

		// Header returns `engine.Header`.
		Header() Header

		// Proto returns the HTTP proto. (HTTP/1.1 etc.)
		Proto() string
		// ProtoMajor() int
		// ProtoMinor() int

		// RemoteAddress returns the client's network address.
		RemoteAddress() string

		// RealIP returns the client's network address based on `X-Forwarded-For`
		// or `X-Real-IP` request header.
		RealIP() string

		// Method returns the request's HTTP function.
		Method() string

		// SetMethod sets the HTTP method of the request.
		SetMethod(string)

		// Body returns request's body.
		Body() io.ReadCloser

		SetBody(io.Reader)

		// FormValue returns the form field value for the provided name.
		FormValue(string) string
		Object() interface{}

		Form() URLValuer
		PostForm() URLValuer

		// MultipartForm returns the multipart form.
		MultipartForm() *multipart.Form

		// IsTLS returns true if HTTP connection is TLS otherwise false.
		IsTLS() bool
		Cookie(string) string
		Referer() string

		// UserAgent returns the client's `User-Agent`.
		UserAgent() string

		// FormFile returns the multipart form file for the provided name.
		FormFile(string) (multipart.File, *multipart.FileHeader, error)

		// Size returns the size of request's body.
		Size() int64

		BasicAuth() (string, string, bool)

		StdRequest() *http.Request
	}

	// Response defines an interface for HTTP response.
	Response interface {
		// Header returns `engine.Header`
		Header() Header

		// WriteHeader sends an HTTP response header with status code.
		WriteHeader(int)

		KeepBody(bool)

		// Write writes the data to the connection as part of an HTTP reply.
		Write([]byte) (int, error)

		// Status returns the HTTP response status.
		Status() int

		// Size returns the number of bytes written to HTTP response.
		Size() int64

		// Committed returns true if HTTP response header is written, otherwise false.
		Committed() bool

		// SetWriter sets the HTTP response writer.
		SetWriter(io.Writer)

		// Write returns the HTTP response writer.
		Writer() io.Writer
		Object() interface{}

		Body() []byte
		Redirect(string, int)
		NotFound()
		SetCookie(*http.Cookie)
		ServeFile(string)
		Stream(func(io.Writer) bool)
		Error(string, ...int)

		StdResponseWriter() http.ResponseWriter
	}

	// Header defines an interface for HTTP header.
	Header interface {
		// Add adds the key, value pair to the header. It appends to any existing values
		// associated with key.
		Add(string, string)

		// Del deletes the values associated with key.
		Del(string)

		// Get gets the first value associated with the given key. If there are
		// no values associated with the key, Get returns "".
		Get(string) string

		// Set sets the header entries associated with key to the single element value.
		// It replaces any existing values associated with key.
		Set(string, string)

		Object() interface{}

		Std() http.Header
	}

	// URLValuer Wrap url.Values
	URLValuer interface {
		Add(string, string)
		Del(string)
		Get(string) string
		Gets(string) []string
		Set(string, string)
		Encode() string
		All() map[string][]string
		Reset(url.Values)
		Merge(url.Values)
	}

	// URL defines an interface for HTTP request url.
	URL interface {
		SetPath(string)
		RawPath() string
		Path() string
		QueryValue(string) string
		QueryValues(string) []string
		Query() url.Values
		RawQuery() string
		SetRawQuery(string)
		String() string
		Object() interface{}
	}

	// Handler defines an interface to server HTTP requests via `ServeHTTP(Request, Response)`
	// function.
	Handler interface {
		ServeHTTP(Request, Response)
	}

	// HandlerFunc is an adapter to allow the use of `func(Request, Response)` as HTTP handlers.
	HandlerFunc func(Request, Response)
)

// ServeHTTP serves HTTP request.
func (h HandlerFunc) ServeHTTP(req Request, res Response) {
	h(req, res)
}
<commit_msg>update<commit_after>package engine

import (
	"context"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/url"

	"github.com/webx-top/echo/logger"
)

type (
	// Engine defines an interface for HTTP server.
	Engine interface {
		SetHandler(Handler)
		SetLogger(logger.Logger)
		Start() error
		Stop() error
		Shutdown(ctx context.Context) error
	}

	// Request defines an interface for HTTP request.
	Request interface {
		// Scheme returns the HTTP protocol scheme, `http` or `https`.
		Scheme() string

		// Host returns HTTP request host. Per RFC 2616, this is either the value of
		// the `Host` header or the host name given in the URL itself.
		Host() string

		// SetHost sets the host of the request.
		SetHost(string)

		// URI returns the unmodified `Request-URI` sent by the client.
		URI() string

		// SetURI sets the URI of the request.
		SetURI(string)

		// URL returns `engine.URL`.
		URL() URL

		// Header returns `engine.Header`.
		Header() Header

		// Proto returns the HTTP proto. (HTTP/1.1 etc.)
		Proto() string
		// ProtoMajor() int
		// ProtoMinor() int

		// RemoteAddress returns the client's network address.
		RemoteAddress() string

		// RealIP returns the client's network address based on `X-Forwarded-For`
		// or `X-Real-IP` request header.
		RealIP() string

		// Method returns the request's HTTP function.
		Method() string

		// SetMethod sets the HTTP method of the request.
		SetMethod(string)

		// Body returns request's body.
		Body() io.ReadCloser

		SetBody(io.Reader)

		// FormValue returns the form field value for the provided name.
		FormValue(string) string
		Object() interface{}

		Form() URLValuer
		PostForm() URLValuer

		// MultipartForm returns the multipart form.
		MultipartForm() *multipart.Form

		// IsTLS returns true if HTTP connection is TLS otherwise false.
		IsTLS() bool
		Cookie(string) string
		Referer() string

		// UserAgent returns the client's `User-Agent`.
		UserAgent() string

		// FormFile returns the multipart form file for the provided name.
		FormFile(string) (multipart.File, *multipart.FileHeader, error)

		// Size returns the size of request's body.
		Size() int64

		BasicAuth() (string, string, bool)

		StdRequest() *http.Request
	}

	// Response defines an interface for HTTP response.
	Response interface {
		// Header returns `engine.Header`
		Header() Header

		// WriteHeader sends an HTTP response header with status code.
		WriteHeader(int)

		KeepBody(bool)

		// Write writes the data to the connection as part of an HTTP reply.
		Write([]byte) (int, error)

		// Status returns the HTTP response status.
		Status() int

		// Size returns the number of bytes written to HTTP response.
		Size() int64

		// Committed returns true if HTTP response header is written, otherwise false.
		Committed() bool

		// SetWriter sets the HTTP response writer.
		SetWriter(io.Writer)

		// Write returns the HTTP response writer.
		Writer() io.Writer
		Object() interface{}

		Hijacker(func(net.Conn)) error
		Body() []byte
		Redirect(string, int)
		NotFound()
		SetCookie(*http.Cookie)
		ServeFile(string)
		Stream(func(io.Writer) bool)
		Error(string, ...int)

		StdResponseWriter() http.ResponseWriter
	}

	// Header defines an interface for HTTP header.
	Header interface {
		// Add adds the key, value pair to the header. It appends to any existing values
		// associated with key.
		Add(string, string)

		// Del deletes the values associated with key.
		Del(string)

		// Get gets the first value associated with the given key. If there are
		// no values associated with the key, Get returns "".
		Get(string) string

		// Set sets the header entries associated with key to the single element value.
		// It replaces any existing values associated with key.
		Set(string, string)

		Object() interface{}

		Std() http.Header
	}

	// URLValuer Wrap url.Values
	URLValuer interface {
		Add(string, string)
		Del(string)
		Get(string) string
		Gets(string) []string
		Set(string, string)
		Encode() string
		All() map[string][]string
		Reset(url.Values)
		Merge(url.Values)
	}

	// URL defines an interface for HTTP request url.
	URL interface {
		SetPath(string)
		RawPath() string
		Path() string
		QueryValue(string) string
		QueryValues(string) []string
		Query() url.Values
		RawQuery() string
		SetRawQuery(string)
		String() string
		Object() interface{}
	}

	// Handler defines an interface to server HTTP requests via `ServeHTTP(Request, Response)`
	// function.
	Handler interface {
		ServeHTTP(Request, Response)
	}

	// HandlerFunc is an adapter to allow the use of `func(Request, Response)` as HTTP handlers.
	HandlerFunc func(Request, Response)
)

// ServeHTTP serves HTTP request.
func (h HandlerFunc) ServeHTTP(req Request, res Response) {
	h(req, res)
}
<|endoftext|>


<commit_before>// Copyright © 2016 Sidharth Kshatriya
//
// 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 engine

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"fmt"
	"github.com/fatih/color"
	"github.com/kr/pty"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	// These strings are not to be changed as these strings are sentinels from the dontbug zend extension
	dontbugZendExtensionLoadedSentinel          = "dontbug zend extension: dontbug.so successfully loaded by PHP"
	dontbugZendXdebugNotLoadedSentinel          = "dontbug zend extension: Xdebug has not been loaded"
	dontbugZendXdebugEntryPointNotFoundSentinel = "dontbug zend extension: Xdebug entrypoint not found"
	// End do not change

	dontbugRRTraceDirSentinel = "rr: Saving execution to trace directory `"

	dontbugNotPatchedXdebugMsg = `Unpatched Xdebug zend extension (xdebug.so) found. See below for more information:
dontbug zend extension currently relies on a patched version of Xdebug to function correctly.
This is a very minor patch and simply makes a single function extern (instead of static) linkage.
It seems you are using the plain vanilla version of Xdebug. Consult documentation on patching Xdebug.
`
)

func getOrCreateDontbugSharePath() string {
	currentUser, err := user.Current()
	fatalIf(err)

	dontbugShareDir := currentUser.HomeDir + "/.local/share/dontbug/"
	mkDirAll(dontbugShareDir)

	return dontbugShareDir
}

func copyAndMakeUniqueDontbugSo(sharedObjectPath, dontbugShareDir string) string {
	uniqueDontbugSoFilename := path.Clean(fmt.Sprintf("%v/at-%v-dontbug.so", dontbugShareDir, time.Now().UnixNano()))
	output, err := exec.Command("cp", sharedObjectPath, uniqueDontbugSoFilename).CombinedOutput()
	if err != nil {
		log.Fatal(output)
	}
	return uniqueDontbugSoFilename
}

// Assumptions:
// - rrPath represents an rr executable that meets dontbug's requirements
// - phpPath represents an php executable that meets dontbug's requirements
// - sharedObject path is the path to xdebug.so that meets dontbug's requirements
// - docrootOrScriptAbsNoSymPath is a valid docroot directory or a php script
func doRecordSession(
	docrootOrScriptAbsNoSymPath,
	sharedObjectPath,
	rrPath,
	phpPath string,
	isCli bool,
	arguments,
	serverListen string,
	serverPort,
	recordPort,
	maxStackDepth int,
	takeSnapshot bool,
	snapShotDir string,
	originalDocrootOrScriptFullPath string,
) {
	newSharedObjectPath := sharedObjectPath
	if takeSnapshot {
		dontbugShareDir := getOrCreateDontbugSharePath()
		newSharedObjectPath = copyAndMakeUniqueDontbugSo(sharedObjectPath, dontbugShareDir)
	}

	// Many of these options are not really necessary to be specified.
	// However, we still do that to override any settings that
	// might be present in user php.ini files and change them
	// to sensible defaults for 'dontbug record'
	rrCmd := []string{
		"record",
		phpPath,
		"-d", "zend_extension=" + newSharedObjectPath,
		"-d", fmt.Sprintf("xdebug.remote_port=%v", recordPort),
		"-d", "xdebug.remote_autostart=1",
		"-d", "xdebug.remote_connect_back=0",
		"-d", "xdebug.remote_enable=1",
		"-d", "xdebug.remote_mode=req",
		"-d", "xdebug.auto_trace=0",
		"-d", "xdebug.trace_enable_trigger=\"\"",
		"-d", "xdebug.coverage_enable=0",
		"-d", "xdebug.extended_info=1",
		"-d", fmt.Sprintf("xdebug.max_nesting_level=%v", maxStackDepth),
		"-d", "xdebug.profiler_enable=0",
		"-d", "xdebug.profiler_enable_trigger=0",
	}

	if isCli {
		arguments = strings.TrimSpace(arguments)
		rrCmd = append(rrCmd, docrootOrScriptAbsNoSymPath)
		if arguments != "" {
			argumentsAr := strings.Split(arguments, " ")
			rrCmd = append(rrCmd, argumentsAr...)
		}
	} else {
		rrCmd = append(
			rrCmd,
			"-S", fmt.Sprintf("%v:%v", serverListen, serverPort),
			"-t", docrootOrScriptAbsNoSymPath)
	}

	Verboseln("dontbug: Issuing command: rr", strings.Join(rrCmd, " "))
	recordSession := exec.Command(rrPath, rrCmd...)

	f, err := pty.Start(recordSession)
	fatalIf(err)

	color.Yellow("dontbug: -- Recording. Ctrl-C to terminate recording if running on the PHP built-in webserver")
	color.Yellow("dontbug: -- Recording. Ctrl-C if running a script or simply wait for it to end")

	rrTraceDir := ""
	go func() {
		wrappedF := bufio.NewReader(f)
		fatalIf(err)

		for {
			line, err := wrappedF.ReadString('\n')
			fmt.Print(line)
			if err == io.EOF {
				return
			} else if err != nil {
				log.Fatal(err)
			}

			if strings.Contains(line, dontbugRRTraceDirSentinel) {
				start := strings.LastIndex(line, "`")
				end := strings.LastIndex(line, "'")
				if start == -1 || end == -1 || start+1 == len(line) {
					log.Fatal("Could not understand rr trace directory message")
				}

				rrTraceDir = line[start+1 : end]
			}

			if strings.Contains(line, dontbugZendXdebugNotLoadedSentinel) ||
				(strings.Contains(line, "Failed loading") && strings.Contains(line, "xdebug.so")) ||
				(strings.Contains(line, "Cannot load Xdebug") && !strings.Contains(line, "Cannot load Xdebug - it was already loaded")) {
				log.Fatal("xdebug zend extension was not loaded. dontbug needs xdebug to work correctly")
			}

			if strings.Contains(line, dontbugZendXdebugEntryPointNotFoundSentinel) {
				log.Fatal(dontbugNotPatchedXdebugMsg)
			}

			if (strings.Contains(line, "Failed loading") && strings.Contains(line, "dontbug.so")) ||
				strings.Contains(line, "Cannot load dontbug") {
				log.Fatal("Could not load dontbug.so")
			}

			if strings.Contains(line, dontbugZendExtensionLoadedSentinel) {
				break
			}
		}

		io.Copy(os.Stdout, f)
	}()

	// Handle a Ctrl+C
	// If we don't do this rr will terminate abruptly and not save the execution traces properly
	c := make(chan os.Signal)
	defer close(c)

	signal.Notify(c, os.Interrupt) // Ctrl+C
	go func() {
		<-c
		color.Yellow("dontbug: Sending a Ctrl+C to recording")
		f.Write([]byte{3}) // Ctrl+C is ASCII code 3
		signal.Stop(c)
	}()

	err = recordSession.Wait()
	fatalIf(err)

	if takeSnapshot {
		if rrTraceDir == "" {
			log.Fatal("Could not detect rr trace dir location")
		}
		createSnapshotMetadata(rrTraceDir, snapShotDir, originalDocrootOrScriptFullPath)
	}
	color.Green("\ndontbug: Closed cleanly. Replay should work properly")
}

func createSnapshotMetadata(rrTraceDir, snapShotDir string, originalDocrootOrScriptFullPath string) {
	fileData := []byte(snapShotDir + ":" + originalDocrootOrScriptFullPath)
	metaDataFilename := rrTraceDir + "/dontbug-snapshot-metadata"
	err := ioutil.WriteFile(metaDataFilename, fileData, 0700)
	if err != nil {
		log.Fatalf("Could not write to %v\n", metaDataFilename)
	}
}

// Here we're basically serving the role of an PHP debugger in an IDE
func startBasicDebuggerClient(recordPort int) {
	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%v", recordPort))
	fatalIf(err)

	Verbosef("Started debug client for recording at 127.0.0.1:%v\n", recordPort)
	go func() {
		for {
			conn, err := listener.Accept()
			fatalIf(err)

			go func(conn net.Conn) {
				buf := make([]byte, 2048)
				seq := 0
				for {
					bytesRead, _ := conn.Read(buf)
					if bytesRead <= 0 {
						return
					}

					nullAt := bytes.IndexByte(buf, byte(0))
					if nullAt == -1 {
						log.Fatal("Could not find length in debugger engine response")
					}

					dataLen, err := strconv.Atoi(string(buf[0:nullAt]))
					fatalIf(err)

					bytesLeft := dataLen - (bytesRead - nullAt - 2)
					if bytesLeft != 0 {
						log.Fatal("There are still some bytes left to receive -- strange")
					}

					seq++

					// Keep running until we are able to record the execution
					runCommand := fmt.Sprintf("run -i %d\x00", seq)
					conn.Write([]byte(runCommand))
				}
			}(conn)
		}
	}()
}

func checkDontbugWasCompiled(extDirAbsPath string) string {
	dlPath := path.Clean(extDirAbsPath + "/modules/dontbug.so")

	// Does the zend extension exist?
	_, err := os.Stat(dlPath)
	if err != nil {
		log.Fatal("Not able to find dontbug.so")
	}

	return dlPath
}

func getAbsNoSymExtDirAndCheckInstallLocation(installLocation string) string {
	if strings.TrimSpace(installLocation) == "" {
		color.Yellow("dontbug: No --install-location specified. Defaulting to $GOPATH/src/github.com/sidkshatriya/dontbug")
		gopath := os.Getenv("GOPATH")
		if gopath == "" {
			log.Fatal("Unable to find environment variable GOPATH. Is go installed properly?")
		}
		installLocation = getAbsNoSymlinkPath(path.Clean(gopath + "/src/github.com/sidkshatriya/dontbug"))
	} else {
		installLocation = getAbsNoSymlinkPath(installLocation)
	}

	color.Green("dontbug: Using --install-location \"%v\"", installLocation)
	extAbsDir := path.Clean(installLocation + "/ext/dontbug")
	_, err := os.Stat(extAbsDir)
	if err != nil {
		log.Fatalf("'%v' does not seem to be a valid install location of dontbug. Error: %v\n", installLocation, err)
	}

	return extAbsDir
}

func DoChecksAndRecord(
	phpExecutable,
	rrExecutable,
	rootDir,
	installLocation,
	docrootOrScriptRelPath string,
	maxStackDepth int,
	isCli bool,
	arguments string,
	recordPort int,
	serverListen string,
	serverPort int,
	takeSnapshot bool,
) {
	rootAbsNoSymDir := getAbsNoSymlinkPath(rootDir)
	extAbsNoSymDir := getAbsNoSymExtDirAndCheckInstallLocation(installLocation)

	docrootOrScriptFullPath := path.Clean(fmt.Sprintf("%v/%v", rootAbsNoSymDir, docrootOrScriptRelPath))

	snapShotDir := ""
	originalDocrootOrScriptFullPath := ""
	if takeSnapshot {
		snapShotDir = doSnapshot(rootAbsNoSymDir)
		originalDocrootOrScriptFullPath = docrootOrScriptFullPath
		docrootOrScriptFullPath = path.Clean(fmt.Sprintf("%v/%v", snapShotDir, docrootOrScriptRelPath))
	}

	docrootOrScriptAbsNoSymPath := getAbsNoSymlinkPath(docrootOrScriptFullPath)

	phpPath := checkPhpExecutable(phpExecutable)
	rrPath := CheckRRExecutable(rrExecutable)

	doGeneration(rootAbsNoSymDir, extAbsNoSymDir, maxStackDepth, phpPath)
	dontbugSharedObjectPath := checkDontbugWasCompiled(extAbsNoSymDir)
	startBasicDebuggerClient(recordPort)
	doRecordSession(
		docrootOrScriptAbsNoSymPath,
		dontbugSharedObjectPath,
		rrPath,
		phpPath,
		isCli,
		arguments,
		serverListen,
		serverPort,
		recordPort,
		maxStackDepth,
		takeSnapshot,
		snapShotDir,
		originalDocrootOrScriptFullPath,
	)
}

func doSnapshot(rootAbsNoSymDir string) string {
	rootAbsNoSymDir = path.Clean(rootAbsNoSymDir) + "/"
	hash := sha1.Sum([]byte(rootAbsNoSymDir))

	sharePath := getOrCreateDontbugSharePath()
	hashx := fmt.Sprintf("%.10x", hash)

	snapShotGroupDir := fmt.Sprintf("%v%v/", sharePath, hashx)
	mkDirAll(snapShotGroupDir)

	matches, err := filepath.Glob(snapShotGroupDir + "snap-*")
	lastSnapExists := false
	lastSnapName := ""
	if len(matches) != 0 {
		lastSnapExists = true
		lastSnapName = matches[len(matches)-1]
		Verbosef("dontbug: Last snapshot was: %v\n", lastSnapName)
	}

	command := []string{}

	// @TODO incomplete?
	common := []string{
		"--exclude=.git",
		"--exclude=.hg",
	}

	snapShotDir := fmt.Sprintf("%vsnap-%v/", snapShotGroupDir, time.Now().UnixNano()/1000000)
	if !lastSnapExists {
		command = []string{
			"rsync",
			"-a",
			rootAbsNoSymDir, snapShotDir,
		}

		Verbosef("dontbug: Creating master snapshot from: %v\n", rootAbsNoSymDir)
	} else {
		command = []string{
			"rsync",
			"-a",
			"--delete",
			fmt.Sprint("--link-dest=../", path.Base(lastSnapName)),
			rootAbsNoSymDir, snapShotDir,
		}

	}

	command = append(command, common...)
	color.Green("dontbug: rsyncing sources and creating a snapshot at: %v", snapShotDir)
	color.Green("dontbug: If this was your second or later snapshot, disk usage should only go up by what was changed from previous snapshot")
	Verboseln("dontbug: Issuing command: ", strings.Join(command, " "))
	outputBytes, err := exec.Command(command[0], command[1:]...).CombinedOutput()
	if err != nil {
		fmt.Println(string(outputBytes))
		log.Fatal(err)
	}

	if VerboseFlag {
		fmt.Println(string(outputBytes))
	}

	return snapShotDir
}
<commit_msg>Set xdebug.remote_host to 127.0.0.1 always during recording<commit_after>// Copyright © 2016 Sidharth Kshatriya
//
// 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 engine

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"fmt"
	"github.com/fatih/color"
	"github.com/kr/pty"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	// These strings are not to be changed as these strings are sentinels from the dontbug zend extension
	dontbugZendExtensionLoadedSentinel          = "dontbug zend extension: dontbug.so successfully loaded by PHP"
	dontbugZendXdebugNotLoadedSentinel          = "dontbug zend extension: Xdebug has not been loaded"
	dontbugZendXdebugEntryPointNotFoundSentinel = "dontbug zend extension: Xdebug entrypoint not found"
	// End do not change

	dontbugRRTraceDirSentinel = "rr: Saving execution to trace directory `"

	dontbugNotPatchedXdebugMsg = `Unpatched Xdebug zend extension (xdebug.so) found. See below for more information:
dontbug zend extension currently relies on a patched version of Xdebug to function correctly.
This is a very minor patch and simply makes a single function extern (instead of static) linkage.
It seems you are using the plain vanilla version of Xdebug. Consult documentation on patching Xdebug.
`
)

func getOrCreateDontbugSharePath() string {
	currentUser, err := user.Current()
	fatalIf(err)

	dontbugShareDir := currentUser.HomeDir + "/.local/share/dontbug/"
	mkDirAll(dontbugShareDir)

	return dontbugShareDir
}

func copyAndMakeUniqueDontbugSo(sharedObjectPath, dontbugShareDir string) string {
	uniqueDontbugSoFilename := path.Clean(fmt.Sprintf("%v/at-%v-dontbug.so", dontbugShareDir, time.Now().UnixNano()))
	output, err := exec.Command("cp", sharedObjectPath, uniqueDontbugSoFilename).CombinedOutput()
	if err != nil {
		log.Fatal(output)
	}
	return uniqueDontbugSoFilename
}

// Assumptions:
// - rrPath represents an rr executable that meets dontbug's requirements
// - phpPath represents an php executable that meets dontbug's requirements
// - sharedObject path is the path to xdebug.so that meets dontbug's requirements
// - docrootOrScriptAbsNoSymPath is a valid docroot directory or a php script
func doRecordSession(
	docrootOrScriptAbsNoSymPath,
	sharedObjectPath,
	rrPath,
	phpPath string,
	isCli bool,
	arguments,
	serverListen string,
	serverPort,
	recordPort,
	maxStackDepth int,
	takeSnapshot bool,
	snapShotDir string,
	originalDocrootOrScriptFullPath string,
) {
	newSharedObjectPath := sharedObjectPath
	if takeSnapshot {
		dontbugShareDir := getOrCreateDontbugSharePath()
		newSharedObjectPath = copyAndMakeUniqueDontbugSo(sharedObjectPath, dontbugShareDir)
	}

	// Many of these options are not really necessary to be specified.
	// However, we still do that to override any settings that
	// might be present in user php.ini files and change them
	// to sensible defaults for 'dontbug record'
	rrCmd := []string{
		"record",
		phpPath,
		"-d", "zend_extension=" + newSharedObjectPath,
		"-d", fmt.Sprintf("xdebug.remote_port=%v", recordPort),
		"-d", "xdebug.remote_autostart=1",
		"-d", "xdebug.remote_host=\"127.0.0.1\"",
		"-d", "xdebug.remote_connect_back=0",
		"-d", "xdebug.remote_enable=1",
		"-d", "xdebug.remote_mode=req",
		"-d", "xdebug.auto_trace=0",
		"-d", "xdebug.trace_enable_trigger=\"\"",
		"-d", "xdebug.coverage_enable=0",
		"-d", "xdebug.extended_info=1",
		"-d", fmt.Sprintf("xdebug.max_nesting_level=%v", maxStackDepth),
		"-d", "xdebug.profiler_enable=0",
		"-d", "xdebug.profiler_enable_trigger=0",
	}

	if isCli {
		arguments = strings.TrimSpace(arguments)
		rrCmd = append(rrCmd, docrootOrScriptAbsNoSymPath)
		if arguments != "" {
			argumentsAr := strings.Split(arguments, " ")
			rrCmd = append(rrCmd, argumentsAr...)
		}
	} else {
		rrCmd = append(
			rrCmd,
			"-S", fmt.Sprintf("%v:%v", serverListen, serverPort),
			"-t", docrootOrScriptAbsNoSymPath)
	}

	Verboseln("dontbug: Issuing command: rr", strings.Join(rrCmd, " "))
	recordSession := exec.Command(rrPath, rrCmd...)

	f, err := pty.Start(recordSession)
	fatalIf(err)

	color.Yellow("dontbug: -- Recording. Ctrl-C to terminate recording if running on the PHP built-in webserver")
	color.Yellow("dontbug: -- Recording. Ctrl-C if running a script or simply wait for it to end")

	rrTraceDir := ""
	go func() {
		wrappedF := bufio.NewReader(f)
		fatalIf(err)

		for {
			line, err := wrappedF.ReadString('\n')
			fmt.Print(line)
			if err == io.EOF {
				return
			} else if err != nil {
				log.Fatal(err)
			}

			if strings.Contains(line, dontbugRRTraceDirSentinel) {
				start := strings.LastIndex(line, "`")
				end := strings.LastIndex(line, "'")
				if start == -1 || end == -1 || start+1 == len(line) {
					log.Fatal("Could not understand rr trace directory message")
				}

				rrTraceDir = line[start+1 : end]
			}

			if strings.Contains(line, dontbugZendXdebugNotLoadedSentinel) ||
				(strings.Contains(line, "Failed loading") && strings.Contains(line, "xdebug.so")) ||
				(strings.Contains(line, "Cannot load Xdebug") && !strings.Contains(line, "Cannot load Xdebug - it was already loaded")) {
				log.Fatal("xdebug zend extension was not loaded. dontbug needs xdebug to work correctly")
			}

			if strings.Contains(line, dontbugZendXdebugEntryPointNotFoundSentinel) {
				log.Fatal(dontbugNotPatchedXdebugMsg)
			}

			if (strings.Contains(line, "Failed loading") && strings.Contains(line, "dontbug.so")) ||
				strings.Contains(line, "Cannot load dontbug") {
				log.Fatal("Could not load dontbug.so")
			}

			if strings.Contains(line, dontbugZendExtensionLoadedSentinel) {
				break
			}
		}

		io.Copy(os.Stdout, f)
	}()

	// Handle a Ctrl+C
	// If we don't do this rr will terminate abruptly and not save the execution traces properly
	c := make(chan os.Signal)
	defer close(c)

	signal.Notify(c, os.Interrupt) // Ctrl+C
	go func() {
		<-c
		color.Yellow("dontbug: Sending a Ctrl+C to recording")
		f.Write([]byte{3}) // Ctrl+C is ASCII code 3
		signal.Stop(c)
	}()

	err = recordSession.Wait()
	fatalIf(err)

	if takeSnapshot {
		if rrTraceDir == "" {
			log.Fatal("Could not detect rr trace dir location")
		}
		createSnapshotMetadata(rrTraceDir, snapShotDir, originalDocrootOrScriptFullPath)
	}
	color.Green("\ndontbug: Closed cleanly. Replay should work properly")
}

func createSnapshotMetadata(rrTraceDir, snapShotDir string, originalDocrootOrScriptFullPath string) {
	fileData := []byte(snapShotDir + ":" + originalDocrootOrScriptFullPath)
	metaDataFilename := rrTraceDir + "/dontbug-snapshot-metadata"
	err := ioutil.WriteFile(metaDataFilename, fileData, 0700)
	if err != nil {
		log.Fatalf("Could not write to %v\n", metaDataFilename)
	}
}

// Here we're basically serving the role of an PHP debugger in an IDE
func startBasicDebuggerClient(recordPort int) {
	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%v", recordPort))
	fatalIf(err)

	Verbosef("Started debug client for recording at 127.0.0.1:%v\n", recordPort)
	go func() {
		for {
			conn, err := listener.Accept()
			fatalIf(err)

			go func(conn net.Conn) {
				buf := make([]byte, 2048)
				seq := 0
				for {
					bytesRead, _ := conn.Read(buf)
					if bytesRead <= 0 {
						return
					}

					nullAt := bytes.IndexByte(buf, byte(0))
					if nullAt == -1 {
						log.Fatal("Could not find length in debugger engine response")
					}

					dataLen, err := strconv.Atoi(string(buf[0:nullAt]))
					fatalIf(err)

					bytesLeft := dataLen - (bytesRead - nullAt - 2)
					if bytesLeft != 0 {
						log.Fatal("There are still some bytes left to receive -- strange")
					}

					seq++

					// Keep running until we are able to record the execution
					runCommand := fmt.Sprintf("run -i %d\x00", seq)
					conn.Write([]byte(runCommand))
				}
			}(conn)
		}
	}()
}

func checkDontbugWasCompiled(extDirAbsPath string) string {
	dlPath := path.Clean(extDirAbsPath + "/modules/dontbug.so")

	// Does the zend extension exist?
	_, err := os.Stat(dlPath)
	if err != nil {
		log.Fatal("Not able to find dontbug.so")
	}

	return dlPath
}

func getAbsNoSymExtDirAndCheckInstallLocation(installLocation string) string {
	if strings.TrimSpace(installLocation) == "" {
		color.Yellow("dontbug: No --install-location specified. Defaulting to $GOPATH/src/github.com/sidkshatriya/dontbug")
		gopath := os.Getenv("GOPATH")
		if gopath == "" {
			log.Fatal("Unable to find environment variable GOPATH. Is go installed properly?")
		}
		installLocation = getAbsNoSymlinkPath(path.Clean(gopath + "/src/github.com/sidkshatriya/dontbug"))
	} else {
		installLocation = getAbsNoSymlinkPath(installLocation)
	}

	color.Green("dontbug: Using --install-location \"%v\"", installLocation)
	extAbsDir := path.Clean(installLocation + "/ext/dontbug")
	_, err := os.Stat(extAbsDir)
	if err != nil {
		log.Fatalf("'%v' does not seem to be a valid install location of dontbug. Error: %v\n", installLocation, err)
	}

	return extAbsDir
}

func DoChecksAndRecord(
	phpExecutable,
	rrExecutable,
	rootDir,
	installLocation,
	docrootOrScriptRelPath string,
	maxStackDepth int,
	isCli bool,
	arguments string,
	recordPort int,
	serverListen string,
	serverPort int,
	takeSnapshot bool,
) {
	rootAbsNoSymDir := getAbsNoSymlinkPath(rootDir)
	extAbsNoSymDir := getAbsNoSymExtDirAndCheckInstallLocation(installLocation)

	docrootOrScriptFullPath := path.Clean(fmt.Sprintf("%v/%v", rootAbsNoSymDir, docrootOrScriptRelPath))

	snapShotDir := ""
	originalDocrootOrScriptFullPath := ""
	if takeSnapshot {
		snapShotDir = doSnapshot(rootAbsNoSymDir)
		originalDocrootOrScriptFullPath = docrootOrScriptFullPath
		docrootOrScriptFullPath = path.Clean(fmt.Sprintf("%v/%v", snapShotDir, docrootOrScriptRelPath))
	}

	docrootOrScriptAbsNoSymPath := getAbsNoSymlinkPath(docrootOrScriptFullPath)

	phpPath := checkPhpExecutable(phpExecutable)
	rrPath := CheckRRExecutable(rrExecutable)

	doGeneration(rootAbsNoSymDir, extAbsNoSymDir, maxStackDepth, phpPath)
	dontbugSharedObjectPath := checkDontbugWasCompiled(extAbsNoSymDir)
	startBasicDebuggerClient(recordPort)
	doRecordSession(
		docrootOrScriptAbsNoSymPath,
		dontbugSharedObjectPath,
		rrPath,
		phpPath,
		isCli,
		arguments,
		serverListen,
		serverPort,
		recordPort,
		maxStackDepth,
		takeSnapshot,
		snapShotDir,
		originalDocrootOrScriptFullPath,
	)
}

func doSnapshot(rootAbsNoSymDir string) string {
	rootAbsNoSymDir = path.Clean(rootAbsNoSymDir) + "/"
	hash := sha1.Sum([]byte(rootAbsNoSymDir))

	sharePath := getOrCreateDontbugSharePath()
	hashx := fmt.Sprintf("%.10x", hash)

	snapShotGroupDir := fmt.Sprintf("%v%v/", sharePath, hashx)
	mkDirAll(snapShotGroupDir)

	matches, err := filepath.Glob(snapShotGroupDir + "snap-*")
	lastSnapExists := false
	lastSnapName := ""
	if len(matches) != 0 {
		lastSnapExists = true
		lastSnapName = matches[len(matches)-1]
		Verbosef("dontbug: Last snapshot was: %v\n", lastSnapName)
	}

	command := []string{}

	// @TODO incomplete?
	common := []string{
		"--exclude=.git",
		"--exclude=.hg",
	}

	snapShotDir := fmt.Sprintf("%vsnap-%v/", snapShotGroupDir, time.Now().UnixNano()/1000000)
	if !lastSnapExists {
		command = []string{
			"rsync",
			"-a",
			rootAbsNoSymDir, snapShotDir,
		}

		Verbosef("dontbug: Creating master snapshot from: %v\n", rootAbsNoSymDir)
	} else {
		command = []string{
			"rsync",
			"-a",
			"--delete",
			fmt.Sprint("--link-dest=../", path.Base(lastSnapName)),
			rootAbsNoSymDir, snapShotDir,
		}

	}

	command = append(command, common...)
	color.Green("dontbug: rsyncing sources and creating a snapshot at: %v", snapShotDir)
	color.Green("dontbug: If this was your second or later snapshot, disk usage should only go up by what was changed from previous snapshot")
	Verboseln("dontbug: Issuing command: ", strings.Join(command, " "))
	outputBytes, err := exec.Command(command[0], command[1:]...).CombinedOutput()
	if err != nil {
		fmt.Println(string(outputBytes))
		log.Fatal(err)
	}

	if VerboseFlag {
		fmt.Println(string(outputBytes))
	}

	return snapShotDir
}
<|endoftext|>


<commit_before>package routers

import (
	"../controllers"
	"github.com/astaxie/beego"
)

func init() {
    beego.Router("/", &controllers.MainController{})
}
<commit_msg>login<commit_after>package routers

import (
	"../controllers"
	"github.com/astaxie/beego"
)

func init() {
    beego.Router("/", &controllers.MainController{})
    beego.Router("/login", &controllers.LoginController{})
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"net/url"
	"os"
	"path/filepath"
	texttemplate "text/template"
	"time"

	"github.com/coreos/go-oidc/key"
	"github.com/coreos/go-oidc/oidc"
	"github.com/coreos/pkg/health"
	"github.com/go-gorp/gorp"

	"github.com/coreos/dex/connector"
	"github.com/coreos/dex/db"
	"github.com/coreos/dex/email"
	sessionmanager "github.com/coreos/dex/session/manager"
	"github.com/coreos/dex/user"
	useremail "github.com/coreos/dex/user/email"
	usermanager "github.com/coreos/dex/user/manager"
)

type ServerConfig struct {
	IssuerURL                string
	IssuerName               string
	IssuerLogoURL            string
	TemplateDir              string
	EmailTemplateDirs        []string
	EmailFromAddress         string
	EmailerConfigFile        string
	StateConfig              StateConfigurer
	EnableRegistration       bool
	EnableClientRegistration bool
}

type StateConfigurer interface {
	Configure(*Server) error
}

type SingleServerConfig struct {
	ClientsFile    string
	ConnectorsFile string
	UsersFile      string
}

type MultiServerConfig struct {
	KeySecrets     [][]byte
	DatabaseConfig db.Config
	UseOldFormat   bool
}

func (cfg *ServerConfig) Server() (*Server, error) {
	iu, err := url.Parse(cfg.IssuerURL)
	if err != nil {
		return nil, err
	}

	tpl, err := getTemplates(cfg.IssuerName, cfg.IssuerLogoURL, cfg.EnableRegistration, cfg.TemplateDir)
	if err != nil {
		return nil, err
	}

	km := key.NewPrivateKeyManager()
	srv := Server{
		IssuerURL:  *iu,
		KeyManager: km,
		Templates:  tpl,

		HealthChecks: []health.Checkable{km},
		Connectors:   []connector.Connector{},

		EnableRegistration:       cfg.EnableRegistration,
		EnableClientRegistration: cfg.EnableClientRegistration,
	}

	err = cfg.StateConfig.Configure(&srv)
	if err != nil {
		return nil, err
	}

	err = setTemplates(&srv, tpl)
	if err != nil {
		return nil, err
	}

	err = setEmailer(&srv, cfg.IssuerName, cfg.EmailFromAddress, cfg.EmailerConfigFile, cfg.EmailTemplateDirs)
	if err != nil {
		return nil, err
	}
	return &srv, nil
}

func (cfg *SingleServerConfig) Configure(srv *Server) error {
	k, err := key.GeneratePrivateKey()
	if err != nil {
		return err
	}

	dbMap := db.NewMemDB()

	ks := key.NewPrivateKeySet([]*key.PrivateKey{k}, time.Now().Add(24*time.Hour))
	kRepo := key.NewPrivateKeySetRepo()
	if err = kRepo.Set(ks); err != nil {
		return err
	}

	clients, err := loadClients(cfg.ClientsFile)
	if err != nil {
		return fmt.Errorf("unable to read clients from file %s: %v", cfg.ClientsFile, err)
	}
	ciRepo, err := db.NewClientIdentityRepoFromClients(dbMap, clients)
	if err != nil {
		return fmt.Errorf("failed to create client identity repo: %v", err)
	}

	f, err := os.Open(cfg.ConnectorsFile)
	if err != nil {
		return fmt.Errorf("opening connectors file: %v", err)
	}
	defer f.Close()
	cfgs, err := connector.ReadConfigs(f)
	if err != nil {
		return fmt.Errorf("decoding connector configs: %v", err)
	}
	cfgRepo := db.NewConnectorConfigRepo(dbMap)
	if err := cfgRepo.Set(cfgs); err != nil {
		return fmt.Errorf("failed to set connectors: %v", err)
	}

	sRepo := db.NewSessionRepo(dbMap)
	skRepo := db.NewSessionKeyRepo(dbMap)
	sm := sessionmanager.NewSessionManager(sRepo, skRepo)

	users, err := loadUsers(cfg.UsersFile)
	if err != nil {
		return fmt.Errorf("unable to read users from file: %v", err)
	}
	userRepo, err := db.NewUserRepoFromUsers(dbMap, users)
	if err != nil {
		return err
	}

	pwiRepo := db.NewPasswordInfoRepo(dbMap)

	refTokRepo := db.NewRefreshTokenRepo(dbMap)

	txnFactory := db.TransactionFactory(dbMap)
	userManager := usermanager.NewUserManager(userRepo, pwiRepo, cfgRepo, txnFactory, usermanager.ManagerOptions{})
	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refTokRepo
	return nil
}

func loadUsers(filepath string) (users []user.UserWithRemoteIdentities, err error) {
	f, err := os.Open(filepath)
	if err != nil {
		return
	}
	defer f.Close()
	err = json.NewDecoder(f).Decode(&users)
	return
}

func loadClients(filepath string) ([]oidc.ClientIdentity, error) {
	f, err := os.Open(filepath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var c []struct {
		ID           string   `json:"id"`
		Secret       string   `json:"secret"`
		RedirectURLs []string `json:"redirectURLs"`
	}
	if err := json.NewDecoder(f).Decode(&c); err != nil {
		return nil, err
	}
	clients := make([]oidc.ClientIdentity, len(c))
	for i, client := range c {
		redirectURIs := make([]url.URL, len(client.RedirectURLs))
		for j, u := range client.RedirectURLs {
			uri, err := url.Parse(u)
			if err != nil {
				return nil, err
			}
			redirectURIs[j] = *uri
		}

		clients[i] = oidc.ClientIdentity{
			Credentials: oidc.ClientCredentials{
				ID:     client.ID,
				Secret: client.Secret,
			},
			Metadata: oidc.ClientMetadata{
				RedirectURIs: redirectURIs,
			},
		}
	}
	return clients, nil
}

func (cfg *MultiServerConfig) Configure(srv *Server) error {
	if len(cfg.KeySecrets) == 0 {
		return errors.New("missing key secret")
	}

	if cfg.DatabaseConfig.DSN == "" {
		return errors.New("missing database connection string")
	}

	dbc, err := db.NewConnection(cfg.DatabaseConfig)
	if err != nil {
		return fmt.Errorf("unable to initialize database connection: %v", err)
	}
	if _, ok := dbc.Dialect.(gorp.PostgresDialect); !ok {
		return errors.New("only postgres backend supported for multi server configurations")
	}

	kRepo, err := db.NewPrivateKeySetRepo(dbc, cfg.UseOldFormat, cfg.KeySecrets...)
	if err != nil {
		return fmt.Errorf("unable to create PrivateKeySetRepo: %v", err)
	}

	ciRepo := db.NewClientIdentityRepo(dbc)
	sRepo := db.NewSessionRepo(dbc)
	skRepo := db.NewSessionKeyRepo(dbc)
	cfgRepo := db.NewConnectorConfigRepo(dbc)
	userRepo := db.NewUserRepo(dbc)
	pwiRepo := db.NewPasswordInfoRepo(dbc)
	userManager := usermanager.NewUserManager(userRepo, pwiRepo, cfgRepo, db.TransactionFactory(dbc), usermanager.ManagerOptions{})
	refreshTokenRepo := db.NewRefreshTokenRepo(dbc)

	sm := sessionmanager.NewSessionManager(sRepo, skRepo)

	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refreshTokenRepo
	return nil
}

func getTemplates(issuerName, issuerLogoURL string,
	enableRegister bool, dir string) (*template.Template, error) {
	tpl := template.New("").Funcs(map[string]interface{}{
		"issuerName": func() string {
			return issuerName
		},
		"issuerLogoURL": func() string {
			return issuerLogoURL
		},
		"enableRegister": func() bool {
			return enableRegister
		},
	})

	return tpl.ParseGlob(dir + "/*.html")
}

func setTemplates(srv *Server, tpls *template.Template) error {
	ltpl, err := findTemplate(LoginPageTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.LoginTemplate = ltpl

	rtpl, err := findTemplate(RegisterTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.RegisterTemplate = rtpl

	vtpl, err := findTemplate(VerifyEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.VerifyEmailTemplate = vtpl

	srtpl, err := findTemplate(SendResetPasswordEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.SendResetPasswordEmailTemplate = srtpl

	rpwtpl, err := findTemplate(ResetPasswordTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.ResetPasswordTemplate = rpwtpl

	return nil
}

func setEmailer(srv *Server, issuerName, fromAddress, emailerConfigFile string, emailTemplateDirs []string) error {

	cfg, err := email.NewEmailerConfigFromFile(emailerConfigFile)
	if err != nil {
		return err
	}

	emailer, err := cfg.Emailer()
	if err != nil {
		return err
	}

	getFileNames := func(dir, ext string) ([]string, error) {
		fns, err := filepath.Glob(dir + "/*." + ext)
		if err != nil {
			return nil, err
		}
		return fns, nil
	}
	getTextFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "txt")
	}
	getHTMLFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "html")
	}

	textTemplates := texttemplate.New("textTemplates")
	htmlTemplates := template.New("htmlTemplates")
	for _, dir := range emailTemplateDirs {
		textFileNames, err := getTextFiles(dir)
		if err != nil {
			return err
		}
		if len(textFileNames) != 0 {
			textTemplates, err = textTemplates.ParseFiles(textFileNames...)
		}
		if err != nil {
			return err
		}

		htmlFileNames, err := getHTMLFiles(dir)
		if err != nil {
			return err
		}
		if len(htmlFileNames) != 0 {
			htmlTemplates, err = htmlTemplates.ParseFiles(htmlFileNames...)
		}
		if err != nil {
			return err
		}
	}
	tMailer := email.NewTemplatizedEmailerFromTemplates(textTemplates, htmlTemplates, emailer)
	tMailer.SetGlobalContext(map[string]interface{}{
		"issuer_name": issuerName,
	})

	ue := useremail.NewUserEmailer(srv.UserRepo,
		srv.PasswordInfoRepo,
		srv.KeyManager.Signer,
		srv.SessionManager.ValidityWindow,
		srv.IssuerURL,
		tMailer,
		fromAddress,
		srv.absURL(httpPathResetPassword),
		srv.absURL(httpPathEmailVerify),
		srv.absURL(httpPathAcceptInvitation),
	)

	srv.UserEmailer = ue
	return nil
}

func findTemplate(name string, tpls *template.Template) (*template.Template, error) {
	tpl := tpls.Lookup(name)
	if tpl == nil {
		return nil, fmt.Errorf("unable to find template: %q", name)
	}
	return tpl, nil
}
<commit_msg>server: add db heatlh checker to server checkers<commit_after>package server

import (
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"net/url"
	"os"
	"path/filepath"
	texttemplate "text/template"
	"time"

	"github.com/coreos/go-oidc/key"
	"github.com/coreos/go-oidc/oidc"
	"github.com/coreos/pkg/health"
	"github.com/go-gorp/gorp"

	"github.com/coreos/dex/connector"
	"github.com/coreos/dex/db"
	"github.com/coreos/dex/email"
	sessionmanager "github.com/coreos/dex/session/manager"
	"github.com/coreos/dex/user"
	useremail "github.com/coreos/dex/user/email"
	usermanager "github.com/coreos/dex/user/manager"
)

type ServerConfig struct {
	IssuerURL                string
	IssuerName               string
	IssuerLogoURL            string
	TemplateDir              string
	EmailTemplateDirs        []string
	EmailFromAddress         string
	EmailerConfigFile        string
	StateConfig              StateConfigurer
	EnableRegistration       bool
	EnableClientRegistration bool
}

type StateConfigurer interface {
	Configure(*Server) error
}

type SingleServerConfig struct {
	ClientsFile    string
	ConnectorsFile string
	UsersFile      string
}

type MultiServerConfig struct {
	KeySecrets     [][]byte
	DatabaseConfig db.Config
	UseOldFormat   bool
}

func (cfg *ServerConfig) Server() (*Server, error) {
	iu, err := url.Parse(cfg.IssuerURL)
	if err != nil {
		return nil, err
	}

	tpl, err := getTemplates(cfg.IssuerName, cfg.IssuerLogoURL, cfg.EnableRegistration, cfg.TemplateDir)
	if err != nil {
		return nil, err
	}

	km := key.NewPrivateKeyManager()
	srv := Server{
		IssuerURL:  *iu,
		KeyManager: km,
		Templates:  tpl,

		HealthChecks: []health.Checkable{km},
		Connectors:   []connector.Connector{},

		EnableRegistration:       cfg.EnableRegistration,
		EnableClientRegistration: cfg.EnableClientRegistration,
	}

	err = cfg.StateConfig.Configure(&srv)
	if err != nil {
		return nil, err
	}

	err = setTemplates(&srv, tpl)
	if err != nil {
		return nil, err
	}

	err = setEmailer(&srv, cfg.IssuerName, cfg.EmailFromAddress, cfg.EmailerConfigFile, cfg.EmailTemplateDirs)
	if err != nil {
		return nil, err
	}
	return &srv, nil
}

func (cfg *SingleServerConfig) Configure(srv *Server) error {
	k, err := key.GeneratePrivateKey()
	if err != nil {
		return err
	}

	dbMap := db.NewMemDB()

	ks := key.NewPrivateKeySet([]*key.PrivateKey{k}, time.Now().Add(24*time.Hour))
	kRepo := key.NewPrivateKeySetRepo()
	if err = kRepo.Set(ks); err != nil {
		return err
	}

	clients, err := loadClients(cfg.ClientsFile)
	if err != nil {
		return fmt.Errorf("unable to read clients from file %s: %v", cfg.ClientsFile, err)
	}
	ciRepo, err := db.NewClientIdentityRepoFromClients(dbMap, clients)
	if err != nil {
		return fmt.Errorf("failed to create client identity repo: %v", err)
	}

	f, err := os.Open(cfg.ConnectorsFile)
	if err != nil {
		return fmt.Errorf("opening connectors file: %v", err)
	}
	defer f.Close()
	cfgs, err := connector.ReadConfigs(f)
	if err != nil {
		return fmt.Errorf("decoding connector configs: %v", err)
	}
	cfgRepo := db.NewConnectorConfigRepo(dbMap)
	if err := cfgRepo.Set(cfgs); err != nil {
		return fmt.Errorf("failed to set connectors: %v", err)
	}

	sRepo := db.NewSessionRepo(dbMap)
	skRepo := db.NewSessionKeyRepo(dbMap)
	sm := sessionmanager.NewSessionManager(sRepo, skRepo)

	users, err := loadUsers(cfg.UsersFile)
	if err != nil {
		return fmt.Errorf("unable to read users from file: %v", err)
	}
	userRepo, err := db.NewUserRepoFromUsers(dbMap, users)
	if err != nil {
		return err
	}

	pwiRepo := db.NewPasswordInfoRepo(dbMap)

	refTokRepo := db.NewRefreshTokenRepo(dbMap)

	txnFactory := db.TransactionFactory(dbMap)
	userManager := usermanager.NewUserManager(userRepo, pwiRepo, cfgRepo, txnFactory, usermanager.ManagerOptions{})
	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refTokRepo
	srv.HealthChecks = append(srv.HealthChecks, db.NewHealthChecker(dbMap))
	return nil
}

func loadUsers(filepath string) (users []user.UserWithRemoteIdentities, err error) {
	f, err := os.Open(filepath)
	if err != nil {
		return
	}
	defer f.Close()
	err = json.NewDecoder(f).Decode(&users)
	return
}

func loadClients(filepath string) ([]oidc.ClientIdentity, error) {
	f, err := os.Open(filepath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var c []struct {
		ID           string   `json:"id"`
		Secret       string   `json:"secret"`
		RedirectURLs []string `json:"redirectURLs"`
	}
	if err := json.NewDecoder(f).Decode(&c); err != nil {
		return nil, err
	}
	clients := make([]oidc.ClientIdentity, len(c))
	for i, client := range c {
		redirectURIs := make([]url.URL, len(client.RedirectURLs))
		for j, u := range client.RedirectURLs {
			uri, err := url.Parse(u)
			if err != nil {
				return nil, err
			}
			redirectURIs[j] = *uri
		}

		clients[i] = oidc.ClientIdentity{
			Credentials: oidc.ClientCredentials{
				ID:     client.ID,
				Secret: client.Secret,
			},
			Metadata: oidc.ClientMetadata{
				RedirectURIs: redirectURIs,
			},
		}
	}
	return clients, nil
}

func (cfg *MultiServerConfig) Configure(srv *Server) error {
	if len(cfg.KeySecrets) == 0 {
		return errors.New("missing key secret")
	}

	if cfg.DatabaseConfig.DSN == "" {
		return errors.New("missing database connection string")
	}

	dbc, err := db.NewConnection(cfg.DatabaseConfig)
	if err != nil {
		return fmt.Errorf("unable to initialize database connection: %v", err)
	}
	if _, ok := dbc.Dialect.(gorp.PostgresDialect); !ok {
		return errors.New("only postgres backend supported for multi server configurations")
	}

	kRepo, err := db.NewPrivateKeySetRepo(dbc, cfg.UseOldFormat, cfg.KeySecrets...)
	if err != nil {
		return fmt.Errorf("unable to create PrivateKeySetRepo: %v", err)
	}

	ciRepo := db.NewClientIdentityRepo(dbc)
	sRepo := db.NewSessionRepo(dbc)
	skRepo := db.NewSessionKeyRepo(dbc)
	cfgRepo := db.NewConnectorConfigRepo(dbc)
	userRepo := db.NewUserRepo(dbc)
	pwiRepo := db.NewPasswordInfoRepo(dbc)
	userManager := usermanager.NewUserManager(userRepo, pwiRepo, cfgRepo, db.TransactionFactory(dbc), usermanager.ManagerOptions{})
	refreshTokenRepo := db.NewRefreshTokenRepo(dbc)

	sm := sessionmanager.NewSessionManager(sRepo, skRepo)

	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refreshTokenRepo
	srv.HealthChecks = append(srv.HealthChecks, db.NewHealthChecker(dbc))
	return nil
}

func getTemplates(issuerName, issuerLogoURL string,
	enableRegister bool, dir string) (*template.Template, error) {
	tpl := template.New("").Funcs(map[string]interface{}{
		"issuerName": func() string {
			return issuerName
		},
		"issuerLogoURL": func() string {
			return issuerLogoURL
		},
		"enableRegister": func() bool {
			return enableRegister
		},
	})

	return tpl.ParseGlob(dir + "/*.html")
}

func setTemplates(srv *Server, tpls *template.Template) error {
	ltpl, err := findTemplate(LoginPageTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.LoginTemplate = ltpl

	rtpl, err := findTemplate(RegisterTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.RegisterTemplate = rtpl

	vtpl, err := findTemplate(VerifyEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.VerifyEmailTemplate = vtpl

	srtpl, err := findTemplate(SendResetPasswordEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.SendResetPasswordEmailTemplate = srtpl

	rpwtpl, err := findTemplate(ResetPasswordTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.ResetPasswordTemplate = rpwtpl

	return nil
}

func setEmailer(srv *Server, issuerName, fromAddress, emailerConfigFile string, emailTemplateDirs []string) error {

	cfg, err := email.NewEmailerConfigFromFile(emailerConfigFile)
	if err != nil {
		return err
	}

	emailer, err := cfg.Emailer()
	if err != nil {
		return err
	}

	getFileNames := func(dir, ext string) ([]string, error) {
		fns, err := filepath.Glob(dir + "/*." + ext)
		if err != nil {
			return nil, err
		}
		return fns, nil
	}
	getTextFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "txt")
	}
	getHTMLFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "html")
	}

	textTemplates := texttemplate.New("textTemplates")
	htmlTemplates := template.New("htmlTemplates")
	for _, dir := range emailTemplateDirs {
		textFileNames, err := getTextFiles(dir)
		if err != nil {
			return err
		}
		if len(textFileNames) != 0 {
			textTemplates, err = textTemplates.ParseFiles(textFileNames...)
		}
		if err != nil {
			return err
		}

		htmlFileNames, err := getHTMLFiles(dir)
		if err != nil {
			return err
		}
		if len(htmlFileNames) != 0 {
			htmlTemplates, err = htmlTemplates.ParseFiles(htmlFileNames...)
		}
		if err != nil {
			return err
		}
	}
	tMailer := email.NewTemplatizedEmailerFromTemplates(textTemplates, htmlTemplates, emailer)
	tMailer.SetGlobalContext(map[string]interface{}{
		"issuer_name": issuerName,
	})

	ue := useremail.NewUserEmailer(srv.UserRepo,
		srv.PasswordInfoRepo,
		srv.KeyManager.Signer,
		srv.SessionManager.ValidityWindow,
		srv.IssuerURL,
		tMailer,
		fromAddress,
		srv.absURL(httpPathResetPassword),
		srv.absURL(httpPathEmailVerify),
		srv.absURL(httpPathAcceptInvitation),
	)

	srv.UserEmailer = ue
	return nil
}

func findTemplate(name string, tpls *template.Template) (*template.Template, error) {
	tpl := tpls.Lookup(name)
	if tpl == nil {
		return nil, fmt.Errorf("unable to find template: %q", name)
	}
	return tpl, nil
}
<|endoftext|>


<commit_before>package soundtrack

import (
	"net/http"

	"github.com/aerogo/aero"
	"github.com/animenotifier/notify.moe/arn"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/server/middleware"
)

// Get track.
func Get(ctx aero.Context) error {
	id := ctx.Get("id")
	track, err := arn.GetSoundTrack(id)
	user := arn.GetUserFromContext(ctx)

	if err != nil {
		return ctx.Error(http.StatusNotFound, "Track not found", err)
	}

	var relatedTracks []*arn.SoundTrack
	for _, anime := range track.Anime() {
		tracks := arn.FilterSoundTracks(func(t *arn.SoundTrack) bool {
			return !t.IsDraft && len(t.Media) > 0 && t.ID != track.ID && arn.Contains(t.Tags, "anime:"+anime.ID)
		})
		relatedTracks = append(relatedTracks, tracks...)
	}

	customCtx := ctx.(*middleware.OpenGraphContext)
	customCtx.OpenGraph = getOpenGraph(track)
	return ctx.HTML(components.SoundTrackPage(track, relatedTracks, user))
}
<commit_msg>Preallocate `relatedTracks`<commit_after>package soundtrack

import (
	"net/http"

	"github.com/aerogo/aero"
	"github.com/animenotifier/notify.moe/arn"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/server/middleware"
)

// Get track.
func Get(ctx aero.Context) error {
	id := ctx.Get("id")
	track, err := arn.GetSoundTrack(id)
	user := arn.GetUserFromContext(ctx)

	if err != nil {
		return ctx.Error(http.StatusNotFound, "Track not found", err)
	}

	relatedTracks := make([]*arn.SoundTrack, 0, 5)
	for _, anime := range track.Anime() {
		tracks := arn.FilterSoundTracks(func(t *arn.SoundTrack) bool {
			return !t.IsDraft && len(t.Media) > 0 && t.ID != track.ID && arn.Contains(t.Tags, "anime:"+anime.ID)
		})
		relatedTracks = append(relatedTracks, tracks...)
	}

	customCtx := ctx.(*middleware.OpenGraphContext)
	customCtx.OpenGraph = getOpenGraph(track)
	return ctx.HTML(components.SoundTrackPage(track, relatedTracks, user))
}
<|endoftext|>