<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) (*http.Response) {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{  // go1.8 deprecated: Use DialContext instead
			Timeout: timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Always close connections to the alternative and production servers.
		DisableKeepAlives: true,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout: timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	UpdateForwardedHeaders(req)
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			_ = handleRequest(alternativeRequest, timeout)
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		body, _ := ioutil.ReadAll(resp.Body)
		w.Write(body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
	           *listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

const XFF_HEADER = "X-Forwarded-For"

func InsertOrExtendXFFHeader(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	remoteIP := request.RemoteAddr[:positionOfColon]
	xff, ok := request.Header[XFF_HEADER]
	if ok {
		// extend
		xff = append(xff, remoteIP)
		request.Header[XFF_HEADER][0] = strings.Join(xff, ", ")
	} else {
		// insert
		request.Header[XFF_HEADER] = []string{remoteIP}
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func InsertOrExtendForwardedHeader(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	remoteIP := request.RemoteAddr[:positionOfColon]
	extension := "for=" + remoteIP
	if header, ok := request.Header[FORWARDED_HEADER]; ok == true {
		// extend
		header = append(header, extension)
		request.Header[FORWARDED_HEADER][0] = strings.Join(header, ", ")
	} else {
		// insert
		request.Header[FORWARDED_HEADER] = []string{extension}
	}
}

func UpdateForwardedHeaders(request *http.Request) {
	InsertOrExtendForwardedHeader(request)
	InsertOrExtendXFFHeader(request)
}
<commit_msg>Unify and simplify code changes<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) (*http.Response) {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{  // go1.8 deprecated: Use DialContext instead
			Timeout: timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Always close connections to the alternative and production servers.
		DisableKeepAlives: true,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout: timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	UpdateForwardedHeaders(req)
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			_ = handleRequest(alternativeRequest, timeout)
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		body, _ := ioutil.ReadAll(resp.Body)
		w.Write(body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
	           *listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

func UpdateForwardedHeaders(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	var remoteIP string
	if positionOfColon != -1 {
		remoteIP = request.RemoteAddr[:positionOfColon]
	} else {
		Logger.Printf("The default format of request.RemoteAddr should be IP:Port but was %s\n", remoteIP)
		remoteIP = request.RemoteAddr
	}
	InsertOrExtendForwardedHeader(request, remoteIP)
	InsertOrExtendXFFHeader(request, remoteIP)
}

const XFF_HEADER = "X-Forwarded-For"

// Implementation according to rfc7239
func InsertOrExtendXFFHeader(request *http.Request, remoteIP string) {
	header := request.Header.Get(XFF_HEADER)
	if header != "" {
		// extend
		request.Header.Set(XFF_HEADER, header+", "+remoteIP)
	} else {
		// insert
		request.Header.Set(XFF_HEADER, remoteIP)
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func InsertOrExtendForwardedHeader(request *http.Request, remoteIP string) {
	extension := "for=" + remoteIP
	header := request.Header.Get(FORWARDED_HEADER)
	if header != "" {
		// extend
		request.Header.Set(FORWARDED_HEADER, header+", "+extension)
	} else {
		// insert
		request.Header.Set(FORWARDED_HEADER, extension)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The go-elog Authors. All rights reserved.
// This code is MIT licensed. See the LICENSE file for more info.

package log

import "text/template"

// funcMap contains the available functions to the log format template.
var (
	funcMap = template.FuncMap{}
	logFmt  = "{{if .Date}}{{.Date}} {{end}}" +
		"{{if .Prefix}}{{.Prefix}} {{end}}" +
		"{{if .LogLabel}}{{.LogLabel}} {{end}}" +
		"{{if .Indent}}{{.Indent}} {{end}}" +
		"{{if .Id}}{{.Id}} {{end}}" +
		"{{if .FileName}}{{.FileName}}: {{end}}" +
		"{{if .FunctionName}}{{.FunctionName}}: {{end}}" +
		"{{if .LineNumber}}Line {{.LineNumber}}: {{end}}" +
		"{{if .Text}}{{.Text}}{{end}}"
)

// format is the possible values that can be used in a log output format
type format struct {
	Prefix       string
	LogLabel     string
	Date         string
	FileName     string
	FunctionName string
	LineNumber   int
	Indent       string
	Id           string
	Text         string
}
<commit_msg>template.go: Remove space after .Indent<commit_after>// Copyright 2013 The go-elog Authors. All rights reserved.
// This code is MIT licensed. See the LICENSE file for more info.

package log

import "text/template"

// funcMap contains the available functions to the log format template.
var (
	funcMap = template.FuncMap{}
	logFmt  = "{{if .Date}}{{.Date}} {{end}}" +
		"{{if .Prefix}}{{.Prefix}} {{end}}" +
		"{{if .LogLabel}}{{.LogLabel}} {{end}}" +
		"{{if .Id}}{{.Id}} {{end}}" +
		"{{if .Indent}}{{.Indent}}{{end}}" +
		"{{if .FileName}}{{.FileName}}: {{end}}" +
		"{{if .FunctionName}}{{.FunctionName}}: {{end}}" +
		"{{if .LineNumber}}Line {{.LineNumber}}: {{end}}" +
		"{{if .Text}}{{.Text}}{{end}}"
)

// format is the possible values that can be used in a log output format
type format struct {
	Prefix       string
	LogLabel     string
	Date         string
	FileName     string
	FunctionName string
	LineNumber   int
	Indent       string
	Id           string
	Text         string
}
<|endoftext|>


<commit_before>package jid

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strings"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
	"github.com/nwidger/jsoncolor"
)

type Terminal struct {
	defaultY   int
	prompt     string
	formatter  *jsoncolor.Formatter
	monochrome bool
	outputArea *[][]termbox.Cell
}

type TerminalDrawAttributes struct {
	Query           string
	Contents        []string
	CandidateIndex  int
	ContentsOffsetY int
	Complete        string
	Candidates      []string
	CursorOffset    int
}

func NewTerminal(prompt string, defaultY int, monochrome bool) *Terminal {
	t := &Terminal{
		prompt:     prompt,
		defaultY:   defaultY,
		monochrome: monochrome,
		outputArea: &[][]termbox.Cell{},
		formatter:  nil,
	}
	if !monochrome {
		t.formatter = t.initColorizeFormatter()
	}
	return t
}

func (t *Terminal) Draw(attr *TerminalDrawAttributes) error {

	query := attr.Query
	complete := attr.Complete
	rows := attr.Contents
	candidates := attr.Candidates
	candidateidx := attr.CandidateIndex
	contentOffsetY := attr.ContentsOffsetY

	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	y := t.defaultY

	t.drawFilterLine(query, complete)

	if len(candidates) > 0 {
		y = t.drawCandidates(0, t.defaultY, candidateidx, candidates)
	}

	cellsArr, err := t.rowsToCells(rows)
	if err != nil {
		return err
	}

	for idx, cells := range cellsArr {
		if i := idx - contentOffsetY; i >= 0 {
			t.drawCells(0, i+y, cells)
		}
	}

	termbox.SetCursor(len(t.prompt)+attr.CursorOffset, 0)

	termbox.Flush()
	return nil
}

func (t *Terminal) drawFilterLine(qs string, complete string) error {
	fs := t.prompt + qs
	cs := complete
	str := fs + cs

	color := termbox.ColorDefault
	backgroundColor := termbox.ColorDefault

	var cells []termbox.Cell
	match := []int{len(fs), len(fs + cs)}

	var c termbox.Attribute
	for i, s := range str {
		c = color
		if i >= match[0] && i < match[1] {
			c = termbox.ColorGreen
		}
		cells = append(cells, termbox.Cell{
			Ch: s,
			Fg: c,
			Bg: backgroundColor,
		})
	}
	t.drawCells(0, 0, cells)
	return nil
}

type termboxSprintfFuncer struct {
	fg         termbox.Attribute
	bg         termbox.Attribute
	outputArea *[][]termbox.Cell
}

func (tsf *termboxSprintfFuncer) SprintfFunc() func(format string, a ...interface{}) string {
	return func(format string, a ...interface{}) string {
		cells := tsf.outputArea
		idx := len(*cells) - 1
		str := fmt.Sprintf(format, a...)
		for _, s := range str {
			if s == '\n' {
				*cells = append(*cells, []termbox.Cell{})
				idx++
				continue
			}
			(*cells)[idx] = append((*cells)[idx], termbox.Cell{
				Ch: s,
				Fg: tsf.fg,
				Bg: tsf.bg,
			})
		}
		return "dummy"
	}
}

func (t *Terminal) initColorizeFormatter() *jsoncolor.Formatter {
	formatter := jsoncolor.NewFormatter()

	regular := &termboxSprintfFuncer{
		fg:         termbox.ColorDefault,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	bold := &termboxSprintfFuncer{
		fg:         termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	blueBold := &termboxSprintfFuncer{
		fg:         termbox.ColorBlue | termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	green := &termboxSprintfFuncer{
		fg:         termbox.ColorGreen,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	blackBold := &termboxSprintfFuncer{
		fg:         termbox.ColorBlack | termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	formatter.SpaceColor = regular
	formatter.CommaColor = bold
	formatter.ColonColor = bold
	formatter.ObjectColor = bold
	formatter.ArrayColor = bold
	formatter.FieldQuoteColor = blueBold
	formatter.FieldColor = blueBold
	formatter.StringQuoteColor = green
	formatter.StringColor = green
	formatter.TrueColor = regular
	formatter.FalseColor = regular
	formatter.NumberColor = regular
	formatter.NullColor = blackBold

	return formatter
}

func (t *Terminal) rowsToCells(rows []string) ([][]termbox.Cell, error) {
	*t.outputArea = [][]termbox.Cell{[]termbox.Cell{}}

	var err error

	if t.formatter != nil {
		err = t.formatter.Format(ioutil.Discard, []byte(strings.Join(rows, "\n")))
	}

	cells := *t.outputArea

	if err != nil || t.monochrome {
		cells = [][]termbox.Cell{}
		for _, row := range rows {
			var cls []termbox.Cell
			for _, char := range row {
				cls = append(cls, termbox.Cell{
					Ch: char,
					Fg: termbox.ColorDefault,
					Bg: termbox.ColorDefault,
				})
			}
			cells = append(cells, cls)
		}
	}

	return cells, nil
}

func (t *Terminal) drawCells(x int, y int, cells []termbox.Cell) {
	i := 0
	for _, c := range cells {
		termbox.SetCell(x+i, y, c.Ch, c.Fg, c.Bg)

		w := runewidth.RuneWidth(c.Ch)
		if w == 0 || w == 2 && runewidth.IsAmbiguousWidth(c.Ch) {
			w = 1
		}

		i += w
	}
}

func (t *Terminal) drawCandidates(x int, y int, index int, candidates []string) int {
	color := termbox.ColorBlack
	backgroundColor := termbox.ColorWhite

	w, _ := termbox.Size()

	ss := candidates[index]
	re := regexp.MustCompile("[[:space:]]" + ss + "[[:space:]]")

	var rows []string
	var str string
	for _, word := range candidates {
		combine := " "
		if l := len(str); l+len(word)+1 >= w {
			rows = append(rows, str+" ")
			str = ""
		}
		str += combine + word
	}
	rows = append(rows, str+" ")

	for i, row := range rows {
		match := re.FindStringIndex(row)
		var c termbox.Attribute
		ii := 0
		for k, s := range row {
			c = color
			backgroundColor = termbox.ColorMagenta
			if match != nil && k >= match[0]+1 && k < match[1]-1 {
				backgroundColor = termbox.ColorWhite
			}
			termbox.SetCell(x+ii, y+i, s, c, backgroundColor)
			w := runewidth.RuneWidth(s)
			if w == 0 || w == 2 && runewidth.IsAmbiguousWidth(s) {
				w = 1
			}
			ii += w
		}
	}
	return y + len(rows)
}
<commit_msg>Improve performance<commit_after>package jid

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strings"

	runewidth "github.com/mattn/go-runewidth"
	termbox "github.com/nsf/termbox-go"
	"github.com/nwidger/jsoncolor"
)

type Terminal struct {
	defaultY   int
	prompt     string
	formatter  *jsoncolor.Formatter
	monochrome bool
	outputArea *[][]termbox.Cell
}

type TerminalDrawAttributes struct {
	Query           string
	Contents        []string
	CandidateIndex  int
	ContentsOffsetY int
	Complete        string
	Candidates      []string
	CursorOffset    int
}

func NewTerminal(prompt string, defaultY int, monochrome bool) *Terminal {
	t := &Terminal{
		prompt:     prompt,
		defaultY:   defaultY,
		monochrome: monochrome,
		outputArea: &[][]termbox.Cell{},
		formatter:  nil,
	}
	if !monochrome {
		t.formatter = t.initColorizeFormatter()
	}
	return t
}

func (t *Terminal) Draw(attr *TerminalDrawAttributes) error {

	query := attr.Query
	complete := attr.Complete
	rows := attr.Contents
	candidates := attr.Candidates
	candidateidx := attr.CandidateIndex
	contentOffsetY := attr.ContentsOffsetY

	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	y := t.defaultY
	_, h := termbox.Size()

	t.drawFilterLine(query, complete)

	if len(candidates) > 0 {
		y = t.drawCandidates(0, t.defaultY, candidateidx, candidates)
	}

	cellsArr, err := t.rowsToCells(rows)
	if err != nil {
		return err
	}

	for idx, cells := range cellsArr {
		i := idx - contentOffsetY
		if i >= 0 {
			t.drawCells(0, i+y, cells)
		}
		if i > h {
			break
		}
	}

	termbox.SetCursor(len(t.prompt)+attr.CursorOffset, 0)

	termbox.Flush()
	return nil
}

func (t *Terminal) drawFilterLine(qs string, complete string) error {
	fs := t.prompt + qs
	cs := complete
	str := fs + cs

	color := termbox.ColorDefault
	backgroundColor := termbox.ColorDefault

	var cells []termbox.Cell
	match := []int{len(fs), len(fs + cs)}

	var c termbox.Attribute
	for i, s := range str {
		c = color
		if i >= match[0] && i < match[1] {
			c = termbox.ColorGreen
		}
		cells = append(cells, termbox.Cell{
			Ch: s,
			Fg: c,
			Bg: backgroundColor,
		})
	}
	t.drawCells(0, 0, cells)
	return nil
}

type termboxSprintfFuncer struct {
	fg         termbox.Attribute
	bg         termbox.Attribute
	outputArea *[][]termbox.Cell
}

func (tsf *termboxSprintfFuncer) SprintfFunc() func(format string, a ...interface{}) string {
	return func(format string, a ...interface{}) string {
		cells := tsf.outputArea
		idx := len(*cells) - 1
		str := fmt.Sprintf(format, a...)
		for _, s := range str {
			if s == '\n' {
				*cells = append(*cells, []termbox.Cell{})
				idx++
				continue
			}
			(*cells)[idx] = append((*cells)[idx], termbox.Cell{
				Ch: s,
				Fg: tsf.fg,
				Bg: tsf.bg,
			})
		}
		return "dummy"
	}
}

func (t *Terminal) initColorizeFormatter() *jsoncolor.Formatter {
	formatter := jsoncolor.NewFormatter()

	regular := &termboxSprintfFuncer{
		fg:         termbox.ColorDefault,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	bold := &termboxSprintfFuncer{
		fg:         termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	blueBold := &termboxSprintfFuncer{
		fg:         termbox.ColorBlue | termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	green := &termboxSprintfFuncer{
		fg:         termbox.ColorGreen,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	blackBold := &termboxSprintfFuncer{
		fg:         termbox.ColorBlack | termbox.AttrBold,
		bg:         termbox.ColorDefault,
		outputArea: t.outputArea,
	}

	formatter.SpaceColor = regular
	formatter.CommaColor = bold
	formatter.ColonColor = bold
	formatter.ObjectColor = bold
	formatter.ArrayColor = bold
	formatter.FieldQuoteColor = blueBold
	formatter.FieldColor = blueBold
	formatter.StringQuoteColor = green
	formatter.StringColor = green
	formatter.TrueColor = regular
	formatter.FalseColor = regular
	formatter.NumberColor = regular
	formatter.NullColor = blackBold

	return formatter
}

func (t *Terminal) rowsToCells(rows []string) ([][]termbox.Cell, error) {
	*t.outputArea = [][]termbox.Cell{[]termbox.Cell{}}

	var err error

	if t.formatter != nil {
		err = t.formatter.Format(ioutil.Discard, []byte(strings.Join(rows, "\n")))
	}

	cells := *t.outputArea

	if err != nil || t.monochrome {
		cells = [][]termbox.Cell{}
		for _, row := range rows {
			var cls []termbox.Cell
			for _, char := range row {
				cls = append(cls, termbox.Cell{
					Ch: char,
					Fg: termbox.ColorDefault,
					Bg: termbox.ColorDefault,
				})
			}
			cells = append(cells, cls)
		}
	}

	return cells, nil
}

func (t *Terminal) drawCells(x int, y int, cells []termbox.Cell) {
	i := 0
	for _, c := range cells {
		termbox.SetCell(x+i, y, c.Ch, c.Fg, c.Bg)

		w := runewidth.RuneWidth(c.Ch)
		if w == 0 || w == 2 && runewidth.IsAmbiguousWidth(c.Ch) {
			w = 1
		}

		i += w
	}
}

func (t *Terminal) drawCandidates(x int, y int, index int, candidates []string) int {
	color := termbox.ColorBlack
	backgroundColor := termbox.ColorWhite

	w, _ := termbox.Size()

	ss := candidates[index]
	re := regexp.MustCompile("[[:space:]]" + ss + "[[:space:]]")

	var rows []string
	var str string
	for _, word := range candidates {
		combine := " "
		if l := len(str); l+len(word)+1 >= w {
			rows = append(rows, str+" ")
			str = ""
		}
		str += combine + word
	}
	rows = append(rows, str+" ")

	for i, row := range rows {
		match := re.FindStringIndex(row)
		var c termbox.Attribute
		ii := 0
		for k, s := range row {
			c = color
			backgroundColor = termbox.ColorMagenta
			if match != nil && k >= match[0]+1 && k < match[1]-1 {
				backgroundColor = termbox.ColorWhite
			}
			termbox.SetCell(x+ii, y+i, s, c, backgroundColor)
			w := runewidth.RuneWidth(s)
			if w == 0 || w == 2 && runewidth.IsAmbiguousWidth(s) {
				w = 1
			}
			ii += w
		}
	}
	return y + len(rows)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/utp"
)

var (
	flClientMode   = flag.Bool("c", false, "client mode")
	flServerMode   = flag.Bool("s", false, "server mode")
	flHost         = flag.String("h", "127.0.0.1", "host")
	flPort         = flag.Int("p", 6001, "port")
	flLen          = flag.Int("l", 1400, "length of data")
	flThreads      = flag.Int("t", 1, "threads")
	flDuration     = flag.Duration("d", time.Second*10, "duration")
	flDurationStat = flag.Duration("ds", time.Second*5, "duration for stats")
)

func main() {
	log.Printf("UTP Benchmark Tool by Artem Andreenko (miolini@gmail.com)")
	flag.Parse()
	if !*flClientMode && !*flServerMode {
		flag.PrintDefaults()
		return
	}
	ts := time.Now()
	wg := sync.WaitGroup{}
	if *flServerMode {
		wg.Add(1)
		go server(&wg, *flHost, *flPort)
	}
	if *flClientMode {
		wg.Add(*flThreads)
		chStat := make(chan int, 100)
		go stat(chStat, *flDurationStat)
		for i := 0; i < *flThreads; i++ {
			go client(&wg, *flHost, *flPort, *flLen, *flDuration, chStat)
		}
	}
	wg.Wait()
	log.Printf("time takes %.2fsec", time.Since(ts).Seconds())
}

func stat(chStat chan int, duration time.Duration) {
	t := time.NewTicker(duration)
	defer t.Stop()
	counter := 0
	for {
		select {
		case n := <-chStat:
			counter += n
		case <-t.C:
			log.Printf("speed %.3f mbit/sec", float64(counter*8)/duration.Seconds()/1024/1024)
			counter = 0
		}
	}
}

func server(wg *sync.WaitGroup, host string, port int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		wg.Done()
	}()
	log.Printf("server listen %s:%d", host, port)
	s, err := utp.NewSocket("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		panic(err)
	}
	defer s.Close()
	for {
		conn, err := s.Accept()
		if err != nil {
			panic(err)
		}
		wg.Add(1)
		go readConn(conn)
	}
}

func readConn(conn net.Conn) {
	defer conn.Close()
	defer log.Printf("client %s disconnected", conn.RemoteAddr().String())
	log.Printf("client %s connected", conn.RemoteAddr().String())
	buf := make([]byte, 4096)
	for {
		_, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("err: %s", err)
			}
			break
		}
	}
}

func client(wg *sync.WaitGroup, host string, port, len int, duration time.Duration, chStat chan int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		log.Printf("disconnected")
		wg.Done()
	}()
	log.Printf("connecting to %s:%d, len %d, duration %s", host, port, len, duration.String())
	conn, err := utp.DialTimeout(fmt.Sprintf("%s:%d", host, port), time.Second)
	if err != nil {
		panic(err)
	}
	defer conn.Close()
	log.Printf("connected")
	buf := bytes.Repeat([]byte("H"), len)
	ts := time.Now()
	for time.Since(ts) < duration {
		n, err := conn.Write(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chStat <- n
	}
}
<commit_msg>add io.Copy for read conn<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/utp"
)

var (
	flClientMode   = flag.Bool("c", false, "client mode")
	flServerMode   = flag.Bool("s", false, "server mode")
	flHost         = flag.String("h", "127.0.0.1", "host")
	flPort         = flag.Int("p", 6001, "port")
	flLen          = flag.Int("l", 1400, "length of data")
	flThreads      = flag.Int("t", 1, "threads")
	flDuration     = flag.Duration("d", time.Second*10, "duration")
	flDurationStat = flag.Duration("ds", time.Second*5, "duration for stats")
)

func main() {
	log.Printf("UTP Benchmark Tool by Artem Andreenko (miolini@gmail.com)")
	flag.Parse()
	if !*flClientMode && !*flServerMode {
		flag.PrintDefaults()
		return
	}
	ts := time.Now()
	wg := sync.WaitGroup{}
	if *flServerMode {
		wg.Add(1)
		go server(&wg, *flHost, *flPort)
	}
	if *flClientMode {
		wg.Add(*flThreads)
		chStat := make(chan int, 100)
		go stat(chStat, *flDurationStat)
		for i := 0; i < *flThreads; i++ {
			go client(&wg, *flHost, *flPort, *flLen, *flDuration, chStat)
		}
	}
	wg.Wait()
	log.Printf("time takes %.2fsec", time.Since(ts).Seconds())
}

func stat(chStat chan int, duration time.Duration) {
	t := time.NewTicker(duration)
	defer t.Stop()
	counter := 0
	for {
		select {
		case n := <-chStat:
			counter += n
		case <-t.C:
			log.Printf("speed %.3f mbit/sec", float64(counter*8)/duration.Seconds()/1024/1024)
			counter = 0
		}
	}
}

func server(wg *sync.WaitGroup, host string, port int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		wg.Done()
	}()
	log.Printf("server listen %s:%d", host, port)
	s, err := utp.NewSocket("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		panic(err)
	}
	defer s.Close()
	for {
		conn, err := s.Accept()
		if err != nil {
			panic(err)
		}
		wg.Add(1)
		go readConn(conn)
	}
}

func readConn(conn net.Conn) {
	defer conn.Close()
	defer log.Printf("client %s disconnected", conn.RemoteAddr().String())
	log.Printf("client %s connected", conn.RemoteAddr().String())
	_, err := io.Copy(ioutil.Discard, conn)
	if err != io.EOF {
		log.Printf("client %s error: %s", conn.RemoteAddr().String(), err)
	}
}

func client(wg *sync.WaitGroup, host string, port, len int, duration time.Duration, chStat chan int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		log.Printf("disconnected")
		wg.Done()
	}()
	log.Printf("connecting to %s:%d, len %d, duration %s", host, port, len, duration.String())
	conn, err := utp.DialTimeout(fmt.Sprintf("%s:%d", host, port), time.Second)
	if err != nil {
		panic(err)
	}
	defer conn.Close()
	log.Printf("connected")
	buf := bytes.Repeat([]byte("H"), len)
	ts := time.Now()
	for time.Since(ts) < duration {
		n, err := conn.Write(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chStat <- n
	}
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

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

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

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

package v2

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/godbus/dbus"

	"github.com/containerd/cgroups/v2/stats"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	cgroupProcs    = "cgroup.procs"
	defaultDirPerm = 0755
)

// defaultFilePerm is a var so that the test framework can change the filemode
// of all files created when the tests are running.  The difference between the
// tests and real world use is that files like "cgroup.procs" will exist when writing
// to a read cgroup filesystem and do not exist prior when running in the tests.
// this is set to a non 0 value in the test code
var defaultFilePerm = os.FileMode(0)

// remove will remove a cgroup path handling EAGAIN and EBUSY errors and
// retrying the remove after a exp timeout
func remove(path string) error {
	var err error
	delay := 10 * time.Millisecond
	for i := 0; i < 5; i++ {
		if i != 0 {
			time.Sleep(delay)
			delay *= 2
		}
		if err = os.RemoveAll(path); err == nil {
			return nil
		}
	}
	return errors.Wrapf(err, "cgroups: unable to remove path %q", path)
}

// parseCgroupProcsFile parses /sys/fs/cgroup/$GROUPPATH/cgroup.procs
func parseCgroupProcsFile(path string) ([]uint64, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var (
		out []uint64
		s   = bufio.NewScanner(f)
	)
	for s.Scan() {
		if t := s.Text(); t != "" {
			pid, err := strconv.ParseUint(t, 10, 0)
			if err != nil {
				return nil, err
			}
			out = append(out, pid)
		}
	}
	return out, nil
}

func parseKV(raw string) (string, interface{}, error) {
	parts := strings.Fields(raw)
	switch len(parts) {
	case 2:
		v, err := parseUint(parts[1], 10, 64)
		if err != nil {
			// if we cannot parse as a uint, parse as a string
			return parts[0], parts[1], nil
		}
		return parts[0], v, nil
	default:
		return "", 0, ErrInvalidFormat
	}
}

func readUint(path string) (uint64, error) {
	v, err := ioutil.ReadFile(path)
	if err != nil {
		return 0, err
	}
	return parseUint(strings.TrimSpace(string(v)), 10, 64)
}

func parseUint(s string, base, bitSize int) (uint64, error) {
	v, err := strconv.ParseUint(s, base, bitSize)
	if err != nil {
		intValue, intErr := strconv.ParseInt(s, base, bitSize)
		// 1. Handle negative values greater than MinInt64 (and)
		// 2. Handle negative values lesser than MinInt64
		if intErr == nil && intValue < 0 {
			return 0, nil
		} else if intErr != nil &&
			intErr.(*strconv.NumError).Err == strconv.ErrRange &&
			intValue < 0 {
			return 0, nil
		}
		return 0, err
	}
	return v, nil
}

// parseCgroupFile parses /proc/PID/cgroup file and return string
func parseCgroupFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return parseCgroupFromReader(f)
}

func parseCgroupFromReader(r io.Reader) (string, error) {
	var (
		s = bufio.NewScanner(r)
	)
	for s.Scan() {
		if err := s.Err(); err != nil {
			return "", err
		}
		var (
			text  = s.Text()
			parts = strings.SplitN(text, ":", 3)
		)
		if len(parts) < 3 {
			return "", fmt.Errorf("invalid cgroup entry: %q", text)
		}
		// text is like "0::/user.slice/user-1001.slice/session-1.scope"
		if parts[0] == "0" && parts[1] == "" {
			return parts[2], nil
		}
	}
	return "", fmt.Errorf("cgroup path not found")
}

// ToResources converts the oci LinuxResources struct into a
// v2 Resources type for use with this package.
//
// converting cgroups configuration from v1 to v2
// ref: https://github.com/containers/crun/blob/master/crun.1.md#cgroup-v2
func ToResources(spec *specs.LinuxResources) *Resources {
	var resources Resources
	if cpu := spec.CPU; cpu != nil {
		resources.CPU = &CPU{
			Cpus: cpu.Cpus,
			Mems: cpu.Mems,
		}
		if shares := cpu.Shares; shares != nil {
			convertedWeight := (1 + ((*shares-2)*9999)/262142)
			resources.CPU.Weight = &convertedWeight
		}
		if period := cpu.Period; period != nil {
			resources.CPU.Max = NewCPUMax(cpu.Quota, period)
		}
	}
	if mem := spec.Memory; mem != nil {
		resources.Memory = &Memory{}
		if swap := mem.Swap; swap != nil {
			resources.Memory.Swap = swap
		}
		if l := mem.Limit; l != nil {
			resources.Memory.Max = l
		}
		if l := mem.Reservation; l != nil {
			resources.Memory.Low = l
		}
	}
	if pids := spec.Pids; pids != nil {
		resources.Pids = &Pids{
			Max: pids.Limit,
		}
	}
	if i := spec.BlockIO; i != nil {
		resources.IO = &IO{}
		if i.Weight != nil {
			resources.IO.BFQ.Weight = *i.Weight
		}
		for t, devices := range map[IOType][]specs.LinuxThrottleDevice{
			ReadBPS:   i.ThrottleReadBpsDevice,
			WriteBPS:  i.ThrottleWriteBpsDevice,
			ReadIOPS:  i.ThrottleReadIOPSDevice,
			WriteIOPS: i.ThrottleWriteIOPSDevice,
		} {
			for _, d := range devices {
				resources.IO.Max = append(resources.IO.Max, Entry{
					Type:  t,
					Major: d.Major,
					Minor: d.Minor,
					Rate:  d.Rate,
				})
			}
		}
	}
	if i := spec.Rdma; i != nil {
		resources.RDMA = &RDMA{}
		for device, value := range spec.Rdma {
			if device != "" && (value.HcaHandles != nil || value.HcaObjects != nil) {
				resources.RDMA.Limit = append(resources.RDMA.Limit, RDMAEntry{
					Device:     device,
					HcaHandles: *value.HcaHandles,
					HcaObjects: *value.HcaObjects,
				})
			}
		}
	}

	return &resources
}

// Gets uint64 parsed content of single value cgroup stat file
func getStatFileContentUint64(filePath string) uint64 {
	contents, err := ioutil.ReadFile(filePath)
	if err != nil {
		return 0
	}
	trimmed := strings.TrimSpace(string(contents))
	if trimmed == "max" {
		return math.MaxUint64
	}

	res, err := parseUint(trimmed, 10, 64)
	if err != nil {
		logrus.Errorf("unable to parse %q as a uint from Cgroup file %q", string(contents), filePath)
		return res
	}

	return res
}

func readIoStats(path string) []*stats.IOEntry {
	// more details on the io.stat file format: https://www.kernel.org/doc/Documentation/cgroup-v2.txt
	var usage []*stats.IOEntry
	fpath := filepath.Join(path, "io.stat")
	currentData, err := ioutil.ReadFile(fpath)
	if err != nil {
		return usage
	}
	entries := strings.Split(string(currentData), "\n")

	for _, entry := range entries {
		parts := strings.Split(entry, " ")
		if len(parts) < 2 {
			continue
		}
		majmin := strings.Split(parts[0], ":")
		if len(majmin) != 2 {
			continue
		}
		major, err := strconv.ParseUint(majmin[0], 10, 0)
		if err != nil {
			return usage
		}
		minor, err := strconv.ParseUint(majmin[1], 10, 0)
		if err != nil {
			return usage
		}
		parts = parts[1:]
		ioEntry := stats.IOEntry{
			Major: major,
			Minor: minor,
		}
		for _, stats := range parts {
			keyPairValue := strings.Split(stats, "=")
			if len(keyPairValue) != 2 {
				continue
			}
			v, err := strconv.ParseUint(keyPairValue[1], 10, 0)
			if err != nil {
				continue
			}
			switch keyPairValue[0] {
			case "rbytes":
				ioEntry.Rbytes = v
			case "wbytes":
				ioEntry.Wbytes = v
			case "rios":
				ioEntry.Rios = v
			case "wios":
				ioEntry.Wios = v
			}
		}
		usage = append(usage, &ioEntry)
	}
	return usage
}

func rdmaStats(filepath string) []*stats.RdmaEntry {
	currentData, err := ioutil.ReadFile(filepath)
	if err != nil {
		return []*stats.RdmaEntry{}
	}
	return toRdmaEntry(strings.Split(string(currentData), "\n"))
}

func parseRdmaKV(raw string, entry *stats.RdmaEntry) {
	var value uint64
	var err error

	parts := strings.Split(raw, "=")
	switch len(parts) {
	case 2:
		if parts[1] == "max" {
			value = math.MaxUint32
		} else {
			value, err = parseUint(parts[1], 10, 32)
			if err != nil {
				return
			}
		}
		if parts[0] == "hca_handle" {
			entry.HcaHandles = uint32(value)
		} else if parts[0] == "hca_object" {
			entry.HcaObjects = uint32(value)
		}
	}
}

func toRdmaEntry(strEntries []string) []*stats.RdmaEntry {
	var rdmaEntries []*stats.RdmaEntry
	for i := range strEntries {
		parts := strings.Fields(strEntries[i])
		switch len(parts) {
		case 3:
			entry := new(stats.RdmaEntry)
			entry.Device = parts[0]
			parseRdmaKV(parts[1], entry)
			parseRdmaKV(parts[2], entry)

			rdmaEntries = append(rdmaEntries, entry)
		default:
			continue
		}
	}
	return rdmaEntries
}

// isUnitExists returns true if the error is that a systemd unit already exists.
func isUnitExists(err error) bool {
	if err != nil {
		if dbusError, ok := err.(dbus.Error); ok {
			return strings.Contains(dbusError.Name, "org.freedesktop.systemd1.UnitExists")
		}
	}
	return false
}

func systemdUnitFromPath(path string) string {
	_, unit := filepath.Split(path)
	return unit
}
<commit_msg>Fix dbus version in utils<commit_after>/*
   Copyright The containerd Authors.

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

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

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

package v2

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/godbus/dbus/v5"

	"github.com/containerd/cgroups/v2/stats"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	cgroupProcs    = "cgroup.procs"
	defaultDirPerm = 0755
)

// defaultFilePerm is a var so that the test framework can change the filemode
// of all files created when the tests are running.  The difference between the
// tests and real world use is that files like "cgroup.procs" will exist when writing
// to a read cgroup filesystem and do not exist prior when running in the tests.
// this is set to a non 0 value in the test code
var defaultFilePerm = os.FileMode(0)

// remove will remove a cgroup path handling EAGAIN and EBUSY errors and
// retrying the remove after a exp timeout
func remove(path string) error {
	var err error
	delay := 10 * time.Millisecond
	for i := 0; i < 5; i++ {
		if i != 0 {
			time.Sleep(delay)
			delay *= 2
		}
		if err = os.RemoveAll(path); err == nil {
			return nil
		}
	}
	return errors.Wrapf(err, "cgroups: unable to remove path %q", path)
}

// parseCgroupProcsFile parses /sys/fs/cgroup/$GROUPPATH/cgroup.procs
func parseCgroupProcsFile(path string) ([]uint64, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var (
		out []uint64
		s   = bufio.NewScanner(f)
	)
	for s.Scan() {
		if t := s.Text(); t != "" {
			pid, err := strconv.ParseUint(t, 10, 0)
			if err != nil {
				return nil, err
			}
			out = append(out, pid)
		}
	}
	return out, nil
}

func parseKV(raw string) (string, interface{}, error) {
	parts := strings.Fields(raw)
	switch len(parts) {
	case 2:
		v, err := parseUint(parts[1], 10, 64)
		if err != nil {
			// if we cannot parse as a uint, parse as a string
			return parts[0], parts[1], nil
		}
		return parts[0], v, nil
	default:
		return "", 0, ErrInvalidFormat
	}
}

func readUint(path string) (uint64, error) {
	v, err := ioutil.ReadFile(path)
	if err != nil {
		return 0, err
	}
	return parseUint(strings.TrimSpace(string(v)), 10, 64)
}

func parseUint(s string, base, bitSize int) (uint64, error) {
	v, err := strconv.ParseUint(s, base, bitSize)
	if err != nil {
		intValue, intErr := strconv.ParseInt(s, base, bitSize)
		// 1. Handle negative values greater than MinInt64 (and)
		// 2. Handle negative values lesser than MinInt64
		if intErr == nil && intValue < 0 {
			return 0, nil
		} else if intErr != nil &&
			intErr.(*strconv.NumError).Err == strconv.ErrRange &&
			intValue < 0 {
			return 0, nil
		}
		return 0, err
	}
	return v, nil
}

// parseCgroupFile parses /proc/PID/cgroup file and return string
func parseCgroupFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return parseCgroupFromReader(f)
}

func parseCgroupFromReader(r io.Reader) (string, error) {
	var (
		s = bufio.NewScanner(r)
	)
	for s.Scan() {
		if err := s.Err(); err != nil {
			return "", err
		}
		var (
			text  = s.Text()
			parts = strings.SplitN(text, ":", 3)
		)
		if len(parts) < 3 {
			return "", fmt.Errorf("invalid cgroup entry: %q", text)
		}
		// text is like "0::/user.slice/user-1001.slice/session-1.scope"
		if parts[0] == "0" && parts[1] == "" {
			return parts[2], nil
		}
	}
	return "", fmt.Errorf("cgroup path not found")
}

// ToResources converts the oci LinuxResources struct into a
// v2 Resources type for use with this package.
//
// converting cgroups configuration from v1 to v2
// ref: https://github.com/containers/crun/blob/master/crun.1.md#cgroup-v2
func ToResources(spec *specs.LinuxResources) *Resources {
	var resources Resources
	if cpu := spec.CPU; cpu != nil {
		resources.CPU = &CPU{
			Cpus: cpu.Cpus,
			Mems: cpu.Mems,
		}
		if shares := cpu.Shares; shares != nil {
			convertedWeight := (1 + ((*shares-2)*9999)/262142)
			resources.CPU.Weight = &convertedWeight
		}
		if period := cpu.Period; period != nil {
			resources.CPU.Max = NewCPUMax(cpu.Quota, period)
		}
	}
	if mem := spec.Memory; mem != nil {
		resources.Memory = &Memory{}
		if swap := mem.Swap; swap != nil {
			resources.Memory.Swap = swap
		}
		if l := mem.Limit; l != nil {
			resources.Memory.Max = l
		}
		if l := mem.Reservation; l != nil {
			resources.Memory.Low = l
		}
	}
	if pids := spec.Pids; pids != nil {
		resources.Pids = &Pids{
			Max: pids.Limit,
		}
	}
	if i := spec.BlockIO; i != nil {
		resources.IO = &IO{}
		if i.Weight != nil {
			resources.IO.BFQ.Weight = *i.Weight
		}
		for t, devices := range map[IOType][]specs.LinuxThrottleDevice{
			ReadBPS:   i.ThrottleReadBpsDevice,
			WriteBPS:  i.ThrottleWriteBpsDevice,
			ReadIOPS:  i.ThrottleReadIOPSDevice,
			WriteIOPS: i.ThrottleWriteIOPSDevice,
		} {
			for _, d := range devices {
				resources.IO.Max = append(resources.IO.Max, Entry{
					Type:  t,
					Major: d.Major,
					Minor: d.Minor,
					Rate:  d.Rate,
				})
			}
		}
	}
	if i := spec.Rdma; i != nil {
		resources.RDMA = &RDMA{}
		for device, value := range spec.Rdma {
			if device != "" && (value.HcaHandles != nil || value.HcaObjects != nil) {
				resources.RDMA.Limit = append(resources.RDMA.Limit, RDMAEntry{
					Device:     device,
					HcaHandles: *value.HcaHandles,
					HcaObjects: *value.HcaObjects,
				})
			}
		}
	}

	return &resources
}

// Gets uint64 parsed content of single value cgroup stat file
func getStatFileContentUint64(filePath string) uint64 {
	contents, err := ioutil.ReadFile(filePath)
	if err != nil {
		return 0
	}
	trimmed := strings.TrimSpace(string(contents))
	if trimmed == "max" {
		return math.MaxUint64
	}

	res, err := parseUint(trimmed, 10, 64)
	if err != nil {
		logrus.Errorf("unable to parse %q as a uint from Cgroup file %q", string(contents), filePath)
		return res
	}

	return res
}

func readIoStats(path string) []*stats.IOEntry {
	// more details on the io.stat file format: https://www.kernel.org/doc/Documentation/cgroup-v2.txt
	var usage []*stats.IOEntry
	fpath := filepath.Join(path, "io.stat")
	currentData, err := ioutil.ReadFile(fpath)
	if err != nil {
		return usage
	}
	entries := strings.Split(string(currentData), "\n")

	for _, entry := range entries {
		parts := strings.Split(entry, " ")
		if len(parts) < 2 {
			continue
		}
		majmin := strings.Split(parts[0], ":")
		if len(majmin) != 2 {
			continue
		}
		major, err := strconv.ParseUint(majmin[0], 10, 0)
		if err != nil {
			return usage
		}
		minor, err := strconv.ParseUint(majmin[1], 10, 0)
		if err != nil {
			return usage
		}
		parts = parts[1:]
		ioEntry := stats.IOEntry{
			Major: major,
			Minor: minor,
		}
		for _, stats := range parts {
			keyPairValue := strings.Split(stats, "=")
			if len(keyPairValue) != 2 {
				continue
			}
			v, err := strconv.ParseUint(keyPairValue[1], 10, 0)
			if err != nil {
				continue
			}
			switch keyPairValue[0] {
			case "rbytes":
				ioEntry.Rbytes = v
			case "wbytes":
				ioEntry.Wbytes = v
			case "rios":
				ioEntry.Rios = v
			case "wios":
				ioEntry.Wios = v
			}
		}
		usage = append(usage, &ioEntry)
	}
	return usage
}

func rdmaStats(filepath string) []*stats.RdmaEntry {
	currentData, err := ioutil.ReadFile(filepath)
	if err != nil {
		return []*stats.RdmaEntry{}
	}
	return toRdmaEntry(strings.Split(string(currentData), "\n"))
}

func parseRdmaKV(raw string, entry *stats.RdmaEntry) {
	var value uint64
	var err error

	parts := strings.Split(raw, "=")
	switch len(parts) {
	case 2:
		if parts[1] == "max" {
			value = math.MaxUint32
		} else {
			value, err = parseUint(parts[1], 10, 32)
			if err != nil {
				return
			}
		}
		if parts[0] == "hca_handle" {
			entry.HcaHandles = uint32(value)
		} else if parts[0] == "hca_object" {
			entry.HcaObjects = uint32(value)
		}
	}
}

func toRdmaEntry(strEntries []string) []*stats.RdmaEntry {
	var rdmaEntries []*stats.RdmaEntry
	for i := range strEntries {
		parts := strings.Fields(strEntries[i])
		switch len(parts) {
		case 3:
			entry := new(stats.RdmaEntry)
			entry.Device = parts[0]
			parseRdmaKV(parts[1], entry)
			parseRdmaKV(parts[2], entry)

			rdmaEntries = append(rdmaEntries, entry)
		default:
			continue
		}
	}
	return rdmaEntries
}

// isUnitExists returns true if the error is that a systemd unit already exists.
func isUnitExists(err error) bool {
	if err != nil {
		if dbusError, ok := err.(dbus.Error); ok {
			return strings.Contains(dbusError.Name, "org.freedesktop.systemd1.UnitExists")
		}
	}
	return false
}

func systemdUnitFromPath(path string) string {
	_, unit := filepath.Split(path)
	return unit
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	fdk "github.com/fnproject/fdk-go"
)

const (
	// InvalidResponseStr is a string that isn't one of the 'hot' formats.
	InvalidResponseStr = "Olive oil is a liquid fat obtained from olives...\n"
)

// AppRequest is the body of the input of a function, it can be used to change
// behavior of this function.
type AppRequest struct {
	// if specified we 'sleep' the specified msecs
	SleepTime int `json:"sleepTime,omitempty"`
	// if specified, this is our response http status code
	ResponseCode int `json:"responseCode,omitempty"`
	// if specified, this is our response content-type
	ResponseContentType string `json:"responseContentType,omitempty"`
	// if specified, this is our response content-type.
	// jason doesn't sit with the other kids at school.
	JasonContentType string `json:"jasonContentType,omitempty"`
	// if specified, this is echoed back to client
	EchoContent string `json:"echoContent,omitempty"`
	// verbose mode
	IsDebug bool `json:"isDebug,omitempty"`
	// simulate crash
	IsCrash bool `json:"isCrash,omitempty"`
	// read a file from disk
	ReadFile string `json:"readFile,omitempty"`
	// fill created with with zero bytes of specified size
	ReadFileSize int `json:"readFileSize,omitempty"`
	// create a file on disk
	CreateFile string `json:"createFile,omitempty"`
	// fill created with with zero bytes of specified size
	CreateFileSize int `json:"createFileSize,omitempty"`
	// allocate RAM and hold until next request
	AllocateMemory int `json:"allocateMemory,omitempty"`
	// leak RAM forever
	LeakMemory int `json:"leakMemory,omitempty"`
	// duplicate trailer if > 0
	TrailerRepeat int `json:"trailerRepeat,omitempty"`
	// corrupt http or json
	InvalidResponse bool `json:"invalidResponse,omitempty"`
	// if specified we 'sleep' the specified msecs *after* processing request
	PostSleepTime int `json:"postSleepTime,omitempty"`
	// spit this out in stdout after processing each request
	PostOutGarbage string `json:"postOutGarbage,omitempty"`
	// spit this out in stderr after processing each request
	PostErrGarbage string `json:"postErrGarbage,omitempty"`
	// test empty body
	IsEmptyBody bool `json:"isEmptyBody,omitempty"`
	// test headers that come into function
	ExpectHeaders map[string][]string `json:"expectHeaders,omitempty"`
	// send some headers out explicitly
	ReturnHeaders map[string][]string `json:"returnHeaders,omitempty"`

	// TODO: simulate slow read/slow write
	// TODO: simulate partial IO write/read
	// TODO: simulate high cpu usage (async and sync)
	// TODO: simulate large body upload/download
	// TODO: infinite loop
}

// Leaks are ever growing memory leak chunks
var Leaks []*[]byte

// Hold is memory to hold on to at every request, new requests overwrite it.
var Hold []byte

// AppResponse is the output of this function, in JSON
type AppResponse struct {
	Request AppRequest        `json:"request"`
	Headers http.Header       `json:"header"`
	Config  map[string]string `json:"config"`
	Data    map[string]string `json:"data"`
	Trailer []string          `json:"trailer"`
}

func init() {
	Leaks = make([]*[]byte, 0, 0)
}

func getTotalLeaks() int {
	total := 0
	for idx := range Leaks {
		total += len(*(Leaks[idx]))
	}
	return total
}

// AppHandler is the fdk.Handler used by this package
func AppHandler(ctx context.Context, in io.Reader, out io.Writer) {
	req, resp := processRequest(ctx, in)

	if req.InvalidResponse {
		_, err := io.Copy(out, strings.NewReader(InvalidResponseStr))
		if err != nil {
			log.Printf("io copy error %v", err)
			panic(err.Error())
		}
	}

	finalizeRequest(out, req, resp)
	err := postProcessRequest(req, out)
	if err != nil {
		panic(err.Error())
	}
}

func finalizeRequest(out io.Writer, req *AppRequest, resp *AppResponse) {
	// custom content type
	if req.ResponseContentType != "" {
		fdk.SetHeader(out, "Content-Type", req.ResponseContentType)
	}
	// NOTE: don't add 'application/json' explicitly here as an else,
	// we will test that go's auto-detection logic does not fade since
	// some people are relying on it now

	if req.JasonContentType != "" {
		// this will get picked up by our json out handler...
		fdk.SetHeader(out, "Content-Type", req.JasonContentType)
	}

	if req.ReturnHeaders != nil {
		for k, vs := range req.ReturnHeaders {
			for _, v := range vs {
				fdk.AddHeader(out, k, v)
			}
		}
	}

	// custom response code
	if req.ResponseCode != 0 {
		fdk.WriteStatus(out, req.ResponseCode)
	}

	if !req.IsEmptyBody {
		json.NewEncoder(out).Encode(resp)
	}
}

func processRequest(ctx context.Context, in io.Reader) (*AppRequest, *AppResponse) {

	fnctx := fdk.GetContext(ctx)

	var request AppRequest
	json.NewDecoder(in).Decode(&request)

	if request.IsDebug {
		log.Printf("BeginOfLogs")
		log.Printf("Received request %#v", request)
		log.Printf("Received headers %v", fnctx.Header())
		log.Printf("Received config %v", fnctx.Config())
	}

	// simulate load if requested
	if request.SleepTime > 0 {
		if request.IsDebug {
			log.Printf("Sleeping %d", request.SleepTime)
		}
		time.Sleep(time.Duration(request.SleepTime) * time.Millisecond)
	}

	data := make(map[string]string)

	// read a file
	if request.ReadFile != "" {
		if request.IsDebug {
			log.Printf("Reading file %s", request.ReadFile)
		}
		out, err := readFile(request.ReadFile, request.ReadFileSize)
		if err != nil {
			data[request.ReadFile+".read_error"] = err.Error()
		} else {
			data[request.ReadFile+".read_output"] = out
		}
	}

	// create a file
	if request.CreateFile != "" {
		if request.IsDebug {
			log.Printf("Creating file %s (size: %d)", request.CreateFile, request.CreateFileSize)
		}
		err := createFile(request.CreateFile, request.CreateFileSize)
		if err != nil {
			data[request.CreateFile+".create_error"] = err.Error()
		}
	}

	// handle one time alloc request (hold on to the memory until next request)
	if request.AllocateMemory != 0 && request.IsDebug {
		log.Printf("Allocating memory size: %d", request.AllocateMemory)
	}
	Hold = getChunk(request.AllocateMemory)

	// leak memory forever
	if request.LeakMemory != 0 {
		if request.IsDebug {
			log.Printf("Leaking memory size: %d total: %d", request.LeakMemory, getTotalLeaks())
		}
		chunk := getChunk(request.LeakMemory)
		Leaks = append(Leaks, &chunk)
	}

	if request.IsDebug {
		info := getDockerInfo()
		log.Printf("DockerInfo %+v", info)
		data["DockerId"] = info.ID
		data["DockerHostname"] = info.Hostname
	}

	// simulate crash
	if request.IsCrash {
		log.Fatalln("Crash requested")
	}

	if request.ExpectHeaders != nil {
		for name, header := range request.ExpectHeaders {
			if h2 := fnctx.Header().Get(name); header[0] != h2 {
				log.Fatalf("Expected header `%s` to be `%s` but was `%s`",
					name, header[0], h2)
			}
		}
	}

	if _, ok := ctx.Deadline(); !ok {
		// XXX(reed): we should plumb the timeout and test it's approximately right but who has time for that?
		log.Fatalf("fdk should set deadline, go fix fdk-go immediately you")
	}

	resp := AppResponse{
		Data:    data,
		Request: request,
		Headers: fnctx.Header(),
		Config:  fnctx.Config(),
		Trailer: make([]string, 0, request.TrailerRepeat),
	}

	for i := request.TrailerRepeat; i > 0; i-- {
		resp.Trailer = append(resp.Trailer, request.EchoContent)
	}

	// Well, almost true.. If panic/errors, we may print stuff after this
	if request.IsDebug {
		log.Printf("EndOfLogs")
	}
	return &request, &resp
}

func postProcessRequest(request *AppRequest, out io.Writer) error {
	if request == nil {
		return nil
	}

	if request.PostSleepTime > 0 {
		if request.IsDebug {
			log.Printf("PostProcess Sleeping %d", request.PostSleepTime)
		}
		time.Sleep(time.Duration(request.PostSleepTime) * time.Millisecond)
	}

	if request.PostOutGarbage != "" {
		if request.IsDebug {
			log.Printf("PostProcess PostOutGarbage %s", request.PostOutGarbage)
		}

		_, err := io.WriteString(out, request.PostOutGarbage)
		if err != nil {
			log.Printf("PostOutGarbage write string error %v", err)
			return err
		}
	}

	if request.PostErrGarbage != "" {
		log.Printf("PostProcess PostErrGarbage %s", request.PostErrGarbage)
	}

	return nil
}

func main() {
	if os.Getenv("ENABLE_HEADER") != "" {
		log.Printf("Container starting")
	}

	fdk.Handle(fdk.HandlerFunc(AppHandler)) // XXX(reed): can extract & instrument

	if os.Getenv("ENABLE_FOOTER") != "" {
		log.Printf("Container ending")
	}
}

func getChunk(size int) []byte {
	chunk := make([]byte, size)
	// fill it
	for idx := range chunk {
		chunk[idx] = 1
	}
	return chunk
}

func readFile(name string, size int) (string, error) {
	// read the whole file into memory
	out, err := ioutil.ReadFile(name)
	if err != nil {
		return "", err
	}
	// only respond with partion output if requested
	if size > 0 {
		return string(out[:size]), nil
	}
	return string(out), nil
}

func createFile(name string, size int) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}

	if size > 0 {
		// create a 1K block (keep this buffer small to keep
		// memory usage small)
		chunk := make([]byte, 1024)
		for i := 0; i < 1024; i++ {
			chunk[i] = byte(i)
		}

		for size > 0 {
			dlen := size
			if dlen > 1024 {
				dlen = 1024
			}

			_, err := f.Write(chunk[:dlen])
			if err != nil {
				return err
			}

			// slightly modify the chunk to avoid any sparse file possibility
			chunk[0]++
			size = size - dlen
		}
	}
	return nil
}

type dockerInfo struct {
	Hostname string
	ID       string
}

func getDockerInfo() dockerInfo {
	var info dockerInfo

	info.Hostname, _ = os.Hostname()

	// cgroup file has lines such as, where last token is the docker id
	/*
		12:freezer:/docker/610d96c712c6983776f920f2bcf10fae056a6fe5274393c86678ca802d184b0a
	*/
	file, err := os.Open("/proc/self/cgroup")
	if err == nil {
		defer file.Close()
		r := bufio.NewReader(file)
		for {
			line, _, err := r.ReadLine()
			if err != nil {
				break
			}

			tokens := bytes.Split(line, []byte("/"))
			tokLen := len(tokens)
			if tokLen >= 3 && bytes.Compare(tokens[tokLen-2], []byte("docker")) == 0 {
				info.ID = string(tokens[tokLen-1])
				break
			}
		}
	}

	return info
}
<commit_msg>fn: adding pre/post sleep and UDS shutdown in fn-test-utils (#1291)<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	fdk "github.com/fnproject/fdk-go"
)

const (
	// InvalidResponseStr is a string that isn't one of the 'hot' formats.
	InvalidResponseStr = "Olive oil is a liquid fat obtained from olives...\n"
)

// AppRequest is the body of the input of a function, it can be used to change
// behavior of this function.
type AppRequest struct {
	// if specified we 'sleep' the specified msecs
	SleepTime int `json:"sleepTime,omitempty"`
	// if specified, this is our response http status code
	ResponseCode int `json:"responseCode,omitempty"`
	// if specified, this is our response content-type
	ResponseContentType string `json:"responseContentType,omitempty"`
	// if specified, this is our response content-type.
	// jason doesn't sit with the other kids at school.
	JasonContentType string `json:"jasonContentType,omitempty"`
	// if specified, this is echoed back to client
	EchoContent string `json:"echoContent,omitempty"`
	// verbose mode
	IsDebug bool `json:"isDebug,omitempty"`
	// simulate crash
	IsCrash bool `json:"isCrash,omitempty"`
	// shutdown UDS after request
	IsShutdown bool `json:"isShutdown,omitempty"`
	// read a file from disk
	ReadFile string `json:"readFile,omitempty"`
	// fill created with with zero bytes of specified size
	ReadFileSize int `json:"readFileSize,omitempty"`
	// create a file on disk
	CreateFile string `json:"createFile,omitempty"`
	// fill created with with zero bytes of specified size
	CreateFileSize int `json:"createFileSize,omitempty"`
	// allocate RAM and hold until next request
	AllocateMemory int `json:"allocateMemory,omitempty"`
	// leak RAM forever
	LeakMemory int `json:"leakMemory,omitempty"`
	// duplicate trailer if > 0
	TrailerRepeat int `json:"trailerRepeat,omitempty"`
	// corrupt http or json
	InvalidResponse bool `json:"invalidResponse,omitempty"`
	// if specified we 'sleep' the specified msecs *after* processing request
	PostSleepTime int `json:"postSleepTime,omitempty"`
	// spit this out in stdout after processing each request
	PostOutGarbage string `json:"postOutGarbage,omitempty"`
	// spit this out in stderr after processing each request
	PostErrGarbage string `json:"postErrGarbage,omitempty"`
	// test empty body
	IsEmptyBody bool `json:"isEmptyBody,omitempty"`
	// test headers that come into function
	ExpectHeaders map[string][]string `json:"expectHeaders,omitempty"`
	// send some headers out explicitly
	ReturnHeaders map[string][]string `json:"returnHeaders,omitempty"`

	// TODO: simulate slow read/slow write
	// TODO: simulate partial IO write/read
	// TODO: simulate high cpu usage (async and sync)
	// TODO: simulate large body upload/download
	// TODO: infinite loop
}

// Leaks are ever growing memory leak chunks
var Leaks []*[]byte

// Hold is memory to hold on to at every request, new requests overwrite it.
var Hold []byte

var GlobCancel context.CancelFunc

// AppResponse is the output of this function, in JSON
type AppResponse struct {
	Request AppRequest        `json:"request"`
	Headers http.Header       `json:"header"`
	Config  map[string]string `json:"config"`
	Data    map[string]string `json:"data"`
	Trailer []string          `json:"trailer"`
}

func init() {
	Leaks = make([]*[]byte, 0, 0)
}

func getTotalLeaks() int {
	total := 0
	for idx := range Leaks {
		total += len(*(Leaks[idx]))
	}
	return total
}

// AppHandler is the fdk.Handler used by this package
func AppHandler(ctx context.Context, in io.Reader, out io.Writer) {
	req, resp := processRequest(ctx, in)

	if req.InvalidResponse {
		_, err := io.Copy(out, strings.NewReader(InvalidResponseStr))
		if err != nil {
			log.Fatalf("io copy error %v", err)
		}
	}

	finalizeRequest(out, req, resp)
	err := postProcessRequest(req, out)
	if err != nil {
		log.Fatalf("post process error %v", err)
	}
}

func finalizeRequest(out io.Writer, req *AppRequest, resp *AppResponse) {
	// custom content type
	if req.ResponseContentType != "" {
		fdk.SetHeader(out, "Content-Type", req.ResponseContentType)
	}
	// NOTE: don't add 'application/json' explicitly here as an else,
	// we will test that go's auto-detection logic does not fade since
	// some people are relying on it now

	if req.JasonContentType != "" {
		// this will get picked up by our json out handler...
		fdk.SetHeader(out, "Content-Type", req.JasonContentType)
	}

	if req.ReturnHeaders != nil {
		for k, vs := range req.ReturnHeaders {
			for _, v := range vs {
				fdk.AddHeader(out, k, v)
			}
		}
	}

	// custom response code
	if req.ResponseCode != 0 {
		fdk.WriteStatus(out, req.ResponseCode)
	}

	if !req.IsEmptyBody {
		json.NewEncoder(out).Encode(resp)
	}
}

func processRequest(ctx context.Context, in io.Reader) (*AppRequest, *AppResponse) {

	fnctx := fdk.GetContext(ctx)

	var request AppRequest
	json.NewDecoder(in).Decode(&request)

	if request.IsDebug {
		log.Printf("BeginOfLogs")
		log.Printf("Received request %#v", request)
		log.Printf("Received headers %v", fnctx.Header())
		log.Printf("Received config %v", fnctx.Config())
	}

	// simulate load if requested
	if request.SleepTime > 0 {
		if request.IsDebug {
			log.Printf("Sleeping %d", request.SleepTime)
		}
		time.Sleep(time.Duration(request.SleepTime) * time.Millisecond)
	}

	data := make(map[string]string)

	// read a file
	if request.ReadFile != "" {
		if request.IsDebug {
			log.Printf("Reading file %s", request.ReadFile)
		}
		out, err := readFile(request.ReadFile, request.ReadFileSize)
		if err != nil {
			data[request.ReadFile+".read_error"] = err.Error()
		} else {
			data[request.ReadFile+".read_output"] = out
		}
	}

	// create a file
	if request.CreateFile != "" {
		if request.IsDebug {
			log.Printf("Creating file %s (size: %d)", request.CreateFile, request.CreateFileSize)
		}
		err := createFile(request.CreateFile, request.CreateFileSize)
		if err != nil {
			data[request.CreateFile+".create_error"] = err.Error()
		}
	}

	// handle one time alloc request (hold on to the memory until next request)
	if request.AllocateMemory != 0 && request.IsDebug {
		log.Printf("Allocating memory size: %d", request.AllocateMemory)
	}
	Hold = getChunk(request.AllocateMemory)

	// leak memory forever
	if request.LeakMemory != 0 {
		if request.IsDebug {
			log.Printf("Leaking memory size: %d total: %d", request.LeakMemory, getTotalLeaks())
		}
		chunk := getChunk(request.LeakMemory)
		Leaks = append(Leaks, &chunk)
	}

	if request.IsDebug {
		info := getDockerInfo()
		log.Printf("DockerInfo %+v", info)
		data["DockerId"] = info.ID
		data["DockerHostname"] = info.Hostname
	}

	// simulate crash
	if request.IsCrash {
		log.Fatalln("Crash requested")
	}

	if request.ExpectHeaders != nil {
		for name, header := range request.ExpectHeaders {
			if h2 := fnctx.Header().Get(name); header[0] != h2 {
				log.Fatalf("Expected header `%s` to be `%s` but was `%s`",
					name, header[0], h2)
			}
		}
	}

	if _, ok := ctx.Deadline(); !ok {
		// XXX(reed): we should plumb the timeout and test it's approximately right but who has time for that?
		log.Fatalf("fdk should set deadline, go fix fdk-go immediately you")
	}

	resp := AppResponse{
		Data:    data,
		Request: request,
		Headers: fnctx.Header(),
		Config:  fnctx.Config(),
		Trailer: make([]string, 0, request.TrailerRepeat),
	}

	for i := request.TrailerRepeat; i > 0; i-- {
		resp.Trailer = append(resp.Trailer, request.EchoContent)
	}

	// Well, almost true.. If panic/errors, we may print stuff after this
	if request.IsDebug {
		log.Printf("EndOfLogs")
	}
	return &request, &resp
}

func postProcessRequest(request *AppRequest, out io.Writer) error {
	if request == nil {
		return nil
	}

	if request.PostSleepTime > 0 {
		if request.IsDebug {
			log.Printf("PostProcess Sleeping %d", request.PostSleepTime)
		}
		time.Sleep(time.Duration(request.PostSleepTime) * time.Millisecond)
	}

	if request.PostOutGarbage != "" {
		if request.IsDebug {
			log.Printf("PostProcess PostOutGarbage %s", request.PostOutGarbage)
		}

		_, err := io.WriteString(out, request.PostOutGarbage)
		if err != nil {
			log.Printf("PostOutGarbage write string error %v", err)
			return err
		}
	}

	if request.PostErrGarbage != "" {
		log.Printf("PostProcess PostErrGarbage %s", request.PostErrGarbage)
	}

	if request.IsShutdown && GlobCancel != nil {
		log.Printf("PostProcess Shutting down UDS")
		GlobCancel()
	}

	return nil
}

func main() {
	if os.Getenv("ENABLE_HEADER") != "" {
		log.Printf("Container starting")
	}

	// simulate long initialization
	if sleeper := os.Getenv("ENABLE_INIT_DELAY_MSEC"); sleeper != "" {
		log.Printf("Container start sleep %v", sleeper)
		delay, err := strconv.ParseInt(sleeper, 10, 64)
		if err != nil {
			log.Fatalf("cannot parse ENABLE_INIT_DELAY_MSEC %v", err)
		}
		time.Sleep(time.Millisecond * time.Duration(delay))
	}

	ctx, cancel := context.WithCancel(context.Background())
	GlobCancel = cancel
	fdk.HandleContext(ctx, fdk.HandlerFunc(AppHandler)) // XXX(reed): can extract & instrument

	// simulate long exit
	if sleeper := os.Getenv("ENABLE_EXIT_DELAY_MSEC"); sleeper != "" {
		log.Printf("Container end sleep %v", sleeper)
		delay, err := strconv.ParseInt(sleeper, 10, 64)
		if err != nil {
			log.Fatalf("cannot parse ENABLE_EXIT_DELAY_MSEC %v", err)
		}
		time.Sleep(time.Millisecond * time.Duration(delay))
	}

	if os.Getenv("ENABLE_FOOTER") != "" {
		log.Printf("Container ending")
	}
}

func getChunk(size int) []byte {
	chunk := make([]byte, size)
	// fill it
	for idx := range chunk {
		chunk[idx] = 1
	}
	return chunk
}

func readFile(name string, size int) (string, error) {
	// read the whole file into memory
	out, err := ioutil.ReadFile(name)
	if err != nil {
		return "", err
	}
	// only respond with partion output if requested
	if size > 0 {
		return string(out[:size]), nil
	}
	return string(out), nil
}

func createFile(name string, size int) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}

	if size > 0 {
		// create a 1K block (keep this buffer small to keep
		// memory usage small)
		chunk := make([]byte, 1024)
		for i := 0; i < 1024; i++ {
			chunk[i] = byte(i)
		}

		for size > 0 {
			dlen := size
			if dlen > 1024 {
				dlen = 1024
			}

			_, err := f.Write(chunk[:dlen])
			if err != nil {
				return err
			}

			// slightly modify the chunk to avoid any sparse file possibility
			chunk[0]++
			size = size - dlen
		}
	}
	return nil
}

type dockerInfo struct {
	Hostname string
	ID       string
}

func getDockerInfo() dockerInfo {
	var info dockerInfo

	info.Hostname, _ = os.Hostname()

	// cgroup file has lines such as, where last token is the docker id
	/*
		12:freezer:/docker/610d96c712c6983776f920f2bcf10fae056a6fe5274393c86678ca802d184b0a
	*/
	file, err := os.Open("/proc/self/cgroup")
	if err == nil {
		defer file.Close()
		r := bufio.NewReader(file)
		for {
			line, _, err := r.ReadLine()
			if err != nil {
				break
			}

			tokens := bytes.Split(line, []byte("/"))
			tokLen := len(tokens)
			if tokLen >= 3 && bytes.Compare(tokens[tokLen-2], []byte("docker")) == 0 {
				info.ID = string(tokens[tokLen-1])
				break
			}
		}
	}

	return info
}
<|endoftext|>


<commit_before>package whitespace

import (
	"testing"
)

func TestWhitespace(t *testing.T) {

	tests := []struct {
		R rune
		Expected bool
	}{
		{' '  , true},

		{'\a' , false}, // bell (== \007)
		{'\f' ,  true}, // form feed (== \014)
		{'\t' ,  true}, // horizontal tab (== \011)
		{'\n' ,  true}, // newline (== \012)
		{'\r' ,  true}, // carriage return (== \015)
		{'\v' ,  true}, // vertical tab character (== \013)

		{'a'  , false},
		{'A'  , false},
		{'x'  , false},
		{'X'  , false},
		{'z'  , false},
		{'Z'  , false},

		{'!'  , false},
		{'@'  , false},
		{'#'  , false},
		{'$'  , false},
		{'%'  , false},
		{'^'  , false},
		{'&'  , false},
		{'*'  , false},
		{'('  , false},
		{')'  , false},
		{'_'  , false},
		{'+'  , false},

		{'1'  , false},
		{'2'  , false},
		{'3'  , false},
		{'4'  , false},
		{'5'  , false},
		{'6'  , false},
		{'7'  , false},
		{'8'  , false},
		{'9'  , false},
		{'0'  , false},
		{'-'  , false},
		{'='  , false},

		{'\u0009', true}, // horizontal tab
		{'\u000A', true}, // line feed
		{'\u000B', true}, // vertical tab
		{'\u000C', true}, // form feed
		{'\u000D', true}, // carriage return
		{'\u0020', true}, // space
		{'\u0085', true}, // next line
		{'\u00A0', true}, // no-break space
		{'\u1680', true}, // ogham space mark
		{'\u180E', true}, // mongolian vowel separator
		{'\u2000', true}, // en quad
		{'\u2001', true}, // em quad
		{'\u2002', true}, // en space
		{'\u2003', true}, // em space
		{'\u2004', true}, // three-per-em space
		{'\u2005', true}, // four-per-em space
		{'\u2006', true}, // six-per-em space
		{'\u2007', true}, // figure space
		{'\u2008', true}, // punctuation space
		{'\u2009', true}, // thin space
		{'\u200A', true}, // hair space
		{'\u2028', true}, // line separator
		{'\u2029', true}, // paragraph separator
		{'\u202F', true}, // narrow no-break space
		{'\u205F', true}, // medium mathematical space
		{'\u3000', true}, // ideographic space

	}

	for _,datum := range tests {

		expected := datum.Expected

		actual := IsWhitespace(datum.R)

		if expected != actual {
			t.Errorf("For rune [%v] ([%v]) expected it to be [%v] that it is white space, but instead got [%v].", datum.R, string(datum.R), expected, actual)
		}

	} // for

}

<commit_msg>updated style of test<commit_after>package whitespace

import (
	"testing"
)

func TestWhitespace(t *testing.T) {

	tests := []struct {
		R rune
		Expected bool
	}{
		{' '  , true},

		{'\a' , false}, // bell (== \007)
		{'\f' ,  true}, // form feed (== \014)
		{'\t' ,  true}, // horizontal tab (== \011)
		{'\n' ,  true}, // newline (== \012)
		{'\r' ,  true}, // carriage return (== \015)
		{'\v' ,  true}, // vertical tab character (== \013)

		{'a'  , false},
		{'A'  , false},
		{'x'  , false},
		{'X'  , false},
		{'z'  , false},
		{'Z'  , false},

		{'!'  , false},
		{'@'  , false},
		{'#'  , false},
		{'$'  , false},
		{'%'  , false},
		{'^'  , false},
		{'&'  , false},
		{'*'  , false},
		{'('  , false},
		{')'  , false},
		{'_'  , false},
		{'+'  , false},

		{'1'  , false},
		{'2'  , false},
		{'3'  , false},
		{'4'  , false},
		{'5'  , false},
		{'6'  , false},
		{'7'  , false},
		{'8'  , false},
		{'9'  , false},
		{'0'  , false},
		{'-'  , false},
		{'='  , false},

		{'\u0009', true}, // horizontal tab
		{'\u000A', true}, // line feed
		{'\u000B', true}, // vertical tab
		{'\u000C', true}, // form feed
		{'\u000D', true}, // carriage return
		{'\u0020', true}, // space
		{'\u0085', true}, // next line
		{'\u00A0', true}, // no-break space
		{'\u1680', true}, // ogham space mark
		{'\u180E', true}, // mongolian vowel separator
		{'\u2000', true}, // en quad
		{'\u2001', true}, // em quad
		{'\u2002', true}, // en space
		{'\u2003', true}, // em space
		{'\u2004', true}, // three-per-em space
		{'\u2005', true}, // four-per-em space
		{'\u2006', true}, // six-per-em space
		{'\u2007', true}, // figure space
		{'\u2008', true}, // punctuation space
		{'\u2009', true}, // thin space
		{'\u200A', true}, // hair space
		{'\u2028', true}, // line separator
		{'\u2029', true}, // paragraph separator
		{'\u202F', true}, // narrow no-break space
		{'\u205F', true}, // medium mathematical space
		{'\u3000', true}, // ideographic space

	}

	for testNumber, datum := range tests {

		expected := datum.Expected

		actual := IsWhitespace(datum.R)

		if expected != actual {
			t.Errorf("For test #%d, expected whitespace.IsWhitespace(%q = %d) = %t, but actually got whitespace.IsWhitespace(%q = %d) = %v.", testNumber, datum.R, datum.R, expected, datum.R, datum.R, actual)
		}

	}
}

<|endoftext|>


<commit_before>package whoson

import (
	"fmt"
	"net"
	"net/textproto"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type Session struct {
	protocol ProtocolType

	udpconn    *net.UDPConn
	remoteAddr *net.UDPAddr
	b          *Buffer

	tcpserver *TCPServer
	conn      net.Conn
	tp        *textproto.Conn
	tpid      uint

	cmdMethod MethodType
	cmdIp     net.IP
	cmdArgs   string
}

func NewSessionUDP(c *net.UDPConn, r *net.UDPAddr, b *Buffer) *Session {
	return &Session{
		protocol:   pUDP,
		udpconn:    c,
		remoteAddr: r,
		b:          b,
	}
}

func NewSessionTCP(s *TCPServer, c net.Conn) *Session {
	return &Session{
		protocol:  pTCP,
		tcpserver: s,
		conn:      c,
		tp:        textproto.NewConn(c),
	}
}

func (ses *Session) setTpId() {
	ses.tpid = ses.tp.Next()
}

func (ses *Session) methodType(m string) MethodType {
	if v, ok := methodFromString[m]; ok {
		return v
	} else {
		return mUnkownMethod
	}
}

func (ses *Session) parseCmd(line string) error {
	var cmd []string

	if strings.TrimSpace(line) == "" {
		return errors.New("command parse error")
	}

	w := strings.Split(line, " ")
	for i := range w {
		if w[i] != "" {
			cmd = append(cmd, strings.TrimSpace(w[i]))
		}
	}

	ses.cmdMethod = ses.methodType(strings.ToUpper(cmd[0]))
	switch ses.cmdMethod {
	case mLogin, mLogout, mQuery:
		if ses.cmdIp = net.ParseIP(cmd[1]); ses.cmdIp == nil {
			return errors.New("command parse error")
		}
		ses.cmdArgs = strings.Join(cmd[2:], " ")
	case mQuit:
		//pp.Println("Quit")
		ses.cmdArgs = strings.Join(cmd[1:], " ")
	default:
		return errors.New("command not found")
	}
	return nil
}

func (ses *Session) readLine() (string, error) {
	ses.tp.StartRequest(ses.tpid)
	l1, err := ses.tp.ReadLine()
	if err != nil {
		return "", err
	}

	l2, err := ses.tp.ReadLine()
	if err != nil {
		return "", err
	}
	ses.tp.EndRequest(ses.tpid)

	if l1 != "" && l2 == "" {
		return l1, nil
	} else {
		return "", errors.New("session read error")
	}
}

func (ses *Session) sendLine(str string) error {
	var err error
	if ses.protocol == pTCP {
		ses.tp.StartResponse(ses.tpid)
		err = ses.tp.PrintfLine(str + CRLF)
		ses.tp.EndResponse(ses.tpid)
	} else {
		b := []byte(str + CRLF + CRLF)
		_, err = ses.udpconn.WriteToUDP(b, ses.remoteAddr)
	}
	return err
}

func (ses *Session) sendResponsePositive(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rPositive], str))
}

func (ses *Session) sendResponseNegative(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rNegative], str))
}

func (ses *Session) sendResponseBadRequest(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rBadRequest], str))
}

func (ses *Session) startHandler() bool {
	var err error

	if ses.protocol == pTCP {
		ses.setTpId()

		line, err := ses.readLine()
		if err != nil {
			if opError, ok := err.(*net.OpError); ok && opError.Timeout() {
				return false
			}
			if err.Error() == "EOF" {
				return false
			} else {
				ses.sendResponseBadRequest(err.Error())
				return true
			}
		}
		err = ses.parseCmd(line)
		if err != nil {
			ses.sendResponseBadRequest(err.Error())
			return true
		}
	} else {
		err = ses.parseCmd(string(ses.b.buf[:ses.b.count]))
		if err != nil {
			ses.sendResponseBadRequest(err.Error())
			return true
		}
	}

	switch ses.cmdMethod {
	case mLogin:
		ses.methodLogin()
	case mLogout:
		ses.methodLogout()
	case mQuery:
		ses.methodQuery()
	case mQuit:
		ses.methodQuit()
	default:
		err := errors.New("handler error")
		ses.sendResponseBadRequest(err.Error())
	}
	return true
}

func (ses *Session) methodLogin() {
	sd := &StoreData{
		Expire: time.Now().Add(StoreDataExpire),
		IP:     ses.cmdIp,
		Data:   ses.cmdArgs,
	}
	MainStore.Set(sd.Key(), sd)
	ses.sendResponsePositive("LOGIN OK")
}

func (ses *Session) methodLogout() {
	ok := MainStore.Del(ses.cmdIp.String())
	if ok {
		ses.sendResponsePositive("LOGOUT record deleted")
	} else {
		ses.sendResponsePositive("LOGOUT no such record, nothing done")
	}
}

func (ses *Session) methodQuery() {
	sd, err := MainStore.Get(ses.cmdIp.String())
	if err != nil {
		ses.sendResponseNegative("Not Logged in")
	} else {
		ses.sendResponsePositive(sd.Data)
	}
}

func (ses *Session) methodQuit() {
	ses.sendResponsePositive("QUIT OK")
	ses.conn.Close()
}

func (ses *Session) close() {
	if ses.protocol == pUDP {
		ses.b.Free()
	}
}
<commit_msg>Close processing is limited to tcp<commit_after>package whoson

import (
	"fmt"
	"net"
	"net/textproto"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type Session struct {
	protocol ProtocolType

	udpconn    *net.UDPConn
	remoteAddr *net.UDPAddr
	b          *Buffer

	tcpserver *TCPServer
	conn      net.Conn
	tp        *textproto.Conn
	tpid      uint

	cmdMethod MethodType
	cmdIp     net.IP
	cmdArgs   string
}

func NewSessionUDP(c *net.UDPConn, r *net.UDPAddr, b *Buffer) *Session {
	return &Session{
		protocol:   pUDP,
		udpconn:    c,
		remoteAddr: r,
		b:          b,
	}
}

func NewSessionTCP(s *TCPServer, c net.Conn) *Session {
	return &Session{
		protocol:  pTCP,
		tcpserver: s,
		conn:      c,
		tp:        textproto.NewConn(c),
	}
}

func (ses *Session) setTpId() {
	ses.tpid = ses.tp.Next()
}

func (ses *Session) methodType(m string) MethodType {
	if v, ok := methodFromString[m]; ok {
		return v
	} else {
		return mUnkownMethod
	}
}

func (ses *Session) parseCmd(line string) error {
	var cmd []string

	if strings.TrimSpace(line) == "" {
		return errors.New("command parse error")
	}

	w := strings.Split(line, " ")
	for i := range w {
		if w[i] != "" {
			cmd = append(cmd, strings.TrimSpace(w[i]))
		}
	}

	ses.cmdMethod = ses.methodType(strings.ToUpper(cmd[0]))
	switch ses.cmdMethod {
	case mLogin, mLogout, mQuery:
		if ses.cmdIp = net.ParseIP(cmd[1]); ses.cmdIp == nil {
			return errors.New("command parse error")
		}
		ses.cmdArgs = strings.Join(cmd[2:], " ")
	case mQuit:
		//pp.Println("Quit")
		ses.cmdArgs = strings.Join(cmd[1:], " ")
	default:
		return errors.New("command not found")
	}
	return nil
}

func (ses *Session) readLine() (string, error) {
	ses.tp.StartRequest(ses.tpid)
	l1, err := ses.tp.ReadLine()
	if err != nil {
		return "", err
	}

	l2, err := ses.tp.ReadLine()
	if err != nil {
		return "", err
	}
	ses.tp.EndRequest(ses.tpid)

	if l1 != "" && l2 == "" {
		return l1, nil
	} else {
		return "", errors.New("session read error")
	}
}

func (ses *Session) sendLine(str string) error {
	var err error
	if ses.protocol == pTCP {
		ses.tp.StartResponse(ses.tpid)
		err = ses.tp.PrintfLine(str + CRLF)
		ses.tp.EndResponse(ses.tpid)
	} else {
		b := []byte(str + CRLF + CRLF)
		_, err = ses.udpconn.WriteToUDP(b, ses.remoteAddr)
	}
	return err
}

func (ses *Session) sendResponsePositive(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rPositive], str))
}

func (ses *Session) sendResponseNegative(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rNegative], str))
}

func (ses *Session) sendResponseBadRequest(str string) error {
	return ses.sendLine(fmt.Sprintf("%s%s", result[rBadRequest], str))
}

func (ses *Session) startHandler() bool {
	var err error

	if ses.protocol == pTCP {
		ses.setTpId()

		line, err := ses.readLine()
		if err != nil {
			if opError, ok := err.(*net.OpError); ok && opError.Timeout() {
				return false
			}
			if err.Error() == "EOF" {
				return false
			} else {
				ses.sendResponseBadRequest(err.Error())
				return true
			}
		}
		err = ses.parseCmd(line)
		if err != nil {
			ses.sendResponseBadRequest(err.Error())
			return true
		}
	} else {
		err = ses.parseCmd(string(ses.b.buf[:ses.b.count]))
		if err != nil {
			ses.sendResponseBadRequest(err.Error())
			return true
		}
	}

	switch ses.cmdMethod {
	case mLogin:
		ses.methodLogin()
	case mLogout:
		ses.methodLogout()
	case mQuery:
		ses.methodQuery()
	case mQuit:
		ses.methodQuit()
	default:
		err := errors.New("handler error")
		ses.sendResponseBadRequest(err.Error())
	}
	return true
}

func (ses *Session) methodLogin() {
	sd := &StoreData{
		Expire: time.Now().Add(StoreDataExpire),
		IP:     ses.cmdIp,
		Data:   ses.cmdArgs,
	}
	MainStore.Set(sd.Key(), sd)
	ses.sendResponsePositive("LOGIN OK")
}

func (ses *Session) methodLogout() {
	ok := MainStore.Del(ses.cmdIp.String())
	if ok {
		ses.sendResponsePositive("LOGOUT record deleted")
	} else {
		ses.sendResponsePositive("LOGOUT no such record, nothing done")
	}
}

func (ses *Session) methodQuery() {
	sd, err := MainStore.Get(ses.cmdIp.String())
	if err != nil {
		ses.sendResponseNegative("Not Logged in")
	} else {
		ses.sendResponsePositive(sd.Data)
	}
}

func (ses *Session) methodQuit() {
	ses.sendResponsePositive("QUIT OK")
	if ses.protocol == pTCP {
		ses.conn.Close()
	}
}

func (ses *Session) close() {
	if ses.protocol == pUDP {
		ses.b.Free()
	}
}
<|endoftext|>