<commit_before>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/influxdata/influxdb-comparisons/bulk_query_gen/mongodb"
	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	limit          int64
	doLoad         bool
	dbName         string
	documentFormat string
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
)

// Global vars
var (
	batchChan      chan *Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
)

// Magic database constants
const (
	pointCollectionName = "point_data"
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 1024)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// batchPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.Int64Var(&limit, "limit", -1, "Number of items to insert (default unlimited).")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.StringVar(&dbName, "db", "benchmark_db", "Database for influx to use (ignored for ElasticSearch).")
	flag.StringVar(&documentFormat, "document-format", "", "Document format specification. ('simpleArrays' is supported; leave empty for previous behaviour)")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	if documentFormat == mongodb.SimpleArraysFormat {
		log.Printf("Using '%s' document serialization", documentFormat)
	}

	for i := 0; i < workers*batchSize; i++ {
		bufPool.Put(bufPool.New())
	}

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		mustCreateCollections(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		session.SetMode(mgo.Eventual, false)
		session.SetSyncTimeout(180 * time.Second)
		session.SetSocketTimeout(180 * time.Second)

		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*10)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	itemRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d values in %fsec with %d workers (mean values rate %f values/sec, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "MongoDB",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          int(limit),
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, itemRate, bytesRate, took)

		if err != nil {
			log.Fatal(err)
		}
	}
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) (int64, int64) {
	var n int
	var itemsRead, bytesRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	start := time.Now()
	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		if itemsRead == limit {
			break
		}
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer object
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			bytesRead += int64(len(itemBuf))
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}

		_ = start
		//if itemsRead > 0 && itemsRead%100000 == 0 {
		//	_ = start
		//	//took := (time.Now().UnixNano() - start.UnixNano())
		//	//if took >= 1e9 {
		//	//	tookUs := float64(took) / 1e3
		//	//	tookSec := float64(took) / 1e9
		//	//	fmt.Fprintf(os.Stderr, "itemsRead: %d, rate: %.0f/sec, lag: %.2fus/op\n",
		//	//		itemsRead, float64(itemsRead)/tookSec, tookUs/float64(itemsRead))
		//	//}
		//}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead, bytesRead
}

// processBatches reads byte buffers from batchChan, interprets them and writes
// them to the target server. Note that mgo forcibly incurs serialization
// overhead (it always encodes to BSON).
func processBatches(session *mgo.Session) {
	db := session.DB(dbName)

	type Tag struct {
		Key string `bson:"key"`
		Val string `bson:"val"`
	}

	type Field struct {
		Key string `bson:"key"`
		Val interface{} `bson:"val"`
	}

	type Point struct {
		// Use `string` here even though they are really `[]byte`.
		// This is so the mongo data is human-readable.
		MeasurementName string      `bson:"measurement"`
		Timestamp       int64       `bson:"timestamp_ns"`
		Tags            []interface{}  `bson:"tags"`
		Fields          []interface{}  `bson:"fields"`
	}
	pPool := &sync.Pool{New: func() interface{} { return &Point{} }}
	pvs := []interface{}{}

	item := &mongo_serialization.Item{}
	destTag := &mongo_serialization.Tag{}
	destField := &mongo_serialization.Field{}
	collection := db.C(pointCollectionName)
	for batch := range batchChan {
		bulk := collection.Bulk()

		if cap(pvs) < len(*batch) {
			pvs = make([]interface{}, len(*batch))
		}
		pvs = pvs[:len(*batch)]

		for i, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)

			x := pPool.Get().(*Point)

			x.MeasurementName = unsafeBytesToString(item.MeasurementNameBytes())
			x.Timestamp = item.TimestampNanos()

			tagLength := item.TagsLength()
			if cap(x.Tags) < tagLength {
				x.Tags = make([]interface{}, 0, tagLength)
			}
			x.Tags = x.Tags[:tagLength]
			for i := 0; i < tagLength; i++ {
				*destTag = mongo_serialization.Tag{} // clear
				item.Tags(destTag, i)
				tagKey := unsafeBytesToString(destTag.KeyBytes())
				tagValue := unsafeBytesToString(destTag.ValBytes())
				if documentFormat == mongodb.SimpleArraysFormat {
					x.Tags[i] = bson.M{tagKey:tagValue}
				} else {
					x.Tags[i] = &Tag{Key:tagKey,Val:tagValue}
				}
			}

			fieldLength := item.FieldsLength()
			if cap(x.Fields) < fieldLength {
				x.Fields = make([]interface{}, 0, fieldLength)
			}
			x.Fields = x.Fields[:fieldLength]
			for i := 0; i < fieldLength; i++ {
				*destField = mongo_serialization.Field{} // clear
				item.Fields(destField, i)
				fieldKey := unsafeBytesToString(destField.KeyBytes())
				var fieldValue interface{}
				switch destField.ValueType() {
				case mongo_serialization.ValueTypeInt:
					fieldValue = destField.IntValue()
				case mongo_serialization.ValueTypeLong:
					fieldValue = destField.LongValue()
				case mongo_serialization.ValueTypeFloat:
					fieldValue = destField.FloatValue()
				case mongo_serialization.ValueTypeDouble:
					fieldValue = destField.DoubleValue()
				case mongo_serialization.ValueTypeString:
					fieldValue = unsafeBytesToString(destField.StringValueBytes())
				default:
					panic("logic error")
				}
				if documentFormat == mongodb.SimpleArraysFormat {
					x.Fields[i] = bson.M{fieldKey:fieldValue}
				} else {
					x.Fields[i] = &Field{Key:fieldKey,Val:fieldValue}
				}
			}
			pvs[i] = x

		}
		bulk.Insert(pvs...)

		if doLoad {
			_, err := bulk.Run()
			if err != nil {
				log.Fatalf("Bulk err: %s\n", err.Error())
			}

		}

		// cleanup pvs
		for _, x := range pvs {
			p := x.(*Point)
			p.Timestamp = 0
			p.Tags = p.Tags[:0]
			p.Fields = p.Fields[:0]
			pPool.Put(p)
		}

		// cleanup item data
		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	workersGroup.Done()
}

func mustCreateCollections(daemonUrl string) {
	session, err := mgo.Dial(daemonUrl)
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	// collection C: point data
	// from (*mgo.Collection).Create
	cmd := make(bson.D, 0, 4)
	cmd = append(cmd, bson.DocElem{"create", pointCollectionName})

	err = session.DB(dbName).Run(cmd, nil)
	if err != nil {
		log.Fatal(err)
	}

	collection := session.DB(dbName).C("point_data")
	index := mgo.Index{
		Key:        []string{"measurement", "tags", "timestamp_ns"},
		Unique:     false, // Unique does not work on the entire array of tags!
		DropDups:   true,
		Background: false,
		Sparse:     false,
	}
	err = collection.EnsureIndex(index)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>output values/sec at the end<commit_after>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/influxdata/influxdb-comparisons/bulk_query_gen/mongodb"
	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	limit          int64
	doLoad         bool
	dbName         string
	documentFormat string
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
)

// Global vars
var (
	batchChan      chan *Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
	valuesRead     int64
)

// Magic database constants
const (
	pointCollectionName = "point_data"
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 1024)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// batchPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.Int64Var(&limit, "limit", -1, "Number of items to insert (default unlimited).")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.StringVar(&dbName, "db", "benchmark_db", "Database for influx to use (ignored for ElasticSearch).")
	flag.StringVar(&documentFormat, "document-format", "", "Document format specification. ('simpleArrays' is supported; leave empty for previous behaviour)")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	if documentFormat == mongodb.SimpleArraysFormat {
		log.Printf("Using '%s' document serialization", documentFormat)
	}

	for i := 0; i < workers*batchSize; i++ {
		bufPool.Put(bufPool.New())
	}

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		mustCreateCollections(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		session.SetMode(mgo.Eventual, false)
		session.SetSyncTimeout(180 * time.Second)
		session.SetSocketTimeout(180 * time.Second)

		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*10)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	itemRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())
	valuesRate := float64(valuesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean point rate %f/sec, mean value rate %f/s, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemRate, valuesRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "MongoDB",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          int(limit),
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, itemRate, bytesRate, took)

		if err != nil {
			log.Fatal(err)
		}
	}
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) (int64, int64) {
	var n int
	var itemsRead, bytesRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	start := time.Now()
	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		if itemsRead == limit {
			break
		}
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer object
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			bytesRead += int64(len(itemBuf))
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}

		_ = start
		//if itemsRead > 0 && itemsRead%100000 == 0 {
		//	_ = start
		//	//took := (time.Now().UnixNano() - start.UnixNano())
		//	//if took >= 1e9 {
		//	//	tookUs := float64(took) / 1e3
		//	//	tookSec := float64(took) / 1e9
		//	//	fmt.Fprintf(os.Stderr, "itemsRead: %d, rate: %.0f/sec, lag: %.2fus/op\n",
		//	//		itemsRead, float64(itemsRead)/tookSec, tookUs/float64(itemsRead))
		//	//}
		//}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead, bytesRead
}

// processBatches reads byte buffers from batchChan, interprets them and writes
// them to the target server. Note that mgo forcibly incurs serialization
// overhead (it always encodes to BSON).
func processBatches(session *mgo.Session) {
	var workerValuesRead int64

	db := session.DB(dbName)

	type Tag struct {
		Key string `bson:"key"`
		Val string `bson:"val"`
	}

	type Field struct {
		Key string `bson:"key"`
		Val interface{} `bson:"val"`
	}

	type Point struct {
		// Use `string` here even though they are really `[]byte`.
		// This is so the mongo data is human-readable.
		MeasurementName string      `bson:"measurement"`
		Timestamp       int64       `bson:"timestamp_ns"`
		Tags            []interface{}  `bson:"tags"`
		Fields          []interface{}  `bson:"fields"`
	}
	pPool := &sync.Pool{New: func() interface{} { return &Point{} }}
	pvs := []interface{}{}

	item := &mongo_serialization.Item{}
	destTag := &mongo_serialization.Tag{}
	destField := &mongo_serialization.Field{}
	collection := db.C(pointCollectionName)
	for batch := range batchChan {
		bulk := collection.Bulk()

		if cap(pvs) < len(*batch) {
			pvs = make([]interface{}, len(*batch))
		}
		pvs = pvs[:len(*batch)]

		for i, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)

			x := pPool.Get().(*Point)

			x.MeasurementName = unsafeBytesToString(item.MeasurementNameBytes())
			x.Timestamp = item.TimestampNanos()

			tagLength := item.TagsLength()
			if cap(x.Tags) < tagLength {
				x.Tags = make([]interface{}, 0, tagLength)
			}
			x.Tags = x.Tags[:tagLength]
			for i := 0; i < tagLength; i++ {
				*destTag = mongo_serialization.Tag{} // clear
				item.Tags(destTag, i)
				tagKey := unsafeBytesToString(destTag.KeyBytes())
				tagValue := unsafeBytesToString(destTag.ValBytes())
				if documentFormat == mongodb.SimpleArraysFormat {
					x.Tags[i] = bson.M{tagKey:tagValue}
				} else {
					x.Tags[i] = &Tag{Key:tagKey,Val:tagValue}
				}
			}

			fieldLength := item.FieldsLength()
			if cap(x.Fields) < fieldLength {
				x.Fields = make([]interface{}, 0, fieldLength)
			}
			x.Fields = x.Fields[:fieldLength]
			for i := 0; i < fieldLength; i++ {
				*destField = mongo_serialization.Field{} // clear
				item.Fields(destField, i)
				fieldKey := unsafeBytesToString(destField.KeyBytes())
				var fieldValue interface{}
				switch destField.ValueType() {
				case mongo_serialization.ValueTypeInt:
					fieldValue = destField.IntValue()
				case mongo_serialization.ValueTypeLong:
					fieldValue = destField.LongValue()
				case mongo_serialization.ValueTypeFloat:
					fieldValue = destField.FloatValue()
				case mongo_serialization.ValueTypeDouble:
					fieldValue = destField.DoubleValue()
				case mongo_serialization.ValueTypeString:
					fieldValue = unsafeBytesToString(destField.StringValueBytes())
				default:
					panic("logic error")
				}
				if documentFormat == mongodb.SimpleArraysFormat {
					x.Fields[i] = bson.M{fieldKey:fieldValue}
				} else {
					x.Fields[i] = &Field{Key:fieldKey,Val:fieldValue}
				}
			}
			pvs[i] = x
			workerValuesRead += int64(fieldLength)
		}
		bulk.Insert(pvs...)

		if doLoad {
			_, err := bulk.Run()
			if err != nil {
				log.Fatalf("Bulk err: %s\n", err.Error())
			}

		}

		// cleanup pvs
		for _, x := range pvs {
			p := x.(*Point)
			p.Timestamp = 0
			p.Tags = p.Tags[:0]
			p.Fields = p.Fields[:0]
			pPool.Put(p)
		}

		// cleanup item data
		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	atomic.AddInt64(&valuesRead, workerValuesRead)
	workersGroup.Done()
}

func mustCreateCollections(daemonUrl string) {
	session, err := mgo.Dial(daemonUrl)
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	// collection C: point data
	// from (*mgo.Collection).Create
	cmd := make(bson.D, 0, 4)
	cmd = append(cmd, bson.DocElem{"create", pointCollectionName})

	err = session.DB(dbName).Run(cmd, nil)
	if err != nil {
		log.Fatal(err)
	}

	collection := session.DB(dbName).C("point_data")
	index := mgo.Index{
		Key:        []string{"measurement", "tags", "timestamp_ns"},
		Unique:     false, // Unique does not work on the entire array of tags!
		DropDups:   true,
		Background: false,
		Sparse:     false,
	}
	err = collection.EnsureIndex(index)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package cmd

// DO NOT EDIT manually! Generated by hack/update-version.sh
var version = "0.10.3"
<commit_msg>Bump version to 0.10.3-dev<commit_after>package cmd

// DO NOT EDIT manually! Generated by hack/update-version.sh
var version = "0.10.3-dev"
<|endoftext|>


<commit_before>package main

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

	"github.com/Cloud-Foundations/keymaster/lib/instrumentedwriter"
	"github.com/Cloud-Foundations/keymaster/lib/paths"
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
)

func (state *RuntimeState) generateAuthJWT(username string) (string, error) {
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.RS256,
		Key:       state.Signer,
	}, (&jose.SignerOptions{}).WithType("JWT"))
	if err != nil {
		return "", err
	}
	issuer := state.idpGetIssuer()
	now := time.Now().Unix()
	authToken := authInfoJWT{
		Issuer:   issuer,
		Subject:  username,
		Audience: []string{issuer},
		Expiration: now + int64(
			state.Config.Base.WebauthTokenForCliLifetime/time.Second),
		NotBefore: now,
		IssuedAt:  now,
		TokenType: "keymaster_token_auth",
	}
	return jwt.Signed(signer).Claims(authToken).CompactSerialize()
}

func (state *RuntimeState) SendAuthDocumentHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugln(1, "Entered SendAuthDocumentHandler()")
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"Error parsing form")
		return
	}
	authData, err := state.checkAuth(w, r, state.getRequiredWebUIAuthLevel())
	if err != nil {
		state.logger.Debugln(1, err)
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(authData.Username)
	state.logger.Printf("%s requested authentication document export\n",
		authData.Username)
	// Fetch form/query data.
	var portNumber uint64
	var token string
	if val, ok := r.Form["port"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No CLI port number provided")
		state.logger.Printf("SendAuthDocument without port number")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one port number allowed")
			state.logger.Printf("SendAuthDocument with multiple port values")
			return
		}
		if portNumber, err = strconv.ParseUint(val[0], 10, 16); err != nil {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Invalid port number")
			state.logger.Printf("SendAuthDocument with invalid port number")
			return
		}
	}
	if val, ok := r.Form["token"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No token provided")
		state.logger.Printf("SendAuthDocument without token")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one token allowed")
			state.logger.Printf("SendAuthDocument with multiple token values")
			return
		}
		token = val[0]
	}
	authInfo, err := state.getAuthInfoFromJWT(token, "keymaster_token_auth")
	if err != nil {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "Bad token")
		state.logger.Debugln(0, err)
		return
	}
	if authInfo.Username != authData.Username {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "User mismatch")
		state.logger.Printf(
			"SendAuthDocumentHandler: authticated user: %s != token user: %s\n",
			authData.Username, authInfo.Username)
		return
	}
	if time.Until(authInfo.ExpiresAt) < 0 {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "Token expired")
		state.logger.Debugln(0, "token expired")
		return
	}
	var authCookie *http.Cookie
	for _, cookie := range r.Cookies() {
		if cookie.Name != authCookieName {
			continue
		}
		authCookie = cookie
	}
	if authCookie == nil {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No auth_cookie")
		state.logger.Debugln(0, "no auth_cookie")
		return
	}
	http.Redirect(w, r,
		fmt.Sprintf("http://localhost:%d%s?auth_cookie=%s",
			portNumber, paths.ReceiveAuthDocument, authCookie.Value),
		http.StatusPermanentRedirect)
}

func (state *RuntimeState) ShowAuthTokenHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugf(1, "Entered GetAuthTokenHandler(). URL: %v\n", r.URL)
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"Error parsing form")
		return
	}
	authData, err := state.checkAuth(w, r, state.getRequiredWebUIAuthLevel())
	if err != nil {
		state.logger.Debugf(1, "%s", err)
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(authData.Username)
	displayData := authCodePageTemplateData{
		Title:        "Keymaster CLI Token Display",
		AuthUsername: authData.Username,
	}
	token, err := state.generateAuthJWT(authData.Username)
	if err != nil {
		state.logger.Debugf(1, "%s", err)
		displayData.ErrorMessage = "Unable to generate token"
	} else {
		state.logger.Printf(
			"generated webauth CLI token for: %s, lifetime: %s\n",
			authData.Username, state.Config.Base.WebauthTokenForCliLifetime)
		displayData.Token = token
	}
	err = state.htmlTemplate.ExecuteTemplate(w, "authTokenPage", displayData)
	if err != nil {
		logger.Printf("Failed to execute %s", err)
		http.Error(w, "error", http.StatusInternalServerError)
		return
	}
}

func (state *RuntimeState) VerifyAuthTokenHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugf(1, "Entered VerifyAuthTokenHandler(). URL: %v\n", r.URL)
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"Error parsing form")
		return
	}
	// Fetch form/query data.
	var token string
	if val, ok := r.Form["token"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No token provided")
		state.logger.Printf("VerifyAuthToken without token")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one token allowed")
			state.logger.Printf("VerifyAuthToken with multiple token values")
			return
		}
		token = val[0]
	}
	authInfo, err := state.getAuthInfoFromJWT(token, "keymaster_token_auth")
	if err != nil {
		state.writeFailureResponse(w, r, http.StatusNotAcceptable, "Bad token")
		state.logger.Debugln(0, err)
		return
	}
	if time.Until(authInfo.ExpiresAt) < 0 {
		state.writeFailureResponse(w, r, http.StatusGone, "Token expired")
		state.logger.Debugln(0, "token expired")
		return
	}
	w.Write([]byte("OK\n"))
}
<commit_msg>keymasterd: change type of long-lived webauth for CLI token.<commit_after>package main

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

	"github.com/Cloud-Foundations/keymaster/lib/instrumentedwriter"
	"github.com/Cloud-Foundations/keymaster/lib/paths"
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
)

func (state *RuntimeState) generateAuthJWT(username string) (string, error) {
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.RS256,
		Key:       state.Signer,
	}, (&jose.SignerOptions{}).WithType("JWT"))
	if err != nil {
		return "", err
	}
	issuer := state.idpGetIssuer()
	now := time.Now().Unix()
	authToken := authInfoJWT{
		Issuer:   issuer,
		Subject:  username,
		Audience: []string{issuer},
		Expiration: now + int64(
			state.Config.Base.WebauthTokenForCliLifetime/time.Second),
		NotBefore: now,
		IssuedAt:  now,
		TokenType: "keymaster_webauth_for_cli_identity",
	}
	return jwt.Signed(signer).Claims(authToken).CompactSerialize()
}

func (state *RuntimeState) SendAuthDocumentHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugln(1, "Entered SendAuthDocumentHandler()")
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		state.logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusInternalServerError,
			"Error parsing form")
		return
	}
	authData, err := state.checkAuth(w, r, state.getRequiredWebUIAuthLevel())
	if err != nil {
		state.logger.Debugln(1, err)
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(authData.Username)
	state.logger.Printf("%s requested authentication document export\n",
		authData.Username)
	// Fetch form/query data.
	var portNumber uint64
	var token string
	if val, ok := r.Form["port"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No CLI port number provided")
		state.logger.Printf("SendAuthDocument without port number")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one port number allowed")
			state.logger.Printf("SendAuthDocument with multiple port values")
			return
		}
		if portNumber, err = strconv.ParseUint(val[0], 10, 16); err != nil {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Invalid port number")
			state.logger.Printf("SendAuthDocument with invalid port number")
			return
		}
	}
	if val, ok := r.Form["token"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No token provided")
		state.logger.Printf("SendAuthDocument without token")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one token allowed")
			state.logger.Printf("SendAuthDocument with multiple token values")
			return
		}
		token = val[0]
	}
	authInfo, err := state.getAuthInfoFromJWT(token,
		"keymaster_webauth_for_cli_identity")
	if err != nil {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "Bad token")
		state.logger.Debugln(0, err)
		return
	}
	if authInfo.Username != authData.Username {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "User mismatch")
		state.logger.Printf(
			"SendAuthDocumentHandler: authticated user: %s != token user: %s\n",
			authData.Username, authInfo.Username)
		return
	}
	if time.Until(authInfo.ExpiresAt) < 0 {
		state.writeFailureResponse(w, r, http.StatusBadRequest, "Token expired")
		state.logger.Debugln(0, "token expired")
		return
	}
	var authCookie *http.Cookie
	for _, cookie := range r.Cookies() {
		if cookie.Name != authCookieName {
			continue
		}
		authCookie = cookie
	}
	if authCookie == nil {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No auth_cookie")
		state.logger.Debugln(0, "no auth_cookie")
		return
	}
	http.Redirect(w, r,
		fmt.Sprintf("http://localhost:%d%s?auth_cookie=%s",
			portNumber, paths.ReceiveAuthDocument, authCookie.Value),
		http.StatusPermanentRedirect)
}

func (state *RuntimeState) ShowAuthTokenHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugf(1, "Entered GetAuthTokenHandler(). URL: %v\n", r.URL)
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"Error parsing form")
		return
	}
	authData, err := state.checkAuth(w, r, state.getRequiredWebUIAuthLevel())
	if err != nil {
		state.logger.Debugf(1, "%s", err)
		return
	}
	w.(*instrumentedwriter.LoggingWriter).SetUsername(authData.Username)
	displayData := authCodePageTemplateData{
		Title:        "Keymaster CLI Token Display",
		AuthUsername: authData.Username,
	}
	token, err := state.generateAuthJWT(authData.Username)
	if err != nil {
		state.logger.Debugf(1, "%s", err)
		displayData.ErrorMessage = "Unable to generate token"
	} else {
		state.logger.Printf(
			"generated webauth CLI token for: %s, lifetime: %s\n",
			authData.Username, state.Config.Base.WebauthTokenForCliLifetime)
		displayData.Token = token
	}
	err = state.htmlTemplate.ExecuteTemplate(w, "authTokenPage", displayData)
	if err != nil {
		logger.Printf("Failed to execute %s", err)
		http.Error(w, "error", http.StatusInternalServerError)
		return
	}
}

func (state *RuntimeState) VerifyAuthTokenHandler(w http.ResponseWriter,
	r *http.Request) {
	state.logger.Debugf(1, "Entered VerifyAuthTokenHandler(). URL: %v\n", r.URL)
	if state.sendFailureToClientIfLocked(w, r) {
		return
	}
	if r.Method != "GET" && r.Method != "POST" {
		state.writeFailureResponse(w, r, http.StatusMethodNotAllowed, "")
		return
	}
	if err := r.ParseForm(); err != nil {
		logger.Println(err)
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"Error parsing form")
		return
	}
	// Fetch form/query data.
	var token string
	if val, ok := r.Form["token"]; !ok {
		state.writeFailureResponse(w, r, http.StatusBadRequest,
			"No token provided")
		state.logger.Printf("VerifyAuthToken without token")
		return
	} else {
		if len(val) > 1 {
			state.writeFailureResponse(w, r, http.StatusBadRequest,
				"Just one token allowed")
			state.logger.Printf("VerifyAuthToken with multiple token values")
			return
		}
		token = val[0]
	}
	authInfo, err := state.getAuthInfoFromJWT(token,
		"keymaster_webauth_for_cli_identity")
	if err != nil {
		state.writeFailureResponse(w, r, http.StatusNotAcceptable, "Bad token")
		state.logger.Debugln(0, err)
		return
	}
	if time.Until(authInfo.ExpiresAt) < 0 {
		state.writeFailureResponse(w, r, http.StatusGone, "Token expired")
		state.logger.Debugln(0, "token expired")
		return
	}
	w.Write([]byte("OK\n"))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strings"
	"time"

	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/pkg/errors"
)

// FindTeleproxy finds a compatible version of Teleproxy in your PATH
func (d *Daemon) FindTeleproxy() error {
	if len(d.teleproxy) == 0 {
		path, err := exec.LookPath("teleproxy")
		if err != nil {
			return err
		}
		cmd := exec.Command(path, "--version")
		outputBytes, err := cmd.CombinedOutput()
		if err != nil {
			return errors.Wrap(err, "teleproxy --version")
		}
		output := string(outputBytes)
		if !strings.Contains(output, "version 0.6") {
			return fmt.Errorf(
				"required teleproxy 0.6.x not found; found %s in your PATH",
				output,
			)
		}
		d.teleproxy = path
	}
	return nil
}

func (d *Daemon) MakeNetOverride(p *supervisor.Process) error {
	if err := d.FindTeleproxy(); err != nil {
		return err
	}
	netOverride, err := CheckedRetryingCommand(
		p,
		"netOverride",
		[]string{d.teleproxy, "--mode", "intercept"},
		&RunAsInfo{},
		checkNetOverride,
		10*time.Second,
	)
	if err != nil {
		return errors.Wrap(err, "teleproxy initial launch")
	}
	d.network = netOverride
	return nil
}

// checkNetOverride checks the status of teleproxy intercept by doing the
// equivalent of curl http://teleproxy/api/tables/. It's okay to create a new
// client each time because we don't want to reuse connections.
func checkNetOverride() error {
	client := http.Client{Timeout: 3 * time.Second}
	res, err := client.Get(fmt.Sprintf(
		"http://teleproxy%d.cachebust.telepresence.io/api/tables",
		time.Now().Unix(),
	))
	if err != nil {
		return err
	}
	_, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Add comment to satisfy linters<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strings"
	"time"

	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/pkg/errors"
)

// FindTeleproxy finds a compatible version of Teleproxy in your PATH
func (d *Daemon) FindTeleproxy() error {
	if len(d.teleproxy) == 0 {
		path, err := exec.LookPath("teleproxy")
		if err != nil {
			return err
		}
		cmd := exec.Command(path, "--version")
		outputBytes, err := cmd.CombinedOutput()
		if err != nil {
			return errors.Wrap(err, "teleproxy --version")
		}
		output := string(outputBytes)
		if !strings.Contains(output, "version 0.6") {
			return fmt.Errorf(
				"required teleproxy 0.6.x not found; found %s in your PATH",
				output,
			)
		}
		d.teleproxy = path
	}
	return nil
}

// MakeNetOverride sets up the network override resource for the daemon
func (d *Daemon) MakeNetOverride(p *supervisor.Process) error {
	if err := d.FindTeleproxy(); err != nil {
		return err
	}
	netOverride, err := CheckedRetryingCommand(
		p,
		"netOverride",
		[]string{d.teleproxy, "--mode", "intercept"},
		&RunAsInfo{},
		checkNetOverride,
		10*time.Second,
	)
	if err != nil {
		return errors.Wrap(err, "teleproxy initial launch")
	}
	d.network = netOverride
	return nil
}

// checkNetOverride checks the status of teleproxy intercept by doing the
// equivalent of curl http://teleproxy/api/tables/. It's okay to create a new
// client each time because we don't want to reuse connections.
func checkNetOverride() error {
	client := http.Client{Timeout: 3 * time.Second}
	res, err := client.Get(fmt.Sprintf(
		"http://teleproxy%d.cachebust.telepresence.io/api/tables",
		time.Now().Unix(),
	))
	if err != nil {
		return err
	}
	_, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


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

// fixrsdp copies the existing RSDP into the EBDA region in low mem.
//
// This is because LinuxBoot tends to be EFI booted, which places the RSDP
// outside of the low 1M or the EBDA. If Linuxboot legacy boots the following
// operating systems, such as with kexec, they may not have a way to find the
// RSDP afterwards.  All u-root commands that open /dev/mem should also flock
// it to ensure safe, sequential access.
package main

import (
	"bytes"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/acpi"
	"github.com/u-root/u-root/pkg/boot/ebda"
)

func main() {
	// Find the RSDP.
	r, err := acpi.GetRSDP()
	if err != nil {
		log.Fatalf("Unable to find system RSDP, got: %v", err)
	}

	rData := r.AllData()
	rLen := len(rData)

	base := r.RSDPAddr()
	// Check if ACPI rsdp is already in low memory
	if base >= 0xe0000 && base+uint64(rLen) < 0xffff0 {
		log.Printf("RSDP is already in low memory at %#X, no need to fix.", base)
		return
	}

	// Find the EBDA
	f, err := os.OpenFile("/dev/mem", os.O_RDWR, 0)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	e, err := ebda.ReadEBDA(f)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("EBDA starts at %#X, length %#X bytes", e.BaseOffset, e.Length)

	// Scan low 1K of EBDA for an empty spot that is 16 byte aligned
	emptyStart := 0
	for i := 16; i < (1024-rLen) && i < (int(e.Length)-rLen); i += 16 {
		// Check if there's an empty spot to put the RSDP table.
		if bytes.Equal(e.Data[i:i+rLen], make([]byte, rLen)) {
			emptyStart = i
			log.Printf("Found empty space at %#X offset into EBDA, will copy RSDP there", emptyStart)
			break
		}
	}

	if emptyStart == 0 {
		log.Fatal("Unable to find empty space to put RSDP")
	}

	copy(e.Data[emptyStart:emptyStart+rLen], rData)

	if err = ebda.WriteEBDA(e, f); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>fix clean-code<commit_after>// Copyright 2019-2021 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// fixrsdp copies the existing RSDP into the EBDA region in low mem.
//
// This is because LinuxBoot tends to be EFI booted, which places the RSDP
// outside of the low 1M or the EBDA. If Linuxboot legacy boots the following
// operating systems, such as with kexec, they may not have a way to find the
// RSDP afterwards.  All u-root commands that open /dev/mem should also flock
// it to ensure safe, sequential access.
package main

import (
	"bytes"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/acpi"
	"github.com/u-root/u-root/pkg/boot/ebda"
)

func main() {
	// Find the RSDP.
	r, err := acpi.GetRSDP()
	if err != nil {
		log.Fatalf("Unable to find system RSDP, got: %v", err)
	}

	rData := r.AllData()
	rLen := len(rData)

	base := r.RSDPAddr()
	// Check if ACPI rsdp is already in low memory
	if base >= 0xe0000 && base+int64(rLen) < 0xffff0 {
		log.Printf("RSDP is already in low memory at %#X, no need to fix.", base)
		return
	}

	// Find the EBDA
	f, err := os.OpenFile("/dev/mem", os.O_RDWR, 0)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	e, err := ebda.ReadEBDA(f)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("EBDA starts at %#X, length %#X bytes", e.BaseOffset, e.Length)

	// Scan low 1K of EBDA for an empty spot that is 16 byte aligned
	emptyStart := 0
	for i := 16; i < (1024-rLen) && i < (int(e.Length)-rLen); i += 16 {
		// Check if there's an empty spot to put the RSDP table.
		if bytes.Equal(e.Data[i:i+rLen], make([]byte, rLen)) {
			emptyStart = i
			log.Printf("Found empty space at %#X offset into EBDA, will copy RSDP there", emptyStart)
			break
		}
	}

	if emptyStart == 0 {
		log.Fatal("Unable to find empty space to put RSDP")
	}

	copy(e.Data[emptyStart:emptyStart+rLen], rData)

	if err = ebda.WriteEBDA(e, f); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"fmt"
	"golang.org/x/net/ipv4"
)

type TCP_Client_Manager struct {
	//reader *IP_Reader
	readBuffer map[uint16](map[string](chan []byte)) // dst protocol, ip
}

func New_TCP_Client_Manager() (*TCP_Client_Manager, error) {
	/*nr, err := NewNetwork_Reader() // TODO: create a global var for the network reader
	if err != nil {
		return nil, err
	}

	ipr, err := nr.NewIP_Reader("*", TCP_PROTO)
	if err != nil {
		return nil, err
	}*/

	cm := &TCP_Client_Manager{
		//reader: ipr,
		readBuffer: make(map[uint16](map[string](chan []byte))),
	}

	//go cm.readAll()

	return cm, nil
}

/*func (cm *TCP_Client_Manager) readAll() {
	for {
		ip, _, payload, err := cm.reader.ReadFrom()
		if err != nil {
			fmt.Println("TCP readAll error", err)
			continue
		}
		fmt.Println("Recved IP:", ip, "with payload", payload)
	}
}*/

type TCP_Client struct {
	manager   *TCP_Client_Manager
	//writer *IP_Writer
	conn   *ipv4.RawConn
	ipAddress string // destination ip address
	srcIP     string // src ip address
	src, dst  uint16 // ports
}

const (
	TCP_FIN = 0x01
	TCP_SYN = 0x02
	TCP_RSH = 0x04
	TCP_PSH = 0x08
	TCP_ACK = 0x10
	TCP_URG = 0x20
	TCP_ECE = 0x40
	TCP_CWR = 0x80
)

func (x *TCP_Client_Manager) New_TCP_Client(src, dst uint16, dstIP string) (*TCP_Client, error) {
	/*write, err := NewIP_Writer(dstIP, TCP_PROTO)
	if err != nil {
		return nil, err
	}*/

	p, err := net.ListenPacket("ip4:6", dstIP)
	if err != nil {
		fmt.Println(err)
		return nil, err
	}

	r, err := ipv4.NewRawConn(p)
	if err != nil {
		fmt.Println(err)
		return nil, err
	}

	return &TCP_Client{
		src: src,
		dst: dst,
		ipAddress: dstIP,
		srcIP: "127.0.0.1",
		manager: x,
		//writer: write,
		conn: r,
	}, nil
}

func (c *TCP_Client) Connect() error {
	// SYN
	initSeqNumber := uint32(3425) // Can be any number between 0 and 2^32 inclusive, TODO this needs to be changed later, it can't be predictable.
	initAckNumber := uint32(0)    // Always 0 at start
	window := uint16(43690)       // TODO calc using http://ithitman.blogspot.com/2013/02/understanding-tcp-window-window-scaling.html
	headerLen := 40
	SYN := []byte{
		(byte)(c.src >> 8), (byte)(c.src), // Source port in byte
		(byte)(c.dst >> 8), (byte)(c.dst), // Destination port in byte slice
		(byte)(initSeqNumber >> 24), (byte)(initSeqNumber >> 16), (byte)(initSeqNumber >> 8), (byte)(initSeqNumber),
		(byte)(initAckNumber >> 24), (byte)(initAckNumber >> 16), (byte)(initAckNumber >> 8), (byte)(initAckNumber),
		(byte)(
			(headerLen/4) << 4, // Size of header in 32 bit chunks. It is always 5 unless options are used. This is also the data offset.
			// bits 5-7 inclusive are reserved, always 0
			// bit 8 is flag 0(NS flag), set to 0 here because only SYN
		),
		(byte)(// Flags 1-8 inclusive
			// CWR
			// ECE
			// URG
			// ACK
			// PSH
			// RST
			0 | TCP_SYN, //SYN
			// FIN
		),
		(byte)(window >> 8), (byte)(window),
		0, 0, // TODO calc checksum, right now just set to 0
		0, 0, // URG pointer, only matters where URG flag is set.
		0x02, 0x04, 0xff, 0xd7, 0x04, 0x02, 0x08, 0x0a, 0x02, 0x64, 0x80, 0x8b, 0, 0, 0, 0, 0x01, 0x03, 0x03, 0x07,
	}
	checksum := checksum(append(append(append(SYN, net.ParseIP(/*c.writer.src*/c.srcIP)...), net.ParseIP(/*c.writer.dst*/c.ipAddress)...), []byte{byte(TCP_PROTO >> 8), byte(TCP_PROTO), byte(headerLen >> 8), byte(headerLen)}...))
	SYN[16] = byte(checksum >> 8)
	SYN[17] = byte(checksum)

	//c.writer.WriteTo(SYN)
	err := c.conn.WriteTo(&ipv4.Header{
		Version:  ipv4.Version, // protocol version
		Len:      20, // header length
		TOS:      0, // type-of-service (0 is everything normal)
		TotalLen: len(SYN) + 20, // packet total length (octets)
		ID:       0, // identification
		Flags:    ipv4.DontFragment, // flags
		FragOff:  0, // fragment offset
		TTL:      64, // time-to-live (maximum lifespan in seconds)
		Protocol: 6, // next protocol (17 is UDP)
		Checksum: 0, // checksum (apparently autocomputed)
		//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
		Dst: net.ParseIP(c.ipAddress), // destination address
		//Options                         // options, extension headers
	}, SYN, nil)
	fmt.Println("Sent data")
	if err != nil {
		fmt.Println("Raw conn send", err)
	}

	// TODO Wait for SYN + ACK, send back ACK

	return nil
}

func (c *TCP_Client) Send(data []byte) error {
	return nil // TODO: implement TCP_Client send
}

func (c *TCP_Client) Close() error {
	return nil // TODO: free manager read buffer
}
<commit_msg>Random changes with 0x0<commit_after>package main

import (
	"net"
	"fmt"
	"golang.org/x/net/ipv4"
)

type TCP_Client_Manager struct {
	//reader *IP_Reader
	readBuffer map[uint16](map[string](chan []byte)) // dst protocol, ip
}

func New_TCP_Client_Manager() (*TCP_Client_Manager, error) {
	/*nr, err := NewNetwork_Reader() // TODO: create a global var for the network reader
	if err != nil {
		return nil, err
	}

	ipr, err := nr.NewIP_Reader("*", TCP_PROTO)
	if err != nil {
		return nil, err
	}*/

	cm := &TCP_Client_Manager{
		//reader: ipr,
		readBuffer: make(map[uint16](map[string](chan []byte))),
	}

	//go cm.readAll()

	return cm, nil
}

/*func (cm *TCP_Client_Manager) readAll() {
	for {
		ip, _, payload, err := cm.reader.ReadFrom()
		if err != nil {
			fmt.Println("TCP readAll error", err)
			continue
		}
		fmt.Println("Recved IP:", ip, "with payload", payload)
	}
}*/

type TCP_Client struct {
	manager   *TCP_Client_Manager
	//writer *IP_Writer
	conn   *ipv4.RawConn
	ipAddress string // destination ip address
	srcIP     string // src ip address
	src, dst  uint16 // ports
}

const (
	TCP_FIN = 0x01
	TCP_SYN = 0x02
	TCP_RSH = 0x04
	TCP_PSH = 0x08
	TCP_ACK = 0x10
	TCP_URG = 0x20
	TCP_ECE = 0x40
	TCP_CWR = 0x80
)

func (x *TCP_Client_Manager) New_TCP_Client(src, dst uint16, dstIP string) (*TCP_Client, error) {
	/*write, err := NewIP_Writer(dstIP, TCP_PROTO)
	if err != nil {
		return nil, err
	}*/

	p, err := net.ListenPacket("ip4:6", dstIP)
	if err != nil {
		fmt.Println(err)
		return nil, err
	}

	r, err := ipv4.NewRawConn(p)
	if err != nil {
		fmt.Println(err)
		return nil, err
	}

	return &TCP_Client{
		src: src,
		dst: dst,
		ipAddress: dstIP,
		srcIP: "127.0.0.1",
		manager: x,
		//writer: write,
		conn: r,
	}, nil
}

func (c *TCP_Client) Connect() error {
	// SYN
	initSeqNumber := uint32(3425) // Can be any number between 0 and 2^32 inclusive, TODO this needs to be changed later, it can't be predictable.
	initAckNumber := uint32(0)    // Always 0 at start
	window := uint16(43690)       // TODO calc using http://ithitman.blogspot.com/2013/02/understanding-tcp-window-window-scaling.html
	headerLen := 40
	SYN := []byte{
		(byte)(c.src >> 8), (byte)(c.src), // Source port in byte
		(byte)(c.dst >> 8), (byte)(c.dst), // Destination port in byte slice
		(byte)(initSeqNumber >> 24), (byte)(initSeqNumber >> 16), (byte)(initSeqNumber >> 8), (byte)(initSeqNumber),
		(byte)(initAckNumber >> 24), (byte)(initAckNumber >> 16), (byte)(initAckNumber >> 8), (byte)(initAckNumber),
		(byte)(
			(headerLen/4) << 4, // Size of header in 32 bit chunks. It is always 5 unless options are used. This is also the data offset.
			// bits 5-7 inclusive are reserved, always 0
			// bit 8 is flag 0(NS flag), set to 0 here because only SYN
		),
		(byte)(// Flags 1-8 inclusive
			// CWR
			// ECE
			// URG
			// ACK
			// PSH
			// RST
			0 | TCP_SYN, //SYN
			// FIN
		),
		(byte)(window >> 8), (byte)(window),
		0, 0, // TODO calc checksum, right now just set to 0
		0, 0, // URG pointer, only matters where URG flag is set.
		0x02, 0x04, 0xff, 0xd7, 0x04, 0x02, 0x08, 0x0a, 0x02, 0x64, 0x80, 0x8b, 0x0, 0x0, 0x0, 0x0, 0x01, 0x03, 0x03, 0x07,
	}
	checksum := checksum(append(append(append(SYN, net.ParseIP(/*c.writer.src*/c.srcIP)...), net.ParseIP(/*c.writer.dst*/c.ipAddress)...), []byte{byte(TCP_PROTO >> 8), byte(TCP_PROTO), byte(headerLen >> 8), byte(headerLen)}...))
	SYN[16] = byte(checksum >> 8)
	SYN[17] = byte(checksum)

	//c.writer.WriteTo(SYN)
	err := c.conn.WriteTo(&ipv4.Header{
		Version:  ipv4.Version, // protocol version
		Len:      20, // header length
		TOS:      0, // type-of-service (0 is everything normal)
		TotalLen: len(SYN) + 20, // packet total length (octets)
		ID:       0, // identification
		Flags:    ipv4.DontFragment, // flags
		FragOff:  0, // fragment offset
		TTL:      64, // time-to-live (maximum lifespan in seconds)
		Protocol: 6, // next protocol (17 is UDP)
		Checksum: 0, // checksum (apparently autocomputed)
		//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
		Dst: net.ParseIP(c.ipAddress), // destination address
		//Options                         // options, extension headers
	}, SYN, nil)
	fmt.Println("Sent data")
	if err != nil {
		fmt.Println("Raw conn send", err)
	}

	// TODO Wait for SYN + ACK, send back ACK

	return nil
}

func (c *TCP_Client) Send(data []byte) error {
	return nil // TODO: implement TCP_Client send
}

func (c *TCP_Client) Close() error {
	return nil // TODO: free manager read buffer
}
<|endoftext|>


<commit_before>package cidre

import (
	"crypto/sha1"
	"fmt"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// SessionConfig is a configuration object for the SessionMiddleware
type SessionConfig struct {
	// default: gossessionid
	CookieName   string
	CookieDomain string
	// default: false
	CookieSecure  bool
	CookiePath    string
	CookieExpires time.Duration
	// A term used to authenticate the cookie value using HMAC
	Secret string
	// default: "cidre.MemorySessionStore"
	SessionStore string
	// default: 30m
	GcInterval time.Duration
	// default: 30m
	LifeTime time.Duration
}

// Returns a SessionConfig object that has default values set.
// If an 'init' function object argument is not nil, this function
// will call the function with the SessionConfig object.
func DefaultSessionConfig(init ...func(*SessionConfig)) *SessionConfig {
	self := &SessionConfig{
		CookieName:    "gosessionid",
		CookieDomain:  "",
		CookieSecure:  false,
		CookiePath:    "",
		CookieExpires: 0,
		Secret:        "",
		SessionStore:  "cidre.MemorySessionStore",
		GcInterval:    time.Minute * 30,
		LifeTime:      time.Minute * 30,
	}
	if len(init) > 0 {
		init[0](self)
	}
	return self
}

// Middleware for session management.
type SessionMiddleware struct {
	app    *App
	Config *SessionConfig
	Store  SessionStore
}

// Returns a new SessionMiddleware object.
func NewSessionMiddleware(app *App, config *SessionConfig, storeConfig interface{}) *SessionMiddleware {
	self := &SessionMiddleware{app: app, Config: config}
	if len(self.Config.Secret) == 0 {
		panic("Session secret must not be empty.")
	}
	DynamicObjectFactory.Register(MemorySessionStore{})
	store, _ := DynamicObjectFactory.New(self.Config.SessionStore).(SessionStore)
	self.Store = store
	self.Store.Init(self, storeConfig)

	app.Hooks.Add("start_server", func(w http.ResponseWriter, r *http.Request, data interface{}) {
		time.AfterFunc(self.Config.GcInterval, self.Gc)
	})

	return self
}

func (self *SessionMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx := RequestContext(r)
	if !ctx.IsDynamicRoute() {
		ctx.MiddlewareChain.DoNext(w, r)
	} else {
		if !strings.HasPrefix(r.URL.Path, self.Config.CookiePath) {
			return
		}
		func() {
			self.Store.Lock()
			defer self.Store.Unlock()
			signedString, _ := r.Cookie(self.Config.CookieName)
			var session *Session
			if signedString != nil {
				sessionId, err := ValidateSignedString(signedString.Value, self.Config.Secret)
				if err != nil {
					panic(err)
				}
				session = self.Store.Load(sessionId)
			} else {
				session = self.Store.NewSession()
			}
			if session != nil {
				ctx.Session = session
				session.UpdateLastAccessTime()
			}
		}()

		w.(ResponseWriter).Hooks().Add("after_write_header", func(w http.ResponseWriter, rnil *http.Request, datanil interface{}) {
			if strings.Index(r.URL.Path, self.Config.CookiePath) != 0 {
				return
			}
			self.Store.Lock()
			defer self.Store.Unlock()
			cookie := &http.Cookie{
				Domain: self.Config.CookieDomain,
				Secure: self.Config.CookieSecure,
				Path:   self.Config.CookiePath,
			}
			if self.Config.CookieExpires != 0 {
				cookie.Expires = time.Now().Add(self.Config.CookieExpires)
			}
			session := ctx.Session
			if session == nil {
				return
			}
			if session.Killed {
				cookie.MaxAge = -1
				self.Store.Delete(session.Id)
			} else {
				self.Store.Save(session)
			}
			cookie.Name = self.Config.CookieName
			cookie.Value = SignString(session.Id, self.Config.Secret)
			http.SetCookie(w, cookie)
		})

		ctx.MiddlewareChain.DoNext(w, r)
	}

}

func (self *SessionMiddleware) Gc() {
	self.Store.Lock()
	defer self.Store.Unlock()
	self.app.Logger(LogLevelDebug, "Session Gc")
	self.Store.Gc()
	time.AfterFunc(self.Config.GcInterval, self.Gc)
}

// Session value container.
type Session struct {
	Dict
	Killed         bool
	Id             string
	LastAccessTime time.Time
}

const FlashKey = "_flash"

func NewSession(id string) *Session {
	self := &Session{
		Dict:   NewDict(),
		Killed: false, Id: id,
		LastAccessTime: time.Now()}
	self.Set(FlashKey, make(map[string][]string))
	return self
}

func (self *Session) UpdateLastAccessTime() {
	self.LastAccessTime = time.Now()
}

func (self *Session) Kill() {
	self.Killed = true
}

// Adds a flash message to the session
func (self *Session) AddFlash(category string, message string) {
	flash := self.Get(FlashKey).(map[string][]string)
	if _, ok := flash[category]; !ok {
		flash[category] = make([]string, 0, 10)
	}
	flash[category] = append(flash[category], message)
}

// Returns a flash message associated with the given category.
func (self *Session) Flash(category string) []string {
	flash := self.Get(FlashKey).(map[string][]string)
	v, ok := flash[category]
	delete(flash, category)
	if !ok {
		return make([]string, 0, 10)
	}
	return v
}

// Returns a list of flash messages from the session.
//
//     session.AddFlash("info", "info message1")
//     session.AddFlash("info", "info message2")
//     session.AddFlash("error", "error message")
//     messages := session.Flashes()
//     // -> {"info":["info message1", "info message2"], "error":["error message"]}
func (self *Session) Flashes() map[string][]string {
	flash := self.Get(FlashKey).(map[string][]string)
	self.Set(FlashKey, make(map[string][]string))
	return flash
}

// SessionStore is an interface for custom session stores.
// See the MemorySessionStore for examples.
type SessionStore interface {
	Lock()
	Unlock()
	Init(*SessionMiddleware, interface{})
	Exists(string) bool
	NewSession() *Session
	Save(*Session)
	Load(string) *Session
	Delete(string)
	Gc()
	Count() int
}

type MemorySessionStore struct {
	sync.Mutex
	middleware *SessionMiddleware
	store      map[string]*Session
}

func (self *MemorySessionStore) Init(middleware *SessionMiddleware, cfg interface{}) {
	self.middleware = middleware
	self.store = make(map[string]*Session, 30)
}

func (self *MemorySessionStore) NewSessionId() string {
	for true {
		now := time.Now().Unix()
		random := strconv.Itoa(rand.Int())
		sessionId := fmt.Sprintf("%x", sha1.Sum([]byte(fmt.Sprintf("%v%v%v", now, random, self.middleware.Config.Secret))))
		if !self.Exists(sessionId) {
			return sessionId
		}
	}
	return ""
}

func (self *MemorySessionStore) Exists(sessionId string) bool {
	_, ok := self.store[sessionId]
	return ok
}

func (self *MemorySessionStore) NewSession() *Session {
	session := NewSession(self.NewSessionId())
	self.store[session.Id] = session
	return session
}

func (self *MemorySessionStore) Save(*Session) { /* Nothing to do */ }

func (self *MemorySessionStore) Load(sessionId string) *Session {
	session, ok := self.store[sessionId]
	if ok {
		return session
	}
	return self.NewSession()
}

func (self *MemorySessionStore) Delete(sessionId string) {
	delete(self.store, sessionId)
}

func (self *MemorySessionStore) Count() int {
	return len(self.store)
}

func (self *MemorySessionStore) Gc() {
	delkeys := make([]string, 0, len(self.store)/10)
	for k, v := range self.store {
		if (time.Now().Sub(v.LastAccessTime)) > self.middleware.Config.LifeTime {
			delkeys = append(delkeys, k)
		}
	}
	for _, key := range delkeys {
		self.Delete(key)
	}
}
<commit_msg>improve session cookie handling<commit_after>package cidre

import (
	"crypto/sha1"
	"fmt"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// SessionConfig is a configuration object for the SessionMiddleware
type SessionConfig struct {
	// default: gossessionid
	CookieName   string
	CookieDomain string
	// default: false
	CookieSecure  bool
	CookiePath    string
	CookieExpires time.Duration
	// A term used to authenticate the cookie value using HMAC
	Secret string
	// default: "cidre.MemorySessionStore"
	SessionStore string
	// default: 30m
	GcInterval time.Duration
	// default: 30m
	LifeTime time.Duration
}

// Returns a SessionConfig object that has default values set.
// If an 'init' function object argument is not nil, this function
// will call the function with the SessionConfig object.
func DefaultSessionConfig(init ...func(*SessionConfig)) *SessionConfig {
	self := &SessionConfig{
		CookieName:    "gosessionid",
		CookieDomain:  "",
		CookieSecure:  false,
		CookiePath:    "",
		CookieExpires: 0,
		Secret:        "",
		SessionStore:  "cidre.MemorySessionStore",
		GcInterval:    time.Minute * 30,
		LifeTime:      time.Minute * 30,
	}
	if len(init) > 0 {
		init[0](self)
	}
	return self
}

// Middleware for session management.
type SessionMiddleware struct {
	app    *App
	Config *SessionConfig
	Store  SessionStore
}

// Returns a new SessionMiddleware object.
func NewSessionMiddleware(app *App, config *SessionConfig, storeConfig interface{}) *SessionMiddleware {
	self := &SessionMiddleware{app: app, Config: config}
	if len(self.Config.Secret) == 0 {
		panic("Session secret must not be empty.")
	}
	DynamicObjectFactory.Register(MemorySessionStore{})
	store, _ := DynamicObjectFactory.New(self.Config.SessionStore).(SessionStore)
	self.Store = store
	self.Store.Init(self, storeConfig)

	app.Hooks.Add("start_server", func(w http.ResponseWriter, r *http.Request, data interface{}) {
		time.AfterFunc(self.Config.GcInterval, self.Gc)
	})

	return self
}

func (self *SessionMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx := RequestContext(r)
	if !ctx.IsDynamicRoute() {
		ctx.MiddlewareChain.DoNext(w, r)
	} else {
		if !strings.HasPrefix(r.URL.Path, self.Config.CookiePath) {
			return
		}
		func() {
			self.Store.Lock()
			defer self.Store.Unlock()
			signedString, _ := r.Cookie(self.Config.CookieName)
			var session *Session
			if signedString != nil {
				sessionId, err := ValidateSignedString(signedString.Value, self.Config.Secret)
				if err != nil {
					panic(err)
				}
				session = self.Store.Load(sessionId)
			} else {
				session = self.Store.NewSession()
			}
			if session != nil {
				ctx.Session = session
				session.UpdateLastAccessTime()
			}
		}()

		w.(ResponseWriter).Hooks().Add("after_write_header", func(w http.ResponseWriter, rnil *http.Request, datanil interface{}) {
			if strings.Index(r.URL.Path, self.Config.CookiePath) != 0 {
				return
			}
			self.Store.Lock()
			defer self.Store.Unlock()
            domain := self.Config.CookieDomain
            if len(domain) == 0 {
              domain = strings.Split(r.Host,":")[0]
            }
			cookie := &http.Cookie{
				Domain: domain,
				Secure: self.Config.CookieSecure,
				Path:   self.Config.CookiePath,
			}
			if self.Config.CookieExpires != 0 {
				cookie.Expires = time.Now().Add(self.Config.CookieExpires)
			}
			session := ctx.Session
			if session == nil {
				return
			}
			if session.Killed {
				cookie.MaxAge = -1
				self.Store.Delete(session.Id)
			} else {
				self.Store.Save(session)
			}
			cookie.Name = self.Config.CookieName
			cookie.Value = SignString(session.Id, self.Config.Secret)
			http.SetCookie(w, cookie)
		})

		ctx.MiddlewareChain.DoNext(w, r)
	}

}

func (self *SessionMiddleware) Gc() {
	self.Store.Lock()
	defer self.Store.Unlock()
	self.app.Logger(LogLevelDebug, "Session Gc")
	self.Store.Gc()
	time.AfterFunc(self.Config.GcInterval, self.Gc)
}

// Session value container.
type Session struct {
	Dict
	Killed         bool
	Id             string
	LastAccessTime time.Time
}

const FlashKey = "_flash"

func NewSession(id string) *Session {
	self := &Session{
		Dict:   NewDict(),
		Killed: false, Id: id,
		LastAccessTime: time.Now()}
	self.Set(FlashKey, make(map[string][]string))
	return self
}

func (self *Session) UpdateLastAccessTime() {
	self.LastAccessTime = time.Now()
}

func (self *Session) Kill() {
	self.Killed = true
}

// Adds a flash message to the session
func (self *Session) AddFlash(category string, message string) {
	flash := self.Get(FlashKey).(map[string][]string)
	if _, ok := flash[category]; !ok {
		flash[category] = make([]string, 0, 10)
	}
	flash[category] = append(flash[category], message)
}

// Returns a flash message associated with the given category.
func (self *Session) Flash(category string) []string {
	flash := self.Get(FlashKey).(map[string][]string)
	v, ok := flash[category]
	delete(flash, category)
	if !ok {
		return make([]string, 0, 10)
	}
	return v
}

// Returns a list of flash messages from the session.
//
//     session.AddFlash("info", "info message1")
//     session.AddFlash("info", "info message2")
//     session.AddFlash("error", "error message")
//     messages := session.Flashes()
//     // -> {"info":["info message1", "info message2"], "error":["error message"]}
func (self *Session) Flashes() map[string][]string {
	flash := self.Get(FlashKey).(map[string][]string)
	self.Set(FlashKey, make(map[string][]string))
	return flash
}

// SessionStore is an interface for custom session stores.
// See the MemorySessionStore for examples.
type SessionStore interface {
	Lock()
	Unlock()
	Init(*SessionMiddleware, interface{})
	Exists(string) bool
	NewSession() *Session
	Save(*Session)
	Load(string) *Session
	Delete(string)
	Gc()
	Count() int
}

type MemorySessionStore struct {
	sync.Mutex
	middleware *SessionMiddleware
	store      map[string]*Session
}

func (self *MemorySessionStore) Init(middleware *SessionMiddleware, cfg interface{}) {
	self.middleware = middleware
	self.store = make(map[string]*Session, 30)
}

func (self *MemorySessionStore) NewSessionId() string {
	for true {
		now := time.Now().Unix()
		random := strconv.Itoa(rand.Int())
		sessionId := fmt.Sprintf("%x", sha1.Sum([]byte(fmt.Sprintf("%v%v%v", now, random, self.middleware.Config.Secret))))
		if !self.Exists(sessionId) {
			return sessionId
		}
	}
	return ""
}

func (self *MemorySessionStore) Exists(sessionId string) bool {
	_, ok := self.store[sessionId]
	return ok
}

func (self *MemorySessionStore) NewSession() *Session {
	session := NewSession(self.NewSessionId())
	self.store[session.Id] = session
	return session
}

func (self *MemorySessionStore) Save(*Session) { /* Nothing to do */ }

func (self *MemorySessionStore) Load(sessionId string) *Session {
	session, ok := self.store[sessionId]
	if ok {
		return session
	}
	return self.NewSession()
}

func (self *MemorySessionStore) Delete(sessionId string) {
	delete(self.store, sessionId)
}

func (self *MemorySessionStore) Count() int {
	return len(self.store)
}

func (self *MemorySessionStore) Gc() {
	delkeys := make([]string, 0, len(self.store)/10)
	for k, v := range self.store {
		if (time.Now().Sub(v.LastAccessTime)) > self.middleware.Config.LifeTime {
			delkeys = append(delkeys, k)
		}
	}
	for _, key := range delkeys {
		self.Delete(key)
	}
}
<|endoftext|>


<commit_before>package yamux

import (
	"fmt"
	"io"
	"math"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Session is used to wrap a reliable ordered connection and to
// multiplex it into multiple streams.
type Session struct {
	// remoteGoAway indicates the remote side does
	// not want futher connections. Must be first for alignment.
	remoteGoAway int32

	// localGoAway indicates that we should stop
	// accepting futher connections. Must be first for alignment.
	localGoAway int32

	// client is true if we are a client size connection
	client bool

	// config holds our configuration
	config *Config

	// conn is the underlying connection
	conn io.ReadWriteCloser

	// pings is used to track inflight pings
	pings    map[uint32]chan struct{}
	pingID   uint32
	pingLock sync.Mutex

	// nextStreamID is the next stream we should
	// send. This depends if we are a client/server.
	nextStreamID uint32

	// streams maps a stream id to a stream
	streams    map[uint32]*Stream
	streamLock sync.Mutex

	// acceptCh is used to pass ready streams to the client
	acceptCh chan *Stream

	// sendCh is used to mark a stream as ready to send,
	// or to send a header out directly.
	sendCh chan sendReady

	// shutdown is used to safely close a session
	shutdown     bool
	shutdownErr  error
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// sendReady is used to either mark a stream as ready
// or to directly send a header
type sendReady struct {
	Hdr  []byte
	Body io.Reader
	Err  chan error
}

// newSession is used to construct a new session
func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := &Session{
		client:     client,
		config:     config,
		conn:       conn,
		pings:      make(map[uint32]chan struct{}),
		streams:    make(map[uint32]*Stream),
		acceptCh:   make(chan *Stream, config.AcceptBacklog),
		sendCh:     make(chan sendReady, 64),
		shutdownCh: make(chan struct{}),
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recv()
	go s.send()
	if config.EnableKeepAlive {
		go s.keepalive()
	}
	return s
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.shutdownCh:
		return true
	default:
		return false
	}
}

// Open is used to create a new stream
func (s *Session) Open() (*Stream, error) {
	if s.IsClosed() {
		return nil, ErrSessionShutdown
	}
	if atomic.LoadInt32(&s.remoteGoAway) == 1 {
		return nil, ErrRemoteGoAway
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if we've exhaused the streams
	id := s.nextStreamID
	if id >= math.MaxUint32-1 {
		return nil, ErrStreamsExhausted
	}
	s.nextStreamID += 2

	// Register the stream
	stream := newStream(s, id, streamInit)
	s.streams[id] = stream

	// Send the window update to create
	return stream, stream.sendWindowUpdate()
}

// Accept is used to block until the next available stream
// is ready to be accepted.
func (s *Session) Accept() (net.Conn, error) {
	return s.AcceptStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.acceptCh:
		return stream, nil
	case <-s.shutdownCh:
		return nil, s.shutdownErr
	}
}

// Close is used to close the session and all streams.
// Attempts to send a GoAway before closing the connection.
func (s *Session) Close() error {
	s.shutdownLock.Lock()
	defer s.shutdownLock.Unlock()

	if s.shutdown {
		return nil
	}
	s.shutdown = true
	if s.shutdownErr == nil {
		s.shutdownErr = ErrSessionShutdown
	}
	close(s.shutdownCh)
	s.conn.Close()

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	for _, stream := range s.streams {
		stream.forceClose()
	}
	return nil
}

// exitErr is used to handle an error that is causing the
// session to terminate.
func (s *Session) exitErr(err error) {
	s.shutdownErr = err
	s.Close()
}

// GoAway can be used to prevent accepting further
// connections. It does not close the underlying conn.
func (s *Session) GoAway() error {
	return s.waitForSend(s.goAway(goAwayNormal), nil)
}

// goAway is used to send a goAway message
func (s *Session) goAway(reason uint32) header {
	atomic.SwapInt32(&s.localGoAway, 1)
	hdr := header(make([]byte, headerSize))
	hdr.encode(typeGoAway, 0, 0, reason)
	return hdr
}

// Ping is used to measure the RTT response time
func (s *Session) Ping() (time.Duration, error) {
	// Get a channel for the ping
	ch := make(chan struct{})

	// Get a new ping id, mark as pending
	s.pingLock.Lock()
	id := s.pingID
	s.pingID++
	s.pings[id] = ch
	s.pingLock.Unlock()

	// Send the ping request
	hdr := header(make([]byte, headerSize))
	hdr.encode(typePing, flagSYN, 0, id)
	if err := s.waitForSend(hdr, nil); err != nil {
		return 0, err
	}

	// Wait for a response
	start := time.Now()
	select {
	case <-ch:
	case <-s.shutdownCh:
		return 0, ErrSessionShutdown
	}

	// Compute the RTT
	return time.Now().Sub(start), nil
}

// keepalive is a long running goroutine that periodically does
// a ping to keep the connection alive.
func (s *Session) keepalive() {
	for {
		select {
		case <-time.After(s.config.KeepAliveInterval):
			s.Ping()
		case <-s.shutdownCh:
			return
		}
	}
}

// waitForSend waits to send a header, checking for a potential shutdown
func (s *Session) waitForSend(hdr header, body io.Reader) error {
	errCh := make(chan error, 1)
	ready := sendReady{Hdr: hdr, Body: body, Err: errCh}
	select {
	case s.sendCh <- ready:
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
	select {
	case err := <-errCh:
		return err
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// sendNoWait does a send without waiting
func (s *Session) sendNoWait(hdr header) error {
	select {
	case s.sendCh <- sendReady{Hdr: hdr}:
		return nil
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// send is a long running goroutine that sends data
func (s *Session) send() {
	for !s.IsClosed() {
		select {
		case ready := <-s.sendCh:
			// Send a header if ready
			if ready.Hdr != nil {
				sent := 0
				for sent < len(ready.Hdr) {
					n, err := s.conn.Write(ready.Hdr[sent:])
					if err != nil {
						asyncSendErr(ready.Err, err)
						s.exitErr(err)
						return
					}
					sent += n
				}
			}

			// Send data from a body if given
			if ready.Body != nil {
				_, err := io.Copy(s.conn, ready.Body)
				if err != nil {
					asyncSendErr(ready.Err, err)
					s.exitErr(err)
					return
				}
			}

			// No error, successful send
			asyncSendErr(ready.Err, nil)
		case <-s.shutdownCh:
			return
		}
	}
}

// recv is a long running goroutine that accepts new data
func (s *Session) recv() {
	hdr := header(make([]byte, headerSize))
	var handler func(header) error
	for !s.IsClosed() {
		// Read the header
		if _, err := io.ReadFull(s.conn, hdr); err != nil {
			s.exitErr(err)
			return
		}

		// Verify the version
		if hdr.Version() != protoVersion {
			s.exitErr(ErrInvalidVersion)
			return
		}

		// Switch on the type
		switch hdr.MsgType() {
		case typeData:
			handler = s.handleStreamMessage
		case typeWindowUpdate:
			handler = s.handleStreamMessage
		case typeGoAway:
			handler = s.handleGoAway
		case typePing:
			handler = s.handlePing
		default:
			s.exitErr(ErrInvalidMsgType)
			return
		}

		// Invoke the handler
		if err := handler(hdr); err != nil {
			s.exitErr(err)
			return
		}
	}
}

// handleStreamMessage handles either a data or window update frame
func (s *Session) handleStreamMessage(hdr header) error {
	// Check for a new stream creation
	id := hdr.StreamID()
	flags := hdr.Flags()
	if flags&flagSYN == flagSYN {
		if err := s.incomingStream(id); err != nil {
			return err
		}
	}

	// Get the stream
	s.streamLock.Lock()
	stream := s.streams[id]
	s.streamLock.Unlock()

	// Make sure we have a stream
	if stream == nil {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return ErrMissingStream
	}

	// Check if this is a window update
	if hdr.MsgType() == typeWindowUpdate {
		if err := stream.incrSendWindow(hdr, flags); err != nil {
			s.sendNoWait(s.goAway(goAwayProtoErr))
			return err
		}
		return nil
	}

	// Read the new data
	if err := stream.readData(hdr, flags, s.conn); err != nil {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return err
	}
	return nil
}

// handlePing is invokde for a typePing frame
func (s *Session) handlePing(hdr header) error {
	flags := hdr.Flags()
	pingID := hdr.Length()

	// Check if this is a query, respond back
	if flags&flagSYN == flagSYN {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typePing, flagACK, 0, pingID)
		s.sendNoWait(hdr)
		return nil
	}

	// Handle a response
	s.pingLock.Lock()
	ch := s.pings[pingID]
	if ch != nil {
		delete(s.pings, pingID)
		close(ch)
	}
	s.pingLock.Unlock()
	return nil
}

// handleGoAway is invokde for a typeGoAway frame
func (s *Session) handleGoAway(hdr header) error {
	code := hdr.Length()
	switch code {
	case goAwayNormal:
		atomic.SwapInt32(&s.remoteGoAway, 1)
	case goAwayProtoErr:
		return fmt.Errorf("yamux protocol error")
	case goAwayInternalErr:
		return fmt.Errorf("remote yamux internal error")
	default:
		return fmt.Errorf("unexpected go away received")
	}
	return nil
}

// incomingStream is used to create a new incoming stream
func (s *Session) incomingStream(id uint32) error {
	// Reject immediately if we are doing a go away
	if atomic.LoadInt32(&s.localGoAway) == 1 {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typeWindowUpdate, flagRST, id, 0)
		return s.waitForSend(hdr, nil)
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if stream already exists
	if _, ok := s.streams[id]; ok {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return ErrDuplicateStream
	}

	// Register the stream
	stream := newStream(s, id, streamSYNReceived)
	s.streams[id] = stream

	// Check if we've exceeded the backlog
	select {
	case s.acceptCh <- stream:
		return nil
	default:
		// Backlog exceeded! RST the stream
		delete(s.streams, id)
		stream.sendHdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(stream.sendHdr)
	}
	return nil
}

// closeStream is used to close a stream once both sides have
// issued a close.
func (s *Session) closeStream(id uint32, withLock bool) {
	if !withLock {
		s.streamLock.Lock()
		defer s.streamLock.Unlock()
	}
	delete(s.streams, id)
}
<commit_msg>Remove useless var<commit_after>package yamux

import (
	"fmt"
	"io"
	"math"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Session is used to wrap a reliable ordered connection and to
// multiplex it into multiple streams.
type Session struct {
	// remoteGoAway indicates the remote side does
	// not want futher connections. Must be first for alignment.
	remoteGoAway int32

	// localGoAway indicates that we should stop
	// accepting futher connections. Must be first for alignment.
	localGoAway int32

	// config holds our configuration
	config *Config

	// conn is the underlying connection
	conn io.ReadWriteCloser

	// pings is used to track inflight pings
	pings    map[uint32]chan struct{}
	pingID   uint32
	pingLock sync.Mutex

	// nextStreamID is the next stream we should
	// send. This depends if we are a client/server.
	nextStreamID uint32

	// streams maps a stream id to a stream
	streams    map[uint32]*Stream
	streamLock sync.Mutex

	// acceptCh is used to pass ready streams to the client
	acceptCh chan *Stream

	// sendCh is used to mark a stream as ready to send,
	// or to send a header out directly.
	sendCh chan sendReady

	// shutdown is used to safely close a session
	shutdown     bool
	shutdownErr  error
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// sendReady is used to either mark a stream as ready
// or to directly send a header
type sendReady struct {
	Hdr  []byte
	Body io.Reader
	Err  chan error
}

// newSession is used to construct a new session
func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := &Session{
		config:     config,
		conn:       conn,
		pings:      make(map[uint32]chan struct{}),
		streams:    make(map[uint32]*Stream),
		acceptCh:   make(chan *Stream, config.AcceptBacklog),
		sendCh:     make(chan sendReady, 64),
		shutdownCh: make(chan struct{}),
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recv()
	go s.send()
	if config.EnableKeepAlive {
		go s.keepalive()
	}
	return s
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.shutdownCh:
		return true
	default:
		return false
	}
}

// Open is used to create a new stream
func (s *Session) Open() (*Stream, error) {
	if s.IsClosed() {
		return nil, ErrSessionShutdown
	}
	if atomic.LoadInt32(&s.remoteGoAway) == 1 {
		return nil, ErrRemoteGoAway
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if we've exhaused the streams
	id := s.nextStreamID
	if id >= math.MaxUint32-1 {
		return nil, ErrStreamsExhausted
	}
	s.nextStreamID += 2

	// Register the stream
	stream := newStream(s, id, streamInit)
	s.streams[id] = stream

	// Send the window update to create
	return stream, stream.sendWindowUpdate()
}

// Accept is used to block until the next available stream
// is ready to be accepted.
func (s *Session) Accept() (net.Conn, error) {
	return s.AcceptStream()
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.acceptCh:
		return stream, nil
	case <-s.shutdownCh:
		return nil, s.shutdownErr
	}
}

// Close is used to close the session and all streams.
// Attempts to send a GoAway before closing the connection.
func (s *Session) Close() error {
	s.shutdownLock.Lock()
	defer s.shutdownLock.Unlock()

	if s.shutdown {
		return nil
	}
	s.shutdown = true
	if s.shutdownErr == nil {
		s.shutdownErr = ErrSessionShutdown
	}
	close(s.shutdownCh)
	s.conn.Close()

	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	for _, stream := range s.streams {
		stream.forceClose()
	}
	return nil
}

// exitErr is used to handle an error that is causing the
// session to terminate.
func (s *Session) exitErr(err error) {
	s.shutdownErr = err
	s.Close()
}

// GoAway can be used to prevent accepting further
// connections. It does not close the underlying conn.
func (s *Session) GoAway() error {
	return s.waitForSend(s.goAway(goAwayNormal), nil)
}

// goAway is used to send a goAway message
func (s *Session) goAway(reason uint32) header {
	atomic.SwapInt32(&s.localGoAway, 1)
	hdr := header(make([]byte, headerSize))
	hdr.encode(typeGoAway, 0, 0, reason)
	return hdr
}

// Ping is used to measure the RTT response time
func (s *Session) Ping() (time.Duration, error) {
	// Get a channel for the ping
	ch := make(chan struct{})

	// Get a new ping id, mark as pending
	s.pingLock.Lock()
	id := s.pingID
	s.pingID++
	s.pings[id] = ch
	s.pingLock.Unlock()

	// Send the ping request
	hdr := header(make([]byte, headerSize))
	hdr.encode(typePing, flagSYN, 0, id)
	if err := s.waitForSend(hdr, nil); err != nil {
		return 0, err
	}

	// Wait for a response
	start := time.Now()
	select {
	case <-ch:
	case <-s.shutdownCh:
		return 0, ErrSessionShutdown
	}

	// Compute the RTT
	return time.Now().Sub(start), nil
}

// keepalive is a long running goroutine that periodically does
// a ping to keep the connection alive.
func (s *Session) keepalive() {
	for {
		select {
		case <-time.After(s.config.KeepAliveInterval):
			s.Ping()
		case <-s.shutdownCh:
			return
		}
	}
}

// waitForSend waits to send a header, checking for a potential shutdown
func (s *Session) waitForSend(hdr header, body io.Reader) error {
	errCh := make(chan error, 1)
	ready := sendReady{Hdr: hdr, Body: body, Err: errCh}
	select {
	case s.sendCh <- ready:
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
	select {
	case err := <-errCh:
		return err
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// sendNoWait does a send without waiting
func (s *Session) sendNoWait(hdr header) error {
	select {
	case s.sendCh <- sendReady{Hdr: hdr}:
		return nil
	case <-s.shutdownCh:
		return ErrSessionShutdown
	}
}

// send is a long running goroutine that sends data
func (s *Session) send() {
	for !s.IsClosed() {
		select {
		case ready := <-s.sendCh:
			// Send a header if ready
			if ready.Hdr != nil {
				sent := 0
				for sent < len(ready.Hdr) {
					n, err := s.conn.Write(ready.Hdr[sent:])
					if err != nil {
						asyncSendErr(ready.Err, err)
						s.exitErr(err)
						return
					}
					sent += n
				}
			}

			// Send data from a body if given
			if ready.Body != nil {
				_, err := io.Copy(s.conn, ready.Body)
				if err != nil {
					asyncSendErr(ready.Err, err)
					s.exitErr(err)
					return
				}
			}

			// No error, successful send
			asyncSendErr(ready.Err, nil)
		case <-s.shutdownCh:
			return
		}
	}
}

// recv is a long running goroutine that accepts new data
func (s *Session) recv() {
	hdr := header(make([]byte, headerSize))
	var handler func(header) error
	for !s.IsClosed() {
		// Read the header
		if _, err := io.ReadFull(s.conn, hdr); err != nil {
			s.exitErr(err)
			return
		}

		// Verify the version
		if hdr.Version() != protoVersion {
			s.exitErr(ErrInvalidVersion)
			return
		}

		// Switch on the type
		switch hdr.MsgType() {
		case typeData:
			handler = s.handleStreamMessage
		case typeWindowUpdate:
			handler = s.handleStreamMessage
		case typeGoAway:
			handler = s.handleGoAway
		case typePing:
			handler = s.handlePing
		default:
			s.exitErr(ErrInvalidMsgType)
			return
		}

		// Invoke the handler
		if err := handler(hdr); err != nil {
			s.exitErr(err)
			return
		}
	}
}

// handleStreamMessage handles either a data or window update frame
func (s *Session) handleStreamMessage(hdr header) error {
	// Check for a new stream creation
	id := hdr.StreamID()
	flags := hdr.Flags()
	if flags&flagSYN == flagSYN {
		if err := s.incomingStream(id); err != nil {
			return err
		}
	}

	// Get the stream
	s.streamLock.Lock()
	stream := s.streams[id]
	s.streamLock.Unlock()

	// Make sure we have a stream
	if stream == nil {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return ErrMissingStream
	}

	// Check if this is a window update
	if hdr.MsgType() == typeWindowUpdate {
		if err := stream.incrSendWindow(hdr, flags); err != nil {
			s.sendNoWait(s.goAway(goAwayProtoErr))
			return err
		}
		return nil
	}

	// Read the new data
	if err := stream.readData(hdr, flags, s.conn); err != nil {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return err
	}
	return nil
}

// handlePing is invokde for a typePing frame
func (s *Session) handlePing(hdr header) error {
	flags := hdr.Flags()
	pingID := hdr.Length()

	// Check if this is a query, respond back
	if flags&flagSYN == flagSYN {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typePing, flagACK, 0, pingID)
		s.sendNoWait(hdr)
		return nil
	}

	// Handle a response
	s.pingLock.Lock()
	ch := s.pings[pingID]
	if ch != nil {
		delete(s.pings, pingID)
		close(ch)
	}
	s.pingLock.Unlock()
	return nil
}

// handleGoAway is invokde for a typeGoAway frame
func (s *Session) handleGoAway(hdr header) error {
	code := hdr.Length()
	switch code {
	case goAwayNormal:
		atomic.SwapInt32(&s.remoteGoAway, 1)
	case goAwayProtoErr:
		return fmt.Errorf("yamux protocol error")
	case goAwayInternalErr:
		return fmt.Errorf("remote yamux internal error")
	default:
		return fmt.Errorf("unexpected go away received")
	}
	return nil
}

// incomingStream is used to create a new incoming stream
func (s *Session) incomingStream(id uint32) error {
	// Reject immediately if we are doing a go away
	if atomic.LoadInt32(&s.localGoAway) == 1 {
		hdr := header(make([]byte, headerSize))
		hdr.encode(typeWindowUpdate, flagRST, id, 0)
		return s.waitForSend(hdr, nil)
	}

	s.streamLock.Lock()
	defer s.streamLock.Unlock()

	// Check if stream already exists
	if _, ok := s.streams[id]; ok {
		s.sendNoWait(s.goAway(goAwayProtoErr))
		return ErrDuplicateStream
	}

	// Register the stream
	stream := newStream(s, id, streamSYNReceived)
	s.streams[id] = stream

	// Check if we've exceeded the backlog
	select {
	case s.acceptCh <- stream:
		return nil
	default:
		// Backlog exceeded! RST the stream
		delete(s.streams, id)
		stream.sendHdr.encode(typeWindowUpdate, flagRST, id, 0)
		s.sendNoWait(stream.sendHdr)
	}
	return nil
}

// closeStream is used to close a stream once both sides have
// issued a close.
func (s *Session) closeStream(id uint32, withLock bool) {
	if !withLock {
		s.streamLock.Lock()
		defer s.streamLock.Unlock()
	}
	delete(s.streams, id)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

const DEFAULT_REGION = "us-east-1"

func SessionNew(config *Config) *s3.S3 {
    creds := credentials.NewStaticCredentials(config.AccessKey, config.SecretKey, "")

    // By default make sure a region is specified
	return s3.New(session.Must(session.NewSession(&aws.Config{ Credentials: creds, Region: aws.String(DEFAULT_REGION) })))
}

func SessionForBucket(svc *s3.S3, bucket string) (*s3.S3, error) {
    if loc, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: &bucket}); err != nil {
        return nil, err
    } else if (loc.LocationConstraint != nil) {
        return s3.New(session.Must(session.NewSession(&svc.Client.Config, &aws.Config{Region: loc.LocationConstraint}))), nil
    }
    return svc, nil
}
<commit_msg>Another cleanup pass on authentication<commit_after>package main

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// DefaultRegion to use for S3 credential creation
const defaultRegion = "us-east-1"

// SessionNew - Read the config for default credentials, if not provided use environment based variables
func SessionNew(config *Config) *s3.S3 {
	// By default make sure a region is specified, this is required for S3 operations
	var sessionConfig = &aws.Config{Region: aws.String(defaultRegion)}

	if config.AccessKey != "" && config.SecretKey != "" {
		sessionConfig.Credentials = credentials.NewStaticCredentials(config.AccessKey, config.SecretKey, "")
	}

	return s3.New(session.Must(session.NewSession(sessionConfig)))
}

// SessionForBucket - For a given S3 bucket, create an approprate session that references the region
// that this bucket is located in
func SessionForBucket(svc *s3.S3, bucket string) (*s3.S3, error) {
	if loc, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: &bucket}); err != nil {
		return nil, err
	} else if loc.LocationConstraint != nil {
		return s3.New(session.Must(session.NewSession(&svc.Client.Config, &aws.Config{Region: loc.LocationConstraint}))), nil
	}
	return svc, nil
}
<|endoftext|>