<commit_before>// Package local is a file system backed store
package file

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-micro/v2/store"
	bolt "go.etcd.io/bbolt"
)

var (
	// DefaultDatabase is the namespace that the bbolt store
	// will use if no namespace is provided.
	DefaultDatabase = "micro"
	// DefaultTable when none is specified
	DefaultTable = "micro"
	// DefaultDir is the default directory for bbolt files
	DefaultDir = filepath.Join(os.TempDir(), "micro", "store")

	// bucket used for data storage
	dataBucket = "data"
)

// NewStore returns a memory store
func NewStore(opts ...store.Option) store.Store {
	s := &fileStore{
		handles: make(map[string]*fileHandle),
	}
	s.init(opts...)
	return s
}

type fileStore struct {
	options store.Options
	dir     string

	// the database handle
	sync.RWMutex
	handles map[string]*fileHandle
}

type fileHandle struct {
	key string
	db  *bolt.DB
}

// record stored by us
type record struct {
	Key       string
	Value     []byte
	Metadata  map[string]interface{}
	ExpiresAt time.Time
}

func key(database, table string) string {
	return database + ":" + table
}

func (m *fileStore) delete(fd *fileHandle, key string) error {
	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}
		return b.Delete([]byte(key))
	})
}

func (m *fileStore) init(opts ...store.Option) error {
	for _, o := range opts {
		o(&m.options)
	}

	if m.options.Database == "" {
		m.options.Database = DefaultDatabase
	}

	if m.options.Table == "" {
		// bbolt requires bucketname to not be empty
		m.options.Table = DefaultTable
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, m.options.Database)
	// Ignoring this as the folder might exist.
	// Reads/Writes updates will return with sensible error messages
	// about the dir not existing in case this cannot create the path anyway
	os.MkdirAll(dir, 0700)

	return nil
}

func (f *fileStore) getDB(database, table string) (*fileHandle, error) {
	if len(database) == 0 {
		database = f.options.Database
	}
	if len(table) == 0 {
		table = f.options.Table
	}

	k := key(database, table)

	f.RLock()
	fd, ok := f.handles[k]
	f.RUnlock()

	// return the file handle
	if ok {
		return fd, nil
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, database)
	// create the database handle
	fname := table + ".db"
	// make the dir
	os.MkdirAll(dir, 0700)
	// database path
	dbPath := filepath.Join(dir, fname)

	// create new db handle
	db, err := bolt.Open(dbPath, 0700, &bolt.Options{Timeout: 5 * time.Second})
	if err != nil {
		return nil, err
	}

	f.Lock()
	fd = &fileHandle{
		key: k,
		db:  db,
	}
	f.handles[k] = fd
	f.Unlock()

	return fd, nil
}

func (m *fileStore) list(fd *fileHandle, limit, offset uint) []string {
	var allItems []string

	fd.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		// nothing to read
		if b == nil {
			return nil
		}

		// @todo very inefficient
		if err := b.ForEach(func(k, v []byte) error {
			storedRecord := &record{}

			if err := json.Unmarshal(v, storedRecord); err != nil {
				return err
			}

			if !storedRecord.ExpiresAt.IsZero() {
				if storedRecord.ExpiresAt.Before(time.Now()) {
					return nil
				}
			}

			allItems = append(allItems, string(k))

			return nil
		}); err != nil {
			return err
		}

		return nil
	})

	allKeys := make([]string, len(allItems))

	for i, k := range allItems {
		allKeys[i] = k
	}

	if limit != 0 || offset != 0 {
		sort.Slice(allKeys, func(i, j int) bool { return allKeys[i] < allKeys[j] })
		min := func(i, j uint) uint {
			if i < j {
				return i
			}
			return j
		}
		return allKeys[offset:min(limit, uint(len(allKeys)))]
	}

	return allKeys
}

func (m *fileStore) get(fd *fileHandle, k string) (*store.Record, error) {
	var value []byte

	fd.db.View(func(tx *bolt.Tx) error {
		// @todo this is still very experimental...
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}

		value = b.Get([]byte(k))
		return nil
	})

	if value == nil {
		return nil, store.ErrNotFound
	}

	storedRecord := &record{}

	if err := json.Unmarshal(value, storedRecord); err != nil {
		return nil, err
	}

	newRecord := &store.Record{}
	newRecord.Key = storedRecord.Key
	newRecord.Value = storedRecord.Value
	newRecord.Metadata = make(map[string]interface{})

	for k, v := range storedRecord.Metadata {
		newRecord.Metadata[k] = v
	}

	if !storedRecord.ExpiresAt.IsZero() {
		if storedRecord.ExpiresAt.Before(time.Now()) {
			return nil, store.ErrNotFound
		}
		newRecord.Expiry = time.Until(storedRecord.ExpiresAt)
	}

	return newRecord, nil
}

func (m *fileStore) set(fd *fileHandle, r *store.Record) error {
	// copy the incoming record and then
	// convert the expiry in to a hard timestamp
	item := &record{}
	item.Key = r.Key
	item.Value = r.Value
	item.Metadata = make(map[string]interface{})

	if r.Expiry != 0 {
		item.ExpiresAt = time.Now().Add(r.Expiry)
	}

	for k, v := range r.Metadata {
		item.Metadata[k] = v
	}

	// marshal the data
	data, _ := json.Marshal(item)

	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			var err error
			b, err = tx.CreateBucketIfNotExists([]byte(dataBucket))
			if err != nil {
				return err
			}
		}
		return b.Put([]byte(r.Key), data)
	})
}

func (f *fileStore) Close() error {
	f.Lock()
	defer f.Unlock()
	for k, v := range f.handles {
		v.db.Close()
		delete(f.handles, k)
	}
	return nil
}

func (f *fileStore) Init(opts ...store.Option) error {
	return f.init(opts...)
}

func (m *fileStore) Delete(key string, opts ...store.DeleteOption) error {
	var deleteOptions store.DeleteOptions
	for _, o := range opts {
		o(&deleteOptions)
	}

	fd, err := m.getDB(deleteOptions.Database, deleteOptions.Table)
	if err != nil {
		return err
	}

	return m.delete(fd, key)
}

func (m *fileStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var readOpts store.ReadOptions
	for _, o := range opts {
		o(&readOpts)
	}

	fd, err := m.getDB(readOpts.Database, readOpts.Table)
	if err != nil {
		return nil, err
	}

	var keys []string

	// Handle Prefix / suffix
	// TODO: do range scan here rather than listing all keys
	if readOpts.Prefix || readOpts.Suffix {
		// list the keys
		k := m.list(fd, readOpts.Limit, readOpts.Offset)

		// check for prefix and suffix
		for _, v := range k {
			if readOpts.Prefix && !strings.HasPrefix(v, key) {
				continue
			}
			if readOpts.Suffix && !strings.HasSuffix(v, key) {
				continue
			}
			keys = append(keys, v)
		}
	} else {
		keys = []string{key}
	}

	var results []*store.Record

	for _, k := range keys {
		r, err := m.get(fd, k)
		if err != nil {
			return results, err
		}
		results = append(results, r)
	}

	return results, nil
}

func (m *fileStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var writeOpts store.WriteOptions
	for _, o := range opts {
		o(&writeOpts)
	}

	fd, err := m.getDB(writeOpts.Database, writeOpts.Table)
	if err != nil {
		return err
	}

	if len(opts) > 0 {
		// Copy the record before applying options, or the incoming record will be mutated
		newRecord := store.Record{}
		newRecord.Key = r.Key
		newRecord.Value = r.Value
		newRecord.Metadata = make(map[string]interface{})
		newRecord.Expiry = r.Expiry

		if !writeOpts.Expiry.IsZero() {
			newRecord.Expiry = time.Until(writeOpts.Expiry)
		}
		if writeOpts.TTL != 0 {
			newRecord.Expiry = writeOpts.TTL
		}

		for k, v := range r.Metadata {
			newRecord.Metadata[k] = v
		}

		return m.set(fd, &newRecord)
	}

	return m.set(fd, r)
}

func (m *fileStore) Options() store.Options {
	return m.options
}

func (m *fileStore) List(opts ...store.ListOption) ([]string, error) {
	var listOptions store.ListOptions

	for _, o := range opts {
		o(&listOptions)
	}

	fd, err := m.getDB(listOptions.Database, listOptions.Table)
	if err != nil {
		return nil, err
	}

	// TODO apply prefix/suffix in range query
	allKeys := m.list(fd, listOptions.Limit, listOptions.Offset)

	if len(listOptions.Prefix) > 0 {
		var prefixKeys []string
		for _, k := range allKeys {
			if strings.HasPrefix(k, listOptions.Prefix) {
				prefixKeys = append(prefixKeys, k)
			}
		}
		allKeys = prefixKeys
	}

	if len(listOptions.Suffix) > 0 {
		var suffixKeys []string
		for _, k := range allKeys {
			if strings.HasSuffix(k, listOptions.Suffix) {
				suffixKeys = append(suffixKeys, k)
			}
		}
		allKeys = suffixKeys
	}

	return allKeys, nil
}

func (m *fileStore) String() string {
	return "file"
}
<commit_msg>Fix race when opening DB for first time (#1691)<commit_after>// Package local is a file system backed store
package file

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-micro/v2/store"
	bolt "go.etcd.io/bbolt"
)

var (
	// DefaultDatabase is the namespace that the bbolt store
	// will use if no namespace is provided.
	DefaultDatabase = "micro"
	// DefaultTable when none is specified
	DefaultTable = "micro"
	// DefaultDir is the default directory for bbolt files
	DefaultDir = filepath.Join(os.TempDir(), "micro", "store")

	// bucket used for data storage
	dataBucket = "data"
)

// NewStore returns a memory store
func NewStore(opts ...store.Option) store.Store {
	s := &fileStore{
		handles: make(map[string]*fileHandle),
	}
	s.init(opts...)
	return s
}

type fileStore struct {
	options store.Options
	dir     string

	// the database handle
	sync.RWMutex
	handles map[string]*fileHandle
}

type fileHandle struct {
	key string
	db  *bolt.DB
}

// record stored by us
type record struct {
	Key       string
	Value     []byte
	Metadata  map[string]interface{}
	ExpiresAt time.Time
}

func key(database, table string) string {
	return database + ":" + table
}

func (m *fileStore) delete(fd *fileHandle, key string) error {
	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}
		return b.Delete([]byte(key))
	})
}

func (m *fileStore) init(opts ...store.Option) error {
	for _, o := range opts {
		o(&m.options)
	}

	if m.options.Database == "" {
		m.options.Database = DefaultDatabase
	}

	if m.options.Table == "" {
		// bbolt requires bucketname to not be empty
		m.options.Table = DefaultTable
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, m.options.Database)
	// Ignoring this as the folder might exist.
	// Reads/Writes updates will return with sensible error messages
	// about the dir not existing in case this cannot create the path anyway
	os.MkdirAll(dir, 0700)

	return nil
}

func (f *fileStore) getDB(database, table string) (*fileHandle, error) {
	if len(database) == 0 {
		database = f.options.Database
	}
	if len(table) == 0 {
		table = f.options.Table
	}

	k := key(database, table)
	f.RLock()
	fd, ok := f.handles[k]
	f.RUnlock()

	// return the file handle
	if ok {
		return fd, nil
	}

	// double check locking
	f.Lock()
	defer f.Unlock()
	if fd, ok := f.handles[k]; ok {
		return fd, nil
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, database)
	// create the database handle
	fname := table + ".db"
	// make the dir
	os.MkdirAll(dir, 0700)
	// database path
	dbPath := filepath.Join(dir, fname)

	// create new db handle
	// Bolt DB only allows one process to open the file R/W so make sure we're doing this under a lock
	db, err := bolt.Open(dbPath, 0700, &bolt.Options{Timeout: 5 * time.Second})
	if err != nil {
		return nil, err
	}
	fd = &fileHandle{
		key: k,
		db:  db,
	}
	f.handles[k] = fd

	return fd, nil
}

func (m *fileStore) list(fd *fileHandle, limit, offset uint) []string {
	var allItems []string

	fd.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		// nothing to read
		if b == nil {
			return nil
		}

		// @todo very inefficient
		if err := b.ForEach(func(k, v []byte) error {
			storedRecord := &record{}

			if err := json.Unmarshal(v, storedRecord); err != nil {
				return err
			}

			if !storedRecord.ExpiresAt.IsZero() {
				if storedRecord.ExpiresAt.Before(time.Now()) {
					return nil
				}
			}

			allItems = append(allItems, string(k))

			return nil
		}); err != nil {
			return err
		}

		return nil
	})

	allKeys := make([]string, len(allItems))

	for i, k := range allItems {
		allKeys[i] = k
	}

	if limit != 0 || offset != 0 {
		sort.Slice(allKeys, func(i, j int) bool { return allKeys[i] < allKeys[j] })
		min := func(i, j uint) uint {
			if i < j {
				return i
			}
			return j
		}
		return allKeys[offset:min(limit, uint(len(allKeys)))]
	}

	return allKeys
}

func (m *fileStore) get(fd *fileHandle, k string) (*store.Record, error) {
	var value []byte

	fd.db.View(func(tx *bolt.Tx) error {
		// @todo this is still very experimental...
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}

		value = b.Get([]byte(k))
		return nil
	})

	if value == nil {
		return nil, store.ErrNotFound
	}

	storedRecord := &record{}

	if err := json.Unmarshal(value, storedRecord); err != nil {
		return nil, err
	}

	newRecord := &store.Record{}
	newRecord.Key = storedRecord.Key
	newRecord.Value = storedRecord.Value
	newRecord.Metadata = make(map[string]interface{})

	for k, v := range storedRecord.Metadata {
		newRecord.Metadata[k] = v
	}

	if !storedRecord.ExpiresAt.IsZero() {
		if storedRecord.ExpiresAt.Before(time.Now()) {
			return nil, store.ErrNotFound
		}
		newRecord.Expiry = time.Until(storedRecord.ExpiresAt)
	}

	return newRecord, nil
}

func (m *fileStore) set(fd *fileHandle, r *store.Record) error {
	// copy the incoming record and then
	// convert the expiry in to a hard timestamp
	item := &record{}
	item.Key = r.Key
	item.Value = r.Value
	item.Metadata = make(map[string]interface{})

	if r.Expiry != 0 {
		item.ExpiresAt = time.Now().Add(r.Expiry)
	}

	for k, v := range r.Metadata {
		item.Metadata[k] = v
	}

	// marshal the data
	data, _ := json.Marshal(item)

	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			var err error
			b, err = tx.CreateBucketIfNotExists([]byte(dataBucket))
			if err != nil {
				return err
			}
		}
		return b.Put([]byte(r.Key), data)
	})
}

func (f *fileStore) Close() error {
	f.Lock()
	defer f.Unlock()
	for k, v := range f.handles {
		v.db.Close()
		delete(f.handles, k)
	}
	return nil
}

func (f *fileStore) Init(opts ...store.Option) error {
	return f.init(opts...)
}

func (m *fileStore) Delete(key string, opts ...store.DeleteOption) error {
	var deleteOptions store.DeleteOptions
	for _, o := range opts {
		o(&deleteOptions)
	}

	fd, err := m.getDB(deleteOptions.Database, deleteOptions.Table)
	if err != nil {
		return err
	}

	return m.delete(fd, key)
}

func (m *fileStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var readOpts store.ReadOptions
	for _, o := range opts {
		o(&readOpts)
	}

	fd, err := m.getDB(readOpts.Database, readOpts.Table)
	if err != nil {
		return nil, err
	}

	var keys []string

	// Handle Prefix / suffix
	// TODO: do range scan here rather than listing all keys
	if readOpts.Prefix || readOpts.Suffix {
		// list the keys
		k := m.list(fd, readOpts.Limit, readOpts.Offset)

		// check for prefix and suffix
		for _, v := range k {
			if readOpts.Prefix && !strings.HasPrefix(v, key) {
				continue
			}
			if readOpts.Suffix && !strings.HasSuffix(v, key) {
				continue
			}
			keys = append(keys, v)
		}
	} else {
		keys = []string{key}
	}

	var results []*store.Record

	for _, k := range keys {
		r, err := m.get(fd, k)
		if err != nil {
			return results, err
		}
		results = append(results, r)
	}

	return results, nil
}

func (m *fileStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var writeOpts store.WriteOptions
	for _, o := range opts {
		o(&writeOpts)
	}

	fd, err := m.getDB(writeOpts.Database, writeOpts.Table)
	if err != nil {
		return err
	}

	if len(opts) > 0 {
		// Copy the record before applying options, or the incoming record will be mutated
		newRecord := store.Record{}
		newRecord.Key = r.Key
		newRecord.Value = r.Value
		newRecord.Metadata = make(map[string]interface{})
		newRecord.Expiry = r.Expiry

		if !writeOpts.Expiry.IsZero() {
			newRecord.Expiry = time.Until(writeOpts.Expiry)
		}
		if writeOpts.TTL != 0 {
			newRecord.Expiry = writeOpts.TTL
		}

		for k, v := range r.Metadata {
			newRecord.Metadata[k] = v
		}

		return m.set(fd, &newRecord)
	}

	return m.set(fd, r)
}

func (m *fileStore) Options() store.Options {
	return m.options
}

func (m *fileStore) List(opts ...store.ListOption) ([]string, error) {
	var listOptions store.ListOptions

	for _, o := range opts {
		o(&listOptions)
	}

	fd, err := m.getDB(listOptions.Database, listOptions.Table)
	if err != nil {
		return nil, err
	}

	// TODO apply prefix/suffix in range query
	allKeys := m.list(fd, listOptions.Limit, listOptions.Offset)

	if len(listOptions.Prefix) > 0 {
		var prefixKeys []string
		for _, k := range allKeys {
			if strings.HasPrefix(k, listOptions.Prefix) {
				prefixKeys = append(prefixKeys, k)
			}
		}
		allKeys = prefixKeys
	}

	if len(listOptions.Suffix) > 0 {
		var suffixKeys []string
		for _, k := range allKeys {
			if strings.HasSuffix(k, listOptions.Suffix) {
				suffixKeys = append(suffixKeys, k)
			}
		}
		allKeys = suffixKeys
	}

	return allKeys, nil
}

func (m *fileStore) String() string {
	return "file"
}
<|endoftext|>


<commit_before>package stboot

import (
	"archive/zip"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/u-root/u-root/pkg/bootconfig"
)

const (
	signaturesDirName string = "signatures"
	rootCertName      string = "root.cert"
	bootFilesDirName  string = "bootconfig"
)

type BootBall struct {
	Archive        string
	dir            string
	config         *Stconfig
	numBootConfigs int
	bootFiles      [][]string
	rootCert       *x509.CertPool
	signatures     [][]signature
	hashes         [][]byte
	signer         Signer
}

func BootBallFromArchie(archive string) (*BootBall, error) {
	var ball = new(BootBall)

	dir, err := ioutil.TempDir(os.TempDir(), "bootball")
	if err != nil {
		return ball, fmt.Errorf("BootBall: cannot create tmp dir: %v", err)
	}

	err = fromZip(archive, dir)
	if err != nil {
		return ball, fmt.Errorf("BootBall: cannot unzip %s: %v", archive, err)
	}

	cfg, err := getConfig(filepath.Join(dir, ConfigName))
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	cert, err := getRootCert(filepath.Join(dir, signaturesDirName, rootCertName))
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	bootFiles, err := getBootFiles(cfg, dir)
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting boot files faild: %v", err)
	}

	ball.Archive = archive
	ball.dir = dir
	ball.config = cfg
	ball.rootCert = cert
	ball.numBootConfigs = len(ball.config.BootConfigs)
	ball.bootFiles = bootFiles
	ball.signer = sha512PssSigner{}

	return ball, nil
}

func BootBallFromConfig(configFile string) (*BootBall, error) {
	var ball = new(BootBall)

	archive := filepath.Join(filepath.Dir(configFile), BallName)

	cfg, err := getConfig(configFile)
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	dir, err := makeConfigDir(cfg, filepath.Dir(configFile))
	if err != nil {
		return ball, fmt.Errorf("BootBall: creating standard configuration directory faild: %v", err)
	}

	cert, err := getRootCert(filepath.Join(dir, signaturesDirName, rootCertName))
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	bootFiles, err := getBootFiles(cfg, dir)
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting boot files faild: %v", err)
	}

	ball.Archive = archive
	ball.dir = dir
	ball.config = cfg
	ball.rootCert = cert
	ball.numBootConfigs = len(ball.config.BootConfigs)
	ball.bootFiles = bootFiles
	ball.signer = sha512PssSigner{}

	err = ball.getSignatures()
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting signatures: %v", err)
	}

	return ball, nil
}

func (ball *BootBall) Clean() (err error) {
	err = os.RemoveAll(ball.dir)
	if err != nil {
		return
	}
	ball.dir = ""
	return
}

// Pack writes the contents BootBall.Archive
func (ball *BootBall) Pack() (err error) {
	if ball.Archive == "" || ball.dir == "" {
		return errors.New("BootBall.Pack: booball.archive and bootball.dir must be set")
	}
	return toZip(ball.dir, ball.Archive)
}

// Dir returns the directory associated with BootBall
func (ball *BootBall) Dir() (dir string) {
	return ball.dir
}

func (ball *BootBall) GetBootConfigByIndex(index int) (bc *bootconfig.BootConfig, err error) {
	bc, err = ball.config.getBootConfig(index)
	if err != nil {
		return
	}
	bc.SetFilePathsPrefix(ball.dir)
	return
}

func (ball *BootBall) Hash() (err error) {
	ball.hashes = make([][]byte, len(ball.config.BootConfigs))
	for i, files := range ball.bootFiles {
		hash, herr := ball.signer.hash(files...)
		if herr != nil {
			return herr
		}
		ball.hashes[i] = hash
	}
	return
}

func (ball *BootBall) Sign(privKeyFile, certFile string) (err error) {
	err = validateFiles("", privKeyFile, certFile)
	if err != nil {
		return
	}

	buf, err := ioutil.ReadFile(certFile)
	if err != nil {
		return
	}

	cert, err := parseCertificate(buf)
	err = validateCertificate(cert, ball.rootCert)
	if err != nil {
		return
	}

	log.Printf("Signing with: %s", privKeyFile)

	if ball.hashes == nil {
		err = ball.Hash()
		if err != nil {
			return
		}
	}

	sigs := make([]signature, len(ball.config.BootConfigs))
	for i, hash := range ball.hashes {
		s, err := ball.signer.sign(privKeyFile, hash)
		if err != nil {
			return err
		}
		sigs[i].Bytes = s
		sigs[i].Cert = cert
	}

	return writeSignatures(sigs, certFile, ball.dir)
}

func (ball *BootBall) VerifyBootconfigs() (verified []int, err error) {
	verified = make([]int, len(ball.signatures))
	for i := range ball.signatures {
		n, err := ball.VerifyBootconfigByIndex(i)
		if err != nil {
			return verified, err
		}
		verified[i] = n
	}
	return
}

func (ball *BootBall) VerifyBootconfigByIndex(index int) (verified int, err error) {
	if ball.hashes == nil {
		err = ball.Hash()
		if err != nil {
			return
		}
	}

	sigs := ball.signatures[index]
	verified = 0
	for _, sig := range sigs {
		err = validateCertificate(sig.Cert, ball.rootCert)
		if err != nil {
			return
		}
		err = ball.signer.verify(sig, ball.hashes[index])
		if err != nil {
			return
		}
		verified++
	}
	return
}

func getConfig(dest string) (cfg *Stconfig, err error) {
	cfgBytes, err := ioutil.ReadFile(dest)
	if err != nil {
		return
	}
	cfg, err = stconfigFromBytes(cfgBytes)
	if err != nil {
		return
	}
	if !(cfg.IsValid()) {
		return cfg, errors.New("invalid configuration")
	}
	return
}

func getRootCert(dest string) (cert *x509.CertPool, err error) {
	certBytes, err := ioutil.ReadFile(dest)
	if err != nil {
		return
	}
	cert, err = certPool(certBytes)
	if err != nil {
		return
	}
	return
}

func getBootFiles(cfg *Stconfig, prefix string) (bootFiles [][]string, err error) {
	bootFiles = make([][]string, 0)
	for _, bc := range cfg.BootConfigs {
		files := make([]string, 0)
		for _, file := range bc.FileNames() {
			file = filepath.Join(prefix, file)
			files = append(files, file)
			if err = validateFiles("", files...); err != nil {
				return
			}
		}
		bootFiles = append(bootFiles, files)
	}
	return
}

func (ball *BootBall) getSignatures() (err error) {
	ball.signatures = make([][]signature, len(ball.config.BootConfigs))
	path := filepath.Join(ball.dir, signaturesDirName)

	sigPool := make([]signature, 0)
	err = filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		ext := filepath.Ext(info.Name())

		if !info.IsDir() && (ext == ".signature") {
			sigBytes, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}

			dir := filepath.Dir(path)
			index, err := strconv.Atoi(dir[len(dir)-1:])
			if err != nil {
				return err
			}

			certFile := strings.TrimSuffix(path, filepath.Ext(path)) + ".cert"
			certBytes, err := ioutil.ReadFile(certFile)
			if err != nil {
				return err
			}

			cert, err := parseCertificate(certBytes)
			if err != nil {
				return err
			}

			sig := signature{
				Bytes: sigBytes,
				Cert:  cert,
			}
			sigPool = append(sigPool, sig)
			ball.signatures[index] = sigPool
		}
		return nil
	})
	if err != nil {
		return
	}
	return
}

func validateFiles(prefix string, files ...string) (err error) {
	for _, file := range files {
		_, err = os.Stat(filepath.Join(prefix, file))
		if err != nil {
			return
		}
	}
	return
}

func writeSignatures(sigs []signature, certFile, dir string) (err error) {
	for i, sig := range sigs {
		d := fmt.Sprintf("%s%d", bootFilesDirName, i)
		path := filepath.Join(dir, signaturesDirName, d)
		os.Mkdir(path, os.ModePerm)

		id := fmt.Sprintf("%x", sig.Cert.PublicKey)[2:18]
		sigName := fmt.Sprintf("%s.signature", id)
		sigPath := filepath.Join(path, sigName)
		werr := ioutil.WriteFile(sigPath, sig.Bytes, 0644)
		if werr != nil {
			return werr
		}

		certName := fmt.Sprintf("%s.cert", id)
		certPath := filepath.Join(path, certName)
		cerr := copyFile(certFile, certPath)
		if cerr != nil {
			return cerr
		}
	}
	return
}

func makeConfigDir(cfg *Stconfig, origDir string) (dir string, err error) {
	if err = validateFiles(cfg.RootCertPath); err != nil {
		return
	}

	for _, bc := range cfg.BootConfigs {
		if err = validateFiles(origDir, bc.FileNames()...); err != nil {
			return
		}
	}

	dir, err = ioutil.TempDir(os.TempDir(), "bootball")
	if err != nil {
		return
	}

	dstPath := filepath.Join(dir, signaturesDirName, rootCertName)
	srcPath := filepath.Join(origDir, cfg.RootCertPath)
	copyFile(srcPath, dstPath)

	for i, bc := range cfg.BootConfigs {
		dirName := fmt.Sprintf("%s%d", bootFilesDirName, i)
		for _, file := range bc.FileNames() {
			fileName := filepath.Base(file)
			dstPath := filepath.Join(dir, dirName, fileName)
			srcPath := filepath.Join(origDir, file)
			copyFile(srcPath, dstPath)
		}

		bc.ChangeFilePaths(dirName)
		cfg.BootConfigs[i] = bc
	}

	dstPath = filepath.Join(dir, ConfigName)
	bytes, err := cfg.bytes()
	if err != nil {
		return
	}
	ioutil.WriteFile(dstPath, bytes, os.ModePerm)

	return
}

func copyFile(src, dst string) (err error) {
	in, err := os.Open(src)
	if err != nil {
		return
	}
	defer in.Close()

	if err = os.MkdirAll(filepath.Dir(dst), os.ModePerm); err != nil {
		return
	}

	out, err := os.Create(dst)
	if err != nil {
		return
	}
	defer out.Close()

	if _, err = io.Copy(out, in); err != nil {
		return
	}
	err = out.Sync()

	return
}

func toZip(srcDir, dest string) (err error) {
	info, err := os.Stat(srcDir)
	if err != nil {
		return
	}
	if !(info.IsDir()) {
		return fmt.Errorf("%s is not a directory", srcDir)
	}
	archive, err := os.Create(dest)
	if err != nil {
		return
	}
	defer archive.Close()

	z := zip.NewWriter(archive)
	defer z.Close()

	filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// do not include srcDir into archive
		if strings.Compare(info.Name(), filepath.Base(srcDir)) == 0 {
			return nil
		}

		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		// adjust header.Name to preserve folder strulture
		header.Name = strings.TrimPrefix(path, srcDir)

		if info.IsDir() {
			header.Name += "/"
		} else {
			header.Method = zip.Deflate
		}

		writer, err := z.CreateHeader(header)
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = io.Copy(writer, file)
		return err
	})

	return
}

func fromZip(src, destDir string) (err error) {
	z, err := zip.OpenReader(src)
	if err != nil {
		return
	}

	if err = os.MkdirAll(destDir, 0755); err != nil {
		return
	}

	for _, file := range z.File {
		path := filepath.Join(destDir, file.Name)
		if file.FileInfo().IsDir() {
			os.MkdirAll(path, file.Mode())
			continue
		}

		fileReader, err := file.Open()
		if err != nil {
			return err
		}

		targetFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, file.Mode())
		if err != nil {
			return err
		}

		if _, err := io.Copy(targetFile, fileReader); err != nil {
			return err
		}

		fileReader.Close()
		targetFile.Close()
	}

	return
}
<commit_msg>pkg/boot/stboot/bootball.go: add init()<commit_after>package stboot

import (
	"archive/zip"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/u-root/u-root/pkg/bootconfig"
)

const (
	signaturesDirName string = "signatures"
	rootCertName      string = "root.cert"
	bootFilesDirName  string = "bootconfig"
)

type BootBall struct {
	Archive        string
	dir            string
	config         *Stconfig
	numBootConfigs int
	bootFiles      [][]string
	rootCert       *x509.CertPool
	signatures     [][]signature
	hashes         [][]byte
	signer         Signer
}

func BootBallFromArchie(archive string) (*BootBall, error) {
	var ball = new(BootBall)

	dir, err := ioutil.TempDir(os.TempDir(), "bootball")
	if err != nil {
		return ball, fmt.Errorf("BootBall: cannot create tmp dir: %v", err)
	}

	err = fromZip(archive, dir)
	if err != nil {
		return ball, fmt.Errorf("BootBall: cannot unzip %s: %v", archive, err)
	}

	cfg, err := getConfig(filepath.Join(dir, ConfigName))
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	ball.Archive = archive
	ball.dir = dir
	ball.config = cfg

	err = ball.init()
	if err != nil {
		return ball, err
	}

	return ball, nil
}

func BootBallFromConfig(configFile string) (*BootBall, error) {
	var ball = new(BootBall)

	archive := filepath.Join(filepath.Dir(configFile), BallName)

	cfg, err := getConfig(configFile)
	if err != nil {
		return ball, fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	dir, err := makeConfigDir(cfg, filepath.Dir(configFile))
	if err != nil {
		return ball, fmt.Errorf("BootBall: creating standard configuration directory faild: %v", err)
	}

	ball.Archive = archive
	ball.dir = dir
	ball.config = cfg

	err = ball.init()
	if err != nil {
		return ball, err
	}

	return ball, nil
}

func (ball *BootBall) init() (err error) {
	cert, err := getRootCert(filepath.Join(ball.dir, signaturesDirName, rootCertName))
	if err != nil {
		return fmt.Errorf("BootBall: getting configuration faild: %v", err)
	}

	bootFiles, err := getBootFiles(ball.config, ball.dir)
	if err != nil {
		return fmt.Errorf("BootBall: getting boot files faild: %v", err)
	}

	ball.rootCert = cert
	ball.numBootConfigs = len(ball.config.BootConfigs)
	ball.bootFiles = bootFiles
	ball.signer = sha512PssSigner{}

	err = ball.getSignatures()
	if err != nil {
		return fmt.Errorf("BootBall: getting signatures: %v", err)
	}
	return
}

func (ball *BootBall) Clean() (err error) {
	err = os.RemoveAll(ball.dir)
	if err != nil {
		return
	}
	ball.dir = ""
	return
}

// Pack writes the contents BootBall.Archive
func (ball *BootBall) Pack() (err error) {
	if ball.Archive == "" || ball.dir == "" {
		return errors.New("BootBall.Pack: booball.archive and bootball.dir must be set")
	}
	return toZip(ball.dir, ball.Archive)
}

// Dir returns the directory associated with BootBall
func (ball *BootBall) Dir() (dir string) {
	return ball.dir
}

func (ball *BootBall) GetBootConfigByIndex(index int) (bc *bootconfig.BootConfig, err error) {
	bc, err = ball.config.getBootConfig(index)
	if err != nil {
		return
	}
	bc.SetFilePathsPrefix(ball.dir)
	return
}

func (ball *BootBall) Hash() (err error) {
	ball.hashes = make([][]byte, len(ball.config.BootConfigs))
	for i, files := range ball.bootFiles {
		hash, herr := ball.signer.hash(files...)
		if herr != nil {
			return herr
		}
		ball.hashes[i] = hash
	}
	return
}

func (ball *BootBall) Sign(privKeyFile, certFile string) (err error) {
	err = validateFiles("", privKeyFile, certFile)
	if err != nil {
		return
	}

	buf, err := ioutil.ReadFile(certFile)
	if err != nil {
		return
	}

	cert, err := parseCertificate(buf)
	err = validateCertificate(cert, ball.rootCert)
	if err != nil {
		return
	}

	log.Printf("Signing with: %s", privKeyFile)

	if ball.hashes == nil {
		err = ball.Hash()
		if err != nil {
			return
		}
	}

	sigs := make([]signature, len(ball.config.BootConfigs))
	for i, hash := range ball.hashes {
		s, err := ball.signer.sign(privKeyFile, hash)
		if err != nil {
			return err
		}
		sigs[i].Bytes = s
		sigs[i].Cert = cert
	}

	return writeSignatures(sigs, certFile, ball.dir)
}

func (ball *BootBall) VerifyBootconfigs() (verified []int, err error) {
	verified = make([]int, len(ball.signatures))
	for i := range ball.signatures {
		n, err := ball.VerifyBootconfigByIndex(i)
		if err != nil {
			return verified, err
		}
		verified[i] = n
	}
	return
}

func (ball *BootBall) VerifyBootconfigByIndex(index int) (verified int, err error) {
	if ball.hashes == nil {
		err = ball.Hash()
		if err != nil {
			return
		}
	}

	sigs := ball.signatures[index]
	verified = 0
	for _, sig := range sigs {
		err = validateCertificate(sig.Cert, ball.rootCert)
		if err != nil {
			return
		}
		err = ball.signer.verify(sig, ball.hashes[index])
		if err != nil {
			return
		}
		verified++
	}
	return
}

func getConfig(dest string) (cfg *Stconfig, err error) {
	cfgBytes, err := ioutil.ReadFile(dest)
	if err != nil {
		return
	}
	cfg, err = stconfigFromBytes(cfgBytes)
	if err != nil {
		return
	}
	if !(cfg.IsValid()) {
		return cfg, errors.New("invalid configuration")
	}
	return
}

func getRootCert(dest string) (cert *x509.CertPool, err error) {
	certBytes, err := ioutil.ReadFile(dest)
	if err != nil {
		return
	}
	cert, err = certPool(certBytes)
	if err != nil {
		return
	}
	return
}

func getBootFiles(cfg *Stconfig, prefix string) (bootFiles [][]string, err error) {
	bootFiles = make([][]string, 0)
	for _, bc := range cfg.BootConfigs {
		files := make([]string, 0)
		for _, file := range bc.FileNames() {
			file = filepath.Join(prefix, file)
			files = append(files, file)
			if err = validateFiles("", files...); err != nil {
				return
			}
		}
		bootFiles = append(bootFiles, files)
	}
	return
}

func (ball *BootBall) getSignatures() (err error) {
	ball.signatures = make([][]signature, len(ball.config.BootConfigs))
	path := filepath.Join(ball.dir, signaturesDirName)

	sigPool := make([]signature, 0)
	err = filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		ext := filepath.Ext(info.Name())

		if !info.IsDir() && (ext == ".signature") {
			sigBytes, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}

			dir := filepath.Dir(path)
			index, err := strconv.Atoi(dir[len(dir)-1:])
			if err != nil {
				return err
			}

			certFile := strings.TrimSuffix(path, filepath.Ext(path)) + ".cert"
			certBytes, err := ioutil.ReadFile(certFile)
			if err != nil {
				return err
			}

			cert, err := parseCertificate(certBytes)
			if err != nil {
				return err
			}

			sig := signature{
				Bytes: sigBytes,
				Cert:  cert,
			}
			sigPool = append(sigPool, sig)
			ball.signatures[index] = sigPool
		}
		return nil
	})
	if err != nil {
		return
	}
	return
}

func validateFiles(prefix string, files ...string) (err error) {
	for _, file := range files {
		_, err = os.Stat(filepath.Join(prefix, file))
		if err != nil {
			return
		}
	}
	return
}

func writeSignatures(sigs []signature, certFile, dir string) (err error) {
	for i, sig := range sigs {
		d := fmt.Sprintf("%s%d", bootFilesDirName, i)
		path := filepath.Join(dir, signaturesDirName, d)
		os.Mkdir(path, os.ModePerm)

		id := fmt.Sprintf("%x", sig.Cert.PublicKey)[2:18]
		sigName := fmt.Sprintf("%s.signature", id)
		sigPath := filepath.Join(path, sigName)
		werr := ioutil.WriteFile(sigPath, sig.Bytes, 0644)
		if werr != nil {
			return werr
		}

		certName := fmt.Sprintf("%s.cert", id)
		certPath := filepath.Join(path, certName)
		cerr := copyFile(certFile, certPath)
		if cerr != nil {
			return cerr
		}
	}
	return
}

func makeConfigDir(cfg *Stconfig, origDir string) (dir string, err error) {
	if err = validateFiles(cfg.RootCertPath); err != nil {
		return
	}

	for _, bc := range cfg.BootConfigs {
		if err = validateFiles(origDir, bc.FileNames()...); err != nil {
			return
		}
	}

	dir, err = ioutil.TempDir(os.TempDir(), "bootball")
	if err != nil {
		return
	}

	dstPath := filepath.Join(dir, signaturesDirName, rootCertName)
	srcPath := filepath.Join(origDir, cfg.RootCertPath)
	copyFile(srcPath, dstPath)

	for i, bc := range cfg.BootConfigs {
		dirName := fmt.Sprintf("%s%d", bootFilesDirName, i)
		for _, file := range bc.FileNames() {
			fileName := filepath.Base(file)
			dstPath := filepath.Join(dir, dirName, fileName)
			srcPath := filepath.Join(origDir, file)
			copyFile(srcPath, dstPath)
		}

		bc.ChangeFilePaths(dirName)
		cfg.BootConfigs[i] = bc
	}

	dstPath = filepath.Join(dir, ConfigName)
	bytes, err := cfg.bytes()
	if err != nil {
		return
	}
	ioutil.WriteFile(dstPath, bytes, os.ModePerm)

	return
}

func copyFile(src, dst string) (err error) {
	in, err := os.Open(src)
	if err != nil {
		return
	}
	defer in.Close()

	if err = os.MkdirAll(filepath.Dir(dst), os.ModePerm); err != nil {
		return
	}

	out, err := os.Create(dst)
	if err != nil {
		return
	}
	defer out.Close()

	if _, err = io.Copy(out, in); err != nil {
		return
	}
	err = out.Sync()

	return
}

func toZip(srcDir, dest string) (err error) {
	info, err := os.Stat(srcDir)
	if err != nil {
		return
	}
	if !(info.IsDir()) {
		return fmt.Errorf("%s is not a directory", srcDir)
	}
	archive, err := os.Create(dest)
	if err != nil {
		return
	}
	defer archive.Close()

	z := zip.NewWriter(archive)
	defer z.Close()

	filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// do not include srcDir into archive
		if strings.Compare(info.Name(), filepath.Base(srcDir)) == 0 {
			return nil
		}

		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		// adjust header.Name to preserve folder strulture
		header.Name = strings.TrimPrefix(path, srcDir)

		if info.IsDir() {
			header.Name += "/"
		} else {
			header.Method = zip.Deflate
		}

		writer, err := z.CreateHeader(header)
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = io.Copy(writer, file)
		return err
	})

	return
}

func fromZip(src, destDir string) (err error) {
	z, err := zip.OpenReader(src)
	if err != nil {
		return
	}

	if err = os.MkdirAll(destDir, 0755); err != nil {
		return
	}

	for _, file := range z.File {
		path := filepath.Join(destDir, file.Name)
		if file.FileInfo().IsDir() {
			os.MkdirAll(path, file.Mode())
			continue
		}

		fileReader, err := file.Open()
		if err != nil {
			return err
		}

		targetFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, file.Mode())
		if err != nil {
			return err
		}

		if _, err := io.Copy(targetFile, fileReader); err != nil {
			return err
		}

		fileReader.Close()
		targetFile.Close()
	}

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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 handler

import (
	"bufio"
	"context"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"knative.dev/pkg/websocket"
)

type timeToFirstByteTimeoutHandler struct {
	handler http.Handler
	timeout time.Duration
	body    string
}

// NewTimeToFirstByteTimeoutHandler returns a Handler that runs `h` with the
// given timeout in which the first byte of the response must be written.
//
// The new Handler calls h.ServeHTTP to handle each request, but if a
// call runs for longer than its time limit, the handler responds with
// a 504 Gateway Timeout error and the given message in its body.
// (If msg is empty, a suitable default message will be sent.)
// After such a timeout, writes by h to its ResponseWriter will return
// ErrHandlerTimeout.
//
// A panic from the underlying handler is propagated as-is to be able to
// make use of custom panic behavior by HTTP handlers. See
// https://golang.org/pkg/net/http/#Handler.
//
// The implementation is largely inspired by http.TimeoutHandler.
func NewTimeToFirstByteTimeoutHandler(h http.Handler, msg string, timeout time.Duration) http.Handler {
	return &timeToFirstByteTimeoutHandler{
		handler: h,
		body:    msg,
		timeout: timeout,
	}
}

func (h *timeToFirstByteTimeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	// done is closed when h.handler.ServeHTTP completes and contains
	// the panic from h.handler.ServeHTTP if h.handler.ServeHTTP panics.
	done := make(chan interface{})
	tw := &timeoutWriter{w: w}
	go func() {
		defer func() {
			defer close(done)
			if p := recover(); p != nil {
				done <- p
			}
		}()
		h.handler.ServeHTTP(tw, r.WithContext(ctx))
	}()

	timeout := time.NewTimer(h.timeout)
	defer timeout.Stop()
	for {
		select {
		case p, ok := <-done:
			if ok {
				panic(p)
			}
			return
		case <-timeout.C:
			if tw.timeoutAndWriteError(h.body) {
				return
			}
		}
	}
}

// timeoutWriter is a wrapper around an http.ResponseWriter. It guards
// writing an error response to whether or not the underlying writer has
// already been written to.
//
// If the underlying writer has not been written to, an error response is
// returned. If it has already been written to, the error is ignored and
// the response is allowed to continue.
type timeoutWriter struct {
	w http.ResponseWriter

	mu        sync.Mutex
	timedOut  bool
	wroteOnce bool
}

var _ http.Flusher = (*timeoutWriter)(nil)
var _ http.ResponseWriter = (*timeoutWriter)(nil)

func (tw *timeoutWriter) Flush() {
	// The inner handler of timeoutHandler can call Flush at any time including after
	// timeoutHandler.ServeHTTP has returned. Forwarding this call to the inner
	// http.ResponseWriter would lead to a panic in HTTP2. See http2/server.go line 2556.
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return
	}

	tw.w.(http.Flusher).Flush()
}

// Hijack calls Hijack() on the wrapped http.ResponseWriter if it implements
// http.Hijacker interface, which is required for net/http/httputil/reverseproxy
// to handle connection upgrade/switching protocol.  Otherwise returns an error.
func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return websocket.HijackIfPossible(tw.w)
}

func (tw *timeoutWriter) Header() http.Header { return tw.w.Header() }

func (tw *timeoutWriter) Write(p []byte) (int, error) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return 0, http.ErrHandlerTimeout
	}

	tw.wroteOnce = true
	return tw.w.Write(p)
}

func (tw *timeoutWriter) WriteHeader(code int) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return
	}
	tw.wroteOnce = true
	tw.w.WriteHeader(code)
}

// timeoutAndWriteError writes an error to the responsewriter if
// nothing has been written to the writer before. Returns whether
// an error was written or not.
//
// If this writes an error, all subsequent calls to Write will
// result in http.ErrHandlerTimeout.
func (tw *timeoutWriter) timeoutAndWriteError(msg string) bool {
	tw.mu.Lock()
	defer tw.mu.Unlock()

	if !tw.wroteOnce {
		tw.w.WriteHeader(http.StatusGatewayTimeout)
		io.WriteString(tw.w, msg)

		tw.timedOut = true
		return true
	}

	return false
}
<commit_msg>Reduce allocations from the timeout handler by pooling the timers (#11771)<commit_after>/*
Copyright 2020 The Knative Authors

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 handler

import (
	"bufio"
	"context"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"knative.dev/pkg/websocket"
)

type timeToFirstByteTimeoutHandler struct {
	handler http.Handler
	timeout time.Duration
	body    string
}

// NewTimeToFirstByteTimeoutHandler returns a Handler that runs `h` with the
// given timeout in which the first byte of the response must be written.
//
// The new Handler calls h.ServeHTTP to handle each request, but if a
// call runs for longer than its time limit, the handler responds with
// a 504 Gateway Timeout error and the given message in its body.
// (If msg is empty, a suitable default message will be sent.)
// After such a timeout, writes by h to its ResponseWriter will return
// ErrHandlerTimeout.
//
// A panic from the underlying handler is propagated as-is to be able to
// make use of custom panic behavior by HTTP handlers. See
// https://golang.org/pkg/net/http/#Handler.
//
// The implementation is largely inspired by http.TimeoutHandler.
func NewTimeToFirstByteTimeoutHandler(h http.Handler, msg string, timeout time.Duration) http.Handler {
	return &timeToFirstByteTimeoutHandler{
		handler: h,
		body:    msg,
		timeout: timeout,
	}
}

func (h *timeToFirstByteTimeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	// done is closed when h.handler.ServeHTTP completes and contains
	// the panic from h.handler.ServeHTTP if h.handler.ServeHTTP panics.
	done := make(chan interface{})
	tw := &timeoutWriter{w: w}
	go func() {
		defer func() {
			defer close(done)
			if p := recover(); p != nil {
				done <- p
			}
		}()
		h.handler.ServeHTTP(tw, r.WithContext(ctx))
	}()

	timeout := getTimer(h.timeout)
	var timeoutDrained bool
	defer func() { putTimer(timeout, timeoutDrained) }()
	for {
		select {
		case p, ok := <-done:
			if ok {
				panic(p)
			}
			return
		case <-timeout.C:
			timeoutDrained = true
			if tw.timeoutAndWriteError(h.body) {
				return
			}
		}
	}
}

// timeoutWriter is a wrapper around an http.ResponseWriter. It guards
// writing an error response to whether or not the underlying writer has
// already been written to.
//
// If the underlying writer has not been written to, an error response is
// returned. If it has already been written to, the error is ignored and
// the response is allowed to continue.
type timeoutWriter struct {
	w http.ResponseWriter

	mu        sync.Mutex
	timedOut  bool
	wroteOnce bool
}

var _ http.Flusher = (*timeoutWriter)(nil)
var _ http.ResponseWriter = (*timeoutWriter)(nil)

func (tw *timeoutWriter) Flush() {
	// The inner handler of timeoutHandler can call Flush at any time including after
	// timeoutHandler.ServeHTTP has returned. Forwarding this call to the inner
	// http.ResponseWriter would lead to a panic in HTTP2. See http2/server.go line 2556.
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return
	}

	tw.w.(http.Flusher).Flush()
}

// Hijack calls Hijack() on the wrapped http.ResponseWriter if it implements
// http.Hijacker interface, which is required for net/http/httputil/reverseproxy
// to handle connection upgrade/switching protocol.  Otherwise returns an error.
func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return websocket.HijackIfPossible(tw.w)
}

func (tw *timeoutWriter) Header() http.Header { return tw.w.Header() }

func (tw *timeoutWriter) Write(p []byte) (int, error) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return 0, http.ErrHandlerTimeout
	}

	tw.wroteOnce = true
	return tw.w.Write(p)
}

func (tw *timeoutWriter) WriteHeader(code int) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return
	}
	tw.wroteOnce = true
	tw.w.WriteHeader(code)
}

// timeoutAndWriteError writes an error to the responsewriter if
// nothing has been written to the writer before. Returns whether
// an error was written or not.
//
// If this writes an error, all subsequent calls to Write will
// result in http.ErrHandlerTimeout.
func (tw *timeoutWriter) timeoutAndWriteError(msg string) bool {
	tw.mu.Lock()
	defer tw.mu.Unlock()

	if !tw.wroteOnce {
		tw.w.WriteHeader(http.StatusGatewayTimeout)
		io.WriteString(tw.w, msg)

		tw.timedOut = true
		return true
	}

	return false
}

var timerPool sync.Pool

func getTimer(timeout time.Duration) *time.Timer {
	if v := timerPool.Get(); v != nil {
		t := v.(*time.Timer)
		t.Reset(timeout)
		return t
	}
	return time.NewTimer(timeout)
}

func putTimer(t *time.Timer, alreadyDrained bool) {
	if !t.Stop() && !alreadyDrained {
		// Stop told us that we didn't *actually* stop the timer, so it expired. We've
		// also not drained the channel yet, so the expiration raced the inner handler
		// finishing, so we know we *have* to drain here.
		<-t.C
	}
	timerPool.Put(t)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vault

import (
	"crypto/x509"
	"errors"
	"fmt"
	"net/http"
	"path"
	"path/filepath"
	"strings"
	"time"

	vault "github.com/hashicorp/vault/api"
	"github.com/hashicorp/vault/helper/certutil"
	corelisters "k8s.io/client-go/listers/core/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

var _ Interface = &Vault{}

type VaultClientBuilder func(namespace string, secretsLister corelisters.SecretLister,
	issuer v1alpha2.GenericIssuer) (Interface, error)

type Interface interface {
	Sign(csrPEM []byte, duration time.Duration) (certPEM []byte, caPEM []byte, err error)
	Sys() *vault.Sys
}

type Client interface {
	NewRequest(method, requestPath string) *vault.Request
	RawRequest(r *vault.Request) (*vault.Response, error)
	SetToken(v string)
	Token() string
	Sys() *vault.Sys
}

type Vault struct {
	secretsLister corelisters.SecretLister
	issuer        v1alpha2.GenericIssuer
	namespace     string

	client Client
}

func New(namespace string, secretsLister corelisters.SecretLister,
	issuer v1alpha2.GenericIssuer) (Interface, error) {
	v := &Vault{
		secretsLister: secretsLister,
		namespace:     namespace,
		issuer:        issuer,
	}

	cfg, err := v.newConfig()
	if err != nil {
		return nil, err
	}

	client, err := vault.NewClient(cfg)
	if err != nil {
		return nil, fmt.Errorf("error initializing Vault client: %s", err.Error())
	}

	if err := v.setToken(client); err != nil {
		return nil, err
	}

	v.client = client

	return v, nil
}

func (v *Vault) Sign(csrPEM []byte, duration time.Duration) (cert []byte, ca []byte, err error) {
	csr, err := pki.DecodeX509CertificateRequestBytes(csrPEM)
	if err != nil {
		return nil, nil, fmt.Errorf("faild to decode CSR for signing: %s", err)
	}

	parameters := map[string]string{
		"common_name": csr.Subject.CommonName,
		"alt_names":   strings.Join(csr.DNSNames, ","),
		"ip_sans":     strings.Join(pki.IPAddressesToString(csr.IPAddresses), ","),
		"ttl":         duration.String(),
		"csr":         string(csrPEM),

		"exclude_cn_from_sans": "true",
	}

	url := path.Join("/v1", v.issuer.GetSpec().Vault.Path)

	request := v.client.NewRequest("POST", url)

	if err := request.SetJSONBody(parameters); err != nil {
		return nil, nil, fmt.Errorf("failed to build vault request: %s", err)
	}

	resp, err := v.client.RawRequest(request)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to sign certificate by vault: %s", err)
	}

	defer resp.Body.Close()

	vaultResult := certutil.Secret{}
	resp.DecodeJSON(&vaultResult)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to decode response returned by vault: %s", err)
	}

	parsedBundle, err := certutil.ParsePKIMap(vaultResult.Data)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to decode response returned by vault: %s", err)
	}

	bundle, err := parsedBundle.ToCertBundle()
	if err != nil {
		return nil, nil, fmt.Errorf("unable to convert certificate bundle to PEM bundle: %s", err.Error())
	}

	var caPem []byte = nil
	if len(bundle.CAChain) > 0 {
		caPem = []byte(bundle.CAChain[0])
	}

	return []byte(bundle.ToPEMBundle()), caPem, nil
}

func (v *Vault) setToken(client Client) error {
	tokenRef := v.issuer.GetSpec().Vault.Auth.TokenSecretRef
	if tokenRef != nil {
		token, err := v.tokenRef(tokenRef.Name, v.namespace, tokenRef.Key)
		if err != nil {
			return err
		}
		client.SetToken(token)

		return nil
	}

	appRole := v.issuer.GetSpec().Vault.Auth.AppRole
	if appRole != nil {
		token, err := v.requestTokenWithAppRoleRef(client, appRole)
		if err != nil {
			return err
		}
		client.SetToken(token)

		return nil
	}

	kubernetesAuth := v.issuer.GetSpec().Vault.Auth.Kubernetes
	if kubernetesAuth != nil {
		token, err := v.requestTokenWithKubernetesAuth(client, kubernetesAuth)
		if err != nil {
			return fmt.Errorf("error reading Kubernetes service account token from %s: %s", kubernetesAuth.SecretRef.Name, err.Error())
		}
		client.SetToken(token)
		return nil
	}

	return fmt.Errorf("error initializing Vault client: tokenSecretRef, appRoleSecretRef, or Kubernetes auth role not set")
}

func (v *Vault) newConfig() (*vault.Config, error) {
	cfg := vault.DefaultConfig()
	cfg.Address = v.issuer.GetSpec().Vault.Server

	certs := v.issuer.GetSpec().Vault.CABundle
	if len(certs) == 0 {
		return cfg, nil
	}

	caCertPool := x509.NewCertPool()
	ok := caCertPool.AppendCertsFromPEM(certs)
	if ok == false {
		return nil, fmt.Errorf("error loading Vault CA bundle")
	}

	cfg.HttpClient.Transport.(*http.Transport).TLSClientConfig.RootCAs = caCertPool

	return cfg, nil
}

func (v *Vault) tokenRef(name, namespace, key string) (string, error) {
	secret, err := v.secretsLister.Secrets(namespace).Get(name)
	if err != nil {
		return "", err
	}

	if key == "" {
		key = v1alpha2.DefaultVaultTokenAuthSecretKey
	}

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", fmt.Errorf("no data for %q in secret '%s/%s'", key, name, namespace)
	}

	token := string(keyBytes)
	token = strings.TrimSpace(token)

	return token, nil
}

func (v *Vault) appRoleRef(appRole *v1alpha2.VaultAppRole) (roleId, secretId string, err error) {
	roleId = strings.TrimSpace(appRole.RoleId)

	secret, err := v.secretsLister.Secrets(v.namespace).Get(appRole.SecretRef.Name)
	if err != nil {
		return "", "", err
	}

	key := appRole.SecretRef.Key

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", "", fmt.Errorf("no data for %q in secret '%s/%s'", key, v.namespace, appRole.SecretRef.Name)
	}

	secretId = string(keyBytes)
	secretId = strings.TrimSpace(secretId)

	return roleId, secretId, nil
}

func (v *Vault) requestTokenWithAppRoleRef(client Client, appRole *v1alpha2.VaultAppRole) (string, error) {
	roleId, secretId, err := v.appRoleRef(appRole)
	if err != nil {
		return "", err
	}

	parameters := map[string]string{
		"role_id":   roleId,
		"secret_id": secretId,
	}

	authPath := appRole.Path
	if authPath == "" {
		authPath = "approle"
	}

	url := path.Join("/v1", "auth", authPath, "login")

	request := client.NewRequest("POST", url)

	err = request.SetJSONBody(parameters)
	if err != nil {
		return "", fmt.Errorf("error encoding Vault parameters: %s", err.Error())
	}

	resp, err := client.RawRequest(request)
	if err != nil {
		return "", fmt.Errorf("error logging in to Vault server: %s", err.Error())
	}

	defer resp.Body.Close()

	vaultResult := vault.Secret{}
	if err := resp.DecodeJSON(&vaultResult); err != nil {
		return "", fmt.Errorf("unable to decode JSON payload: %s", err.Error())
	}

	token, err := vaultResult.TokenID()
	if err != nil {
		return "", fmt.Errorf("unable to read token: %s", err.Error())
	}

	if token == "" {
		return "", errors.New("no token returned")
	}

	return token, nil
}

func (v *Vault) requestTokenWithKubernetesAuth(client Client, kubernetesAuth *v1alpha2.VaultKubernetesAuth) (string, error) {
	secret, err := v.secretsLister.Secrets(v.namespace).Get(kubernetesAuth.SecretRef.Name)
	if err != nil {
		return "", err
	}

	key := kubernetesAuth.SecretRef.Key
	if key == "" {
		key = v1alpha2.DefaultVaultTokenAuthSecretKey
	}

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", fmt.Errorf("no data for %q in secret '%s/%s'", key, v.namespace, kubernetesAuth.SecretRef.Name)
	}

	jwt := string(keyBytes)

	parameters := map[string]string{
		"role": kubernetesAuth.Role,
		"jwt":  jwt,
	}

	mountPath := kubernetesAuth.Path
	if mountPath == "" {
		mountPath = v1alpha2.DefaultVaultKubernetesAuthMountPath
	}

	url := filepath.Join("/v1", "auth", mountPath, "login")
	request := client.NewRequest("POST", url)
	err = request.SetJSONBody(parameters)
	if err != nil {
		return "", fmt.Errorf("error encoding Vault parameters: %s", err.Error())
	}

	resp, err := client.RawRequest(request)
	if err != nil {
		return "", fmt.Errorf("error calling Vault server: %s", err.Error())
	}

	defer resp.Body.Close()
	vaultResult := vault.Secret{}
	err = resp.DecodeJSON(&vaultResult)
	if err != nil {
		return "", fmt.Errorf("unable to decode JSON payload: %s", err.Error())
	}

	token, err := vaultResult.TokenID()
	if err != nil {
		return "", fmt.Errorf("unable to read token: %s", err.Error())
	}

	return token, nil
}

func (v *Vault) Sys() *vault.Sys {
	return v.client.Sys()
}
<commit_msg>Adds uri_sans to vault parameters<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vault

import (
	"crypto/x509"
	"errors"
	"fmt"
	"net/http"
	"path"
	"path/filepath"
	"strings"
	"time"

	vault "github.com/hashicorp/vault/api"
	"github.com/hashicorp/vault/helper/certutil"
	corelisters "k8s.io/client-go/listers/core/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

var _ Interface = &Vault{}

type VaultClientBuilder func(namespace string, secretsLister corelisters.SecretLister,
	issuer v1alpha2.GenericIssuer) (Interface, error)

type Interface interface {
	Sign(csrPEM []byte, duration time.Duration) (certPEM []byte, caPEM []byte, err error)
	Sys() *vault.Sys
}

type Client interface {
	NewRequest(method, requestPath string) *vault.Request
	RawRequest(r *vault.Request) (*vault.Response, error)
	SetToken(v string)
	Token() string
	Sys() *vault.Sys
}

type Vault struct {
	secretsLister corelisters.SecretLister
	issuer        v1alpha2.GenericIssuer
	namespace     string

	client Client
}

func New(namespace string, secretsLister corelisters.SecretLister,
	issuer v1alpha2.GenericIssuer) (Interface, error) {
	v := &Vault{
		secretsLister: secretsLister,
		namespace:     namespace,
		issuer:        issuer,
	}

	cfg, err := v.newConfig()
	if err != nil {
		return nil, err
	}

	client, err := vault.NewClient(cfg)
	if err != nil {
		return nil, fmt.Errorf("error initializing Vault client: %s", err.Error())
	}

	if err := v.setToken(client); err != nil {
		return nil, err
	}

	v.client = client

	return v, nil
}

func (v *Vault) Sign(csrPEM []byte, duration time.Duration) (cert []byte, ca []byte, err error) {
	csr, err := pki.DecodeX509CertificateRequestBytes(csrPEM)
	if err != nil {
		return nil, nil, fmt.Errorf("faild to decode CSR for signing: %s", err)
	}

	parameters := map[string]string{
		"common_name": csr.Subject.CommonName,
		"alt_names":   strings.Join(csr.DNSNames, ","),
		"ip_sans":     strings.Join(pki.IPAddressesToString(csr.IPAddresses), ","),
		"uri_sans":    strings.Join(pki.URLsToString(csr.URIs), ","),
		"ttl":         duration.String(),
		"csr":         string(csrPEM),

		"exclude_cn_from_sans": "true",
	}

	url := path.Join("/v1", v.issuer.GetSpec().Vault.Path)

	request := v.client.NewRequest("POST", url)

	if err := request.SetJSONBody(parameters); err != nil {
		return nil, nil, fmt.Errorf("failed to build vault request: %s", err)
	}

	resp, err := v.client.RawRequest(request)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to sign certificate by vault: %s", err)
	}

	defer resp.Body.Close()

	vaultResult := certutil.Secret{}
	resp.DecodeJSON(&vaultResult)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to decode response returned by vault: %s", err)
	}

	parsedBundle, err := certutil.ParsePKIMap(vaultResult.Data)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to decode response returned by vault: %s", err)
	}

	bundle, err := parsedBundle.ToCertBundle()
	if err != nil {
		return nil, nil, fmt.Errorf("unable to convert certificate bundle to PEM bundle: %s", err.Error())
	}

	var caPem []byte = nil
	if len(bundle.CAChain) > 0 {
		caPem = []byte(bundle.CAChain[0])
	}

	return []byte(bundle.ToPEMBundle()), caPem, nil
}

func (v *Vault) setToken(client Client) error {
	tokenRef := v.issuer.GetSpec().Vault.Auth.TokenSecretRef
	if tokenRef != nil {
		token, err := v.tokenRef(tokenRef.Name, v.namespace, tokenRef.Key)
		if err != nil {
			return err
		}
		client.SetToken(token)

		return nil
	}

	appRole := v.issuer.GetSpec().Vault.Auth.AppRole
	if appRole != nil {
		token, err := v.requestTokenWithAppRoleRef(client, appRole)
		if err != nil {
			return err
		}
		client.SetToken(token)

		return nil
	}

	kubernetesAuth := v.issuer.GetSpec().Vault.Auth.Kubernetes
	if kubernetesAuth != nil {
		token, err := v.requestTokenWithKubernetesAuth(client, kubernetesAuth)
		if err != nil {
			return fmt.Errorf("error reading Kubernetes service account token from %s: %s", kubernetesAuth.SecretRef.Name, err.Error())
		}
		client.SetToken(token)
		return nil
	}

	return fmt.Errorf("error initializing Vault client: tokenSecretRef, appRoleSecretRef, or Kubernetes auth role not set")
}

func (v *Vault) newConfig() (*vault.Config, error) {
	cfg := vault.DefaultConfig()
	cfg.Address = v.issuer.GetSpec().Vault.Server

	certs := v.issuer.GetSpec().Vault.CABundle
	if len(certs) == 0 {
		return cfg, nil
	}

	caCertPool := x509.NewCertPool()
	ok := caCertPool.AppendCertsFromPEM(certs)
	if ok == false {
		return nil, fmt.Errorf("error loading Vault CA bundle")
	}

	cfg.HttpClient.Transport.(*http.Transport).TLSClientConfig.RootCAs = caCertPool

	return cfg, nil
}

func (v *Vault) tokenRef(name, namespace, key string) (string, error) {
	secret, err := v.secretsLister.Secrets(namespace).Get(name)
	if err != nil {
		return "", err
	}

	if key == "" {
		key = v1alpha2.DefaultVaultTokenAuthSecretKey
	}

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", fmt.Errorf("no data for %q in secret '%s/%s'", key, name, namespace)
	}

	token := string(keyBytes)
	token = strings.TrimSpace(token)

	return token, nil
}

func (v *Vault) appRoleRef(appRole *v1alpha2.VaultAppRole) (roleId, secretId string, err error) {
	roleId = strings.TrimSpace(appRole.RoleId)

	secret, err := v.secretsLister.Secrets(v.namespace).Get(appRole.SecretRef.Name)
	if err != nil {
		return "", "", err
	}

	key := appRole.SecretRef.Key

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", "", fmt.Errorf("no data for %q in secret '%s/%s'", key, v.namespace, appRole.SecretRef.Name)
	}

	secretId = string(keyBytes)
	secretId = strings.TrimSpace(secretId)

	return roleId, secretId, nil
}

func (v *Vault) requestTokenWithAppRoleRef(client Client, appRole *v1alpha2.VaultAppRole) (string, error) {
	roleId, secretId, err := v.appRoleRef(appRole)
	if err != nil {
		return "", err
	}

	parameters := map[string]string{
		"role_id":   roleId,
		"secret_id": secretId,
	}

	authPath := appRole.Path
	if authPath == "" {
		authPath = "approle"
	}

	url := path.Join("/v1", "auth", authPath, "login")

	request := client.NewRequest("POST", url)

	err = request.SetJSONBody(parameters)
	if err != nil {
		return "", fmt.Errorf("error encoding Vault parameters: %s", err.Error())
	}

	resp, err := client.RawRequest(request)
	if err != nil {
		return "", fmt.Errorf("error logging in to Vault server: %s", err.Error())
	}

	defer resp.Body.Close()

	vaultResult := vault.Secret{}
	if err := resp.DecodeJSON(&vaultResult); err != nil {
		return "", fmt.Errorf("unable to decode JSON payload: %s", err.Error())
	}

	token, err := vaultResult.TokenID()
	if err != nil {
		return "", fmt.Errorf("unable to read token: %s", err.Error())
	}

	if token == "" {
		return "", errors.New("no token returned")
	}

	return token, nil
}

func (v *Vault) requestTokenWithKubernetesAuth(client Client, kubernetesAuth *v1alpha2.VaultKubernetesAuth) (string, error) {
	secret, err := v.secretsLister.Secrets(v.namespace).Get(kubernetesAuth.SecretRef.Name)
	if err != nil {
		return "", err
	}

	key := kubernetesAuth.SecretRef.Key
	if key == "" {
		key = v1alpha2.DefaultVaultTokenAuthSecretKey
	}

	keyBytes, ok := secret.Data[key]
	if !ok {
		return "", fmt.Errorf("no data for %q in secret '%s/%s'", key, v.namespace, kubernetesAuth.SecretRef.Name)
	}

	jwt := string(keyBytes)

	parameters := map[string]string{
		"role": kubernetesAuth.Role,
		"jwt":  jwt,
	}

	mountPath := kubernetesAuth.Path
	if mountPath == "" {
		mountPath = v1alpha2.DefaultVaultKubernetesAuthMountPath
	}

	url := filepath.Join("/v1", "auth", mountPath, "login")
	request := client.NewRequest("POST", url)
	err = request.SetJSONBody(parameters)
	if err != nil {
		return "", fmt.Errorf("error encoding Vault parameters: %s", err.Error())
	}

	resp, err := client.RawRequest(request)
	if err != nil {
		return "", fmt.Errorf("error calling Vault server: %s", err.Error())
	}

	defer resp.Body.Close()
	vaultResult := vault.Secret{}
	err = resp.DecodeJSON(&vaultResult)
	if err != nil {
		return "", fmt.Errorf("unable to decode JSON payload: %s", err.Error())
	}

	token, err := vaultResult.TokenID()
	if err != nil {
		return "", fmt.Errorf("unable to read token: %s", err.Error())
	}

	return token, nil
}

func (v *Vault) Sys() *vault.Sys {
	return v.client.Sys()
}
<|endoftext|>


<commit_before>// Package i18n provides internationalization utilites.
package i18n

import (
	"fmt"
	"html/template"
	"strings"

	"golang.org/x/text/language"
)

// Translations is a map of translations associated to keys.
type Translations map[string]string

// Locales is a map of locales and their translations.
type Locales map[language.Tag]Translations

// Has tells if locale lt exists in the locales map.
func (l Locales) Has(lt language.Tag) bool {
	_, ok := (l)[lt]
	return ok
}

// A Translator contains the current locale with its translations and provides functions to get them.
type Translator struct {
	Locale       language.Tag
	Translations Translations
}

// TnPlaceholder is the placeholder replaced by n in a translation, when using the Tn function.
const TnPlaceholder = "{{.n}}"

// Tn returns the translation associated to key, for the client locale.
// If the translation defines plural forms (zero, one, other), it uses the most appropriate.
// All TnPlaceholder in the translation are replaced with number n.
// When translation is not found, an empty string is returned.
func (t *Translator) Tn(key string, n interface{}, a ...interface{}) (s string) {
	if n == 0 {
		if v, ok := t.Translations[key+"Zero"]; ok {
			s = fmt.Sprintf(v, a...)
		}
	} else if n == 1 {
		if v, ok := t.Translations[key+"One"]; ok {
			s = fmt.Sprintf(v, a...)
		}
	} else if v, ok := t.Translations[key+"Other"]; ok {
		s = fmt.Sprintf(v, a...)
	} else if v, ok := t.Translations[key]; ok {
		s = fmt.Sprintf(v, a...)
	}
	s = strings.Replace(s, TnPlaceholder, Fmtn(t.Locale, n), -1)
	if s == "" {
		s = fmt.Sprintf("[%s]", key)
	}
	return
}

// T returns the translation associated to key, for the client locale.
func (t *Translator) T(key string, a ...interface{}) string {
	return t.Tn(key, -1, a...)
}

// TnHTML works like Tn but returns an HTML unescaped translation.
// An "nl2br" function is applied to the result.
func (t *Translator) TnHTML(key string, n interface{}, a ...interface{}) template.HTML {
	return template.HTML(strings.Replace(t.Tn(key, n, a...), "\n", "
", -1))
}

// THTML works like T but returns an HTML unescaped translation.
// An "nl2br" function is applied to the result.
func (t *Translator) THTML(key string, a ...interface{}) template.HTML {
	return t.TnHTML(key, -1, a...)
}

// CleanAcceptLanguage parses, cleans and returns the contents of a Accept-Language header.
// If an error is encountered, the returned string is the same as received.
func CleanAcceptLanguage(s string) (string, error) {
	tt, q, err := language.ParseAcceptLanguage(s)
	if err != nil {
		return s, err
	}

	s = ""
	for i := 0; i < len(tt); i++ {
		if i > 0 {
			s += ","
		}
		s += fmt.Sprintf("%s;q=%g", tt[i].String(), q[i])
	}
	return s, nil
}
<commit_msg>Fix typo<commit_after>// Package i18n provides internationalization utilities.
package i18n

import (
	"fmt"
	"html/template"
	"strings"

	"golang.org/x/text/language"
)

// Translations is a map of translations associated to keys.
type Translations map[string]string

// Locales is a map of locales and their translations.
type Locales map[language.Tag]Translations

// Has tells if locale lt exists in the locales map.
func (l Locales) Has(lt language.Tag) bool {
	_, ok := (l)[lt]
	return ok
}

// A Translator contains the current locale with its translations and provides functions to get them.
type Translator struct {
	Locale       language.Tag
	Translations Translations
}

// TnPlaceholder is the placeholder replaced by n in a translation, when using the Tn function.
const TnPlaceholder = "{{.n}}"

// Tn returns the translation associated to key, for the client locale.
// If the translation defines plural forms (zero, one, other), it uses the most appropriate.
// All TnPlaceholder in the translation are replaced with number n.
// When translation is not found, an empty string is returned.
func (t *Translator) Tn(key string, n interface{}, a ...interface{}) (s string) {
	if n == 0 {
		if v, ok := t.Translations[key+"Zero"]; ok {
			s = fmt.Sprintf(v, a...)
		}
	} else if n == 1 {
		if v, ok := t.Translations[key+"One"]; ok {
			s = fmt.Sprintf(v, a...)
		}
	} else if v, ok := t.Translations[key+"Other"]; ok {
		s = fmt.Sprintf(v, a...)
	} else if v, ok := t.Translations[key]; ok {
		s = fmt.Sprintf(v, a...)
	}
	s = strings.Replace(s, TnPlaceholder, Fmtn(t.Locale, n), -1)
	if s == "" {
		s = fmt.Sprintf("[%s]", key)
	}
	return
}

// T returns the translation associated to key, for the client locale.
func (t *Translator) T(key string, a ...interface{}) string {
	return t.Tn(key, -1, a...)
}

// TnHTML works like Tn but returns an HTML unescaped translation.
// An "nl2br" function is applied to the result.
func (t *Translator) TnHTML(key string, n interface{}, a ...interface{}) template.HTML {
	return template.HTML(strings.Replace(t.Tn(key, n, a...), "\n", "
", -1))
}

// THTML works like T but returns an HTML unescaped translation.
// An "nl2br" function is applied to the result.
func (t *Translator) THTML(key string, a ...interface{}) template.HTML {
	return t.TnHTML(key, -1, a...)
}

// CleanAcceptLanguage parses, cleans and returns the contents of a Accept-Language header.
// If an error is encountered, the returned string is the same as received.
func CleanAcceptLanguage(s string) (string, error) {
	tt, q, err := language.ParseAcceptLanguage(s)
	if err != nil {
		return s, err
	}

	s = ""
	for i := 0; i < len(tt); i++ {
		if i > 0 {
			s += ","
		}
		s += fmt.Sprintf("%s;q=%g", tt[i].String(), q[i])
	}
	return s, nil
}
<|endoftext|>


<commit_before>package mapconv

import (
	"testing"

	"github.com/stretchr/testify/require"
)

type dummyTagged struct {
	A          string `mapconv:"ValueA.A"`
	B          string `mapconv:"ValueA.ValueB.B"`
	C          string `mapconv:"ValueA.ValueB.ValueC.C"`
	NoTag      string
	unexported string
}

type dummyNaked struct {
	ValueA *struct {
		A      string
		ValueB *struct {
			B      string
			ValueC *struct {
				C string
			}
		}
	}
	NoTag      string
	unexported string
}

func TestToNaked(t *testing.T) {

	expects := []struct {
		tagged *dummyTagged
		naked  *dummyNaked
		err    error
	}{
		{
			tagged: &dummyTagged{
				A:          "A",
				B:          "B",
				C:          "C",
				NoTag:      "NoTag",
				unexported: "unexported",
			},
			naked: &dummyNaked{
				ValueA: &struct {
					A      string
					ValueB *struct {
						B      string
						ValueC *struct {
							C string
						}
					}
				}{
					A: "A",
					ValueB: &struct {
						B      string
						ValueC *struct {
							C string
						}
					}{
						B: "B",
						ValueC: &struct {
							C string
						}{
							C: "C",
						},
					},
				},
				NoTag: "NoTag",
			},
		},
	}

	for _, expect := range expects {
		naked := &dummyNaked{}
		err := ToNaked(expect.tagged, naked)
		require.Equal(t, expect.err, err)
		if err == nil {
			require.Equal(t, expect.naked, naked)
		}
	}

}

func TestFromNaked(t *testing.T) {

	expects := []struct {
		tagged *dummyTagged
		naked  *dummyNaked
		err    error
	}{
		{
			tagged: &dummyTagged{
				A:     "A",
				B:     "B",
				C:     "C",
				NoTag: "NoTag",
			},
			naked: &dummyNaked{
				ValueA: &struct {
					A      string
					ValueB *struct {
						B      string
						ValueC *struct {
							C string
						}
					}
				}{
					A: "A",
					ValueB: &struct {
						B      string
						ValueC *struct {
							C string
						}
					}{
						B: "B",
						ValueC: &struct {
							C string
						}{
							C: "C",
						},
					},
				},
				NoTag: "NoTag",
			},
		},
	}

	for _, expect := range expects {
		tagged := &dummyTagged{}
		err := FromNaked(expect.naked, tagged)
		require.Equal(t, expect.err, err)
		if err == nil {
			require.Equal(t, expect.tagged, tagged)
		}
	}

}
<commit_msg>Add testcase to mapconv<commit_after>package mapconv

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

type dummyTagged struct {
	A          string `mapconv:"ValueA.A"`
	B          string `mapconv:"ValueA.ValueB.B"`
	C          string `mapconv:"ValueA.ValueB.ValueC.C"`
	Pointer    *time.Time
	Slice      []string
	NoTag      string
	unexported string
}

type dummyNaked struct {
	ValueA *struct {
		A      string
		ValueB *struct {
			B      string
			ValueC *struct {
				C string
			}
		}
	}
	Pointer    *time.Time
	Slice      []string
	NoTag      string
	unexported string
}

func TestToNaked(t *testing.T) {
	zeroTime := time.Unix(0, 0)
	expects := []struct {
		tagged *dummyTagged
		naked  *dummyNaked
		err    error
	}{
		{
			tagged: &dummyTagged{
				A:          "A",
				B:          "B",
				C:          "C",
				Pointer:    &zeroTime,
				Slice:      []string{"a", "b", "c"},
				NoTag:      "NoTag",
				unexported: "unexported",
			},
			naked: &dummyNaked{
				ValueA: &struct {
					A      string
					ValueB *struct {
						B      string
						ValueC *struct {
							C string
						}
					}
				}{
					A: "A",
					ValueB: &struct {
						B      string
						ValueC *struct {
							C string
						}
					}{
						B: "B",
						ValueC: &struct {
							C string
						}{
							C: "C",
						},
					},
				},
				Pointer: &zeroTime,
				Slice:   []string{"a", "b", "c"},
				NoTag:   "NoTag",
			},
		},
	}

	for _, expect := range expects {
		naked := &dummyNaked{}
		err := ToNaked(expect.tagged, naked)
		require.Equal(t, expect.err, err)
		if err == nil {
			require.Equal(t, expect.naked, naked)
		}
	}

}

func TestFromNaked(t *testing.T) {

	expects := []struct {
		tagged *dummyTagged
		naked  *dummyNaked
		err    error
	}{
		{
			tagged: &dummyTagged{
				A:     "A",
				B:     "B",
				C:     "C",
				NoTag: "NoTag",
			},
			naked: &dummyNaked{
				ValueA: &struct {
					A      string
					ValueB *struct {
						B      string
						ValueC *struct {
							C string
						}
					}
				}{
					A: "A",
					ValueB: &struct {
						B      string
						ValueC *struct {
							C string
						}
					}{
						B: "B",
						ValueC: &struct {
							C string
						}{
							C: "C",
						},
					},
				},
				NoTag: "NoTag",
			},
		},
	}

	for _, expect := range expects {
		tagged := &dummyTagged{}
		err := FromNaked(expect.naked, tagged)
		require.Equal(t, expect.err, err)
		if err == nil {
			require.Equal(t, expect.tagged, tagged)
		}
	}

}
<|endoftext|>


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

package playground

import (
        "log"
	"net/http"
)

func client(r *http.Request) *http.Client {
        return http.DefaultClient
}

func report(r *http.Request, err error) {
	log.Printf("%v", err)
}
<commit_msg>remove more appengine code<commit_after><|endoftext|>


<commit_before>/*
Copyright The Helm Authors.
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 repotest

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"helm.sh/helm/v3/internal/tlsutil"

	"sigs.k8s.io/yaml"

	"helm.sh/helm/v3/pkg/repo"
)

// NewTempServerWithCleanup creates a server inside of a temp dir.
//
// If the passed in string is not "", it will be treated as a shell glob, and files
// will be copied from that path to the server's docroot.
//
// The caller is responsible for stopping the server.
// The temp dir will be removed by testing package automatically when test finished.
func NewTempServerWithCleanup(t *testing.T, glob string) (*Server, error) {
	srv, err := NewTempServer(glob)
	t.Cleanup(func() { os.RemoveAll(srv.docroot) })
	return srv, err
}

// NewTempServer creates a server inside of a temp dir.
//
// If the passed in string is not "", it will be treated as a shell glob, and files
// will be copied from that path to the server's docroot.
//
// The caller is responsible for destroying the temp directory as well as stopping
// the server.
func NewTempServer(glob string) (*Server, error) {
	tdir, err := ioutil.TempDir("", "helm-repotest-")
	if err != nil {
		return nil, err
	}
	srv := NewServer(tdir)

	if glob != "" {
		if _, err := srv.CopyCharts(glob); err != nil {
			srv.Stop()
			return srv, err
		}
	}

	return srv, nil
}

// NewServer creates a repository server for testing.
//
// docroot should be a temp dir managed by the caller.
//
// This will start the server, serving files off of the docroot.
//
// Use CopyCharts to move charts into the repository and then index them
// for service.
func NewServer(docroot string) *Server {
	root, err := filepath.Abs(docroot)
	if err != nil {
		panic(err)
	}
	srv := &Server{
		docroot: root,
	}
	srv.Start()
	// Add the testing repository as the only repo.
	if err := setTestingRepository(srv.URL(), filepath.Join(root, "repositories.yaml")); err != nil {
		panic(err)
	}
	return srv
}

// Server is an implementation of a repository server for testing.
type Server struct {
	docroot    string
	srv        *httptest.Server
	middleware http.HandlerFunc
}

// WithMiddleware injects middleware in front of the server. This can be used to inject
// additional functionality like layering in an authentication frontend.
func (s *Server) WithMiddleware(middleware http.HandlerFunc) {
	s.middleware = middleware
}

// Root gets the docroot for the server.
func (s *Server) Root() string {
	return s.docroot
}

// CopyCharts takes a glob expression and copies those charts to the server root.
func (s *Server) CopyCharts(origin string) ([]string, error) {
	files, err := filepath.Glob(origin)
	if err != nil {
		return []string{}, err
	}
	copied := make([]string, len(files))
	for i, f := range files {
		base := filepath.Base(f)
		newname := filepath.Join(s.docroot, base)
		data, err := ioutil.ReadFile(f)
		if err != nil {
			return []string{}, err
		}
		if err := ioutil.WriteFile(newname, data, 0644); err != nil {
			return []string{}, err
		}
		copied[i] = newname
	}

	err = s.CreateIndex()
	return copied, err
}

// CreateIndex will read docroot and generate an index.yaml file.
func (s *Server) CreateIndex() error {
	// generate the index
	index, err := repo.IndexDirectory(s.docroot, s.URL())
	if err != nil {
		return err
	}

	d, err := yaml.Marshal(index)
	if err != nil {
		return err
	}

	ifile := filepath.Join(s.docroot, "index.yaml")
	return ioutil.WriteFile(ifile, d, 0644)
}

func (s *Server) Start() {
	s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.middleware != nil {
			s.middleware.ServeHTTP(w, r)
		}
		http.FileServer(http.Dir(s.docroot)).ServeHTTP(w, r)
	}))
}

func (s *Server) StartTLS() {
	cd := "../../testdata"
	ca, pub, priv := filepath.Join(cd, "rootca.crt"), filepath.Join(cd, "crt.pem"), filepath.Join(cd, "key.pem")

	s.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.middleware != nil {
			s.middleware.ServeHTTP(w, r)
		}
		http.FileServer(http.Dir(s.Root())).ServeHTTP(w, r)
	}))
	tlsConf, err := tlsutil.NewClientTLS(pub, priv, ca)
	if err != nil {
		panic(err)
	}
	tlsConf.BuildNameToCertificate()
	tlsConf.ServerName = "helm.sh"
	s.srv.TLS = tlsConf
	s.srv.StartTLS()

	// Set up repositories config with ca file
	repoConfig := filepath.Join(s.Root(), "repositories.yaml")

	r := repo.NewFile()
	r.Add(&repo.Entry{
		Name:   "test",
		URL:    s.URL(),
		CAFile: filepath.Join("../../testdata", "rootca.crt"),
	})

	if err := r.WriteFile(repoConfig, 0644); err != nil {
		panic(err)
	}
}

// Stop stops the server and closes all connections.
//
// It should be called explicitly.
func (s *Server) Stop() {
	s.srv.Close()
}

// URL returns the URL of the server.
//
// Example:
//	http://localhost:1776
func (s *Server) URL() string {
	return s.srv.URL
}

// LinkIndices links the index created with CreateIndex and makes a symbolic link to the cache index.
//
// This makes it possible to simulate a local cache of a repository.
func (s *Server) LinkIndices() error {
	lstart := filepath.Join(s.docroot, "index.yaml")
	ldest := filepath.Join(s.docroot, "test-index.yaml")
	return os.Symlink(lstart, ldest)
}

// setTestingRepository sets up a testing repository.yaml with only the given URL.
func setTestingRepository(url, fname string) error {
	r := repo.NewFile()
	r.Add(&repo.Entry{
		Name: "test",
		URL:  url,
	})
	return r.WriteFile(fname, 0644)
}
<commit_msg>mark NewTempServer as Deprecated<commit_after>/*
Copyright The Helm Authors.
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 repotest

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"helm.sh/helm/v3/internal/tlsutil"

	"sigs.k8s.io/yaml"

	"helm.sh/helm/v3/pkg/repo"
)

// NewTempServerWithCleanup creates a server inside of a temp dir.
//
// If the passed in string is not "", it will be treated as a shell glob, and files
// will be copied from that path to the server's docroot.
//
// The caller is responsible for stopping the server.
// The temp dir will be removed by testing package automatically when test finished.
func NewTempServerWithCleanup(t *testing.T, glob string) (*Server, error) {
	srv, err := NewTempServer(glob)
	t.Cleanup(func() { os.RemoveAll(srv.docroot) })
	return srv, err
}

// NewTempServer creates a server inside of a temp dir.
//
// If the passed in string is not "", it will be treated as a shell glob, and files
// will be copied from that path to the server's docroot.
//
// The caller is responsible for destroying the temp directory as well as stopping
// the server.
//
// Deprecated: use NewTempServerWithCleanup
func NewTempServer(glob string) (*Server, error) {
	tdir, err := ioutil.TempDir("", "helm-repotest-")
	if err != nil {
		return nil, err
	}
	srv := NewServer(tdir)

	if glob != "" {
		if _, err := srv.CopyCharts(glob); err != nil {
			srv.Stop()
			return srv, err
		}
	}

	return srv, nil
}

// NewServer creates a repository server for testing.
//
// docroot should be a temp dir managed by the caller.
//
// This will start the server, serving files off of the docroot.
//
// Use CopyCharts to move charts into the repository and then index them
// for service.
func NewServer(docroot string) *Server {
	root, err := filepath.Abs(docroot)
	if err != nil {
		panic(err)
	}
	srv := &Server{
		docroot: root,
	}
	srv.Start()
	// Add the testing repository as the only repo.
	if err := setTestingRepository(srv.URL(), filepath.Join(root, "repositories.yaml")); err != nil {
		panic(err)
	}
	return srv
}

// Server is an implementation of a repository server for testing.
type Server struct {
	docroot    string
	srv        *httptest.Server
	middleware http.HandlerFunc
}

// WithMiddleware injects middleware in front of the server. This can be used to inject
// additional functionality like layering in an authentication frontend.
func (s *Server) WithMiddleware(middleware http.HandlerFunc) {
	s.middleware = middleware
}

// Root gets the docroot for the server.
func (s *Server) Root() string {
	return s.docroot
}

// CopyCharts takes a glob expression and copies those charts to the server root.
func (s *Server) CopyCharts(origin string) ([]string, error) {
	files, err := filepath.Glob(origin)
	if err != nil {
		return []string{}, err
	}
	copied := make([]string, len(files))
	for i, f := range files {
		base := filepath.Base(f)
		newname := filepath.Join(s.docroot, base)
		data, err := ioutil.ReadFile(f)
		if err != nil {
			return []string{}, err
		}
		if err := ioutil.WriteFile(newname, data, 0644); err != nil {
			return []string{}, err
		}
		copied[i] = newname
	}

	err = s.CreateIndex()
	return copied, err
}

// CreateIndex will read docroot and generate an index.yaml file.
func (s *Server) CreateIndex() error {
	// generate the index
	index, err := repo.IndexDirectory(s.docroot, s.URL())
	if err != nil {
		return err
	}

	d, err := yaml.Marshal(index)
	if err != nil {
		return err
	}

	ifile := filepath.Join(s.docroot, "index.yaml")
	return ioutil.WriteFile(ifile, d, 0644)
}

func (s *Server) Start() {
	s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.middleware != nil {
			s.middleware.ServeHTTP(w, r)
		}
		http.FileServer(http.Dir(s.docroot)).ServeHTTP(w, r)
	}))
}

func (s *Server) StartTLS() {
	cd := "../../testdata"
	ca, pub, priv := filepath.Join(cd, "rootca.crt"), filepath.Join(cd, "crt.pem"), filepath.Join(cd, "key.pem")

	s.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.middleware != nil {
			s.middleware.ServeHTTP(w, r)
		}
		http.FileServer(http.Dir(s.Root())).ServeHTTP(w, r)
	}))
	tlsConf, err := tlsutil.NewClientTLS(pub, priv, ca)
	if err != nil {
		panic(err)
	}
	tlsConf.BuildNameToCertificate()
	tlsConf.ServerName = "helm.sh"
	s.srv.TLS = tlsConf
	s.srv.StartTLS()

	// Set up repositories config with ca file
	repoConfig := filepath.Join(s.Root(), "repositories.yaml")

	r := repo.NewFile()
	r.Add(&repo.Entry{
		Name:   "test",
		URL:    s.URL(),
		CAFile: filepath.Join("../../testdata", "rootca.crt"),
	})

	if err := r.WriteFile(repoConfig, 0644); err != nil {
		panic(err)
	}
}

// Stop stops the server and closes all connections.
//
// It should be called explicitly.
func (s *Server) Stop() {
	s.srv.Close()
}

// URL returns the URL of the server.
//
// Example:
//	http://localhost:1776
func (s *Server) URL() string {
	return s.srv.URL
}

// LinkIndices links the index created with CreateIndex and makes a symbolic link to the cache index.
//
// This makes it possible to simulate a local cache of a repository.
func (s *Server) LinkIndices() error {
	lstart := filepath.Join(s.docroot, "index.yaml")
	ldest := filepath.Join(s.docroot, "test-index.yaml")
	return os.Symlink(lstart, ldest)
}

// setTestingRepository sets up a testing repository.yaml with only the given URL.
func setTestingRepository(url, fname string) error {
	r := repo.NewFile()
	r.Add(&repo.Entry{
		Name: "test",
		URL:  url,
	})
	return r.WriteFile(fname, 0644)
}
<|endoftext|>


<commit_before>package dufu

// IPv4
type IPv4 []byte

// Version returns IPv4 version, it's always 4.
func (b IPv4) Version() uint8 { return uint8(b[0] >> 4) }

// HeaderLength returns IPv4 headerlength, it's unit is 4.
func (b IPv4) HeaderLength() uint8 { return uint8(b[0]) & 0xf }

// ToS returns IPv4 type of service.
func (b IPv4) ToS() uint8 { return uint8(b[1]) }

// TotalLength returns the total length of the IPv4 header and payload.
func (b IPv4) TotalLength() uint16 { return binary.BigEndian.Uint16(b[2:4]) }

// Identification returns identification of the IPv4 header.
func (b IPv4) Identification() uint16 { return binary.BigEndian.Uint16(b[4:6]) }

// flagsOffset returns flags and offset filed of the header.
func (b IPv4) flagsOffset() uint16 { return binary.BigEnddian.Uint16(b[6:8]) }

// Flags returns segmentation flags of the header.
func (b IPv4) Flags() { return uint8(b.flagsOffset() >> 13) }

// FragmentOffset returns offset of the IP fragment.
func (b IPv4) FragmentOffset() { return uint16(b.flagsOffset()) }

// TTL returns the time to live.
func (b IPv4) TTL() uint8 { return uint8(b[8]) }

// Protocol returns the protocol used in the data portion.
func (b IPv4) Protocol() uint8 { return uint8(b[9]) }

// Checksum returns header checksum.
func (b IPv4) Checksum() uint16 { return binary.BigEndian.Uint16(b[10:]) }

// SourceAddress returns IP source address.
func (b IPv4) SourceAddress() []byte { return b[12:16] }

// DestinationAddress returns IP destination address.
func (b IPv4) DestinationAddress() []byte { return b[16:20] }
func (b IPv4) Message() []byte            { return b[b.HeaderLength()*4:] }

// IPHandle handles IP packet.
func IPRcv(l2l *L2Layer, packet []byte) {
	ip := IPv4(packet)
	if ip.Protocol() == ICMPv4ProtocolNumber {
		ICMPRcv(ip)
	}
}

func IPSend(packet []byte){

}
<commit_msg>ipv4: support checksum calculation<commit_after>package dufu

import (
	"encoding/binary"
	"fmt"
)

const (
	IPv4ProtocolNumber = 0x0800
	IPv6ProtocolNumber = 0x86DD
)

// IPv4
type IPv4 []byte

// Version returns IPv4 version, it's always 4.
func (b IPv4) Version() uint8     { return uint8(b[0] >> 4) }
func (b IPv4) SetVersion(v uint8) { b[0] = (b[0] & 0x0f) | (v << 4) }

// HeaderLength returns IPv4 headerlength, it's unit is 4.
func (b IPv4) HeaderLength() uint8     { return uint8(b[0]) & 0xf }
func (b IPv4) SetHeaderLength(l uint8) { b[0] = (b[0] & 0xf0) | (l & 0xf) }

// ToS returns IPv4 type of service.
func (b IPv4) ToS() uint8       { return uint8(b[1]) }
func (b IPv4) SetToS(tos uint8) { b[1] = tos; return }

// TotalLength returns the total length of the IPv4 header and payload.
func (b IPv4) TotalLength() uint16        { return binary.BigEndian.Uint16(b[2:4]) }
func (b IPv4) SetTotalLength(totl uint16) { binary.BigEndian.PutUint16(b[2:4], totl); return }

// Identification returns identification of the IPv4 header.
func (b IPv4) Identification() uint16 { return binary.BigEndian.Uint16(b[4:6]) }

// SetIdentification sets packet ID for this IPv4 packet.
func (b IPv4) SetIdentification(id uint16) { binary.BigEndian.PutUint16(b[4:6], id) }

// flagsOffset returns flags and offset filed of the header.
func (b IPv4) FlagsOffset() uint16 { return binary.BigEndian.Uint16(b[6:8]) }

// Flags returns segmentation flags of the header.
func (b IPv4) Flags() uint8 { return uint8(b.FlagsOffset() >> 13) }

// FragmentOffset returns offset of the IP fragment.
func (b IPv4) FragmentOffset() uint16 { return uint16(b.FlagsOffset()) }

// TTL returns the time to live.
func (b IPv4) TTL() uint8       { return uint8(b[8]) }
func (b IPv4) SetTTL(ttl uint8) { b[8] = byte(ttl); return }

// Protocol returns the protocol used in the data portion.
func (b IPv4) Protocol() uint8     { return uint8(b[9]) }
func (b IPv4) SetProtocol(p uint8) { b[9] = byte(p) }

// Checksum returns header checksum.
func (b IPv4) Checksum() uint16       { return binary.BigEndian.Uint16(b[10:12]) }
func (b IPv4) SetChecksum(cks uint16) { binary.BigEndian.PutUint16(b[10:12], cks) }

// SourceAddress returns IP source address.
func (b IPv4) SourceAddress() []byte { return b[12:16] }

// DestinationAddress returns IP destination address.
func (b IPv4) DestinationAddress() []byte { return b[16:20] }
func (b IPv4) Message() []byte            { return b[b.HeaderLength()*4:] }

// L3Layer is a IP layer.
type L3Layer struct {
	*L2Layer
	// TODO: preprendable buffer
}

// IPHandle handles IP packet.
func (l3 *L3Layer) IPRcv(l2l *L2Layer, skb *SkBuff) {
	ip := IPv4(skb.Data())
	fmt.Printf("IP Protocol 0x%x\n", ip.Protocol())
	fmt.Printf("Header Length %d\n",ip.HeaderLength())
	fmt.Printf("Header Checksum 0x%x\n", ip.Checksum())
	fmt.Printf("IPv4 Checksum 0x%x\n", ipv4Checksum(ip[:ip.HeaderLength()*4]))
	if ip.Protocol() == ICMPv4ProtocolNumber {
		ICMPRcv(l3, ip)
	}
}

func (l3 *L3Layer) IPSend(packet []byte) {
	// assume up layer has set the destination address and source address.
	ip := IPv4(packet)
	ip.SetVersion(4)
	ip.SetTotalLength(uint16(len(packet)))
	ip.SetIdentification(uint16(IpSelectIdent(ip)))
}

// ipv4Checksum get buf checksum based on rfc1071
func ipv4Checksum(buf []byte) uint16 {
	var sum uint32
	count := len(buf)
	for i := 0; count > 1; i += 2 {
		sum += uint32(buf[i])<<8 + uint32(buf[i+1])
		count -= 2
	}
	if count > 0 {
		sum += uint32(buf[len(buf)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = sum&0xffff + (sum >> 16)
	}
	return ^uint16(sum)
}
<|endoftext|>