<commit_before>// Package llog is a generic logging library used by leven labs. The log methods
// come in different severities: Debug, Info, Warn, Error, and Fatal.
//
// The log methods take in a string describing the error, and a set of key/value
// pairs giving the specific context around the error. The string is intended to
// always be the same no matter what, while the key/value pairs give information
// like which userID the error happened to, or any other relevant dynamic
// information.
//
// By default logs will be output to Stdout, without a timestamp attached to
// them, and only showing entries of level Info or above. All of these can be
// configured.
//
// All public functions in this package are thread-safe and can be called at any
// time. The public variables in this package are NOT thread-safe and should
// only be modified before any logging takes place
//
// Examples:
//
//	Info("Something important has occured")
//	Error("Could not open file", llog.KV{"filename": filename, "err": err})
//
package llog

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Out is the io.Writer all log entries will be written to. It can be changed to
// anything you like, but the change should happen before any logging occurs. If
// an error occurs while writing to Out the entry will be written to Stdout
// instead
var Out io.Writer = os.Stdout
var defaultOut = os.Stdout

// DisplayTimestamp determines whether or not a timestamp is displayed in the
// log messages. By default one is not displayed. This can be changed by it
// should only be changed before any logging occurs
var DisplayTimestamp bool

// Level describes the severity of a particular log message
type Level int

// All defined log levels
const (
	DebugLevel Level = iota
	InfoLevel
	WarnLevel
	ErrorLevel
	FatalLevel
)

func (l Level) String() string {
	switch l {
	case DebugLevel:
		return "DEBUG"
	case InfoLevel:
		return "INFO"
	case WarnLevel:
		return "WARN"
	case ErrorLevel:
		return "ERROR"
	case FatalLevel:
		return "FATAL"
	}
	return "unknown level"
}

var currLevel = InfoLevel
var currLevelLock sync.RWMutex

// GetLevel returns the current log level
func GetLevel() Level {
	currLevelLock.RLock()
	defer currLevelLock.RUnlock()
	return currLevel
}

// SetLevel sets the current minimum log level which will be written to Out
func SetLevel(l Level) {
	currLevelLock.Lock()
	defer currLevelLock.Unlock()
	currLevel = l
}

// SetLevelFromString attempts to interpret the given string as a log level and
// sets the current log level to that. If the string can't be interpreted an
// error is returned and the log level remains what it was
func SetLevelFromString(ls string) error {
	switch strings.ToUpper(ls) {
	case "DEBUG":
		SetLevel(DebugLevel)
	case "INFO":
		SetLevel(InfoLevel)
	case "WARN":
		SetLevel(WarnLevel)
	case "ERROR":
		SetLevel(ErrorLevel)
	case "FATAL":
		SetLevel(FatalLevel)
	default:
		return fmt.Errorf("unknown log level %q", ls)
	}

	return nil
}

// KV is used to provide context to a log entry in the form of a dynamic set of
// key/value pairs which can be different for every entry.
type KV map[string]interface{}

type entry struct {
	level   Level
	msg     string
	kv      KV
	blockCh chan struct{} // can be nil
}

var (
	prefix         = []byte("~ ")
	separator      = []byte(" --")
	separatorSpace = append(separator, ' ')
	tsPrefix       = []byte("[")
	tsSuffix       = []byte("] ")
	space          = []byte(" ")
	equals         = []byte("=")
	newline        = []byte("\n")
)

func writeHelper(b []byte, w io.Writer, lastErr error) error {
	if lastErr != nil {
		return lastErr
	}
	_, lastErr = w.Write(b)
	return lastErr
}

func (e entry) printOut(w io.Writer, displayTS bool) error {
	var err error
	err = writeHelper(prefix, w, err)
	if displayTS {
		err = writeHelper(tsPrefix, w, err)
		err = writeHelper([]byte(time.Now().String()), w, err)
		err = writeHelper(tsSuffix, w, err)
	}
	err = writeHelper([]byte(e.level.String()), w, err)
	err = writeHelper(separatorSpace, w, err)
	err = writeHelper([]byte(e.msg), w, err)
	if len(e.kv) > 0 {
		err = writeHelper(separator, w, err)
		for k, v := range e.kv {
			err = writeHelper(space, w, err)
			err = writeHelper([]byte(k), w, err)
			err = writeHelper(equals, w, err)
			vstr := fmt.Sprint(v)
			err = writeHelper([]byte(strconv.QuoteToASCII(vstr)), w, err)
		}
	}
	err = writeHelper(newline, w, err)

	return err
}

type syncer interface {
	Sync()
}

type flusher interface {
	Flush()
}

var entryCh = make(chan entry)

func init() {
	go func() {
		for e := range entryCh {
			err := e.printOut(Out, DisplayTimestamp)

			// If we couldn't write the entry to Out we write an error to that
			// effect to Stdout, then try to write the original entry as well
			if err != nil && Out != defaultOut {
				erre := entry{
					level: ErrorLevel,
					msg:   "Could not write to error Out",
					kv: KV{
						"err": err,
					},
				}
				erre.printOut(defaultOut, DisplayTimestamp)
				e.printOut(defaultOut, DisplayTimestamp)
			}

			// If the error level is fatal this is the last entry we should ever
			// write. We do want to attempt to flush Out though, in case it's
			// buffered, otherwise exiting now will cause the fatal message to
			// never be shown. We try to cast to either an interface with a Sync
			// or a Flush command as a form of ghetto reflection, to see if the
			// writer has either, and use one if found.
			if e.level == FatalLevel {
				if so, ok := Out.(syncer); ok {
					so.Sync()
				} else if fo, ok := Out.(flusher); ok {
					fo.Flush()
				}
			}

			if e.blockCh != nil {
				close(e.blockCh)
			}
		}
	}()
}

func kvNormalize(kv []KV) KV {
	if len(kv) > 0 {
		kvCopy := make(KV, len(kv[0]))
		for k, v := range kv[0] {
			kvCopy[k] = v
		}
		return kvCopy
	}
	return nil
}

func logEntry(l Level, msg string, kv []KV, blockCh chan struct{}) {
	if l >= GetLevel() {
		entryCh <- entry{
			level:   l,
			msg:     msg,
			kv:      kvNormalize(kv),
			blockCh: blockCh,
		}
	}
}

// LogFunc is the function signature used by the different log functions (Debug,
// Info, Warn, Error, and Fatal). It's useful for writing wrapper functions
type LogFunc func(string, ...KV)

// Debug writes a Debug message to Out, with an optional set of key/value pairs
func Debug(msg string, kv ...KV) {
	logEntry(DebugLevel, msg, kv, nil)
}

// Info writes an Info message to Out, with an optional set of key/value pairs
func Info(msg string, kv ...KV) {
	logEntry(InfoLevel, msg, kv, nil)
}

// Warn writes a Warn message to Out, with an optional set of key/value pairs
func Warn(msg string, kv ...KV) {
	logEntry(WarnLevel, msg, kv, nil)
}

// Error writes an Error message to Out, with an optional set of key/value pairs
func Error(msg string, kv ...KV) {
	logEntry(ErrorLevel, msg, kv, nil)
}

// Fatal writes a Fatal message to Out, with an optional set of key/value pairs.
// Once written the process will be exited with an exit code of 1
func Fatal(msg string, kv ...KV) {
	blockCh := make(chan struct{})
	logEntry(FatalLevel, msg, kv, blockCh)
	<-blockCh
	os.Exit(1)
}
<commit_msg>sort the kv pairs by key name<commit_after>// Package llog is a generic logging library used by leven labs. The log methods
// come in different severities: Debug, Info, Warn, Error, and Fatal.
//
// The log methods take in a string describing the error, and a set of key/value
// pairs giving the specific context around the error. The string is intended to
// always be the same no matter what, while the key/value pairs give information
// like which userID the error happened to, or any other relevant dynamic
// information.
//
// By default logs will be output to Stdout, without a timestamp attached to
// them, and only showing entries of level Info or above. All of these can be
// configured.
//
// All public functions in this package are thread-safe and can be called at any
// time. The public variables in this package are NOT thread-safe and should
// only be modified before any logging takes place
//
// Examples:
//
//	Info("Something important has occured")
//	Error("Could not open file", llog.KV{"filename": filename, "err": err})
//
package llog

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Out is the io.Writer all log entries will be written to. It can be changed to
// anything you like, but the change should happen before any logging occurs. If
// an error occurs while writing to Out the entry will be written to Stdout
// instead
var Out io.Writer = os.Stdout
var defaultOut = os.Stdout

// DisplayTimestamp determines whether or not a timestamp is displayed in the
// log messages. By default one is not displayed. This can be changed by it
// should only be changed before any logging occurs
var DisplayTimestamp bool

// Level describes the severity of a particular log message
type Level int

// All defined log levels
const (
	DebugLevel Level = iota
	InfoLevel
	WarnLevel
	ErrorLevel
	FatalLevel
)

func (l Level) String() string {
	switch l {
	case DebugLevel:
		return "DEBUG"
	case InfoLevel:
		return "INFO"
	case WarnLevel:
		return "WARN"
	case ErrorLevel:
		return "ERROR"
	case FatalLevel:
		return "FATAL"
	}
	return "unknown level"
}

var currLevel = InfoLevel
var currLevelLock sync.RWMutex

// GetLevel returns the current log level
func GetLevel() Level {
	currLevelLock.RLock()
	defer currLevelLock.RUnlock()
	return currLevel
}

// SetLevel sets the current minimum log level which will be written to Out
func SetLevel(l Level) {
	currLevelLock.Lock()
	defer currLevelLock.Unlock()
	currLevel = l
}

// SetLevelFromString attempts to interpret the given string as a log level and
// sets the current log level to that. If the string can't be interpreted an
// error is returned and the log level remains what it was
func SetLevelFromString(ls string) error {
	switch strings.ToUpper(ls) {
	case "DEBUG":
		SetLevel(DebugLevel)
	case "INFO":
		SetLevel(InfoLevel)
	case "WARN":
		SetLevel(WarnLevel)
	case "ERROR":
		SetLevel(ErrorLevel)
	case "FATAL":
		SetLevel(FatalLevel)
	default:
		return fmt.Errorf("unknown log level %q", ls)
	}

	return nil
}

// KV is used to provide context to a log entry in the form of a dynamic set of
// key/value pairs which can be different for every entry.
type KV map[string]interface{}

type kvE struct {
	K string
	V interface{}
}

type kvL []kvE

func (l kvL) Len() int {
	return len(l)
}

func (l kvL) Less(i, j int) bool {
	return l[i].K < l[j].K
}

func (l kvL) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func kvTo(kv KV) kvL {
	kvl := make(kvL, 0, len(kv))
	for k, v := range kv {
		kvl = append(kvl, kvE{K: k, V: v})
	}
	sort.Sort(kvl)
	return kvl
}

type entry struct {
	level   Level
	msg     string
	kv      kvL
	blockCh chan struct{} // can be nil
}

var (
	prefix         = []byte("~ ")
	separator      = []byte(" --")
	separatorSpace = append(separator, ' ')
	tsPrefix       = []byte("[")
	tsSuffix       = []byte("] ")
	space          = []byte(" ")
	equals         = []byte("=")
	newline        = []byte("\n")
)

func writeHelper(b []byte, w io.Writer, lastErr error) error {
	if lastErr != nil {
		return lastErr
	}
	_, lastErr = w.Write(b)
	return lastErr
}

func (e entry) printOut(w io.Writer, displayTS bool) error {
	var err error
	err = writeHelper(prefix, w, err)
	if displayTS {
		err = writeHelper(tsPrefix, w, err)
		err = writeHelper([]byte(time.Now().String()), w, err)
		err = writeHelper(tsSuffix, w, err)
	}
	err = writeHelper([]byte(e.level.String()), w, err)
	err = writeHelper(separatorSpace, w, err)
	err = writeHelper([]byte(e.msg), w, err)
	if len(e.kv) > 0 {
		err = writeHelper(separator, w, err)
		for _, kve := range e.kv {
			err = writeHelper(space, w, err)
			err = writeHelper([]byte(kve.K), w, err)
			err = writeHelper(equals, w, err)
			vstr := fmt.Sprint(kve.V)
			err = writeHelper([]byte(strconv.QuoteToASCII(vstr)), w, err)
		}
	}
	err = writeHelper(newline, w, err)

	return err
}

type syncer interface {
	Sync()
}

type flusher interface {
	Flush()
}

var entryCh = make(chan entry)

func init() {
	go func() {
		for e := range entryCh {
			err := e.printOut(Out, DisplayTimestamp)

			// If we couldn't write the entry to Out we write an error to that
			// effect to Stdout, then try to write the original entry as well
			if err != nil && Out != defaultOut {
				erre := entry{
					level: ErrorLevel,
					msg:   "Could not write to error Out",
					kv: kvTo(KV{
						"err": err,
					}),
				}
				erre.printOut(defaultOut, DisplayTimestamp)
				e.printOut(defaultOut, DisplayTimestamp)
			}

			// If the error level is fatal this is the last entry we should ever
			// write. We do want to attempt to flush Out though, in case it's
			// buffered, otherwise exiting now will cause the fatal message to
			// never be shown. We try to cast to either an interface with a Sync
			// or a Flush command as a form of ghetto reflection, to see if the
			// writer has either, and use one if found.
			if e.level == FatalLevel {
				if so, ok := Out.(syncer); ok {
					so.Sync()
				} else if fo, ok := Out.(flusher); ok {
					fo.Flush()
				}
			}

			if e.blockCh != nil {
				close(e.blockCh)
			}
		}
	}()
}

func kvNormalize(kv []KV) kvL {
	if len(kv) > 0 {
		return kvTo(kv[0])
	}
	return nil
}

func logEntry(l Level, msg string, kv []KV, blockCh chan struct{}) {
	if l >= GetLevel() {
		entryCh <- entry{
			level:   l,
			msg:     msg,
			kv:      kvNormalize(kv),
			blockCh: blockCh,
		}
	}
}

// LogFunc is the function signature used by the different log functions (Debug,
// Info, Warn, Error, and Fatal). It's useful for writing wrapper functions
type LogFunc func(string, ...KV)

// Debug writes a Debug message to Out, with an optional set of key/value pairs
func Debug(msg string, kv ...KV) {
	logEntry(DebugLevel, msg, kv, nil)
}

// Info writes an Info message to Out, with an optional set of key/value pairs
func Info(msg string, kv ...KV) {
	logEntry(InfoLevel, msg, kv, nil)
}

// Warn writes a Warn message to Out, with an optional set of key/value pairs
func Warn(msg string, kv ...KV) {
	logEntry(WarnLevel, msg, kv, nil)
}

// Error writes an Error message to Out, with an optional set of key/value pairs
func Error(msg string, kv ...KV) {
	logEntry(ErrorLevel, msg, kv, nil)
}

// Fatal writes a Fatal message to Out, with an optional set of key/value pairs.
// Once written the process will be exited with an exit code of 1
func Fatal(msg string, kv ...KV) {
	blockCh := make(chan struct{})
	logEntry(FatalLevel, msg, kv, blockCh)
	<-blockCh
	os.Exit(1)
}
<|endoftext|>


<commit_before>// Copyright ©2012 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

// Package llrb implements a Left-Leaning Red Black tree as described in
//  http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
//  http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java
//  http://www.teachsolaisgames.com/articles/balanced_left_leaning.html
package llrb

const (
	TD234 = iota
	BU23
)

// Operation mode of the LLRB tree. Currently only BU23 tests as correct.
const Mode = BU23

// A Comparable is a type that can be inserted into a Tree or used as a range
// or equality query on the tree,
type Comparable interface {
	// Compare returns a value indicating the sort order relationship between the
	// receiver and the parameter.
	//
	// Given c = a.Compare(b):
	//  c < 0 if a r < b;
	//  c == 0 if a == b; and
	//  c > 0 if a > b.
	//
	Compare(Comparable) int
}

// A Color represents the color of a Node.
type Color bool

// String returns a string representation of a Color.
func (c Color) String() string {
	if c {
		return "Black"
	}
	return "Red"
}

const (
	// Red as false give us the defined behaviour that new nodes are red. Although this
	// is incorrect for the root node, that is resolved on the first insertion.
	Red   Color = false
	Black Color = true
)

// A Node represents a node in the LLRB tree.
type Node struct {
	Elem        Comparable
	Left, Right *Node
	Color       Color
}

// A Tree represent the root node of an LLRB tree. Public methods of nodes are exposed
// through this type.
type Tree Node

// Helper methods

// color returns the effect color of a Node. A nil node returns black.
func (self *Node) color() Color {
	if self == nil {
		return Black
	}
	return self.Color
}

// (a,c)b -rotL-> ((a,)b,)c
func (self *Node) rotateLeft() (root *Node) {
	// Assumes: self has two children.
	root = self.Right
	self.Right = root.Left
	root.Left = self
	root.Color = self.Color
	self.Color = Red
	return
}

// (a,c)b -rotR-> (,(,c)b)a
func (self *Node) rotateRight() (root *Node) {
	// Assumes: self has two children.
	root = self.Left
	self.Left = root.Right
	root.Right = self
	root.Color = self.Color
	self.Color = Red
	return
}

// (aR,cR)bB -flipC-> (aB,cB)bR | (aB,cB)bR -flipC-> (aR,cR)bB
func (self *Node) flipColors() {
	// Assumes: self has two children.
	self.Color = !self.Color
	self.Left.Color = !self.Left.Color
	self.Right.Color = !self.Right.Color
}

// fixUp ensures that black link balance is correct, that red nodes lean left,
// and that 4 nodes are split.
func (self *Node) fixUp() *Node {
	if self.Right.color() == Red {
		self = self.rotateLeft()
	}
	if self.Left.color() == Red && self.Left.Left.color() == Red {
		self = self.rotateRight()
	}
	if self.Left.color() == Red && self.Right.color() == Red {
		self.flipColors()
	}
	return self
}

func (self *Node) moveRedLeft() *Node {
	self.flipColors()
	if self.Right.Left.color() == Red {
		self.Right = self.Right.rotateRight()
		self = self.rotateLeft()
		self.flipColors()
	}
	return self
}

func (self *Node) moveRedRight() *Node {
	self.flipColors()
	if self.Left.Left.color() == Red {
		self = self.rotateRight()
		self.flipColors()
	}
	return self
}

// Get returns the first match of q in the Tree. If insertion without
// replacement is used, this is probably not what you want.
func (self *Tree) Get(q Comparable) Comparable {
	return (*Node)(self).search(q).Elem
}

func (self *Node) search(q Comparable) (n *Node) {
	n = self
	for n != nil {
		switch c := q.Compare(n.Elem); {
		case c == 0:
			return n
		case c < 0:
			n = n.Left
		default:
			n = n.Right
		}
	}

	return
}

// Insert inserts the Comparable e into the Tree at the first match found
// with e or when a nil node is reached. Insertion without replacement can
// specified by ensuring that e.Compare() never returns 0. If insert without
// replacement is performed, a distinct query Comparable must be used that
// does return 0 for a Compare() call.
func (self *Tree) Insert(e Comparable) (root *Tree) {
	root = (*Tree)((*Node)(self).insert(e))
	root.Color = Black
	return
}

func (self *Node) insert(e Comparable) (root *Node) {
	if self == nil {
		return &Node{Elem: e}
	} else if self.Elem == nil {
		self.Elem = e
		return self
	}

	if Mode == TD234 {
		if self.Left.color() == Red && self.Right.color() == Red {
			self.flipColors()
		}
	}

	switch c := e.Compare(self.Elem); {
	case c == 0:
		self.Elem = e
	case c < 0:
		self.Left = self.Left.insert(e)
	default:
		self.Right = self.Right.insert(e)
	}

	if self.Right.color() == Red && self.Left.color() == Black {
		self = self.rotateLeft()
	}
	if self.Left.color() == Red && self.Left.Left.color() == Red {
		self = self.rotateRight()
	}

	if Mode == BU23 {
		if self.Left.color() == Red && self.Right.color() == Red {
			self.flipColors()
		}
	}

	return self
}

// DeleteMin deletes the node with the minimum value in the tree.
func (self *Tree) DeleteMin() (root *Tree) {
	root = (*Tree)((*Node)(self).deleteMin())
	root.Color = Black
	return
}

func (self *Node) deleteMin() *Node {
	if self.Left == nil {
		return nil
	}
	if self.Left.color() == Black && self.Left.Left.color() == Black {
		self = self.moveRedLeft()
	}
	self.Left = self.Left.deleteMin()
	return self.fixUp()
}

// DeleteMax deletes the node with the maximum value in the tree.
func (self *Tree) DeleteMax() (root *Tree) {
	root = (*Tree)((*Node)(self).deleteMax())
	root.Color = Black
	return
}

func (self *Node) deleteMax() *Node {
	if self.Left != nil && self.Left.color() == Red {
		self = self.rotateRight()
	}
	if self.Right == nil {
		return nil
	}
	if self.Right.color() == Black && self.Right.Left.color() == Black {
		self = self.moveRedRight()
	}
	self.Right = self.Right.deleteMax()
	return self.fixUp()
}

// Delete deletes the first node found that matches e according to Compare().
func (self *Tree) Delete(e Comparable) (root *Tree) {
	root = (*Tree)((*Node)(self).delete(e))
	if root == nil {
		return
	}
	root.Color = Black
	return
}

func (self *Node) delete(e Comparable) (root *Node) {
	if self == nil {
		return
	}
	if e.Compare(self.Elem) < 0 {
		if self.Left.color() == Black && self.Left != nil && self.Left.Left.color() == Black {
			self = self.moveRedLeft()
		}
		self.Left = self.Left.delete(e)
	} else {
		if self.Left.color() == Red {
			self = self.rotateRight()
		}
		if e.Compare(self.Elem) == 0 && self.Right == nil {
			return nil
		}
		if self.Right.color() == Black && self.Right != nil && self.Right.Left.color() == Black {
			self = self.moveRedRight()
		}
		if e.Compare(self.Elem) == 0 {
			self.Elem = self.Right.min().Elem
			self.Right = self.Right.deleteMin()
		} else {
			self.Right = self.Right.delete(e)
		}
	}
	return self.fixUp()
}

// Return the minimum value stored in the tree. This will be the left-most value if
// insertion without replacement is allowed.
func (self *Tree) Min() Comparable {
	return (*Node)(self).min().Elem
}

func (self *Node) min() *Node {
	if self.Left == nil {
		return self
	}
	return self.Left.min()
}

// Return the maximum value stored in the tree. This will be the right-most value if
// insertion without replacement is allowed.
func (self *Tree) Max() Comparable {
	return (*Node)(self).max().Elem
}

func (self *Node) max() *Node {
	if self.Right == nil {
		return self
	}
	return self.Right.max()
}

// An Operation is a function that operates on a Comparable.
type Operation func(Comparable)

// Do performs fn on all values stored in the tree. If fn alters the values of
// stored values, future tree operation behaviors are undefined.
func (self *Tree) Do(fn Operation) {
	(*Node)(self).do(fn)
}

func (self *Node) do(fn Operation) {
	if self == nil {
		return
	}
	self.Left.do(fn)
	fn(self.Elem)
	self.Right.do(fn)
}

// DoRange performs fn on all values stored in the tree between from and to. If from is
// less than to, the operations are performed from left to right. If from is greater than
// to then the operations are performed from right to left. If fn alters the values of
// stored values, future tree operation behaviors are undefined.
func (self *Tree) DoRange(fn Operation, from, to Comparable) {
	switch order := from.Compare(to); {
	case order < 0:
		(*Node)(self).doRange(fn, from, to)
	case order > 0:
		(*Node)(self).doRangeReverse(fn, from, to)
	default:
		(*Node)(self).doMatch(fn, from)
	}
}

func (self *Node) doRange(fn Operation, from, to Comparable) {
	if self == nil {
		return
	}
	fc, tc := from.Compare(self.Elem), to.Compare(self.Elem)
	if fc <= 0 {
		self.Left.doRange(fn, from, to)
	}
	if fc <= 0 && tc >= 0 {
		fn(self.Elem)
	}
	if tc >= 0 {
		self.Right.doRange(fn, from, to)
	}
}

func (self *Node) doRangeReverse(fn Operation, from, to Comparable) {
	if self == nil {
		return
	}
	fc, tc := from.Compare(self.Elem), to.Compare(self.Elem)
	if tc >= 0 {
		self.Right.doRangeReverse(fn, from, to)
	}
	if fc <= 0 && tc >= 0 {
		fn(self.Elem)
	}
	if fc <= 0 {
		self.Left.doRangeReverse(fn, from, to)
	}
}

// DoMatch performs fn on all values stored in the tree that match q according to Compare.
// If fn alters the values of stored values, future tree operation behaviors are undefined.
func (self *Tree) DoMatching(fn Operation, q Comparable) {
	(*Node)(self).doMatch(fn, q)
}

func (self *Node) doMatch(fn Operation, q Comparable) {
	if self == nil {
		return
	}
	c := q.Compare(self.Elem)
	if c <= 0 {
		self.Left.doMatch(fn, q)
	}
	if c == 0 {
		fn(self.Elem)
	}
	if c >= 0 {
		self.Right.doMatch(fn, q)
	}
}
<commit_msg>Check for nil nodes with min/max methods<commit_after>// Copyright ©2012 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

// Package llrb implements a Left-Leaning Red Black tree as described in
//  http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
//  http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java
//  http://www.teachsolaisgames.com/articles/balanced_left_leaning.html
package llrb

const (
	TD234 = iota
	BU23
)

// Operation mode of the LLRB tree. Currently only BU23 tests as correct.
const Mode = BU23

// A Comparable is a type that can be inserted into a Tree or used as a range
// or equality query on the tree,
type Comparable interface {
	// Compare returns a value indicating the sort order relationship between the
	// receiver and the parameter.
	//
	// Given c = a.Compare(b):
	//  c < 0 if a r < b;
	//  c == 0 if a == b; and
	//  c > 0 if a > b.
	//
	Compare(Comparable) int
}

// A Color represents the color of a Node.
type Color bool

// String returns a string representation of a Color.
func (c Color) String() string {
	if c {
		return "Black"
	}
	return "Red"
}

const (
	// Red as false give us the defined behaviour that new nodes are red. Although this
	// is incorrect for the root node, that is resolved on the first insertion.
	Red   Color = false
	Black Color = true
)

// A Node represents a node in the LLRB tree.
type Node struct {
	Elem        Comparable
	Left, Right *Node
	Color       Color
}

// A Tree represent the root node of an LLRB tree. Public methods of nodes are exposed
// through this type.
type Tree Node

// Helper methods

// color returns the effect color of a Node. A nil node returns black.
func (self *Node) color() Color {
	if self == nil {
		return Black
	}
	return self.Color
}

// (a,c)b -rotL-> ((a,)b,)c
func (self *Node) rotateLeft() (root *Node) {
	// Assumes: self has two children.
	root = self.Right
	self.Right = root.Left
	root.Left = self
	root.Color = self.Color
	self.Color = Red
	return
}

// (a,c)b -rotR-> (,(,c)b)a
func (self *Node) rotateRight() (root *Node) {
	// Assumes: self has two children.
	root = self.Left
	self.Left = root.Right
	root.Right = self
	root.Color = self.Color
	self.Color = Red
	return
}

// (aR,cR)bB -flipC-> (aB,cB)bR | (aB,cB)bR -flipC-> (aR,cR)bB
func (self *Node) flipColors() {
	// Assumes: self has two children.
	self.Color = !self.Color
	self.Left.Color = !self.Left.Color
	self.Right.Color = !self.Right.Color
}

// fixUp ensures that black link balance is correct, that red nodes lean left,
// and that 4 nodes are split.
func (self *Node) fixUp() *Node {
	if self.Right.color() == Red {
		self = self.rotateLeft()
	}
	if self.Left.color() == Red && self.Left.Left.color() == Red {
		self = self.rotateRight()
	}
	if self.Left.color() == Red && self.Right.color() == Red {
		self.flipColors()
	}
	return self
}

func (self *Node) moveRedLeft() *Node {
	self.flipColors()
	if self.Right.Left.color() == Red {
		self.Right = self.Right.rotateRight()
		self = self.rotateLeft()
		self.flipColors()
	}
	return self
}

func (self *Node) moveRedRight() *Node {
	self.flipColors()
	if self.Left.Left.color() == Red {
		self = self.rotateRight()
		self.flipColors()
	}
	return self
}

// Get returns the first match of q in the Tree. If insertion without
// replacement is used, this is probably not what you want.
func (self *Tree) Get(q Comparable) Comparable {
	return (*Node)(self).search(q).Elem
}

func (self *Node) search(q Comparable) (n *Node) {
	n = self
	for n != nil {
		switch c := q.Compare(n.Elem); {
		case c == 0:
			return n
		case c < 0:
			n = n.Left
		default:
			n = n.Right
		}
	}

	return
}

// Insert inserts the Comparable e into the Tree at the first match found
// with e or when a nil node is reached. Insertion without replacement can
// specified by ensuring that e.Compare() never returns 0. If insert without
// replacement is performed, a distinct query Comparable must be used that
// does return 0 for a Compare() call.
func (self *Tree) Insert(e Comparable) (root *Tree) {
	root = (*Tree)((*Node)(self).insert(e))
	root.Color = Black
	return
}

func (self *Node) insert(e Comparable) (root *Node) {
	if self == nil {
		return &Node{Elem: e}
	} else if self.Elem == nil {
		self.Elem = e
		return self
	}

	if Mode == TD234 {
		if self.Left.color() == Red && self.Right.color() == Red {
			self.flipColors()
		}
	}

	switch c := e.Compare(self.Elem); {
	case c == 0:
		self.Elem = e
	case c < 0:
		self.Left = self.Left.insert(e)
	default:
		self.Right = self.Right.insert(e)
	}

	if self.Right.color() == Red && self.Left.color() == Black {
		self = self.rotateLeft()
	}
	if self.Left.color() == Red && self.Left.Left.color() == Red {
		self = self.rotateRight()
	}

	if Mode == BU23 {
		if self.Left.color() == Red && self.Right.color() == Red {
			self.flipColors()
		}
	}

	return self
}

// DeleteMin deletes the node with the minimum value in the tree.
func (self *Tree) DeleteMin() (root *Tree) {
	root = (*Tree)((*Node)(self).deleteMin())
	root.Color = Black
	return
}

func (self *Node) deleteMin() *Node {
	if self.Left == nil {
		return nil
	}
	if self.Left.color() == Black && self.Left.Left.color() == Black {
		self = self.moveRedLeft()
	}
	self.Left = self.Left.deleteMin()
	return self.fixUp()
}

// DeleteMax deletes the node with the maximum value in the tree.
func (self *Tree) DeleteMax() (root *Tree) {
	root = (*Tree)((*Node)(self).deleteMax())
	root.Color = Black
	return
}

func (self *Node) deleteMax() *Node {
	if self.Left != nil && self.Left.color() == Red {
		self = self.rotateRight()
	}
	if self.Right == nil {
		return nil
	}
	if self.Right.color() == Black && self.Right.Left.color() == Black {
		self = self.moveRedRight()
	}
	self.Right = self.Right.deleteMax()
	return self.fixUp()
}

// Delete deletes the first node found that matches e according to Compare().
func (self *Tree) Delete(e Comparable) (root *Tree) {
	root = (*Tree)((*Node)(self).delete(e))
	if root == nil {
		return
	}
	root.Color = Black
	return
}

func (self *Node) delete(e Comparable) (root *Node) {
	if self == nil {
		return
	}
	if e.Compare(self.Elem) < 0 {
		if self.Left.color() == Black && self.Left != nil && self.Left.Left.color() == Black {
			self = self.moveRedLeft()
		}
		self.Left = self.Left.delete(e)
	} else {
		if self.Left.color() == Red {
			self = self.rotateRight()
		}
		if e.Compare(self.Elem) == 0 && self.Right == nil {
			return nil
		}
		if self.Right.color() == Black && self.Right != nil && self.Right.Left.color() == Black {
			self = self.moveRedRight()
		}
		if e.Compare(self.Elem) == 0 {
			self.Elem = self.Right.min().Elem
			self.Right = self.Right.deleteMin()
		} else {
			self.Right = self.Right.delete(e)
		}
	}
	return self.fixUp()
}

// Return the minimum value stored in the tree. This will be the left-most value if
// insertion without replacement is allowed.
func (self *Tree) Min() Comparable {
	n := (*Node)(self).min()
	if n == nil {
		return nil
	}
	return n.Elem
}

func (self *Node) min() *Node {
	if self == nil {
		return nil
	}
	if self.Left == nil {
		return self
	}
	return self.Left.min()
}

// Return the maximum value stored in the tree. This will be the right-most value if
// insertion without replacement is allowed.
func (self *Tree) Max() Comparable {
	n := (*Node)(self).max()
	if n == nil {
		return nil
	}
	return n.Elem
}

func (self *Node) max() *Node {
	if self == nil {
		return nil
	}
	if self.Right == nil {
		return self
	}
	return self.Right.max()
}

// An Operation is a function that operates on a Comparable.
type Operation func(Comparable)

// Do performs fn on all values stored in the tree. If fn alters the values of
// stored values, future tree operation behaviors are undefined.
func (self *Tree) Do(fn Operation) {
	(*Node)(self).do(fn)
}

func (self *Node) do(fn Operation) {
	if self == nil {
		return
	}
	self.Left.do(fn)
	fn(self.Elem)
	self.Right.do(fn)
}

// DoRange performs fn on all values stored in the tree between from and to. If from is
// less than to, the operations are performed from left to right. If from is greater than
// to then the operations are performed from right to left. If fn alters the values of
// stored values, future tree operation behaviors are undefined.
func (self *Tree) DoRange(fn Operation, from, to Comparable) {
	switch order := from.Compare(to); {
	case order < 0:
		(*Node)(self).doRange(fn, from, to)
	case order > 0:
		(*Node)(self).doRangeReverse(fn, from, to)
	default:
		(*Node)(self).doMatch(fn, from)
	}
}

func (self *Node) doRange(fn Operation, from, to Comparable) {
	if self == nil {
		return
	}
	fc, tc := from.Compare(self.Elem), to.Compare(self.Elem)
	if fc <= 0 {
		self.Left.doRange(fn, from, to)
	}
	if fc <= 0 && tc >= 0 {
		fn(self.Elem)
	}
	if tc >= 0 {
		self.Right.doRange(fn, from, to)
	}
}

func (self *Node) doRangeReverse(fn Operation, from, to Comparable) {
	if self == nil {
		return
	}
	fc, tc := from.Compare(self.Elem), to.Compare(self.Elem)
	if tc >= 0 {
		self.Right.doRangeReverse(fn, from, to)
	}
	if fc <= 0 && tc >= 0 {
		fn(self.Elem)
	}
	if fc <= 0 {
		self.Left.doRangeReverse(fn, from, to)
	}
}

// DoMatch performs fn on all values stored in the tree that match q according to Compare.
// If fn alters the values of stored values, future tree operation behaviors are undefined.
func (self *Tree) DoMatching(fn Operation, q Comparable) {
	(*Node)(self).doMatch(fn, q)
}

func (self *Node) doMatch(fn Operation, q Comparable) {
	if self == nil {
		return
	}
	c := q.Compare(self.Elem)
	if c <= 0 {
		self.Left.doMatch(fn, q)
	}
	if c == 0 {
		fn(self.Elem)
	}
	if c >= 0 {
		self.Right.doMatch(fn, q)
	}
}
<|endoftext|>


<commit_before>// Command pubsub is an example of a fanout exchange with dynamic reliable
// membership, reading from stdin, writing to stdout.
//
// This example shows how to implement reconnect logic independent from a
// publish/subscribe loop with bridges to application types.

package main

import (
	"bufio"
	"crypto/sha1"
	"flag"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

var url = flag.String("url", "amqp:///", "AMQP url for both the publisher and subscriber")

// exchange binds the publishers to the subscribers
const exchange = "pubsub"

// message is the application type for a message.  This can contain identity,
// or a reference to the recevier chan for further demuxing.
type message []byte

// session composes an amqp.Connection with an amqp.Channel
type session struct {
	*amqp.Connection
	*amqp.Channel
}

// Close tears the connection down, taking the channel with it.
func (s session) Close() error {
	if s.Connection == nil {
		return nil
	}
	return s.Connection.Close()
}

// redial continually connects to the URL, exiting the program when no longer possible
func redial(ctx context.Context, url string) chan chan session {
	sessions := make(chan chan session)

	go func() {
		sess := make(chan session)
		defer close(sessions)

		for {
			select {
			case sessions <- sess:
			case <-ctx.Done():
				log.Println("shutting down session factory")
				return
			}

			conn, err := amqp.Dial(url)
			if err != nil {
				log.Fatalf("cannot (re)dial: %v: %q", err, url)
			}

			ch, err := conn.Channel()
			if err != nil {
				log.Fatalf("cannot create channel: %v", err)
			}

			if err := ch.ExchangeDeclare(exchange, "fanout", false, true, false, false, nil); err != nil {
				log.Fatalf("cannot declare fanout exchange: %v", err)
			}

			select {
			case sess <- session{conn, ch}:
			case <-ctx.Done():
				log.Println("shutting down new session")
				return
			}
		}
	}()

	return sessions
}

// publish publishes messages to a reconnecting session to a fanout exchange.
// It receives from the application specific source of messages.
func publish(sessions chan chan session, messages <-chan message) {
	var (
		running   bool
		reading   = messages
		pending   = make(chan message, 1)
		ack, nack = make(chan uint64), make(chan uint64)
	)

	for session := range sessions {
		pub := <-session

		// publisher confirms for this channel/connection
		if err := pub.Confirm(false); err != nil {
			log.Printf("publisher confirms not supported")
			close(ack) // confirms not supported, simulate by always acking
		} else {
			pub.NotifyConfirm(ack, nack)
		}

		log.Printf("publishing...")

		for {
			var body message
			select {
			case <-ack:
				reading = messages

			case id := <-nack:
				log.Printf("nack message %d, body: %q", id, string(body))
				reading = messages

			case body = <-pending:
				routingKey := "ignored for fanout exchanges, application dependent for other exchanges"
				err := pub.Publish(exchange, routingKey, false, false, amqp.Publishing{
					Body: body,
				})
				// Retry failed delivery on the next session
				if err != nil {
					pending <- body
					pub.Close()
					break
				}

			case body, running = <-reading:
				// all messages consumed
				if !running {
					return
				}
				// work on pending delivery until ack'd
				pending <- body
				reading = nil
			}
		}
	}
}

// identity returns the same host/process unique string for the lifetime of
// this process so that subscriber reconnections reuse the same queue name.
func identity() string {
	hostname, err := os.Hostname()
	h := sha1.New()
	fmt.Fprint(h, hostname)
	fmt.Fprint(h, err)
	fmt.Fprint(h, os.Getpid())
	return fmt.Sprintf("%x", h.Sum(nil))
}

// subscribe consumes deliveries from an exclusive queue from a fanout exchange and sends to the application specific messages chan.
func subscribe(sessions chan chan session, messages chan<- message) {
	queue := identity()

	for session := range sessions {
		sub := <-session

		if _, err := sub.QueueDeclare(queue, false, true, true, false, nil); err != nil {
			log.Printf("cannot consume from exclusive queue: %q, %v", queue, err)
			return
		}

		routingKey := "application specific routing key for fancy toplogies"
		if err := sub.QueueBind(queue, routingKey, exchange, false, nil); err != nil {
			log.Printf("cannot consume without a binding to exchange: %q, %v", exchange, err)
			return
		}

		deliveries, err := sub.Consume(queue, "", false, true, false, false, nil)
		if err != nil {
			log.Printf("cannot consume from: %q, %v", queue, err)
			return
		}

		log.Printf("subscribed...")

		for msg := range deliveries {
			messages <- message(msg.Body)
			sub.Ack(msg.DeliveryTag, false)
		}
	}
}

// read is this application's translation to the message format, scanning from
// stdin.
func read(r io.Reader) <-chan message {
	lines := make(chan message)
	go func() {
		defer close(lines)
		scan := bufio.NewScanner(r)
		for scan.Scan() {
			lines <- message(scan.Bytes())
		}
	}()
	return lines
}

// write is this application's subscriber of application messages, printing to
// stdout.
func write(w io.Writer) chan<- message {
	lines := make(chan message)
	go func() {
		for line := range lines {
			fmt.Fprintln(w, string(line))
		}
	}()
	return lines
}

func main() {
	flag.Parse()

	ctx, done := context.WithCancel(context.Background())

	go func() {
		publish(redial(ctx, *url), read(os.Stdin))
		done()
	}()

	go func() {
		subscribe(redial(ctx, *url), write(os.Stdout))
		done()
	}()

	<-ctx.Done()
}
<commit_msg>Update pubsub.go<commit_after>// Command pubsub is an example of a fanout exchange with dynamic reliable
// membership, reading from stdin, writing to stdout.
//
// This example shows how to implement reconnect logic independent from a
// publish/subscribe loop with bridges to application types.

package main

import (
	"bufio"
	"crypto/sha1"
	"flag"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

var url = flag.String("url", "amqp:///", "AMQP url for both the publisher and subscriber")

// exchange binds the publishers to the subscribers
const exchange = "pubsub"

// message is the application type for a message.  This can contain identity,
// or a reference to the recevier chan for further demuxing.
type message []byte

// session composes an amqp.Connection with an amqp.Channel
type session struct {
	*amqp.Connection
	*amqp.Channel
}

// Close tears the connection down, taking the channel with it.
func (s session) Close() error {
	if s.Connection == nil {
		return nil
	}
	return s.Connection.Close()
}

// redial continually connects to the URL, exiting the program when no longer possible
func redial(ctx context.Context, url string) chan chan session {
	sessions := make(chan chan session)

	go func() {
		sess := make(chan session)
		defer close(sessions)

		for {
			select {
			case sessions <- sess:
			case <-ctx.Done():
				log.Println("shutting down session factory")
				return
			}

			conn, err := amqp.Dial(url)
			if err != nil {
				log.Fatalf("cannot (re)dial: %v: %q", err, url)
			}

			ch, err := conn.Channel()
			if err != nil {
				log.Fatalf("cannot create channel: %v", err)
			}

			if err := ch.ExchangeDeclare(exchange, "fanout", false, true, false, false, nil); err != nil {
				log.Fatalf("cannot declare fanout exchange: %v", err)
			}

			select {
			case sess <- session{conn, ch}:
			case <-ctx.Done():
				log.Println("shutting down new session")
				return
			}
		}
	}()

	return sessions
}

// publish publishes messages to a reconnecting session to a fanout exchange.
// It receives from the application specific source of messages.
func publish(sessions chan chan session, messages <-chan message) {
	var (
		running   bool
		reading   = messages
		pending   = make(chan message, 1)
		ack, nack = make(chan uint64), make(chan uint64)
	)

	for session := range sessions {
		pub := <-session

		// publisher confirms for this channel/connection
		if err := pub.Confirm(false); err != nil {
			log.Printf("publisher confirms not supported")
			close(ack) // confirms not supported, simulate by always acking
		} else {
			pub.NotifyConfirm(ack, nack)
		}

		log.Printf("publishing...")

PUBLISH:
		for {
			var body message
			select {
			case <-ack:
				reading = messages

			case id := <-nack:
				log.Printf("nack message %d, body: %q", id, string(body))
				reading = messages

			case body = <-pending:
				routingKey := "ignored for fanout exchanges, application dependent for other exchanges"
				err := pub.Publish(exchange, routingKey, false, false, amqp.Publishing{
					Body: body,
				})
				// Retry failed delivery on the next session
				if err != nil {
					pending <- body
					pub.Close()
					break PUBLISH
				}

			case body, running = <-reading:
				// all messages consumed
				if !running {
					return
				}
				// work on pending delivery until ack'd
				pending <- body
				reading = nil
			}
		}
	}
}

// identity returns the same host/process unique string for the lifetime of
// this process so that subscriber reconnections reuse the same queue name.
func identity() string {
	hostname, err := os.Hostname()
	h := sha1.New()
	fmt.Fprint(h, hostname)
	fmt.Fprint(h, err)
	fmt.Fprint(h, os.Getpid())
	return fmt.Sprintf("%x", h.Sum(nil))
}

// subscribe consumes deliveries from an exclusive queue from a fanout exchange and sends to the application specific messages chan.
func subscribe(sessions chan chan session, messages chan<- message) {
	queue := identity()

	for session := range sessions {
		sub := <-session

		if _, err := sub.QueueDeclare(queue, false, true, true, false, nil); err != nil {
			log.Printf("cannot consume from exclusive queue: %q, %v", queue, err)
			return
		}

		routingKey := "application specific routing key for fancy toplogies"
		if err := sub.QueueBind(queue, routingKey, exchange, false, nil); err != nil {
			log.Printf("cannot consume without a binding to exchange: %q, %v", exchange, err)
			return
		}

		deliveries, err := sub.Consume(queue, "", false, true, false, false, nil)
		if err != nil {
			log.Printf("cannot consume from: %q, %v", queue, err)
			return
		}

		log.Printf("subscribed...")

		for msg := range deliveries {
			messages <- message(msg.Body)
			sub.Ack(msg.DeliveryTag, false)
		}
	}
}

// read is this application's translation to the message format, scanning from
// stdin.
func read(r io.Reader) <-chan message {
	lines := make(chan message)
	go func() {
		defer close(lines)
		scan := bufio.NewScanner(r)
		for scan.Scan() {
			lines <- message(scan.Bytes())
		}
	}()
	return lines
}

// write is this application's subscriber of application messages, printing to
// stdout.
func write(w io.Writer) chan<- message {
	lines := make(chan message)
	go func() {
		for line := range lines {
			fmt.Fprintln(w, string(line))
		}
	}()
	return lines
}

func main() {
	flag.Parse()

	ctx, done := context.WithCancel(context.Background())

	go func() {
		publish(redial(ctx, *url), read(os.Stdin))
		done()
	}()

	go func() {
		subscribe(redial(ctx, *url), write(os.Stdout))
		done()
	}()

	<-ctx.Done()
}
<|endoftext|>


<commit_before>// Copyright 2015 The LUCI 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 isolate

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/flag/stringmapflag"
)

// IsolatedGenJSONVersion is used in the batcharchive json format.
//
// TODO(tandrii): Migrate to batch_archive.go.
const IsolatedGenJSONVersion = 1

// ValidVariable is the regexp of valid isolate variable name.
const ValidVariable = "[A-Za-z_][A-Za-z_0-9]*"

var validVariableMatcher = regexp.MustCompile(ValidVariable)
var variableSubstitutionMatcher = regexp.MustCompile("<\\(" + ValidVariable + "\\)")

// IsValidVariable returns true if the variable is a valid symbol name.
func IsValidVariable(variable string) bool {
	return validVariableMatcher.MatchString(variable)
}

// ArchiveOptions for archiving trees.
type ArchiveOptions struct {
	Isolate             string              `json:"isolate"`
	IgnoredPathFilterRe string              `json:"ignored_path_filter_re"`
	PathVariables       stringmapflag.Value `json:"path_variables"`
	ConfigVariables     stringmapflag.Value `json:"config_variables"`
	AllowMissingFileDir bool                `json:"allow_missing_file_dir"`
}

// Init initializes with non-nil values.
func (a *ArchiveOptions) Init() {
	a.PathVariables = map[string]string{}
	if runtime.GOOS == "windows" {
		a.PathVariables["EXECUTABLE_SUFFIX"] = ".exe"
	} else {
		a.PathVariables["EXECUTABLE_SUFFIX"] = ""
	}
	a.ConfigVariables = map[string]string{}
}

func genExtensionsRegex(exts ...string) string {
	if len(exts) == 0 {
		return ""
	}
	res := make([]string, len(exts))
	for i, e := range exts {
		res[i] = `(\.` + e + `)`
	}
	return "((" + strings.Join(res, "|") + ")$)"
}

func genDirectoriesRegex(dirs ...string) string {
	if len(dirs) == 0 {
		return ""
	}
	res := make([]string, len(dirs))
	for i, d := range dirs {
		res[i] = "(" + d + ")"
	}
	// #Backslashes: https://stackoverflow.com/a/4025505/12003165
	return `((^|[\\/])(` + strings.Join(res, "|") + `)([\\/]|$))`
}

// PostProcess post-processes the flags to fix any compatibility issue.
func (a *ArchiveOptions) PostProcess(cwd string) {
	if a.IgnoredPathFilterRe == "" {
		// Set default ignored paths regexp
		// .swp are vim files
		a.IgnoredPathFilterRe = genExtensionsRegex("pyc", "swp") + "|" + genDirectoriesRegex(`\.git`, `\.hg`, `\.svn`)
	}
	if !filepath.IsAbs(a.Isolate) {
		a.Isolate = filepath.Join(cwd, a.Isolate)
	}
	a.Isolate = filepath.Clean(a.Isolate)

	for k, v := range a.PathVariables {
		// This is due to a Windows + GYP specific issue, where double-quoted paths
		// would get mangled in a way that cannot be resolved unless a space is
		// injected.
		a.PathVariables[k] = strings.TrimSpace(v)
	}
}

// ReplaceVariables replaces any occurrences of '<(FOO)' in 'str' with the
// corresponding variable from 'opts'.
//
// If any substitution refers to a variable that is missing, the returned error will
// refer to the first such variable. In the case of errors, the returned string will
// still contain a valid result for any non-missing substitutions.
func ReplaceVariables(str string, opts *ArchiveOptions) (string, error) {
	var err error
	subst := variableSubstitutionMatcher.ReplaceAllStringFunc(str,
		func(match string) string {
			varName := match[2 : len(match)-1]
			if v, ok := opts.PathVariables[varName]; ok {
				return v
			}
			if v, ok := opts.ConfigVariables[varName]; ok {
				return v
			}
			if err == nil {
				err = errors.New("no value for variable '" + varName + "'")
			}
			return match
		})
	return subst, err
}

func processDependencies(deps []string, isolateDir string, opts *ArchiveOptions) ([]string, string, error) {
	// Expand variables in the deps, and convert each path to an absolute form.
	for i := range deps {
		dep, err := ReplaceVariables(deps[i], opts)
		if err != nil {
			return nil, "", err
		}
		deps[i] = filepath.Join(isolateDir, dep)
	}

	// Find the root directory of all the files (the root might be above isolateDir).
	rootDir := isolateDir
	resultDeps := make([]string, 0, len(deps))
	for _, dep := range deps {
		// Check if the dep is outside isolateDir.
		info, err := os.Stat(dep)
		if err != nil {
			if !opts.AllowMissingFileDir {
				return nil, "", errors.Annotate(err, "faled to call Stat for %s", dep).Err()
			}
			log.Printf("Ignore missing dep: %s, err: %v", dep, err)
			continue
		}
		base := filepath.Dir(dep)
		if info.IsDir() {
			base = dep
			// Downstream expects the dependency of a directory to always end
			// with '/', but filepath.Join() removes that, so we add it back.
			dep += osPathSeparator
		}
		resultDeps = append(resultDeps, dep)
		for {
			rel, err := filepath.Rel(rootDir, base)
			if err != nil {
				return nil, "", errors.Annotate(err, "failed to call filepath.Rel(%s, %s)", rootDir, base).Err()
			}
			if !strings.HasPrefix(rel, "..") {
				break
			}
			newRootDir := filepath.Dir(rootDir)
			if newRootDir == rootDir {
				return nil, "", errors.New("failed to find root dir")
			}
			rootDir = newRootDir
		}
	}
	if rootDir != isolateDir {
		log.Printf("Root: %s", rootDir)
	}
	return resultDeps, rootDir, nil
}

// ProcessIsolate parses an isolate file, returning the list of dependencies
func ProcessIsolate(opts *ArchiveOptions) ([]string, string, error) {
	content, err := ioutil.ReadFile(opts.Isolate)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to read file: %s", opts.Isolate).Err()
	}
	deps, isolateDir, err := LoadIsolateForConfig(filepath.Dir(opts.Isolate), content, opts.ConfigVariables)
	if err != nil {
		return nil, "", err
	}

	deps, rootDir, err := processDependencies(deps, isolateDir, opts)
	if err != nil {
		return nil, "", err
	}
	return deps, rootDir, nil
}

// ProcessIsolateForCAS works similarly to ProcessIsolate. However, it is
// simpler in that it returns a list of dependency *relative* paths and the
// root directory, which are the necessary input to upload to RBE-CAS.
//
// TODO(crbug.com/1193375): remove after migrating to RBE's cas package.
func ProcessIsolateForCAS(opts *ArchiveOptions) ([]string, string, error) {
	content, err := ioutil.ReadFile(opts.Isolate)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to call ReadFile: %s", opts.Isolate).Err()
	}
	deps, isolateDir, err := LoadIsolateForConfig(filepath.Dir(opts.Isolate), content, opts.ConfigVariables)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to call LoadIsolateForConfig: %s", opts.Isolate).Err()
	}

	deps, rootDir, err := processDependencies(deps, isolateDir, opts)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to casll processDependencies: %s", opts.Isolate).Err()
	}

	relDeps := make([]string, len(deps))
	for i, dep := range deps {
		rel, err := filepath.Rel(rootDir, dep)
		if err != nil {
			return nil, "", errors.Annotate(err, "failed to call filepath.Rel(%s, %s)", rootDir, dep).Err()
		}
		if strings.HasSuffix(dep, osPathSeparator) && !strings.HasSuffix(rel, osPathSeparator) {
			// Make it consistent with the isolated format such that directory paths must end with osPathSeparator.
			rel += osPathSeparator
		}
		relDeps[i] = rel
	}
	return relDeps, rootDir, nil
}
<commit_msg>isolate: Fix typo.<commit_after>// Copyright 2015 The LUCI 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 isolate

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/flag/stringmapflag"
)

// IsolatedGenJSONVersion is used in the batcharchive json format.
//
// TODO(tandrii): Migrate to batch_archive.go.
const IsolatedGenJSONVersion = 1

// ValidVariable is the regexp of valid isolate variable name.
const ValidVariable = "[A-Za-z_][A-Za-z_0-9]*"

var validVariableMatcher = regexp.MustCompile(ValidVariable)
var variableSubstitutionMatcher = regexp.MustCompile("<\\(" + ValidVariable + "\\)")

// IsValidVariable returns true if the variable is a valid symbol name.
func IsValidVariable(variable string) bool {
	return validVariableMatcher.MatchString(variable)
}

// ArchiveOptions for archiving trees.
type ArchiveOptions struct {
	Isolate             string              `json:"isolate"`
	IgnoredPathFilterRe string              `json:"ignored_path_filter_re"`
	PathVariables       stringmapflag.Value `json:"path_variables"`
	ConfigVariables     stringmapflag.Value `json:"config_variables"`
	AllowMissingFileDir bool                `json:"allow_missing_file_dir"`
}

// Init initializes with non-nil values.
func (a *ArchiveOptions) Init() {
	a.PathVariables = map[string]string{}
	if runtime.GOOS == "windows" {
		a.PathVariables["EXECUTABLE_SUFFIX"] = ".exe"
	} else {
		a.PathVariables["EXECUTABLE_SUFFIX"] = ""
	}
	a.ConfigVariables = map[string]string{}
}

func genExtensionsRegex(exts ...string) string {
	if len(exts) == 0 {
		return ""
	}
	res := make([]string, len(exts))
	for i, e := range exts {
		res[i] = `(\.` + e + `)`
	}
	return "((" + strings.Join(res, "|") + ")$)"
}

func genDirectoriesRegex(dirs ...string) string {
	if len(dirs) == 0 {
		return ""
	}
	res := make([]string, len(dirs))
	for i, d := range dirs {
		res[i] = "(" + d + ")"
	}
	// #Backslashes: https://stackoverflow.com/a/4025505/12003165
	return `((^|[\\/])(` + strings.Join(res, "|") + `)([\\/]|$))`
}

// PostProcess post-processes the flags to fix any compatibility issue.
func (a *ArchiveOptions) PostProcess(cwd string) {
	if a.IgnoredPathFilterRe == "" {
		// Set default ignored paths regexp
		// .swp are vim files
		a.IgnoredPathFilterRe = genExtensionsRegex("pyc", "swp") + "|" + genDirectoriesRegex(`\.git`, `\.hg`, `\.svn`)
	}
	if !filepath.IsAbs(a.Isolate) {
		a.Isolate = filepath.Join(cwd, a.Isolate)
	}
	a.Isolate = filepath.Clean(a.Isolate)

	for k, v := range a.PathVariables {
		// This is due to a Windows + GYP specific issue, where double-quoted paths
		// would get mangled in a way that cannot be resolved unless a space is
		// injected.
		a.PathVariables[k] = strings.TrimSpace(v)
	}
}

// ReplaceVariables replaces any occurrences of '<(FOO)' in 'str' with the
// corresponding variable from 'opts'.
//
// If any substitution refers to a variable that is missing, the returned error will
// refer to the first such variable. In the case of errors, the returned string will
// still contain a valid result for any non-missing substitutions.
func ReplaceVariables(str string, opts *ArchiveOptions) (string, error) {
	var err error
	subst := variableSubstitutionMatcher.ReplaceAllStringFunc(str,
		func(match string) string {
			varName := match[2 : len(match)-1]
			if v, ok := opts.PathVariables[varName]; ok {
				return v
			}
			if v, ok := opts.ConfigVariables[varName]; ok {
				return v
			}
			if err == nil {
				err = errors.New("no value for variable '" + varName + "'")
			}
			return match
		})
	return subst, err
}

func processDependencies(deps []string, isolateDir string, opts *ArchiveOptions) ([]string, string, error) {
	// Expand variables in the deps, and convert each path to an absolute form.
	for i := range deps {
		dep, err := ReplaceVariables(deps[i], opts)
		if err != nil {
			return nil, "", err
		}
		deps[i] = filepath.Join(isolateDir, dep)
	}

	// Find the root directory of all the files (the root might be above isolateDir).
	rootDir := isolateDir
	resultDeps := make([]string, 0, len(deps))
	for _, dep := range deps {
		// Check if the dep is outside isolateDir.
		info, err := os.Stat(dep)
		if err != nil {
			if !opts.AllowMissingFileDir {
				return nil, "", errors.Annotate(err, "failed to call Stat for %s", dep).Err()
			}
			log.Printf("Ignore missing dep: %s, err: %v", dep, err)
			continue
		}
		base := filepath.Dir(dep)
		if info.IsDir() {
			base = dep
			// Downstream expects the dependency of a directory to always end
			// with '/', but filepath.Join() removes that, so we add it back.
			dep += osPathSeparator
		}
		resultDeps = append(resultDeps, dep)
		for {
			rel, err := filepath.Rel(rootDir, base)
			if err != nil {
				return nil, "", errors.Annotate(err, "failed to call filepath.Rel(%s, %s)", rootDir, base).Err()
			}
			if !strings.HasPrefix(rel, "..") {
				break
			}
			newRootDir := filepath.Dir(rootDir)
			if newRootDir == rootDir {
				return nil, "", errors.New("failed to find root dir")
			}
			rootDir = newRootDir
		}
	}
	if rootDir != isolateDir {
		log.Printf("Root: %s", rootDir)
	}
	return resultDeps, rootDir, nil
}

// ProcessIsolate parses an isolate file, returning the list of dependencies
func ProcessIsolate(opts *ArchiveOptions) ([]string, string, error) {
	content, err := ioutil.ReadFile(opts.Isolate)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to read file: %s", opts.Isolate).Err()
	}
	deps, isolateDir, err := LoadIsolateForConfig(filepath.Dir(opts.Isolate), content, opts.ConfigVariables)
	if err != nil {
		return nil, "", err
	}

	deps, rootDir, err := processDependencies(deps, isolateDir, opts)
	if err != nil {
		return nil, "", err
	}
	return deps, rootDir, nil
}

// ProcessIsolateForCAS works similarly to ProcessIsolate. However, it is
// simpler in that it returns a list of dependency *relative* paths and the
// root directory, which are the necessary input to upload to RBE-CAS.
//
// TODO(crbug.com/1193375): remove after migrating to RBE's cas package.
func ProcessIsolateForCAS(opts *ArchiveOptions) ([]string, string, error) {
	content, err := ioutil.ReadFile(opts.Isolate)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to call ReadFile: %s", opts.Isolate).Err()
	}
	deps, isolateDir, err := LoadIsolateForConfig(filepath.Dir(opts.Isolate), content, opts.ConfigVariables)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to call LoadIsolateForConfig: %s", opts.Isolate).Err()
	}

	deps, rootDir, err := processDependencies(deps, isolateDir, opts)
	if err != nil {
		return nil, "", errors.Annotate(err, "failed to casll processDependencies: %s", opts.Isolate).Err()
	}

	relDeps := make([]string, len(deps))
	for i, dep := range deps {
		rel, err := filepath.Rel(rootDir, dep)
		if err != nil {
			return nil, "", errors.Annotate(err, "failed to call filepath.Rel(%s, %s)", rootDir, dep).Err()
		}
		if strings.HasSuffix(dep, osPathSeparator) && !strings.HasSuffix(rel, osPathSeparator) {
			// Make it consistent with the isolated format such that directory paths must end with osPathSeparator.
			rel += osPathSeparator
		}
		relDeps[i] = rel
	}
	return relDeps, rootDir, nil
}
<|endoftext|>


<commit_before>package pnet

import (
	"crypto/cipher"
	"crypto/rand"

	salsa20 "github.com/davidlazar/go-crypto/salsa20"
	mpool "github.com/jbenet/go-msgio/mpool"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	tconn "github.com/libp2p/go-libp2p-transport"
)

// we are using buffer pool as user needs their slice back
// so we can't do XOR cripter in place
var (
	bufPool = mpool.ByteSlicePool

	errShortNonce  = ipnet.NewError("could not read full nonce")
	errInsecureNil = ipnet.NewError("insecure is nil")
	errPSKNil      = ipnet.NewError("pre-shread key is nil")
)

type pskConn struct {
	tconn.Conn
	psk *[32]byte

	writeS20 cipher.Stream
	readS20  cipher.Stream
}

func (c *pskConn) Read(out []byte) (int, error) {
	if c.readS20 == nil {
		nonce := make([]byte, 24)
		n, err := c.Conn.Read(nonce)
		if err != nil {
			return 0, err
		}
		if n != 24 {
			return 0, errShortNonce
		}
		c.readS20 = salsa20.New(c.psk, nonce)
	}

	maxn := uint32(len(out))
	in := bufPool.Get(maxn).([]byte) // get buffer
	defer bufPool.Put(maxn, in)      // put the buffer back

	in = in[:maxn]            // truncate to required length
	n, err := c.Conn.Read(in) // read to in
	if err != nil {
		return 0, err
	}

	c.readS20.XORKeyStream(out[:n], in[:n]) // decrypt to out buffer

	return n, nil
}

func (c *pskConn) Write(in []byte) (int, error) {
	if c.writeS20 == nil {
		nonce := make([]byte, 24)
		_, err := rand.Read(nonce)
		if err != nil {
			return 0, err
		}
		_, err = c.Conn.Write(nonce)
		if err != nil {
			return 0, err
		}

		c.writeS20 = salsa20.New(c.psk, nonce)
	}
	n := uint32(len(in))
	out := bufPool.Get(n).([]byte) // get buffer
	defer bufPool.Put(n, out)      // put the buffer back

	out = out[:n]                    // truncate to required length
	c.writeS20.XORKeyStream(out, in) // encrypt

	return c.Conn.Write(out) // send
}

var _ tconn.Conn = (*pskConn)(nil)

func newPSKConn(psk *[32]byte, insecure tconn.Conn) (tconn.Conn, error) {
	if insecure == nil {
		return nil, errInsecureNil
	}
	if psk == nil {
		return nil, errPSKNil
	}
	return &pskConn{
		Conn: insecure,
		psk:  psk,
	}, nil
}
<commit_msg>read nonces, even if they are fragmented<commit_after>package pnet

import (
	"crypto/cipher"
	"crypto/rand"
	"io"

	salsa20 "github.com/davidlazar/go-crypto/salsa20"
	mpool "github.com/jbenet/go-msgio/mpool"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	tconn "github.com/libp2p/go-libp2p-transport"
)

// we are using buffer pool as user needs their slice back
// so we can't do XOR cripter in place
var (
	bufPool = mpool.ByteSlicePool

	errShortNonce  = ipnet.NewError("could not read full nonce")
	errInsecureNil = ipnet.NewError("insecure is nil")
	errPSKNil      = ipnet.NewError("pre-shread key is nil")
)

type pskConn struct {
	tconn.Conn
	psk *[32]byte

	writeS20 cipher.Stream
	readS20  cipher.Stream
}

func (c *pskConn) Read(out []byte) (int, error) {
	if c.readS20 == nil {
		nonce := make([]byte, 24)
		_, err := io.ReadFull(c.Conn, nonce)
		if err != nil {
			return 0, errShortNonce
		}
		c.readS20 = salsa20.New(c.psk, nonce)
	}

	maxn := uint32(len(out))
	in := bufPool.Get(maxn).([]byte) // get buffer
	defer bufPool.Put(maxn, in)      // put the buffer back

	in = in[:maxn]            // truncate to required length
	n, err := c.Conn.Read(in) // read to in
	if err != nil {
		return 0, err
	}

	c.readS20.XORKeyStream(out[:n], in[:n]) // decrypt to out buffer

	return n, nil
}

func (c *pskConn) Write(in []byte) (int, error) {
	if c.writeS20 == nil {
		nonce := make([]byte, 24)
		_, err := rand.Read(nonce)
		if err != nil {
			return 0, err
		}
		_, err = c.Conn.Write(nonce)
		if err != nil {
			return 0, err
		}

		c.writeS20 = salsa20.New(c.psk, nonce)
	}
	n := uint32(len(in))
	out := bufPool.Get(n).([]byte) // get buffer
	defer bufPool.Put(n, out)      // put the buffer back

	out = out[:n]                    // truncate to required length
	c.writeS20.XORKeyStream(out, in) // encrypt

	return c.Conn.Write(out) // send
}

var _ tconn.Conn = (*pskConn)(nil)

func newPSKConn(psk *[32]byte, insecure tconn.Conn) (tconn.Conn, error) {
	if insecure == nil {
		return nil, errInsecureNil
	}
	if psk == nil {
		return nil, errPSKNil
	}
	return &pskConn{
		Conn: insecure,
		psk:  psk,
	}, nil
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"errors"
	"net"
	"os"
	"runtime"
	"time"

	logging "github.com/ipfs/go-log"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/transport"
	tptu "github.com/libp2p/go-libp2p-transport-upgrader"
	rtpt "github.com/libp2p/go-reuseport-transport"

	ma "github.com/multiformats/go-multiaddr"
	mafmt "github.com/multiformats/go-multiaddr-fmt"
	manet "github.com/multiformats/go-multiaddr/net"
)

// DefaultConnectTimeout is the (default) maximum amount of time the TCP
// transport will spend on the initial TCP connect before giving up.
var DefaultConnectTimeout = 5 * time.Second

var log = logging.Logger("tcp-tpt")

const keepAlivePeriod = 30 * time.Second

type canKeepAlive interface {
	SetKeepAlive(bool) error
	SetKeepAlivePeriod(time.Duration) error
}

var _ canKeepAlive = &net.TCPConn{}

func tryKeepAlive(conn net.Conn, keepAlive bool) {
	keepAliveConn, ok := conn.(canKeepAlive)
	if !ok {
		log.Errorf("Can't set TCP keepalives.")
		return
	}
	if err := keepAliveConn.SetKeepAlive(keepAlive); err != nil {
		// Sometimes we seem to get "invalid argument" results from this function on Darwin.
		// This might be due to a closed connection, but I can't reproduce that on Linux.
		//
		// But there's nothing we can do about invalid arguments, so we'll drop this to a
		// debug.
		if errors.Is(err, os.ErrInvalid) {
			log.Debugw("failed to enable TCP keepalive", "error", err)
		} else {
			log.Errorw("failed to enable TCP keepalive", "error", err)
		}
		return
	}

	if runtime.GOOS != "openbsd" {
		if err := keepAliveConn.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
			log.Errorw("failed set keepalive period", "error", err)
		}
	}
}

// try to set linger on the connection, if possible.
func tryLinger(conn net.Conn, sec int) {
	type canLinger interface {
		SetLinger(int) error
	}

	if lingerConn, ok := conn.(canLinger); ok {
		_ = lingerConn.SetLinger(sec)
	}
}

type tcpListener struct {
	manet.Listener
	sec int
}

func (ll *tcpListener) Accept() (manet.Conn, error) {
	c, err := ll.Listener.Accept()
	if err != nil {
		return nil, err
	}
	tryLinger(c, ll.sec)
	tryKeepAlive(c, true)
	return c, nil
}

// TcpTransport is the TCP transport.
type TcpTransport struct {
	// Connection upgrader for upgrading insecure stream connections to
	// secure multiplex connections.
	Upgrader *tptu.Upgrader

	// Explicitly disable reuseport.
	DisableReuseport bool

	// TCP connect timeout
	ConnectTimeout time.Duration

	reuse rtpt.Transport
}

var _ transport.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport(upgrader *tptu.Upgrader) *TcpTransport {
	return &TcpTransport{Upgrader: upgrader, ConnectTimeout: DefaultConnectTimeout}
}

var dialMatcher = mafmt.And(mafmt.IP, mafmt.Base(ma.P_TCP))

// CanDial returns true if this transport believes it can dial the given
// multiaddr.
func (t *TcpTransport) CanDial(addr ma.Multiaddr) bool {
	return dialMatcher.Matches(addr)
}

func (t *TcpTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	// Apply the deadline iff applicable
	if t.ConnectTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, t.ConnectTimeout)
		defer cancel()
	}

	if t.UseReuseport() {
		return t.reuse.DialContext(ctx, raddr)
	}
	var d manet.Dialer
	return d.DialContext(ctx, raddr)
}

// Dial dials the peer at the remote address.
func (t *TcpTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (transport.CapableConn, error) {
	conn, err := t.maDial(ctx, raddr)
	if err != nil {
		return nil, err
	}
	// Set linger to 0 so we never get stuck in the TIME-WAIT state. When
	// linger is 0, connections are _reset_ instead of closed with a FIN.
	// This means we can immediately reuse the 5-tuple and reconnect.
	tryLinger(conn, 0)
	tryKeepAlive(conn, true)
	c, err := newTracingConn(conn, true)
	if err != nil {
		return nil, err
	}
	return t.Upgrader.UpgradeOutbound(ctx, t, c, p)
}

// UseReuseport returns true if reuseport is enabled and available.
func (t *TcpTransport) UseReuseport() bool {
	return !t.DisableReuseport && ReuseportIsAvailable()
}

func (t *TcpTransport) maListen(laddr ma.Multiaddr) (manet.Listener, error) {
	if t.UseReuseport() {
		return t.reuse.Listen(laddr)
	}
	return manet.Listen(laddr)
}

// Listen listens on the given multiaddr.
func (t *TcpTransport) Listen(laddr ma.Multiaddr) (transport.Listener, error) {
	list, err := t.maListen(laddr)
	if err != nil {
		return nil, err
	}
	list = newTracingListener(&tcpListener{list, 0})
	return t.Upgrader.UpgradeListener(t, list), nil
}

// Protocols returns the list of terminal protocols this transport can dial.
func (t *TcpTransport) Protocols() []int {
	return []int{ma.P_TCP}
}

// Proxy always returns false for the TCP transport.
func (t *TcpTransport) Proxy() bool {
	return false
}

func (t *TcpTransport) String() string {
	return "TCP"
}
<commit_msg>chore: update go-log to v2<commit_after>package tcp

import (
	"context"
	"errors"
	"net"
	"os"
	"runtime"
	"time"

	logging "github.com/ipfs/go-log/v2"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/transport"
	tptu "github.com/libp2p/go-libp2p-transport-upgrader"
	rtpt "github.com/libp2p/go-reuseport-transport"

	ma "github.com/multiformats/go-multiaddr"
	mafmt "github.com/multiformats/go-multiaddr-fmt"
	manet "github.com/multiformats/go-multiaddr/net"
)

// DefaultConnectTimeout is the (default) maximum amount of time the TCP
// transport will spend on the initial TCP connect before giving up.
var DefaultConnectTimeout = 5 * time.Second

var log = logging.Logger("tcp-tpt")

const keepAlivePeriod = 30 * time.Second

type canKeepAlive interface {
	SetKeepAlive(bool) error
	SetKeepAlivePeriod(time.Duration) error
}

var _ canKeepAlive = &net.TCPConn{}

func tryKeepAlive(conn net.Conn, keepAlive bool) {
	keepAliveConn, ok := conn.(canKeepAlive)
	if !ok {
		log.Errorf("Can't set TCP keepalives.")
		return
	}
	if err := keepAliveConn.SetKeepAlive(keepAlive); err != nil {
		// Sometimes we seem to get "invalid argument" results from this function on Darwin.
		// This might be due to a closed connection, but I can't reproduce that on Linux.
		//
		// But there's nothing we can do about invalid arguments, so we'll drop this to a
		// debug.
		if errors.Is(err, os.ErrInvalid) {
			log.Debugw("failed to enable TCP keepalive", "error", err)
		} else {
			log.Errorw("failed to enable TCP keepalive", "error", err)
		}
		return
	}

	if runtime.GOOS != "openbsd" {
		if err := keepAliveConn.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
			log.Errorw("failed set keepalive period", "error", err)
		}
	}
}

// try to set linger on the connection, if possible.
func tryLinger(conn net.Conn, sec int) {
	type canLinger interface {
		SetLinger(int) error
	}

	if lingerConn, ok := conn.(canLinger); ok {
		_ = lingerConn.SetLinger(sec)
	}
}

type tcpListener struct {
	manet.Listener
	sec int
}

func (ll *tcpListener) Accept() (manet.Conn, error) {
	c, err := ll.Listener.Accept()
	if err != nil {
		return nil, err
	}
	tryLinger(c, ll.sec)
	tryKeepAlive(c, true)
	return c, nil
}

// TcpTransport is the TCP transport.
type TcpTransport struct {
	// Connection upgrader for upgrading insecure stream connections to
	// secure multiplex connections.
	Upgrader *tptu.Upgrader

	// Explicitly disable reuseport.
	DisableReuseport bool

	// TCP connect timeout
	ConnectTimeout time.Duration

	reuse rtpt.Transport
}

var _ transport.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport(upgrader *tptu.Upgrader) *TcpTransport {
	return &TcpTransport{Upgrader: upgrader, ConnectTimeout: DefaultConnectTimeout}
}

var dialMatcher = mafmt.And(mafmt.IP, mafmt.Base(ma.P_TCP))

// CanDial returns true if this transport believes it can dial the given
// multiaddr.
func (t *TcpTransport) CanDial(addr ma.Multiaddr) bool {
	return dialMatcher.Matches(addr)
}

func (t *TcpTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	// Apply the deadline iff applicable
	if t.ConnectTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, t.ConnectTimeout)
		defer cancel()
	}

	if t.UseReuseport() {
		return t.reuse.DialContext(ctx, raddr)
	}
	var d manet.Dialer
	return d.DialContext(ctx, raddr)
}

// Dial dials the peer at the remote address.
func (t *TcpTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (transport.CapableConn, error) {
	conn, err := t.maDial(ctx, raddr)
	if err != nil {
		return nil, err
	}
	// Set linger to 0 so we never get stuck in the TIME-WAIT state. When
	// linger is 0, connections are _reset_ instead of closed with a FIN.
	// This means we can immediately reuse the 5-tuple and reconnect.
	tryLinger(conn, 0)
	tryKeepAlive(conn, true)
	c, err := newTracingConn(conn, true)
	if err != nil {
		return nil, err
	}
	return t.Upgrader.UpgradeOutbound(ctx, t, c, p)
}

// UseReuseport returns true if reuseport is enabled and available.
func (t *TcpTransport) UseReuseport() bool {
	return !t.DisableReuseport && ReuseportIsAvailable()
}

func (t *TcpTransport) maListen(laddr ma.Multiaddr) (manet.Listener, error) {
	if t.UseReuseport() {
		return t.reuse.Listen(laddr)
	}
	return manet.Listen(laddr)
}

// Listen listens on the given multiaddr.
func (t *TcpTransport) Listen(laddr ma.Multiaddr) (transport.Listener, error) {
	list, err := t.maListen(laddr)
	if err != nil {
		return nil, err
	}
	list = newTracingListener(&tcpListener{list, 0})
	return t.Upgrader.UpgradeListener(t, list), nil
}

// Protocols returns the list of terminal protocols this transport can dial.
func (t *TcpTransport) Protocols() []int {
	return []int{ma.P_TCP}
}

// Proxy always returns false for the TCP transport.
func (t *TcpTransport) Proxy() bool {
	return false
}

func (t *TcpTransport) String() string {
	return "TCP"
}
<|endoftext|>


<commit_before>// Copyright 2022 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.

package linux

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/boot/bzimage"
	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/uio"
)

const (
	DEFAULT_INITRD_ADDR_MAX  uint = 0x37FFFFFF
	DEFAULT_BZIMAGE_ADDR_MAX uint = 0x37FFFFFF
	bootParams                    = "/sys/kernel/boot_params/data"
)

const defaultPurgatory = "default"

var (
	// Debug is called to print out verbose debug info.
	//
	// Set this to appropriate output stream for display
	// of useful debug info.
	Debug        = log.Printf // func(string, ...interface{}) {}
	curPurgatory = Purgatories[defaultPurgatory]
)

// KexecLoad loads a bzImage-formated Linux kernel file as the to-be-kexeced
// kernel with the given ramfs file and cmdline string.
//
// It uses the kexec_load system call.
func KexecLoad(kernel, ramfs *os.File, cmdline string) error {
	bzimage.Debug = Debug

	// A collection of vars used for processing the kernel for kexec
	var err error
	// bzimage is the deserialized bzImage from the kernel
	// io.ReaderAt.
	var bzimg bzimage.BzImage
	// kmem is a struct holding kexec segments.
	//
	// It has routines to work with physical memory
	// ranges.
	var kmem *kexec.Memory
	// TODO(10000TB): construct default params in go.
	//
	// boot_params directory is x86 specific. So for now, following code only
	// works on x86.
	// https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-kernel-boot_params
	bp, err := ioutil.ReadFile("/sys/kernel/boot_params/data")
	if err != nil {
		return fmt.Errorf("reading boot_param data: %w", err)
	}
	var lp = &bzimage.LinuxParams{}
	if err := lp.UnmarshalBinary(bp); err != nil {
		return fmt.Errorf("unmarshaling header: %w", err)
	}

	kb, err := uio.ReadAll(kernel)
	if err != nil {
		return fmt.Errorf("reading Linux kernel into memory: %w", err)
	}
	if err := bzimg.UnmarshalBinary(kb); err != nil {
		return fmt.Errorf("parsing bzImage Linux kernel: %w", err)
	}

	if len(bzimg.KernelCode) < 1024 {
		return fmt.Errorf("kernel code size smaller than 1024 bytes: %d", len(bzimg.KernelCode))
	}

	kelf, err := bzimg.ELF()
	if err != nil {
		return fmt.Errorf("getting ELF from bzImage: %w", err)
	}
	kernelEntry := uintptr(kelf.Entry)
	Debug("kernelEntry: %v", kernelEntry)

	// Prepare segments.
	kmem = &kexec.Memory{}
	Debug("Try parsing memory map...")
	// TODO(10000TB): refactor this call into initialization of
	// kexec.Memory, as it does not depend on specific boot.
	if err := kmem.ParseMemoryMap(); err != nil {
		return fmt.Errorf("parse memory map: %v", err)
	}

	var relocatableKernel bool
	if bzimg.Header.Protocolversion < 0x0205 {
		return fmt.Errorf("bzImage boot protocol earlier thatn 2.05 is not supported currently: %v", bzimg.Header.Protocolversion)
	}
	relocatableKernel = bzimg.Header.RelocatableKernel != 0
	// Only protected mode is currently supported.
	// In protected mode, kernel need be relocatable, or it will need to fall
	// to real mode executing.
	if !relocatableKernel {
		return errors.New("non-relocateable Kernels are not supported")
	}
	if _, err := kmem.LoadElfSegments(bytes.NewReader(bzimg.KernelCode)); err != nil {
		return fmt.Errorf("loading kernel ELF segments: %w", err)
	}

	var ramfsRange kexec.Range
	if ramfs != nil {
		ramfsContents, err := ioutil.ReadAll(ramfs)
		if err != nil {
			return fmt.Errorf("unable to read initramfs: %w", err)
		}
		if ramfsRange, err = kmem.AddKexecSegment(ramfsContents); err != nil {
			return fmt.Errorf("add initramfs segment: %w", err)
		}
		Debug("Added %d byte initramfs at %s", len(ramfsContents), ramfsRange)
		lp.Initrdstart = uint32(ramfsRange.Start)
		lp.Initrdsize = uint32(ramfsRange.Size)
	}

	Debug("Kernel cmdline to append: %s", cmdline)
	if len(cmdline) > 0 {
		var cmdlineRange kexec.Range
		Debug("Add cmdline: %s", cmdline)
		cmdlineBytes := []byte(cmdline)
		if cmdlineRange, err = kmem.AddKexecSegment(cmdlineBytes); err != nil {
			return fmt.Errorf("add cmdline segment: %v", err)
		}
		Debug("Added %d byte of cmdline at %s", len(cmdlineBytes), cmdlineRange)
		lp.CLPtr = uint32(cmdlineRange.Start)      // 2.02+
		lp.CmdLineSize = uint32(cmdlineRange.Size) // 2.06+
	}

	// The kernel is a bzImage kernel if the protocol >= 2.00 and the 0x01
	// bit (LOAD_HIGH) in the loadflags field is set.
	// TODO(10000TB): check on loadflags.
	linuxParam, err := lp.MarshalBinary()
	if err != nil {
		return fmt.Errorf("re-marshaling header: %w", err)
	}

	// TODO(10000TB): Free mem hole start end aligns by
	// max(16, pagesize).
	//
	// Push alignment logic to kexec memory functions, e.g. a similar
	// function to FindSpace.
	setupRange, err := kmem.AddPhysSegment(
		linuxParam,
		kexec.RangeFromInterval(
			uintptr(0x90000),
			uintptr(len(linuxParam)),
		),
		// TODO(10000TB): evaluate if we need to provide  option to
		// reserve from end.
		//
		// Our go code defaults to pick up a mem block of requested
		// size from beginning, e.g.
		//
		//   [Range.Start, Range.Start+memsz)
		//
		// Kexec userspace use the range from end, e.g.
		//
		//   [Range.end-memsz+1, Range.end)
		//
	)
	if err != nil {
		return fmt.Errorf("add real mode data and cmdline: %v", err)
	}

	Debug("Loaded real mode data and cmdline at: %v", setupRange)

	// Verify purgatory loads higher than the parameters.
	// TODO(10000TB): if rel_addr < setupRange.Start then return error.

	// Load purgatory.
	purgatoryEntry, err := kexec.PurgeLoad(kmem, curPurgatory.Code, kernelEntry, setupRange.Start)
	if err != nil {
		return err
	}
	Debug("purgatory entry: %v", purgatoryEntry)

	// Load it.
	if err := kexec.Load(purgatoryEntry, kmem.Segments, 0); err != nil {
		return fmt.Errorf("Linux kexec: %w", err)
	}
	return nil
}

// SelectPurgatory picks a purgatory, returning an error if none is found
func SelectPurgatory(name string) error {
	p, ok := Purgatories[name]
	if !ok {
		var s []string
		for i := range Purgatories {
			s = append(s, i)
		}
		return fmt.Errorf("%s: no such purgatory, try one of %v", name, s)

	}
	curPurgatory = p
	return nil
}
<commit_msg>boot/linux: cmdline must be null-terminated.<commit_after>// Copyright 2022 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.

package linux

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/boot/bzimage"
	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/uio"
)

const (
	DEFAULT_INITRD_ADDR_MAX  uint = 0x37FFFFFF
	DEFAULT_BZIMAGE_ADDR_MAX uint = 0x37FFFFFF
	bootParams                    = "/sys/kernel/boot_params/data"
)

const defaultPurgatory = "default"

var (
	// Debug is called to print out verbose debug info.
	//
	// Set this to appropriate output stream for display
	// of useful debug info.
	Debug        = log.Printf // func(string, ...interface{}) {}
	curPurgatory = Purgatories[defaultPurgatory]
)

// KexecLoad loads a bzImage-formated Linux kernel file as the to-be-kexeced
// kernel with the given ramfs file and cmdline string.
//
// It uses the kexec_load system call.
func KexecLoad(kernel, ramfs *os.File, cmdline string) error {
	bzimage.Debug = Debug

	// A collection of vars used for processing the kernel for kexec
	var err error
	// bzimage is the deserialized bzImage from the kernel
	// io.ReaderAt.
	var bzimg bzimage.BzImage
	// kmem is a struct holding kexec segments.
	//
	// It has routines to work with physical memory
	// ranges.
	var kmem *kexec.Memory
	// TODO(10000TB): construct default params in go.
	//
	// boot_params directory is x86 specific. So for now, following code only
	// works on x86.
	// https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-kernel-boot_params
	bp, err := ioutil.ReadFile("/sys/kernel/boot_params/data")
	if err != nil {
		return fmt.Errorf("reading boot_param data: %w", err)
	}
	var lp = &bzimage.LinuxParams{}
	if err := lp.UnmarshalBinary(bp); err != nil {
		return fmt.Errorf("unmarshaling header: %w", err)
	}

	kb, err := uio.ReadAll(kernel)
	if err != nil {
		return fmt.Errorf("reading Linux kernel into memory: %w", err)
	}
	if err := bzimg.UnmarshalBinary(kb); err != nil {
		return fmt.Errorf("parsing bzImage Linux kernel: %w", err)
	}

	if len(bzimg.KernelCode) < 1024 {
		return fmt.Errorf("kernel code size smaller than 1024 bytes: %d", len(bzimg.KernelCode))
	}

	kelf, err := bzimg.ELF()
	if err != nil {
		return fmt.Errorf("getting ELF from bzImage: %w", err)
	}
	kernelEntry := uintptr(kelf.Entry)
	Debug("kernelEntry: %v", kernelEntry)

	// Prepare segments.
	kmem = &kexec.Memory{}
	Debug("Try parsing memory map...")
	// TODO(10000TB): refactor this call into initialization of
	// kexec.Memory, as it does not depend on specific boot.
	if err := kmem.ParseMemoryMap(); err != nil {
		return fmt.Errorf("parse memory map: %v", err)
	}

	var relocatableKernel bool
	if bzimg.Header.Protocolversion < 0x0205 {
		return fmt.Errorf("bzImage boot protocol earlier thatn 2.05 is not supported currently: %v", bzimg.Header.Protocolversion)
	}
	relocatableKernel = bzimg.Header.RelocatableKernel != 0
	// Only protected mode is currently supported.
	// In protected mode, kernel need be relocatable, or it will need to fall
	// to real mode executing.
	if !relocatableKernel {
		return errors.New("non-relocateable Kernels are not supported")
	}
	if _, err := kmem.LoadElfSegments(bytes.NewReader(bzimg.KernelCode)); err != nil {
		return fmt.Errorf("loading kernel ELF segments: %w", err)
	}

	var ramfsRange kexec.Range
	if ramfs != nil {
		ramfsContents, err := ioutil.ReadAll(ramfs)
		if err != nil {
			return fmt.Errorf("unable to read initramfs: %w", err)
		}
		if ramfsRange, err = kmem.AddKexecSegment(ramfsContents); err != nil {
			return fmt.Errorf("add initramfs segment: %w", err)
		}
		Debug("Added %d byte initramfs at %s", len(ramfsContents), ramfsRange)
		lp.Initrdstart = uint32(ramfsRange.Start)
		lp.Initrdsize = uint32(ramfsRange.Size)
	}

	Debug("Kernel cmdline to append: %s", cmdline)
	if len(cmdline) > 0 {
		var cmdlineRange kexec.Range
		Debug("Add cmdline: %s", cmdline)

		// Cmdline must be null-terminated.
		cmdlineBytes := []byte(cmdline + "\x00")
		if cmdlineRange, err = kmem.AddKexecSegment(cmdlineBytes); err != nil {
			return fmt.Errorf("add cmdline segment: %v", err)
		}
		Debug("Added %d byte of cmdline at %s", len(cmdlineBytes), cmdlineRange)
		lp.CLPtr = uint32(cmdlineRange.Start)      // 2.02+
		lp.CmdLineSize = uint32(cmdlineRange.Size) // 2.06+
	}

	// The kernel is a bzImage kernel if the protocol >= 2.00 and the 0x01
	// bit (LOAD_HIGH) in the loadflags field is set.
	// TODO(10000TB): check on loadflags.
	linuxParam, err := lp.MarshalBinary()
	if err != nil {
		return fmt.Errorf("re-marshaling header: %w", err)
	}

	// TODO(10000TB): Free mem hole start end aligns by
	// max(16, pagesize).
	//
	// Push alignment logic to kexec memory functions, e.g. a similar
	// function to FindSpace.
	setupRange, err := kmem.AddPhysSegment(
		linuxParam,
		kexec.RangeFromInterval(
			uintptr(0x90000),
			uintptr(len(linuxParam)),
		),
		// TODO(10000TB): evaluate if we need to provide  option to
		// reserve from end.
		//
		// Our go code defaults to pick up a mem block of requested
		// size from beginning, e.g.
		//
		//   [Range.Start, Range.Start+memsz)
		//
		// Kexec userspace use the range from end, e.g.
		//
		//   [Range.end-memsz+1, Range.end)
		//
	)
	if err != nil {
		return fmt.Errorf("add real mode data and cmdline: %v", err)
	}

	Debug("Loaded real mode data and cmdline at: %v", setupRange)

	// Verify purgatory loads higher than the parameters.
	// TODO(10000TB): if rel_addr < setupRange.Start then return error.

	// Load purgatory.
	purgatoryEntry, err := kexec.PurgeLoad(kmem, curPurgatory.Code, kernelEntry, setupRange.Start)
	if err != nil {
		return err
	}
	Debug("purgatory entry: %v", purgatoryEntry)

	// Load it.
	if err := kexec.Load(purgatoryEntry, kmem.Segments, 0); err != nil {
		return fmt.Errorf("Linux kexec: %w", err)
	}
	return nil
}

// SelectPurgatory picks a purgatory, returning an error if none is found
func SelectPurgatory(name string) error {
	p, ok := Purgatories[name]
	if !ok {
		var s []string
		for i := range Purgatories {
			s = append(s, i)
		}
		return fmt.Errorf("%s: no such purgatory, try one of %v", name, s)

	}
	curPurgatory = p
	return nil
}
<|endoftext|>