<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"flag"
	"fmt"
	"hash/crc32"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/tarm/serial"
	"gopkg.in/readline.v1"
)

var (
	rlInstance *readline.Instance
	conn       io.ReadWriter
	serIn      = make(chan []byte)
	outBound   = make(chan string)
	progress   = make(chan bool, 1)
	incLevel   = make(chan int)

	port    = flag.String("p", "", "serial port (required: /dev/tty* or COM*)")
	baud    = flag.Int("b", 115200, "baud rate")
	upload  = flag.String("u", "", "upload the specified firmware, then quit")
	expand  = flag.String("e", "", "expand specified file to stdout, then quit")
	verbose = flag.Bool("v", false, "verbose output, for debugging only")
	capture = flag.String("c", "", "a file where captured output is appended")
	timeout = flag.Duration("t", 500*time.Millisecond, "serial echo timeout")
)

func main() {
	flag.Parse()
	var err error

	// expansion does not use the serial port, it just expands include lines
	if *expand != "" {
		expandFile()
		return
	}

	if *port == "" || flag.NArg() > 0 {
		flag.PrintDefaults()
		os.Exit(1)
	}
	conn, err = connect(*port)
	check(err)
	//defer conn.Close()
	fmt.Println("Connected to:", *port)

	go serialInput() // feed the serIn channel

	// firmware upload uses serial in a different way, needs to quit when done
	if *upload != "" {
		firmwareUpload()
		return
	}

	rlInstance, err = readline.NewEx(&readline.Config{
		UniqueEditLine: true,
	})
	check(err)
	defer rlInstance.Close()

	go serialExchange()

	//conn.Flush()
	readWithTimeout() // get rid of partial pending data
	for {
		line, err := rlInstance.Readline()
		if err != nil { // io.EOF, readline.ErrInterrupt
			break
		}
		parseAndSend(line)
		if strings.HasPrefix(line, "include ") {
			fmt.Println("\\ done.")
		}
	}
}

// isNetPort returns true if the argument is of the form "...:<N>"
func isNetPort(s string) bool {
	if n := strings.Index(s, ":"); n > 0 {
		p, e := strconv.Atoi(s[n+1:])
		return e == nil && p > 0
	}
	return false
}

// connect to either a net port for telnet use, or to a serial device
func connect(name string) (io.ReadWriter, error) {
	if isNetPort(name) {
		return net.Dial("tcp", name)
	}

	config := serial.Config{Name: *port, Baud: *baud}
	if *upload != "" {
		config.Parity = serial.ParityEven
	}
	return serial.OpenPort(&config)
}

func parseAndSend(line string) {
	if strings.HasPrefix(line, "include ") {
		for _, fname := range strings.Split(line[8:], " ") {
			doInclude(fname)
		}
	} else {
		outBound <- line
		<-progress
	}
}

func check(err error) {
	if err != nil {
		if rlInstance != nil {
			rlInstance.Close()
		}
		log.Fatal(err)
	}
}

func serialInput() {
	var f *os.File
	if *capture != "" {
		var err error
		opts := os.O_WRONLY | os.O_APPEND | os.O_CREATE
		f, err = os.OpenFile(*capture, opts, 0666)
		check(err)
		defer f.Close()
	}
	for {
		buf := make([]byte, 600)
		n, err := conn.Read(buf)
		check(err)
		if n == 0 {
			close(serIn)
			return
		}
		if f != nil {
			f.Write(buf[:n])
		}
		serIn <- buf[:n]
	}
}

func readWithTimeout() []byte {
	select {
	case data := <-serIn:
		return data
	case <-time.After(*timeout):
		return nil
	}
}

func serialExchange() {
	depth := 0
	for {
		select {

		case data := <-serIn:
			if len(data) == 0 {
				return
			}
			fmt.Print(string(data))

		case line := <-outBound:
			immediate := depth == 0
			// the task here is to omit "normal" output for included lines,
			// i.e. lines which only generate an echo, a space, and " ok.\n"
			// everything else should be echoed in full, including the input
			if len(line) > 0 {
				serialSend(line)
				prefix, matched := expectEcho(line, false, func(s string) {
					fmt.Print(s) // called to flush pending serial input lines
				})
				fmt.Print(prefix)
				if matched && immediate {
					fmt.Print(line)
					line = ""
				}
			}
			// now that the echo is done, send a CR and wait for the prompt
			serialSend("\r")
			prompt := " ok.\n"
			prefix, matched := expectEcho(prompt, immediate, func(s string) {
				fmt.Print(line + s) // show original command first
				line = ""
			})
			if !matched {
				prompt = ""
			}
			if immediate || prefix != " " || !matched {
				fmt.Print(line + prefix + prompt)
			}
			// signal to sender that this request has been processed
			progress <- matched

		case n := <-incLevel:
			depth += n
		}
	}
}

func expectEcho(match string, immed bool, flusher func(string)) (string, bool) {
	var collected []byte
	for {
		data := readWithTimeout()
		collected = append(collected, data...)
		if bytes.HasSuffix(collected, []byte(match)) {
			bytesBefore := len(collected) - len(match)
			return string(collected[:bytesBefore]), true
		}
		if immed || len(data) == 0 {
			return string(collected), false
		}
		if n := bytes.LastIndexByte(collected, '\n'); n >= 0 {
			flusher(string(collected[:n+1]))
			collected = collected[n+1:]
		}
	}
}

func serialSend(data string) {
	_, err := conn.Write([]byte(data))
	check(err)
}

func doInclude(fname string) {
	if fname == "" {
		return // silently ignore empty files
	}

	incLevel <- +1
	defer func() { incLevel <- -1 }()

	lineNum := 0
	fmt.Printf("\\       >>> include %s\n", fname)
	defer func() {
		fmt.Printf("\\       <<<<<<<<<<< %s (%d lines)\n", fname, lineNum)
	}()

	f, err := os.Open(fname)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		lineNum++

		s := strings.TrimLeft(line, " ")
		if s == "" || strings.HasPrefix(s, "\\ ") {
			if len(*expand) == 0 {
				continue // don't send empty or comment-only lines
			}
		}

		parseAndSend(line)
	}
}

func expandFile() {
	go func() {
		for line := range outBound {
			fmt.Println(line)
			progress <- true
		}
	}()

	go func() {
		for range incLevel {
		}
	}()

	for _, fname := range strings.Split(*expand, ",") {
		doInclude(fname)
	}
}

func hexToBin(data []byte) []byte {
	var bin []byte
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasSuffix(line, "\r") {
			line = line[:len(line)-1]
		}
		if len(line) == 0 {
			continue
		}
		if line[0] != ':' || len(line) < 11 {
			fmt.Println("Not ihex format:", line)
			os.Exit(1)
		}
		bytes, err := hex.DecodeString(line[1:])
		check(err)
		if bytes[3] != 0x00 {
			continue
		}
		offset := (int(bytes[1]) << 8) + int(bytes[2])
		length := bytes[0]
		for offset > len(bin) {
			bin = append(bin, 0xFF)
		}
		bin = append(bin, bytes[4:4+length]...)
	}
	return bin
}

func firmwareUpload() {
	f, err := os.Open(*upload)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	check(err)

	// convert to binary if first bytes look like they are in hex format
	tag := ""
	if len(data) > 11 && data[0] == ':' {
		_, err = hex.DecodeString(string(data[1:11]))
		if err == nil {
			data = hexToBin(data)
			tag = " (converted from Intel HEX)"
		}
	}

	fmt.Printf("        File: %s\n", *upload)
	fmt.Printf("       Count: %d bytes%s\n", len(data), tag)
	fmt.Printf("    Checksum: %08x hex\n", crc32.ChecksumIEEE(data))

	uploadSTM32(data)
}
<commit_msg>add "-x ..." remote execution option<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"flag"
	"fmt"
	"hash/crc32"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/tarm/serial"
	"gopkg.in/readline.v1"
)

var (
	rlInstance *readline.Instance
	conn       io.ReadWriter
	serIn      = make(chan []byte)
	outBound   = make(chan string)
	progress   = make(chan bool, 1)
	incLevel   = make(chan int)

	port    = flag.String("p", "", "serial port (required: /dev/tty* or COM*)")
	exe     = flag.Bool("x", false, "executable (consumes all remaining args)")
	baud    = flag.Int("b", 115200, "baud rate")
	upload  = flag.String("u", "", "upload the specified firmware, then quit")
	expand  = flag.String("e", "", "expand specified file to stdout, then quit")
	verbose = flag.Bool("v", false, "verbose output, for debugging only")
	capture = flag.String("c", "", "a file where captured output is appended")
	timeout = flag.Duration("t", 500*time.Millisecond, "serial echo timeout")
)

func main() {
	flag.Parse()
	var err error

	// expansion does not use the serial port, it just expands include lines
	if *expand != "" {
		expandFile()
		return
	}

	if *exe && flag.NArg() > 0 {
		conn, err = launch(flag.Args())
		*port = "(" + strings.Join(flag.Args(), " ") + ")"
	} else if *port == "" || flag.NArg() > 0 {
		flag.PrintDefaults()
		os.Exit(1)
		conn, err = connect(*port)
	}
	check(err)
	//defer conn.Close()
	fmt.Println("Connected to:", *port)

	go serialInput() // feed the serIn channel

	// firmware upload uses serial in a different way, needs to quit when done
	if *upload != "" {
		firmwareUpload()
		return
	}

	rlInstance, err = readline.NewEx(&readline.Config{
		UniqueEditLine: true,
	})
	check(err)
	defer rlInstance.Close()

	go serialExchange()

	//conn.Flush()
	readWithTimeout() // get rid of partial pending data
	for {
		line, err := rlInstance.Readline()
		if err != nil { // io.EOF, readline.ErrInterrupt
			break
		}
		parseAndSend(line)
		if strings.HasPrefix(line, "include ") {
			fmt.Println("\\ done.")
		}
	}
}

// isNetPort returns true if the argument is of the form "...:<N>"
func isNetPort(s string) bool {
	if n := strings.Index(s, ":"); n > 0 {
		p, e := strconv.Atoi(s[n+1:])
		return e == nil && p > 0
	}
	return false
}

// connect to either a net port for telnet use, or to a serial device
func connect(name string) (io.ReadWriter, error) {
	if isNetPort(name) {
		return net.Dial("tcp", name)
	}

	config := serial.Config{Name: *port, Baud: *baud}
	if *upload != "" {
		config.Parity = serial.ParityEven
	}
	return serial.OpenPort(&config)
}

// launch an executable and talk to it via pipes
func launch(args []string) (io.ReadWriter, error) {
	cmd := exec.Command(args[0], args[1:]...)
	var cmdio exePipe
	var err error
	cmdio.to, err = cmd.StdinPipe()
	if err == nil {
		cmdio.from, err = cmd.StdoutPipe()
		if err == nil {
			err = cmd.Start()
		}
	}
	return &cmdio, err
}

type exePipe struct {
	from io.Reader
	to   io.Writer
}

func (e *exePipe) Read(p []byte) (n int, err error) {
	return e.from.Read(p)
}

func (e *exePipe) Write(p []byte) (n int, err error) {
	return e.to.Write(p)
}

func parseAndSend(line string) {
	if strings.HasPrefix(line, "include ") {
		for _, fname := range strings.Split(line[8:], " ") {
			doInclude(fname)
		}
	} else {
		outBound <- line
		<-progress
	}
}

func check(err error) {
	if err != nil {
		if rlInstance != nil {
			rlInstance.Close()
		}
		log.Fatal(err)
	}
}

func serialInput() {
	var f *os.File
	if *capture != "" {
		var err error
		opts := os.O_WRONLY | os.O_APPEND | os.O_CREATE
		f, err = os.OpenFile(*capture, opts, 0666)
		check(err)
		defer f.Close()
	}
	for {
		buf := make([]byte, 600)
		n, err := conn.Read(buf)
		check(err)
		if n == 0 {
			close(serIn)
			return
		}
		if f != nil {
			f.Write(buf[:n])
		}
		serIn <- buf[:n]
	}
}

func readWithTimeout() []byte {
	select {
	case data := <-serIn:
		return data
	case <-time.After(*timeout):
		return nil
	}
}

func serialExchange() {
	depth := 0
	for {
		select {

		case data := <-serIn:
			if len(data) == 0 {
				return
			}
			fmt.Print(string(data))

		case line := <-outBound:
			immediate := depth == 0
			// the task here is to omit "normal" output for included lines,
			// i.e. lines which only generate an echo, a space, and " ok.\n"
			// everything else should be echoed in full, including the input
			if len(line) > 0 {
				serialSend(line)
				prefix, matched := expectEcho(line, false, func(s string) {
					fmt.Print(s) // called to flush pending serial input lines
				})
				fmt.Print(prefix)
				if matched && immediate {
					fmt.Print(line)
					line = ""
				}
			}
			// now that the echo is done, send a CR and wait for the prompt
			serialSend("\r")
			prompt := " ok.\n"
			prefix, matched := expectEcho(prompt, immediate, func(s string) {
				fmt.Print(line + s) // show original command first
				line = ""
			})
			if !matched {
				prompt = ""
			}
			if immediate || prefix != " " || !matched {
				fmt.Print(line + prefix + prompt)
			}
			// signal to sender that this request has been processed
			progress <- matched

		case n := <-incLevel:
			depth += n
		}
	}
}

func expectEcho(match string, immed bool, flusher func(string)) (string, bool) {
	var collected []byte
	for {
		data := readWithTimeout()
		collected = append(collected, data...)
		if bytes.HasSuffix(collected, []byte(match)) {
			bytesBefore := len(collected) - len(match)
			return string(collected[:bytesBefore]), true
		}
		if immed || len(data) == 0 {
			return string(collected), false
		}
		if n := bytes.LastIndexByte(collected, '\n'); n >= 0 {
			flusher(string(collected[:n+1]))
			collected = collected[n+1:]
		}
	}
}

func serialSend(data string) {
	_, err := conn.Write([]byte(data))
	check(err)
}

func doInclude(fname string) {
	if fname == "" {
		return // silently ignore empty files
	}

	incLevel <- +1
	defer func() { incLevel <- -1 }()

	lineNum := 0
	fmt.Printf("\\       >>> include %s\n", fname)
	defer func() {
		fmt.Printf("\\       <<<<<<<<<<< %s (%d lines)\n", fname, lineNum)
	}()

	f, err := os.Open(fname)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		lineNum++

		s := strings.TrimLeft(line, " ")
		if s == "" || strings.HasPrefix(s, "\\ ") {
			if len(*expand) == 0 {
				continue // don't send empty or comment-only lines
			}
		}

		parseAndSend(line)
	}
}

func expandFile() {
	go func() {
		for line := range outBound {
			fmt.Println(line)
			progress <- true
		}
	}()

	go func() {
		for range incLevel {
		}
	}()

	for _, fname := range strings.Split(*expand, ",") {
		doInclude(fname)
	}
}

func hexToBin(data []byte) []byte {
	var bin []byte
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasSuffix(line, "\r") {
			line = line[:len(line)-1]
		}
		if len(line) == 0 {
			continue
		}
		if line[0] != ':' || len(line) < 11 {
			fmt.Println("Not ihex format:", line)
			os.Exit(1)
		}
		bytes, err := hex.DecodeString(line[1:])
		check(err)
		if bytes[3] != 0x00 {
			continue
		}
		offset := (int(bytes[1]) << 8) + int(bytes[2])
		length := bytes[0]
		for offset > len(bin) {
			bin = append(bin, 0xFF)
		}
		bin = append(bin, bytes[4:4+length]...)
	}
	return bin
}

func firmwareUpload() {
	f, err := os.Open(*upload)
	if err != nil {
		fmt.Println(err)
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	check(err)

	// convert to binary if first bytes look like they are in hex format
	tag := ""
	if len(data) > 11 && data[0] == ':' {
		_, err = hex.DecodeString(string(data[1:11]))
		if err == nil {
			data = hexToBin(data)
			tag = " (converted from Intel HEX)"
		}
	}

	fmt.Printf("        File: %s\n", *upload)
	fmt.Printf("       Count: %d bytes%s\n", len(data), tag)
	fmt.Printf("    Checksum: %08x hex\n", crc32.ChecksumIEEE(data))

	uploadSTM32(data)
}
<|endoftext|>


<commit_before>package gorethink

import (
	"code.google.com/p/goprotobuf/proto"
	"fmt"
	p "github.com/dancannon/gorethink/ql2"
	"strconv"
	"strings"
)

type termsList []RqlTerm
type termsObj map[string]RqlTerm
type RqlTerm struct {
	name     string
	termType p.Term_TermType
	data     interface{}
	args     []RqlTerm
	optArgs  map[string]RqlTerm
}

// build takes the query tree and turns it into a protobuf term tree.
func (t RqlTerm) build() *p.Term {
	switch t.termType {
	case p.Term_DATUM:
		datum, err := constructDatum(t)
		if err != nil {
			panic(err)
		}
		return datum
	default:
		args := []*p.Term{}
		optArgs := []*p.Term_AssocPair{}
		term := &p.Term{
			Type: t.termType.Enum(),
		}

		for _, v := range t.args {
			args = append(args, v.build())
		}

		for k, v := range t.optArgs {
			optArgs = append(optArgs, &p.Term_AssocPair{
				Key: proto.String(k),
				Val: v.build(),
			})
		}

		term.Args = args
		term.Optargs = optArgs

		return term
	}
}

// compose returns a string representation of the query tree
func (t RqlTerm) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.RqlTerm) r.RqlTerm { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}

	default:
		if t.name != "" {
			return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
		} else {
			return fmt.Sprintf("(%s)", strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
		}
	}
}

type WriteResponse struct {
	Inserted      int
	Errors        int
	Updated       int
	Unchanged     int
	Replaced      int
	Deleted       int
	GeneratedKeys []string    `gorethink:"generated_keys"`
	FirstError    string      `gorethink:"first_error"` // populated if Errors > 0
	NewValue      interface{} `gorethink:"new_val"`
	OldValue      interface{} `gorethink:"old_val"`
}

// Run runs a query using the given connection.
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//	for rows.Next() {
//		doc := MyDocumentType{}
//		err := r.Scan(&doc)
//		    // Do something with row
//	}
func (t RqlTerm) Run(s *Session, args ...interface{}) (*ResultRows, error) {
	argm := optArgsToMap([]string{"db", "use_outdated", "noreply", "time_format"}, args)
	return s.startQuery(t, argm)
}

// Run runs a query using the given connection but unlike Run returns ResultRow.
// This function should be used if your query only returns a single row.
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	row, err := query.RunRow(sess, "use_outdated", true)
//	if err != nil {
//		// error
//	}
//	if row.IsNil() {
//		// nothing was found
//	}
//	err = row.Scan(&doc)
func (t RqlTerm) RunRow(s *Session, args ...interface{}) (*ResultRow, error) {
	rows, err := t.Run(s, args...)
	if err == nil {
		defer rows.Close()
		if !rows.Next() {
			err = RqlDriverError{"No rows in the result set"}
		}
	}
	return &ResultRow{rows: rows, err: err}, err
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponse. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess, "noreply", true)
func (t RqlTerm) RunWrite(s *Session, args ...interface{}) (WriteResponse, error) {
	var response WriteResponse
	row, err := t.RunRow(s, args...)
	if err == nil {
		err = row.Scan(&response)
	}
	return response, err
}

// Exec runs the query but does not return the result (It also automatically sets
// the noreply option).
//
// Optional arguments :
// "db", "use_outdated" (defaults to false) and "time_format".
func (t RqlTerm) Exec(s *Session, args ...interface{}) error {
	// Ensure that noreply is set to true
	args = append(args, "noreply")
	args = append(args, true)

	_, err := t.Run(s, args...)
	return err
}
<commit_msg>Added Created field so that RunWrite will create an accurate WriteResponse as a return.<commit_after>package gorethink

import (
	"code.google.com/p/goprotobuf/proto"
	"fmt"
	p "github.com/dancannon/gorethink/ql2"
	"strconv"
	"strings"
)

type termsList []RqlTerm
type termsObj map[string]RqlTerm
type RqlTerm struct {
	name     string
	termType p.Term_TermType
	data     interface{}
	args     []RqlTerm
	optArgs  map[string]RqlTerm
}

// build takes the query tree and turns it into a protobuf term tree.
func (t RqlTerm) build() *p.Term {
	switch t.termType {
	case p.Term_DATUM:
		datum, err := constructDatum(t)
		if err != nil {
			panic(err)
		}
		return datum
	default:
		args := []*p.Term{}
		optArgs := []*p.Term_AssocPair{}
		term := &p.Term{
			Type: t.termType.Enum(),
		}

		for _, v := range t.args {
			args = append(args, v.build())
		}

		for k, v := range t.optArgs {
			optArgs = append(optArgs, &p.Term_AssocPair{
				Key: proto.String(k),
				Val: v.build(),
			})
		}

		term.Args = args
		term.Optargs = optArgs

		return term
	}
}

// compose returns a string representation of the query tree
func (t RqlTerm) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.RqlTerm) r.RqlTerm { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}

	default:
		if t.name != "" {
			return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
		} else {
			return fmt.Sprintf("(%s)", strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
		}
	}
}

type WriteResponse struct {
	Inserted      int
	Errors        int
	Updated       int
	Unchanged     int
	Replaced      int
	Deleted       int
	Created       int
	GeneratedKeys []string    `gorethink:"generated_keys"`
	FirstError    string      `gorethink:"first_error"` // populated if Errors > 0
	NewValue      interface{} `gorethink:"new_val"`
	OldValue      interface{} `gorethink:"old_val"`
}

// Run runs a query using the given connection.
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//	for rows.Next() {
//		doc := MyDocumentType{}
//		err := r.Scan(&doc)
//		    // Do something with row
//	}
func (t RqlTerm) Run(s *Session, args ...interface{}) (*ResultRows, error) {
	argm := optArgsToMap([]string{"db", "use_outdated", "noreply", "time_format"}, args)
	return s.startQuery(t, argm)
}

// Run runs a query using the given connection but unlike Run returns ResultRow.
// This function should be used if your query only returns a single row.
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	row, err := query.RunRow(sess, "use_outdated", true)
//	if err != nil {
//		// error
//	}
//	if row.IsNil() {
//		// nothing was found
//	}
//	err = row.Scan(&doc)
func (t RqlTerm) RunRow(s *Session, args ...interface{}) (*ResultRow, error) {
	rows, err := t.Run(s, args...)
	if err == nil {
		defer rows.Close()
		if !rows.Next() {
			err = RqlDriverError{"No rows in the result set"}
		}
	}
	return &ResultRow{rows: rows, err: err}, err
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponse. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess, "noreply", true)
func (t RqlTerm) RunWrite(s *Session, args ...interface{}) (WriteResponse, error) {
	var response WriteResponse
	row, err := t.RunRow(s, args...)
	if err == nil {
		err = row.Scan(&response)
	}
	return response, err
}

// Exec runs the query but does not return the result (It also automatically sets
// the noreply option).
//
// Optional arguments :
// "db", "use_outdated" (defaults to false) and "time_format".
func (t RqlTerm) Exec(s *Session, args ...interface{}) error {
	// Ensure that noreply is set to true
	args = append(args, "noreply")
	args = append(args, true)

	_, err := t.Run(s, args...)
	return err
}
<|endoftext|>


<commit_before><commit_msg>Add test<commit_after><|endoftext|>


<commit_before><commit_msg>Added LinearSpace function<commit_after><|endoftext|>


<commit_before>package log

import (
	"context"
	"crypto/rand"
	"fmt"
	"net/http"
	"strings"
	"time"
)

// RequestID is but a simple token for tracing requests.
type RequestID struct {
	id []byte
}

// String returns a string formatting for the request id.
func (r *RequestID) String() string {
	return fmt.Sprintf("%X-%X", r.id[0:4], r.id[4:8])
}

// NewRequestID returns a new random request id.
func newRequestID() *RequestID {
	r := &RequestID{
		id: make([]byte, 8),
	}
	rand.Read(r.id)
	return r
}

type ctxKey struct{}

// Trace retreives the request id from a request as a string.
func Trace(r *http.Request) string {
	rid, ok := r.Context().Value(&ctxKey{}).(*RequestID)
	if ok {
		return rid.String()
	}
	return ""
}

// GetRequestID retreives the request id from a request.
func GetRequestID(r *http.Request) *RequestID {
	return r.Context().Value(&ctxKey{}).(*RequestID)
}

// SetRequestID saves the request id in the request context.
func SetRequestID(r *http.Request, rid *RequestID) *http.Request {
	ctx := r.Context()
	ctx = context.WithValue(ctx, &ctxKey{}, rid)
	return r.WithContext(ctx)
}

// Middleware adds a logging wrapper and request tracing to requests.
func Middleware(h http.HandlerFunc) http.HandlerFunc {

	return func(w http.ResponseWriter, r *http.Request) {
		requestID := newRequestID()
		r = SetRequestID(r, requestID) // Sets on context for handlers

		level := LevelInfo

		// For assets etc, use level debug as they clutter up logs
		if r.URL.Path == "/favicon.ico" || strings.HasPrefix(r.URL.Path, "/assets") {
			level = LevelDebug
		}

		Log(Values{
			MessageKey: "<- Request",
			"method":   r.Method,
			URLKey:     r.RequestURI,
			"len":      r.ContentLength,
			IPKey:      r.RemoteAddr,
			TraceKey:   requestID.String(),
			LevelKey:   level,
		})

		start := time.Now()
		h(w, r)

		Time(start, Values{
			MessageKey: "-> Response",
			URLKey:     r.RequestURI,
			TraceKey:   requestID.String(),
			LevelKey:   level,
		})
	}

}
<commit_msg>Updated display of trace variables<commit_after>package log

import (
	"context"
	"crypto/rand"
	"fmt"
	"net/http"
	"strings"
	"time"
)

// RequestID is but a simple token for tracing requests.
type RequestID struct {
	id []byte
}

// String returns a string formatting for the request id.
func (r *RequestID) String() string {
	return fmt.Sprintf("%X-%X-%X-%X", r.id[0:2], r.id[2:4], r.id[4:6], r.id[6:8])
}

// NewRequestID returns a new random request id.
func newRequestID() *RequestID {
	r := &RequestID{
		id: make([]byte, 8),
	}
	rand.Read(r.id)
	return r
}

type ctxKey struct{}

// Trace retreives the request id from a request as a string.
func Trace(r *http.Request) string {
	rid, ok := r.Context().Value(&ctxKey{}).(*RequestID)
	if ok {
		return rid.String()
	}
	return ""
}

// GetRequestID retreives the request id from a request.
func GetRequestID(r *http.Request) *RequestID {
	return r.Context().Value(&ctxKey{}).(*RequestID)
}

// SetRequestID saves the request id in the request context.
func SetRequestID(r *http.Request, rid *RequestID) *http.Request {
	ctx := r.Context()
	ctx = context.WithValue(ctx, &ctxKey{}, rid)
	return r.WithContext(ctx)
}

// Middleware adds a logging wrapper and request tracing to requests.
func Middleware(h http.HandlerFunc) http.HandlerFunc {

	return func(w http.ResponseWriter, r *http.Request) {
		requestID := newRequestID()
		r = SetRequestID(r, requestID) // Sets on context for handlers

		level := LevelInfo

		// For assets etc, use level debug as they clutter up logs
		if r.URL.Path == "/favicon.ico" || strings.HasPrefix(r.URL.Path, "/assets") {
			level = LevelDebug
		}

		Log(Values{
			MessageKey: "<- Request",
			"method":   r.Method,
			URLKey:     r.RequestURI,
			"len":      r.ContentLength,
			IPKey:      r.RemoteAddr,
			TraceKey:   requestID.String(),
			LevelKey:   level,
		})

		start := time.Now()
		h(w, r)

		Time(start, Values{
			MessageKey: "-> Response",
			URLKey:     r.RequestURI,
			TraceKey:   requestID.String(),
			LevelKey:   level,
		})
	}

}
<|endoftext|>


<commit_before>// Copyright 2013, Cong Ding. All rights reserved.
//
// 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.
//
// author: Cong Ding <dinggnu@gmail.com>
//
package logging

import (
	"fmt"
	"os"
	"runtime"
	"sync/atomic"
	"time"
)

type log struct {
	level           Level
	name            string
	nextSeqid       uint64
	levelno         int
	levelname       string
	pathname        string
	filename        string
	module          string
	lineno          int
	funcName        string
	created         int64
	asctime         string
	msecs           int64
	relativeCreated int64
	thread          int
	threadName      string
	process         int
	message         string
	timestamp       int64
}

type field func(*logging, *log) interface{}

var fields = map[string]field{
	"name":            (*logging).lname,
	"nextSeqid":       (*logging).nextSeqid,
	"levelno":         (*logging).levelno,
	"levelname":       (*logging).levelname,
	"pathname":        (*logging).pathname,
	"filename":        (*logging).filename,
	"module":          (*logging).module,
	"lineno":          (*logging).lineno,
	"funcName":        (*logging).funcName,
	"created":         (*logging).created,
	"asctime":         (*logging).asctime,
	"msecs":           (*logging).msecs,
	"relativeCreated": (*logging).relativeCreated,
	"thread":          (*logging).thread,
	"threadName":      (*logging).threadName,
	"process":         (*logging).process,
	"message":         (*logging).message,
	"timestamp":       (*logging).timestamp,
}

const errString = "???"

// GetGoId returns the id of goroutine, which is defined in ./get_go_id.c
func GetGoId() int32

// generate the runtime information, including pathname, function name,
// filename, line number.
func genRuntime(l *log) {
	calldepth := 5
	pc, file, line, ok := runtime.Caller(calldepth)
	if ok {
		// generate short filename
		short := file
		for i := len(file) - 1; i > 0; i-- {
			if file[i] == '/' {
				short = file[i+1:]
				break
			}
		}

		// generate short function name
		fname := runtime.FuncForPC(pc).Name()
		fshort := fname
		for i := len(fname) - 1; i > 0; i-- {
			if fname[i] == '.' {
				fshort = fname[i+1:]
				break
			}
		}

		l.pathname = file
		l.funcName = fshort
		l.filename = short
		l.lineno = line
	} else {
		l.pathname = errString
		l.funcName = errString
		l.filename = errString
		l.lineno = 0
	}
}

func (logger *logging) lname(l *log) interface{} {
	return logger.name
}

func (logger *logging) nextSeqid(l *log) interface{} {
	return atomic.AddUint64(&(logger.seqid), 1)
}

func (logger *logging) levelno(l *log) interface{} {
	return int(l.level)
}

func (logger *logging) levelname(l *log) interface{} {
	return levelNames[l.level]
}

func (logger *logging) pathname(l *log) interface{} {
	if l.pathname == "" {
		genRuntime(l)
	}
	return l.pathname
}

func (logger *logging) filename(l *log) interface{} {
	if l.filename == "" {
		genRuntime(l)
	}
	return l.filename
}

func (logger *logging) module(l *log) interface{} {
	return ""
}

func (logger *logging) lineno(l *log) interface{} {
	if l.lineno == 0 {
		genRuntime(l)
	}
	return l.lineno
}

func (logger *logging) funcName(l *log) interface{} {
	if l.funcName == "" {
		genRuntime(l)
	}
	return l.funcName
}

func (logger *logging) created(l *log) interface{} {
	return logger.startTime
}

func (logger *logging) asctime(l *log) interface{} {
	if l.asctime == "" {
		l.asctime = time.Now().String()
	}
	return l.asctime
}

func (logger *logging) msecs(l *log) interface{} {
	if l.msecs == 0 {
		l.msecs = logger.startTime % 1000
	}
	return l.msecs
}

func (logger *logging) timestamp(l *log) interface{} {
	if l.timestamp == 0 {
		l.timestamp = time.Now().UnixNano()
	}
	return l.timestamp
}

func (logger *logging) relativeCreated(l *log) interface{} {
	if l.relativeCreated == 0 {
		l.relativeCreated = time.Now().UnixNano() - logger.startTime
	}
	return l.relativeCreated
}

func (logger *logging) thread(l *log) interface{} {
	if l.thread == 0 {
		l.thread = int(GetGoId())
	}
	return l.thread
}

func (logger *logging) threadName(l *log) interface{} {
	if l.threadName == "" {
		l.threadName = fmt.Sprintf("Thread-%d", GetGoId())
	}
	return l.threadName
}

// Process ID
func (logger *logging) process(l *log) interface{} {
	if l.process == 0 {
		l.process = os.Getpid()
	}
	return l.process
}

func (logger *logging) message(l *log) interface{} {
	return l.message
}
<commit_msg>use path.Base instead of parsing string<commit_after>// Copyright 2013, Cong Ding. All rights reserved.
//
// 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.
//
// author: Cong Ding <dinggnu@gmail.com>
//
package logging

import (
	"fmt"
	"os"
	"path"
	"runtime"
	"sync/atomic"
	"time"
)

type log struct {
	level           Level
	name            string
	nextSeqid       uint64
	levelno         int
	levelname       string
	pathname        string
	filename        string
	module          string
	lineno          int
	funcName        string
	created         int64
	asctime         string
	msecs           int64
	relativeCreated int64
	thread          int
	threadName      string
	process         int
	message         string
	timestamp       int64
}

type field func(*logging, *log) interface{}

var fields = map[string]field{
	"name":            (*logging).lname,
	"nextSeqid":       (*logging).nextSeqid,
	"levelno":         (*logging).levelno,
	"levelname":       (*logging).levelname,
	"pathname":        (*logging).pathname,
	"filename":        (*logging).filename,
	"module":          (*logging).module,
	"lineno":          (*logging).lineno,
	"funcName":        (*logging).funcName,
	"created":         (*logging).created,
	"asctime":         (*logging).asctime,
	"msecs":           (*logging).msecs,
	"relativeCreated": (*logging).relativeCreated,
	"thread":          (*logging).thread,
	"threadName":      (*logging).threadName,
	"process":         (*logging).process,
	"message":         (*logging).message,
	"timestamp":       (*logging).timestamp,
}

const errString = "???"

// GetGoId returns the id of goroutine, which is defined in ./get_go_id.c
func GetGoId() int32

// generate the runtime information, including pathname, function name,
// filename, line number.
func genRuntime(l *log) {
	calldepth := 5
	pc, file, line, ok := runtime.Caller(calldepth)
	if ok {
		// generate short filename

		// generate short function name
		fname := runtime.FuncForPC(pc).Name()
		fshort := fname
		for i := len(fname) - 1; i > 0; i-- {
			if fname[i] == '.' {
				fshort = fname[i+1:]
				break
			}
		}

		l.pathname = file
		l.funcName = fshort
		l.filename = path.Base(file)
		l.lineno = line
	} else {
		l.pathname = errString
		l.funcName = errString
		l.filename = errString
		l.lineno = 0
	}
}

func (logger *logging) lname(l *log) interface{} {
	return logger.name
}

func (logger *logging) nextSeqid(l *log) interface{} {
	return atomic.AddUint64(&(logger.seqid), 1)
}

func (logger *logging) levelno(l *log) interface{} {
	return int(l.level)
}

func (logger *logging) levelname(l *log) interface{} {
	return levelNames[l.level]
}

func (logger *logging) pathname(l *log) interface{} {
	if l.pathname == "" {
		genRuntime(l)
	}
	return l.pathname
}

func (logger *logging) filename(l *log) interface{} {
	if l.filename == "" {
		genRuntime(l)
	}
	return l.filename
}

func (logger *logging) module(l *log) interface{} {
	return ""
}

func (logger *logging) lineno(l *log) interface{} {
	if l.lineno == 0 {
		genRuntime(l)
	}
	return l.lineno
}

func (logger *logging) funcName(l *log) interface{} {
	if l.funcName == "" {
		genRuntime(l)
	}
	return l.funcName
}

func (logger *logging) created(l *log) interface{} {
	return logger.startTime
}

func (logger *logging) asctime(l *log) interface{} {
	if l.asctime == "" {
		l.asctime = time.Now().String()
	}
	return l.asctime
}

func (logger *logging) msecs(l *log) interface{} {
	if l.msecs == 0 {
		l.msecs = logger.startTime % 1000
	}
	return l.msecs
}

func (logger *logging) timestamp(l *log) interface{} {
	if l.timestamp == 0 {
		l.timestamp = time.Now().UnixNano()
	}
	return l.timestamp
}

func (logger *logging) relativeCreated(l *log) interface{} {
	if l.relativeCreated == 0 {
		l.relativeCreated = time.Now().UnixNano() - logger.startTime
	}
	return l.relativeCreated
}

func (logger *logging) thread(l *log) interface{} {
	if l.thread == 0 {
		l.thread = int(GetGoId())
	}
	return l.thread
}

func (logger *logging) threadName(l *log) interface{} {
	if l.threadName == "" {
		l.threadName = fmt.Sprintf("Thread-%d", GetGoId())
	}
	return l.threadName
}

// Process ID
func (logger *logging) process(l *log) interface{} {
	if l.process == 0 {
		l.process = os.Getpid()
	}
	return l.process
}

func (logger *logging) message(l *log) interface{} {
	return l.message
}
<|endoftext|>


<commit_before>package logging

import (
	"encoding/json"
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/Sirupsen/logrus/formatters/logstash"
	"net/http"
	"os"
	"strings"
	"time"
)

var Logger *logrus.Logger

// The of cookies which should not be logged
var AccessLogCookiesBlacklist = []string{}

var LifecycleEnvVars = []string{"BUILD_NUMBER", "BUILD_HASH", "BUILD_DATE"}

func init() {
	Set("info", false)
}

// Set creates a new Logger with the matching specification
func Set(level string, textLogging bool) error {
	l, err := logrus.ParseLevel(level)
	if err != nil {
		return err
	}

	logger := logrus.New()
	if textLogging {
		logger.Formatter = &logrus.TextFormatter{}
	} else {
		logger.Formatter = &logstash.LogstashFormatter{TimestampFormat: time.RFC3339Nano}
	}
	logger.Level = l
	Logger = logger
	return nil
}

// Access logs an access entry with call duration and status code
func Access(r *http.Request, start time.Time, statusCode int) {
	e := access(r, start, statusCode, nil)

	var msg string
	if len(r.URL.RawQuery) == 0 {
		msg = fmt.Sprintf("%v ->%v %v", statusCode, r.Method, r.URL.Path)
	} else {
		msg = fmt.Sprintf("%v ->%v %v?...", statusCode, r.Method, r.URL.Path)
	}

	if statusCode >= 200 && statusCode <= 399 {
		e.Info(msg)
	} else if statusCode >= 400 && statusCode <= 499 {
		e.Warn(msg)
	} else {
		e.Error(msg)
	}
}

// AccessError logs an error while accessing
func AccessError(r *http.Request, start time.Time, err error) {
	e := access(r, start, 0, err)
	e.Errorf("ERROR ->%v %v", r.Method, r.URL.Path)
}

func access(r *http.Request, start time.Time, statusCode int, err error) *logrus.Entry {
	url := r.URL.Path
	if r.URL.RawQuery != "" {
		url += "?" + r.URL.RawQuery
	}
	fields := logrus.Fields{
		"type":       "access",
		"@timestamp": start,
		"remote_ip":  getRemoteIp(r),
		"host":       r.Host,
		"url":        url,
		"method":     r.Method,
		"proto":      r.Proto,
		"duration":   time.Since(start).Nanoseconds() / 1000000,
	}

	if statusCode != 0 {
		fields["response_status"] = statusCode
	}

	if err != nil {
		fields[logrus.ErrorKey] = err.Error()
	}

	setCorrelationIds(fields, r.Header)

	cookies := map[string]string{}
	for _, c := range r.Cookies() {
		if !contains(AccessLogCookiesBlacklist, c.Name) {
			cookies[c.Name] = c.Value
		}
	}
	if len(cookies) > 0 {
		fields["cookies"] = cookies
	}

	return Logger.WithFields(fields)
}

// Call logs the result of an outgoing call
func Call(r *http.Request, resp *http.Response, start time.Time, err error) {
	url := r.URL.Path
	if r.URL.RawQuery != "" {
		url += "?" + r.URL.RawQuery
	}
	fields := logrus.Fields{
		"type":       "call",
		"@timestamp": start,
		"host":       r.Host,
		"url":        url,
		"full_url":   r.URL.String(),
		"method":     r.Method,
		"duration":   time.Since(start).Nanoseconds() / 1000000,
	}

	setCorrelationIds(fields, r.Header)

	if err != nil {
		fields[logrus.ErrorKey] = err.Error()
		Logger.WithFields(fields).Error(err)
		return
	}

	if resp != nil {
		fields["response_status"] = resp.StatusCode
		fields["content_type"] = resp.Header.Get("Content-Type")
		e := Logger.WithFields(fields)
		msg := fmt.Sprintf("%v %v-> %v", resp.StatusCode, r.Method, r.URL.String())

		if resp.StatusCode >= 200 && resp.StatusCode <= 399 {
			e.Info(msg)
		} else if resp.StatusCode >= 400 && resp.StatusCode <= 499 {
			e.Warn(msg)
		} else {
			e.Error(msg)
		}
		return
	}

	Logger.WithFields(fields).Warn("call, but no response given")
}

// Cacheinfo logs the hit information a accessing a ressource
func Cacheinfo(url string, hit bool) {
	var msg string
	if hit {
		msg = fmt.Sprintf("cache hit: %v", url)
	} else {
		msg = fmt.Sprintf("cache miss: %v", url)
	}
	Logger.WithFields(
		logrus.Fields{
			"type": "cacheinfo",
			"url":  url,
			"hit":  hit,
		}).
		Debug(msg)
}

// Return a log entry for application logs,
// prefilled with the correlation ids out of the supplied request.
func Application(h http.Header) *logrus.Entry {
	fields := logrus.Fields{
		"type": "application",
	}
	setCorrelationIds(fields, h)
	return Logger.WithFields(fields)
}

// LifecycleStart logs the start of an application
// with the configuration struct or map as paramter.
func LifecycleStart(appName string, args interface{}) {
	fields := logrus.Fields{}

	jsonString, err := json.Marshal(args)
	if err == nil {
		err := json.Unmarshal(jsonString, &fields)
		if err != nil {
			fields["parse_error"] = err.Error()
		}
	}
	fields["type"] = "lifecycle"
	fields["event"] = "start"
	for _, env := range LifecycleEnvVars {
		if os.Getenv(env) != "" {
			fields[strings.ToLower(env)] = os.Getenv(env)
		}
	}

	Logger.WithFields(fields).Infof("starting application: %v", appName)
}

// LifecycleStop logs the stop of an application
func LifecycleStop(appName string, signal os.Signal, err error) {
	fields := logrus.Fields{
		"type":  "lifecycle",
		"event": "stop",
	}
	if signal != nil {
		fields["signal"] = signal.String()
	}

	if os.Getenv("BUILD_NUMBER") != "" {
		fields["build_number"] = os.Getenv("BUILD_NUMBER")
	}

	if err != nil {
		Logger.WithFields(fields).
			WithError(err).
			Errorf("stopping application: %v (%v)", appName, err)
	} else {
		Logger.WithFields(fields).Infof("stopping application: %v (%v)", appName, signal.String())
	}
}

func getRemoteIp(r *http.Request) string {
	if r.Header.Get("X-Cluster-Client-Ip") != "" {
		return r.Header.Get("X-Cluster-Client-Ip")
	}
	if r.Header.Get("X-Real-Ip") != "" {
		return r.Header.Get("X-Real-Ip")
	}
	return strings.Split(r.RemoteAddr, ":")[0]
}

func setCorrelationIds(fields logrus.Fields, h http.Header) {
	correlationId := GetCorrelationId(h)
	if correlationId != "" {
		fields["correlation_id"] = correlationId
	}

	userCorrelationId := GetUserCorrelationId(h)
	if userCorrelationId != "" {
		fields["user_correlation_id"] = userCorrelationId
	}
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
<commit_msg>moved import of logstash formater<commit_after>package logging

import (
	"encoding/json"
	"fmt"
	"github.com/Sirupsen/logrus"
	logstash "github.com/bshuster-repo/logrus-logstash-hook"
	"net/http"
	"os"
	"strings"
	"time"
)

var Logger *logrus.Logger

// The of cookies which should not be logged
var AccessLogCookiesBlacklist = []string{}

var LifecycleEnvVars = []string{"BUILD_NUMBER", "BUILD_HASH", "BUILD_DATE"}

func init() {
	Set("info", false)
}

// Set creates a new Logger with the matching specification
func Set(level string, textLogging bool) error {
	l, err := logrus.ParseLevel(level)
	if err != nil {
		return err
	}

	logger := logrus.New()
	if textLogging {
		logger.Formatter = &logrus.TextFormatter{}
	} else {
		logger.Formatter = &logstash.LogstashFormatter{TimestampFormat: time.RFC3339Nano}
	}
	logger.Level = l
	Logger = logger
	return nil
}

// Access logs an access entry with call duration and status code
func Access(r *http.Request, start time.Time, statusCode int) {
	e := access(r, start, statusCode, nil)

	var msg string
	if len(r.URL.RawQuery) == 0 {
		msg = fmt.Sprintf("%v ->%v %v", statusCode, r.Method, r.URL.Path)
	} else {
		msg = fmt.Sprintf("%v ->%v %v?...", statusCode, r.Method, r.URL.Path)
	}

	if statusCode >= 200 && statusCode <= 399 {
		e.Info(msg)
	} else if statusCode >= 400 && statusCode <= 499 {
		e.Warn(msg)
	} else {
		e.Error(msg)
	}
}

// AccessError logs an error while accessing
func AccessError(r *http.Request, start time.Time, err error) {
	e := access(r, start, 0, err)
	e.Errorf("ERROR ->%v %v", r.Method, r.URL.Path)
}

func access(r *http.Request, start time.Time, statusCode int, err error) *logrus.Entry {
	url := r.URL.Path
	if r.URL.RawQuery != "" {
		url += "?" + r.URL.RawQuery
	}
	fields := logrus.Fields{
		"type":       "access",
		"@timestamp": start,
		"remote_ip":  getRemoteIp(r),
		"host":       r.Host,
		"url":        url,
		"method":     r.Method,
		"proto":      r.Proto,
		"duration":   time.Since(start).Nanoseconds() / 1000000,
	}

	if statusCode != 0 {
		fields["response_status"] = statusCode
	}

	if err != nil {
		fields[logrus.ErrorKey] = err.Error()
	}

	setCorrelationIds(fields, r.Header)

	cookies := map[string]string{}
	for _, c := range r.Cookies() {
		if !contains(AccessLogCookiesBlacklist, c.Name) {
			cookies[c.Name] = c.Value
		}
	}
	if len(cookies) > 0 {
		fields["cookies"] = cookies
	}

	return Logger.WithFields(fields)
}

// Call logs the result of an outgoing call
func Call(r *http.Request, resp *http.Response, start time.Time, err error) {
	url := r.URL.Path
	if r.URL.RawQuery != "" {
		url += "?" + r.URL.RawQuery
	}
	fields := logrus.Fields{
		"type":       "call",
		"@timestamp": start,
		"host":       r.Host,
		"url":        url,
		"full_url":   r.URL.String(),
		"method":     r.Method,
		"duration":   time.Since(start).Nanoseconds() / 1000000,
	}

	setCorrelationIds(fields, r.Header)

	if err != nil {
		fields[logrus.ErrorKey] = err.Error()
		Logger.WithFields(fields).Error(err)
		return
	}

	if resp != nil {
		fields["response_status"] = resp.StatusCode
		fields["content_type"] = resp.Header.Get("Content-Type")
		e := Logger.WithFields(fields)
		msg := fmt.Sprintf("%v %v-> %v", resp.StatusCode, r.Method, r.URL.String())

		if resp.StatusCode >= 200 && resp.StatusCode <= 399 {
			e.Info(msg)
		} else if resp.StatusCode >= 400 && resp.StatusCode <= 499 {
			e.Warn(msg)
		} else {
			e.Error(msg)
		}
		return
	}

	Logger.WithFields(fields).Warn("call, but no response given")
}

// Cacheinfo logs the hit information a accessing a ressource
func Cacheinfo(url string, hit bool) {
	var msg string
	if hit {
		msg = fmt.Sprintf("cache hit: %v", url)
	} else {
		msg = fmt.Sprintf("cache miss: %v", url)
	}
	Logger.WithFields(
		logrus.Fields{
			"type": "cacheinfo",
			"url":  url,
			"hit":  hit,
		}).
		Debug(msg)
}

// Return a log entry for application logs,
// prefilled with the correlation ids out of the supplied request.
func Application(h http.Header) *logrus.Entry {
	fields := logrus.Fields{
		"type": "application",
	}
	setCorrelationIds(fields, h)
	return Logger.WithFields(fields)
}

// LifecycleStart logs the start of an application
// with the configuration struct or map as paramter.
func LifecycleStart(appName string, args interface{}) {
	fields := logrus.Fields{}

	jsonString, err := json.Marshal(args)
	if err == nil {
		err := json.Unmarshal(jsonString, &fields)
		if err != nil {
			fields["parse_error"] = err.Error()
		}
	}
	fields["type"] = "lifecycle"
	fields["event"] = "start"
	for _, env := range LifecycleEnvVars {
		if os.Getenv(env) != "" {
			fields[strings.ToLower(env)] = os.Getenv(env)
		}
	}

	Logger.WithFields(fields).Infof("starting application: %v", appName)
}

// LifecycleStop logs the stop of an application
func LifecycleStop(appName string, signal os.Signal, err error) {
	fields := logrus.Fields{
		"type":  "lifecycle",
		"event": "stop",
	}
	if signal != nil {
		fields["signal"] = signal.String()
	}

	if os.Getenv("BUILD_NUMBER") != "" {
		fields["build_number"] = os.Getenv("BUILD_NUMBER")
	}

	if err != nil {
		Logger.WithFields(fields).
			WithError(err).
			Errorf("stopping application: %v (%v)", appName, err)
	} else {
		Logger.WithFields(fields).Infof("stopping application: %v (%v)", appName, signal.String())
	}
}

func getRemoteIp(r *http.Request) string {
	if r.Header.Get("X-Cluster-Client-Ip") != "" {
		return r.Header.Get("X-Cluster-Client-Ip")
	}
	if r.Header.Get("X-Real-Ip") != "" {
		return r.Header.Get("X-Real-Ip")
	}
	return strings.Split(r.RemoteAddr, ":")[0]
}

func setCorrelationIds(fields logrus.Fields, h http.Header) {
	correlationId := GetCorrelationId(h)
	if correlationId != "" {
		fields["correlation_id"] = correlationId
	}

	userCorrelationId := GetUserCorrelationId(h)
	if userCorrelationId != "" {
		fields["user_correlation_id"] = userCorrelationId
	}
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package util

import (
	"crypto/md5"
	"io"
)

// big endian

func BytesToUint64(b []byte) (v uint64) {
	length := uint(len(b))
	for i := uint(0); i < length-1; i++ {
		v += uint64(b[i])
		v <<= 8
	}
	v += uint64(b[length-1])
	return
}
func BytesToUint32(b []byte) (v uint32) {
	length := uint(len(b))
	for i := uint(0); i < length-1; i++ {
		v += uint32(b[i])
		v <<= 8
	}
	v += uint32(b[length-1])
	return
}
func BytesToUint16(b []byte) (v uint16) {
	v += uint16(b[0])
	v <<= 8
	v += uint16(b[1])
	return
}
func Uint64toBytes(b []byte, v uint64) {
	for i := uint(0); i < 8; i++ {
		b[7-i] = byte(v >> (i * 8))
	}
}
func Uint32toBytes(b []byte, v uint32) {
	for i := uint(0); i < 4; i++ {
		b[3-i] = byte(v >> (i * 8))
	}
}
func Uint16toBytes(b []byte, v uint16) {
	b[0] = byte(v >> 8)
	b[1] = byte(v)
}
func Uint8toBytes(b []byte, v uint8) {
	b[0] = byte(v)
}

// returns a 64 bit big int
func HashStringToLong(dir string) (v int64) {
	h := md5.New()
	io.WriteString(h, dir)

	b := h.Sum(nil)

	v += int64(b[0])
	v <<= 8
	v += int64(b[1])
	v <<= 8
	v += int64(b[2])
	v <<= 8
	v += int64(b[3])
	v <<= 8
	v += int64(b[4])
	v <<= 8
	v += int64(b[5])
	v <<= 8
	v += int64(b[6])
	v <<= 8
	v += int64(b[7])

	return
}

func HashToInt32(data []byte) (v int32) {
	h := md5.New()
	h.Write(data)

	b := h.Sum(nil)

	v += int32(b[0])
	v <<= 8
	v += int32(b[1])
	v <<= 8
	v += int32(b[2])
	v <<= 8
	v += int32(b[3])

	return
}
<commit_msg>add util for md5<commit_after>package util

import (
	"crypto/md5"
	"fmt"
	"io"
)

// big endian

func BytesToUint64(b []byte) (v uint64) {
	length := uint(len(b))
	for i := uint(0); i < length-1; i++ {
		v += uint64(b[i])
		v <<= 8
	}
	v += uint64(b[length-1])
	return
}
func BytesToUint32(b []byte) (v uint32) {
	length := uint(len(b))
	for i := uint(0); i < length-1; i++ {
		v += uint32(b[i])
		v <<= 8
	}
	v += uint32(b[length-1])
	return
}
func BytesToUint16(b []byte) (v uint16) {
	v += uint16(b[0])
	v <<= 8
	v += uint16(b[1])
	return
}
func Uint64toBytes(b []byte, v uint64) {
	for i := uint(0); i < 8; i++ {
		b[7-i] = byte(v >> (i * 8))
	}
}
func Uint32toBytes(b []byte, v uint32) {
	for i := uint(0); i < 4; i++ {
		b[3-i] = byte(v >> (i * 8))
	}
}
func Uint16toBytes(b []byte, v uint16) {
	b[0] = byte(v >> 8)
	b[1] = byte(v)
}
func Uint8toBytes(b []byte, v uint8) {
	b[0] = byte(v)
}

// returns a 64 bit big int
func HashStringToLong(dir string) (v int64) {
	h := md5.New()
	io.WriteString(h, dir)

	b := h.Sum(nil)

	v += int64(b[0])
	v <<= 8
	v += int64(b[1])
	v <<= 8
	v += int64(b[2])
	v <<= 8
	v += int64(b[3])
	v <<= 8
	v += int64(b[4])
	v <<= 8
	v += int64(b[5])
	v <<= 8
	v += int64(b[6])
	v <<= 8
	v += int64(b[7])

	return
}

func HashToInt32(data []byte) (v int32) {
	h := md5.New()
	h.Write(data)

	b := h.Sum(nil)

	v += int32(b[0])
	v <<= 8
	v += int32(b[1])
	v <<= 8
	v += int32(b[2])
	v <<= 8
	v += int32(b[3])

	return
}

func Md5(data []byte) string {
	hash := md5.New()
	hash.Write(data)
	return fmt.Sprintf("%x", hash.Sum(nil))
}
<|endoftext|>


<commit_before>package models

import (
	"time"

	"github.com/go-ozzo/ozzo-validation"
	"github.com/jinzhu/gorm"
)

type TaskModelInterface interface {
	FindAll() ([]Task, error)
	FindByID(id int) (*Task, error)
	Save(model *Task) error
	Update(model *Task) error
	Delete(model Task) error
}

type taskModel struct {
	session *gorm.DB
}

type Task struct {
	ID          uint       `json:"id"`
	CreatedAt   time.Time  `json:"createdAt"`
	UpdatedAt   time.Time  `json:"updatedAt"`
	DeletedAt   *time.Time `json:"-"`
	Description string     `json:"description"`
	Completed   bool       `json:"completed"`
}

func NewTaskModel(db *gorm.DB) *taskModel {
	return &taskModel{db}
}

func (t Task) Validate() error {
	return validation.ValidateStruct(
		&t,
		validation.Field(
			&t.Description,
			validation.Required,
			validation.Length(1, 140)),
	)
}

func (t taskModel) FindAll() ([]Task, error) {
	tasks := []Task{}
	err := t.session.Find(&tasks)
	if err.Error != nil {
		return nil, err.Error
	}
	return tasks, nil
}

func (t taskModel) FindByID(id int) (*Task, error) {
	task := &Task{}
	err := t.session.First(task, id)
	if err.Error != nil {
		return nil, err.Error
	}
	return task, nil
}

func (t taskModel) Save(model *Task) error {
	err := t.session.Create(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}

func (t taskModel) Update(model *Task) error {
	err := t.session.Update(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}

func (t taskModel) Delete(model Task) error {
	err := t.session.Delete(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}
<commit_msg>[fix] task model delete method argument and update method<commit_after>package models

import (
	"time"

	"github.com/go-ozzo/ozzo-validation"
	"github.com/jinzhu/gorm"
)

type TaskModelInterface interface {
	FindAll() ([]Task, error)
	FindByID(id int) (*Task, error)
	Save(model *Task) error
	Update(model *Task) error
	Delete(model *Task) error
}

type taskModel struct {
	session *gorm.DB
}

type Task struct {
	ID          uint       `json:"id"`
	CreatedAt   time.Time  `json:"createdAt"`
	UpdatedAt   time.Time  `json:"updatedAt"`
	DeletedAt   *time.Time `json:"-"`
	Description string     `json:"description"`
	Completed   bool       `json:"completed"`
}

func NewTaskModel(db *gorm.DB) *taskModel {
	return &taskModel{db}
}

func (t Task) Validate() error {
	return validation.ValidateStruct(
		&t,
		validation.Field(
			&t.Description,
			validation.Required,
			validation.Length(1, 140)),
	)
}

func (t taskModel) FindAll() ([]Task, error) {
	tasks := []Task{}
	err := t.session.Find(&tasks)
	if err.Error != nil {
		return nil, err.Error
	}
	return tasks, nil
}

func (t taskModel) FindByID(id int) (*Task, error) {
	task := &Task{}
	err := t.session.First(task, id)
	if err.Error != nil {
		return nil, err.Error
	}
	return task, nil
}

func (t taskModel) Save(model *Task) error {
	err := t.session.Create(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}

func (t taskModel) Update(model *Task) error {
	err := t.session.Save(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}

func (t taskModel) Delete(model *Task) error {
	err := t.session.Delete(model)
	if err.Error != nil {
		return err.Error
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"time"
	"github.com/golang/glog"
	"reflect"

	consulapi "github.com/hashicorp/consul/api"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcache "k8s.io/kubernetes/pkg/client/cache"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
)

var (
	argConsulAgent   = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile   = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
	argDryRun        = flag.Bool("dryrun", false, "Runs without connecting to consul")
	argChecks        = flag.Bool("checks", false, "Adds TCP service checks for each TCP Service")
	argConsulSync    = flag.Int("consul-sync", 0, "Set >0 to run a consul sync loop. Useful for consul restarts")
	argKubeSync    = flag.Int("kube-sync", 0, "Set >0 to run a kube sync loop. Useful for kube sanity checks if a notifcation is missed")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

func Contains(s []string, e string) bool {
	for _, i := range s {
		if i == e {
			return true
		}
	}
	return false
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}

	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}

	glog.Info("Parsed Master URL:", parsedUrl.String())
	return parsedUrl.String(), nil
}

func createConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
		glog.Infof("Error parsing Consul url")
		return nil, err
	}

	if consulAgentUrl.Host != "" {
		consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
		glog.Infof("Error creating Consul client")
		return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func createKubeClient() (*kclient.Client, error) {
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}

	overrides := &kclientcmd.ConfigOverrides{}
	overrides.ClusterInfo.Server = masterUrl

	rules := kclientcmd.NewDefaultClientConfigLoadingRules()
	kubeConfig, err := kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig()

	kubeConfig.Host = masterUrl
	if err != nil {
		glog.Error("Error creating Kube Config", err)
		return nil, err
	}
	glog.Infof("Using %s for kubernetes master", kubeConfig.Host)
	glog.Infof("Using kubernetes API %s", kubeConfig.Version)
	return kclient.New(kubeConfig)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to nodes.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func nodeReady(node *kapi.Node) bool {
	for  i := range node.Status.Conditions {
		if node.Status.Conditions[i].Type == kapi.NodeReady {
			return node.Status.Conditions[i].Status == kapi.ConditionTrue
		}
	}

	glog.Error("NodeReady condition is missing from node: ", node.Name)
	return false
}

//TODO Make Generic with Service
func sendNodeWork(action KubeWorkAction, queue chan<- KubeWork, oldObject,newObject interface{}) {
	if node, ok := newObject.(*kapi.Node); ok {
		glog.V(4).Info("Node Action: ", action, " for node ", node.Name)
		work := KubeWork {
			Action: action,
			Node: node,
		}

		if action == KubeWorkUpdateNode{
			if oldNode, ok := oldObject.(*kapi.Node); ok {
				if nodeReady(node) != nodeReady(oldNode) {
					glog.V(4).Info("Ready state change. Old:", nodeReady(oldNode), " New: ", nodeReady(node))
					queue <- work
				}
			}
		} else {
				queue <-work
		}
	}
}

//TODO Make Generic with Node
func sendServiceWork(action KubeWorkAction, queue chan<- KubeWork, serviceObj interface{}) {
	if service, ok := serviceObj.(*kapi.Service); ok {
		glog.V(4).Info("Service Action: ", action, " for service ", service.Name)
		queue <- KubeWork{
			Action: action,
			Service: service,
		}
	}
}

//TODO Combine with watchForServices
func watchForNodes(kubeClient *kclient.Client, queue chan<- KubeWork) {
	_, nodeController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    func(obj interface{}) {
				sendNodeWork(KubeWorkAddNode, queue, nil, obj)
			},
			DeleteFunc: func(obj interface{}) {
				sendNodeWork(KubeWorkRemoveNode, queue, nil, obj)
			},
			UpdateFunc: func(newObj,oldObj interface{}) {
				sendNodeWork(KubeWorkUpdateNode, queue, oldObj, newObj)
			},
		},
	)
	go nodeController.Run(util.NeverStop)
}

//TODO Combine with watchForNodes
func watchForServices(kubeClient *kclient.Client, queue chan<- KubeWork) {
	_, nodeController := kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    func(obj interface{}) {
				sendServiceWork(KubeWorkAddService, queue, obj)
			},
			DeleteFunc: func(obj interface{}) {
				sendServiceWork(KubeWorkRemoveService, queue, obj)
			},
			UpdateFunc: func(newObj,oldObj interface{}) {
				if reflect.DeepEqual(newObj,oldObj) == false {
					sendServiceWork(KubeWorkUpdateService, queue, newObj)
				}
			},
		},
	)
	go nodeController.Run(util.NeverStop)
}

func main() {
	flag.Parse()

	//Attempt to create Kube Client
	kubeClient, err := createKubeClient()
	if err != nil {
		glog.Fatal("Could not connect to Kube Master", err)
	}

	if _, err := kubeClient.ServerVersion(); err != nil {
		glog.Fatal("Could not connect to Kube Master", err)
	} else {
		glog.Info("Connected to K8S API Server")
	}
	//Attempt to create Consul Client (All ways needed so that channels are not blocked)
	var consulClient *consulapi.Client
	if *argDryRun == false {
		consulClient, err = createConsulClient(*argConsulAgent)
	}

	//Do System Setup stuff (create channels?)
	kubeWorkQueue := make(chan KubeWork)
	consulWorkQueue := make(chan ConsulWork)
	go RunBookKeeper(kubeWorkQueue, consulWorkQueue, kubeClient)
	//Launch KubeLoops
	watchForNodes(kubeClient, kubeWorkQueue)
	watchForServices(kubeClient, kubeWorkQueue)
	//Launch Consul Loops
	go RunConsulWorker(consulWorkQueue, consulClient)

	//Launch Sync Loops (if needed)
	if *argConsulSync > 0 {
		glog.Info("Running Consul Sync loop every: ", *argConsulSync, " seconds")
		syncer := time.NewTicker(time.Second * time.Duration(*argConsulSync))

	  go func() {
	    for t := range syncer.C {
	      glog.V(4).Info("Consul Sync request at:", t)
	      consulWorkQueue <- ConsulWork {
					Action: ConsulWorkSyncDNS,
				}
	    }
	  }()
	}

	if *argKubeSync > 0 {
		glog.Info("Running Kube Sync loop every: ", *argConsulSync, " seconds")
		syncer := time.NewTicker(time.Second * time.Duration(*argKubeSync))

	  go func() {
	    for t := range syncer.C {
	      glog.V(4).Info("Kube Sync request at:", t)
	      kubeWorkQueue <- KubeWork {
					Action: KubeWorkSync,
				}
	    }
	  }()
	}

	// Prevent exit
	select {}
}
<commit_msg>Adding node label selector support<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"reflect"
	"time"

	"github.com/golang/glog"

	consulapi "github.com/hashicorp/consul/api"

	kapi "k8s.io/kubernetes/pkg/api"
	kcache "k8s.io/kubernetes/pkg/client/cache"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util"
)

var (
	argConsulAgent   = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile   = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
	argDryRun        = flag.Bool("dryrun", false, "Runs without connecting to consul")
	argChecks        = flag.Bool("checks", false, "Adds TCP service checks for each TCP Service")
	argConsulSync    = flag.Int("consul-sync", 0, "Set >0 to run a consul sync loop. Useful for consul restarts")
	argKubeSync      = flag.Int("kube-sync", 0, "Set >0 to run a kube sync loop. Useful for kube sanity checks if a notifcation is missed")
	argNodeSelector  = flag.String("node-selector", "", "Node label selector to use for proxies. Leave blank for all")

	nodeSelector = labels.Everything()
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

func Contains(s []string, e string) bool {
	for _, i := range s {
		if i == e {
			return true
		}
	}
	return false
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}

	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}

	glog.Info("Parsed Master URL:", parsedUrl.String())
	return parsedUrl.String(), nil
}

func createConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
		glog.Infof("Error parsing Consul url")
		return nil, err
	}

	if consulAgentUrl.Host != "" {
		consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
		glog.Infof("Error creating Consul client")
		return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func createKubeClient() (*kclient.Client, error) {
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}

	overrides := &kclientcmd.ConfigOverrides{}
	overrides.ClusterInfo.Server = masterUrl

	rules := kclientcmd.NewDefaultClientConfigLoadingRules()
	kubeConfig, err := kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig()

	kubeConfig.Host = masterUrl
	if err != nil {
		glog.Error("Error creating Kube Config", err)
		return nil, err
	}
	glog.Infof("Using %s for kubernetes master", kubeConfig.Host)
	glog.Infof("Using kubernetes API %s", kubeConfig.Version)
	return kclient.New(kubeConfig)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to nodes.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func nodeReady(node *kapi.Node) bool {
	for i := range node.Status.Conditions {
		if node.Status.Conditions[i].Type == kapi.NodeReady {
			return node.Status.Conditions[i].Status == kapi.ConditionTrue
		}
	}

	glog.Error("NodeReady condition is missing from node: ", node.Name)
	return false
}

//TODO Make Generic with Service
func sendNodeWork(action KubeWorkAction, queue chan<- KubeWork, oldObject, newObject interface{}) {
	if node, ok := newObject.(*kapi.Node); ok {
		if nodeSelector.Matches(labels.Set(node.Labels)) == false {
			glog.V(2).Infof("Ignoring node %s due to label selectors", node.Name)
			return
		}

		glog.V(4).Info("Node Action: ", action, " for node ", node.Name)
		work := KubeWork{
			Action: action,
			Node:   node,
		}

		if action == KubeWorkUpdateNode {
			if oldNode, ok := oldObject.(*kapi.Node); ok {
				if nodeReady(node) != nodeReady(oldNode) {
					glog.V(4).Info("Ready state change. Old:", nodeReady(oldNode), " New: ", nodeReady(node))
					queue <- work
				}
			}
		} else {
			queue <- work
		}
	}
}

//TODO Make Generic with Node
func sendServiceWork(action KubeWorkAction, queue chan<- KubeWork, serviceObj interface{}) {
	if service, ok := serviceObj.(*kapi.Service); ok {
		glog.V(4).Info("Service Action: ", action, " for service ", service.Name)
		queue <- KubeWork{
			Action:  action,
			Service: service,
		}
	}
}

//TODO Combine with watchForServices
func watchForNodes(kubeClient *kclient.Client, queue chan<- KubeWork) {
	_, nodeController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				sendNodeWork(KubeWorkAddNode, queue, nil, obj)
			},
			DeleteFunc: func(obj interface{}) {
				sendNodeWork(KubeWorkRemoveNode, queue, nil, obj)
			},
			UpdateFunc: func(newObj, oldObj interface{}) {
				sendNodeWork(KubeWorkUpdateNode, queue, oldObj, newObj)
			},
		},
	)
	go nodeController.Run(util.NeverStop)
}

//TODO Combine with watchForNodes
func watchForServices(kubeClient *kclient.Client, queue chan<- KubeWork) {
	_, nodeController := kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				sendServiceWork(KubeWorkAddService, queue, obj)
			},
			DeleteFunc: func(obj interface{}) {
				sendServiceWork(KubeWorkRemoveService, queue, obj)
			},
			UpdateFunc: func(newObj, oldObj interface{}) {
				if reflect.DeepEqual(newObj, oldObj) == false {
					sendServiceWork(KubeWorkUpdateService, queue, newObj)
				}
			},
		},
	)
	go nodeController.Run(util.NeverStop)
}

func main() {
	flag.Parse()

	//Attempt to create Kube Client
	kubeClient, err := createKubeClient()
	if err != nil {
		glog.Fatal("Could not connect to Kube Master", err)
	}

	if _, err := kubeClient.ServerVersion(); err != nil {
		glog.Fatal("Could not connect to Kube Master", err)
	} else {
		glog.Info("Connected to K8S API Server")
	}

	nodeSelector, err = labels.Parse(*argNodeSelector)
	if err != nil {
		glog.Fatal("Could not parse node label selector. Error:", err)
	}

	//Attempt to create Consul Client (All ways needed so that channels are not blocked)
	var consulClient *consulapi.Client
	if *argDryRun == false {
		consulClient, err = createConsulClient(*argConsulAgent)
	}

	//Do System Setup stuff (create channels?)
	kubeWorkQueue := make(chan KubeWork)
	consulWorkQueue := make(chan ConsulWork)
	go RunBookKeeper(kubeWorkQueue, consulWorkQueue, kubeClient)
	//Launch KubeLoops
	watchForNodes(kubeClient, kubeWorkQueue)
	watchForServices(kubeClient, kubeWorkQueue)
	//Launch Consul Loops
	go RunConsulWorker(consulWorkQueue, consulClient)

	//Launch Sync Loops (if needed)
	if *argConsulSync > 0 {
		glog.Info("Running Consul Sync loop every: ", *argConsulSync, " seconds")
		syncer := time.NewTicker(time.Second * time.Duration(*argConsulSync))

		go func() {
			for t := range syncer.C {
				glog.V(4).Info("Consul Sync request at:", t)
				consulWorkQueue <- ConsulWork{
					Action: ConsulWorkSyncDNS,
				}
			}
		}()
	}

	if *argKubeSync > 0 {
		glog.Info("Running Kube Sync loop every: ", *argConsulSync, " seconds")
		syncer := time.NewTicker(time.Second * time.Duration(*argKubeSync))

		go func() {
			for t := range syncer.C {
				glog.V(4).Info("Kube Sync request at:", t)
				kubeWorkQueue <- KubeWork{
					Action: KubeWorkSync,
				}
			}
		}()
	}

	// Prevent exit
	select {}
}
<|endoftext|>


<commit_before>package nagiosplugin

import (
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"
)

// Range is a combination of a lower boundary, an upper boundary
// and a flag for inverted (@) range semantics.
type Range struct {
	Start         float64
	End           float64
	AlertOnInside bool
}

// Returns a new range object and nil if the given
// range definition was valid or nil and an error.
func ParseRange(rangeStr string) (*Range, error) {
	// Set Defaults
	t := &Range{
		Start:         0,
		End:           math.Inf(1),
		AlertOnInside: false,
	}
	// Remove leading and trailing whitespaces
	rangeStr = strings.Trim(rangeStr, " \n\r")

	// Check for inverted semantics
	if rangeStr[0] == '@' {
		t.AlertOnInside = true
		rangeStr = rangeStr[1:]
	}

	// parse lower limit
	endPos := strings.Index(rangeStr, ":")
	if endPos > -1 {
		if rangeStr[0] == '~' {
			t.Start = math.Inf(-1)
		} else {
			min, err := strconv.ParseFloat(rangeStr[0:endPos], 64)
			if err != nil {
				fmt.Println("failed to parse lower limit:", err)
			} else {
				t.Start = min
			}
		}
		rangeStr = rangeStr[endPos+1:]
	}

	// parse upper limit
	if len(rangeStr) > 0 {
		max, err := strconv.ParseFloat(rangeStr, 64)
		if err != nil {
			fmt.Println("failed to parse upper limit:", err)
		} else {
			t.End = max
		}
	}

	if t.End < t.Start {
		return nil, errors.New("Invalid range definition. min <= max violated!")
	}

	// OK
	return t, nil
}

// Returs true if an alert should be raised based on the range
func (r *Range) Check(value float64) bool {
	var no bool = false
	var yes bool = true

	if r.AlertOnInside {
		no = true
		yes = false
	}
	// see https://www.monitoring-plugins.org/doc/guidelines.html#THRESHOLDFORMAT
	if r.Start <= value && value <= r.End {
		return no
	} else {
		return yes
	}
	return yes
}

// CheckInt is a convenience method which does an unchecked type
// conversion from an int to a float64
func (r *Range) CheckInt(val int) bool {
	return r.Check(float64(val))
}

// CheckUint64 is a convenience method which does an unchecked type
// conversion from an uint64 to a float64
func (r *Range) CheckUint64(val uint64) bool {
	return r.Check(float64(val))
}
<commit_msg>Minor doc fixes<commit_after>package nagiosplugin

import (
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"
)

// Range is a combination of a lower boundary, an upper boundary
// and a flag for inverted (@) range semantics.
type Range struct {
	Start         float64
	End           float64
	AlertOnInside bool
}

// Returns a new range object and nil if the given
// range definition was valid or nil and an error.
func ParseRange(rangeStr string) (*Range, error) {
	// Set Defaults
	t := &Range{
		Start:         0,
		End:           math.Inf(1),
		AlertOnInside: false,
	}
	// Remove leading and trailing whitespaces
	rangeStr = strings.Trim(rangeStr, " \n\r")

	// Check for inverted semantics
	if rangeStr[0] == '@' {
		t.AlertOnInside = true
		rangeStr = rangeStr[1:]
	}

	// parse lower limit
	endPos := strings.Index(rangeStr, ":")
	if endPos > -1 {
		if rangeStr[0] == '~' {
			t.Start = math.Inf(-1)
		} else {
			min, err := strconv.ParseFloat(rangeStr[0:endPos], 64)
			if err != nil {
				fmt.Println("failed to parse lower limit:", err)
			} else {
				t.Start = min
			}
		}
		rangeStr = rangeStr[endPos+1:]
	}

	// parse upper limit
	if len(rangeStr) > 0 {
		max, err := strconv.ParseFloat(rangeStr, 64)
		if err != nil {
			fmt.Println("failed to parse upper limit:", err)
		} else {
			t.End = max
		}
	}

	if t.End < t.Start {
		return nil, errors.New("Invalid range definition. min <= max violated!")
	}

	// OK
	return t, nil
}

// Returns true if an alert should be raised based on the range.
func (r *Range) Check(value float64) bool {
	var no bool = false
	var yes bool = true

	if r.AlertOnInside {
		no = true
		yes = false
	}
	// see https://www.monitoring-plugins.org/doc/guidelines.html#THRESHOLDFORMAT
	if r.Start <= value && value <= r.End {
		return no
	} else {
		return yes
	}
	return yes
}

// CheckInt is a convenience method which does an unchecked type
// conversion from an int to a float64.
func (r *Range) CheckInt(val int) bool {
	return r.Check(float64(val))
}

// CheckUint64 is a convenience method which does an unchecked type
// conversion from an uint64 to a float64.
func (r *Range) CheckUint64(val uint64) bool {
	return r.Check(float64(val))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.6.1"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	info("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil || !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local")
		return
	}

	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(120 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileKey := path.Join(usr.HomeDir, ".lightsockskey")
	data, err := ioutil.ReadFile(fileKey)
	if err != nil {
		fmt.Printf("Failed to load key file: %v\n", err)
		os.Exit(1)
	}
	s := strings.TrimSpace(string(data))
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<commit_msg>increased conn timeout<commit_after>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.6.2"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	info("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil || !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local")
		return
	}

	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(3600 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileKey := path.Join(usr.HomeDir, ".lightsockskey")
	data, err := ioutil.ReadFile(fileKey)
	if err != nil {
		fmt.Printf("Failed to load key file: %v\n", err)
		os.Exit(1)
	}
	s := strings.TrimSpace(string(data))
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<|endoftext|>