<commit_before>// A simple package for xsd validation, uses libxml2.
//
// The rationale behind this package is to preload xsd files and use their in-memory structure to validate incoming xml documents in a concurrent environment, eg. the post bodys of xml service endpoints, and return useful error messages when appropriate.
//
// This package is part of a rewrite for an online service I'm currently working on. As there are limited resources on how to handle xsd validation in go this might be useful for somebody.
//
// libxml2-dev is needed, below an example how to install the latest sources (Ubuntu, change prefix according to where libs and include files are located):
//  curl -sL ftp://xmlsoft.org/libxml2/libxml2-2.9.5.tar.gz | tar -xzf -
//  cd ./libxml2-2.9.5/
//  ./configure --prefix=/usr  --enable-static --with-threads --with-history
//  make
//  sudo make install
package xsdvalidate

import "C"
import (
	"errors"
	"log"
)

// The type for parser/validation options.
type Options int16

// The parser options, ParsErrVerbose will slow down parsing considerably!
const (
	ParsErrDefault Options = 1 << iota // Default parser error output
	ParsErrVerbose                     // Verbose parser error output, considerably slower!
)

// Validation options for possible future enhancements.
const (
	ValidErrDefault Options = 1 << iota // Default validation error output
)

// Initializes the libxml2 parser, suggested for multithreading, see http://xmlsoft.org/threads.html.
func Init() {
	log.Printf("Initializing the libxml2 parser")
	libXml2Init()
}

// Cleans up the libxml2 parser, use this when application ends or libxml2 is not needed anymore.
func Cleanup() {
	log.Printf("Cleaning up the libxml2 parser")
	libXml2Cleanup()
}

// Initialize the xml handler struct.
// Always use the Free() method when done using this handler or memory will be leaking.
// The go garbage collector will not collect the allocated resources.
func NewXmlHandlerMem(inXml []byte, options Options) (*XmlHandler, error) {
	xPtr, err := parseXmlMem(inXml, options)
	return &XmlHandler{xPtr}, err
}

// Initialize the xml handler struct.
// Always use Free() method when done using this handler or memory will be leaking.
// The go garbage collector will not collect the allocated resources.
func NewXsdHandlerUrl(url string, options Options) (*XsdHandler, error) {
	sPtr, err := parseUrlSchema(url, options)
	return &XsdHandler{sPtr}, err
}

// The validation method validates an xmlHandler against an xsdHandler and returns the libxml2 validation error text.
// Both xmlHandler and xsdHandler have to be created first with the appropriate New... functions.
func (xsdHandler *XsdHandler) Validate(xmlHandler *XmlHandler, options Options) error {
	if xsdHandler == nil || xsdHandler.schemaPtr == nil {
		return errors.New("Xsd handler not properly initialized, use 'New...'")

	}
	if xsdHandler == nil || xmlHandler.docPtr == nil {
		return errors.New("Xml handler not properly initialized, use 'New...'")
	}
	return validateWithXsd(xmlHandler, xsdHandler)

}

// Frees the schemaPtr, call this when this handler is not needed anymore.
func (xsdHandler *XsdHandler) Free() {
	freeSchemaPtr(xsdHandler)
}

// Frees the xml docPtr, call this when this handler is not needed anymore.
func (xmlHandler *XmlHandler) Free() {
	freeDocPtr(xmlHandler)
}
<commit_msg>Typo fix<commit_after>// A simple package for xsd validation, uses libxml2.
//
// The rationale behind this package is to preload xsd files and use their in-memory structure to validate incoming xml documents in a concurrent environment, eg. the post bodys of xml service endpoints, and return useful error messages when appropriate.
//
// This package is part of a rewrite for an online service I'm currently working on. As there are limited resources on how to handle xsd validation in go this might be useful for somebody.
//
// libxml2-dev is needed, below an example how to install the latest sources (Ubuntu, change prefix according to where libs and include files are located):
//  curl -sL ftp://xmlsoft.org/libxml2/libxml2-2.9.5.tar.gz | tar -xzf -
//  cd ./libxml2-2.9.5/
//  ./configure --prefix=/usr  --enable-static --with-threads --with-history
//  make
//  sudo make install
package xsdvalidate

import "C"
import (
	"errors"
	"log"
)

// The type for parser/validation options.
type Options int16

// The parser options, ParsErrVerbose will slow down parsing considerably!
const (
	ParsErrDefault Options = 1 << iota // Default parser error output
	ParsErrVerbose                     // Verbose parser error output, considerably slower!
)

// Validation options for possible future enhancements.
const (
	ValidErrDefault Options = 1 << iota // Default validation error output
)

// Initializes the libxml2 parser, suggested for multithreading, see http://xmlsoft.org/threads.html.
func Init() {
	log.Printf("Initializing the libxml2 parser")
	libXml2Init()
}

// Cleans up the libxml2 parser, use this when application ends or libxml2 is not needed anymore.
func Cleanup() {
	log.Printf("Cleaning up the libxml2 parser")
	libXml2Cleanup()
}

// Initialize the xml handler struct.
// Always use the Free() method when done using this handler or memory will be leaking.
// The go garbage collector will not collect the allocated resources.
func NewXmlHandlerMem(inXml []byte, options Options) (*XmlHandler, error) {
	xPtr, err := parseXmlMem(inXml, options)
	return &XmlHandler{xPtr}, err
}

// Initialize the xml handler struct.
// Always use Free() method when done using this handler or memory will be leaking.
// The go garbage collector will not collect the allocated resources.
func NewXsdHandlerUrl(url string, options Options) (*XsdHandler, error) {
	sPtr, err := parseUrlSchema(url, options)
	return &XsdHandler{sPtr}, err
}

// The validation method validates an xmlHandler against an xsdHandler and returns the libxml2 validation error text.
// Both xmlHandler and xsdHandler have to be created first with the appropriate New... functions.
func (xsdHandler *XsdHandler) Validate(xmlHandler *XmlHandler, options Options) error {
	if xsdHandler == nil || xsdHandler.schemaPtr == nil {
		return errors.New("Xsd handler not properly initialized, use 'New...'")

	}
	if xmlHandler == nil || xmlHandler.docPtr == nil {
		return errors.New("Xml handler not properly initialized, use 'New...'")
	}
	return validateWithXsd(xmlHandler, xsdHandler)

}

// Frees the schemaPtr, call this when this handler is not needed anymore.
func (xsdHandler *XsdHandler) Free() {
	freeSchemaPtr(xsdHandler)
}

// Frees the xml docPtr, call this when this handler is not needed anymore.
func (xmlHandler *XmlHandler) Free() {
	freeDocPtr(xmlHandler)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"

	"github.com/valyala/fasthttp"

	openrtb "gopkg.in/bsm/openrtb.v2"
)

var openRTBVersion = "2.2"
var l = logger{}

func fastHandleAuctions(ctx *fasthttp.RequestCtx, agents []Agent) {
	var (
		ok    = true
		tmpOk = true
	)

	// enc := json.NewEncoder(w)
	// body, _ := ioutil.ReadAll(r.Body)
	// fmt.Println(string(body))
	var req *openrtb.BidRequest
	err := json.Unmarshal(ctx.PostBody(), &req)
	// req, err := openrtb.ParseRequest(r.Body)

	if err != nil {
		log.Println("ERROR", err.Error())
		ctx.SetStatusCode(fasthttp.StatusNoContent)
		return
	}

	// log.Println("INFO Received bid request", req.ID)

	ids := ExternalIdsFromRequest(req)
	res := EmptyResponseWithOneSeat(req)

	for _, agent := range agents {
		res, tmpOk = agent.DoBid(req, res, ids)
		ok = tmpOk || ok

		if tmpOk {
			agent.BidIncoming()
		}
	}

	if ok {
		ld := logData{AuctionData: req, BidData: res}
		go l.log(ld)

		ctx.Response.Header.Set("Content-type", "application/json")
		ctx.Response.Header.Set("x-openrtb-version", openRTBVersion)
		ctx.SetStatusCode(http.StatusOK)

		bytes, _ := json.Marshal(res)
		ctx.SetBody(bytes)

		print(string(bytes))

		return
	}
	log.Println("No bid.")
	ctx.SetStatusCode(204)
}

func errorMux(ctx *fasthttp.RequestCtx) {
	s := string(ctx.Request.Header.Header()[:])
	log.Println("string is", s)

	// for name, headers := range ctx.Request.Header.Header() {
	// 	name = strings.ToLower(name)
	// 	for _, h := range headers {
	// 		log.Println(name, h)
	// 		// 	request = append(request, fmt.Sprintf(“%v: %v”, name, h))
	// 	}
	// }

	ctx.SetStatusCode(http.StatusOK)
	BidEvent()
}

func winMux(ctx *fasthttp.RequestCtx) {
	// log.Println(ctx.PostBody())
	ctx.SetStatusCode(fasthttp.StatusOK)
	BidWin()
}

func eventMux(ctx *fasthttp.RequestCtx) {
	// var f interface{}

	// s := string(ctx.Request.Header.Header()[:])
	// s := string(ctx.Request.Body()[:])
	// log.Println("string is", s)

	// for name, headers := range ctx.Request.Header.Header() {
	// 	name = strings.ToLower(name)
	// 	for _, h := range headers {
	// 		log.Println(name, h)
	// 		// 	request = append(request, fmt.Sprintf(“%v: %v”, name, h))
	// 	}
	// }

	log.Println("Event!!!!!")
	ctx.SetStatusCode(http.StatusOK)
	BidEvent()
}
<commit_msg>commenting out logger temporarily.<commit_after>package main

import (
	"encoding/json"
	"log"
	"net/http"

	"github.com/valyala/fasthttp"

	openrtb "gopkg.in/bsm/openrtb.v2"
)

var openRTBVersion = "2.2"
var l = logger{}

func fastHandleAuctions(ctx *fasthttp.RequestCtx, agents []Agent) {
	var (
		ok    = true
		tmpOk = true
	)

	// enc := json.NewEncoder(w)
	// body, _ := ioutil.ReadAll(r.Body)
	// fmt.Println(string(body))
	var req *openrtb.BidRequest
	err := json.Unmarshal(ctx.PostBody(), &req)
	// req, err := openrtb.ParseRequest(r.Body)

	if err != nil {
		log.Println("ERROR", err.Error())
		ctx.SetStatusCode(fasthttp.StatusNoContent)
		return
	}

	// log.Println("INFO Received bid request", req.ID)

	ids := ExternalIdsFromRequest(req)
	res := EmptyResponseWithOneSeat(req)

	for _, agent := range agents {
		res, tmpOk = agent.DoBid(req, res, ids)
		ok = tmpOk || ok

		if tmpOk {
			agent.BidIncoming()
		}
	}

	if ok {
		// ld := logData{AuctionData: req, BidData: res}
		// go l.log(ld)

		ctx.Response.Header.Set("Content-type", "application/json")
		ctx.Response.Header.Set("x-openrtb-version", openRTBVersion)
		ctx.SetStatusCode(http.StatusOK)

		bytes, _ := json.Marshal(res)
		ctx.SetBody(bytes)

		print(string(bytes))

		return
	}
	log.Println("No bid.")
	ctx.SetStatusCode(204)
}

func errorMux(ctx *fasthttp.RequestCtx) {
	s := string(ctx.Request.Header.Header()[:])
	log.Println("string is", s)

	// for name, headers := range ctx.Request.Header.Header() {
	// 	name = strings.ToLower(name)
	// 	for _, h := range headers {
	// 		log.Println(name, h)
	// 		// 	request = append(request, fmt.Sprintf(“%v: %v”, name, h))
	// 	}
	// }

	ctx.SetStatusCode(http.StatusOK)
	BidEvent()
}

func winMux(ctx *fasthttp.RequestCtx) {
	// log.Println(ctx.PostBody())
	ctx.SetStatusCode(fasthttp.StatusOK)
	BidWin()
}

func eventMux(ctx *fasthttp.RequestCtx) {
	// var f interface{}

	// s := string(ctx.Request.Header.Header()[:])
	// s := string(ctx.Request.Body()[:])
	// log.Println("string is", s)

	// for name, headers := range ctx.Request.Header.Header() {
	// 	name = strings.ToLower(name)
	// 	for _, h := range headers {
	// 		log.Println(name, h)
	// 		// 	request = append(request, fmt.Sprintf(“%v: %v”, name, h))
	// 	}
	// }

	log.Println("Event!!!!!")
	ctx.SetStatusCode(http.StatusOK)
	BidEvent()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/chrissnell/syslog"
)

type handler struct {
	// To simplify implementation of our handler we embed helper
	// syslog.BaseHandler struct.
	*syslog.BaseHandler
}

type LogentriesHostEntity struct {
	Response  string `json:"response"`
	Host      Host   `json:"host"`
	Host_key  string `json:"host_key"`
	Worker    string `json:"worker"`
	Agent_key string `json:"agent_key"`
}

type Host struct {
	C        float64 `json:"c"`
	Name     string  `json:"name"`
	Distver  string  `json:"distver"`
	Hostname string  `json:"hostname"`
	Object   string  `json:"object"`
	Distname string  `json:"distname"`
	Key      string  `json:"key"`
}

type LogentriesLogEntity struct {
	Response string `json:"response"`
	Log_key  string `json:"log_key"`
	Log      Log    `json:"log"`
}

type Log struct {
	Token     string  `json:"token"`
	Created   float64 `json:"created"`
	Name      string  `json:"name`
	Retention float64 `json:"retention"`
	Filename  string  `json:"filename"`
	Object    string  `json:"object"`
	Type      string  `json:"type"`
	Key       string  `json:"key"`
	Follow    string  `json:"folow"`
}

type LogLine struct {
	Line  syslog.Message
	Token string
}

var (
	logconsumerPtr        *string
	logentriesAPIKeyPtr   *string
	listenAddrPtr         *string
	logentities           = make(map[string]LogentriesLogEntity)
	hostentities          = make(map[string]LogentriesHostEntity)
	tokenchan             = make(chan string)
	logentities_filename  = "logentries-logentities.gob"
	hostentities_filename = "logentries-hostentities.gob"
)

func newHandler() *handler {
	msg := make(chan syslog.Message)
	// Filter function name set to nil to disable filtering
	h := handler{syslog.NewBaseHandler(5, nil, false)}
	go h.mainLoop(msg)
	go ProcessLogMessage(msg)
	return &h
}

func (h *handler) mainLoop(msg chan syslog.Message) {
	for {
		m := h.Get()
		if m == nil {
			break
		}
		msg <- *m
	}
	fmt.Println("Exit handler")
	h.End()
}

func ProcessLogMessage(msg chan syslog.Message) {
	tokenfetchdone := make(chan bool, 1)
	logentrieschan := make(chan LogLine)
	lh := make(chan struct{ host, log string })

	var logline LogLine

	for m := range msg {
		if m.Hostname == "" {
			m.Hostname = "NONE"
		}
		go GetTokenForLog(tokenfetchdone, lh)
		lh <- struct{ host, log string }{m.Hostname, m.Tag}
		token := <-tokenchan
		<-tokenfetchdone

		logline.Token = token
		logline.Line = m

		go SendLogMessages(logentrieschan)
		logentrieschan <- logline

	}
}

func GetTokenForLog(tokenfetchdone chan bool, lh chan struct{ host, log string }) {
	select {
	case lht, msg_ok := <-lh:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {

			var hostentity LogentriesHostEntity
			var logentity LogentriesLogEntity

			l := strings.Join([]string{lht.host, lht.log}, "::")

			hostentity = hostentities[lht.host]
			if hostentity.Host.Key == "" {
				hostentity = RegisterNewHost(lht.host)

				// Store our new host token in our map and sync it to disk
				hostentities[lht.host] = hostentity
				err := SyncHostEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			}

			logentity = logentities[l]
			if logentity.Log.Token == "" {
				logentity := RegisterNewLog(hostentity, lht.log)
				logentities[l] = logentity
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
				err := SyncLogEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			} else {
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
			}
		}
	}
}

func DialLogEntries() (conn net.Conn, err error) {
	for {
		conn, err = net.Dial("tcp", *logconsumerPtr)
		if err == nil {
			return conn, err
		} else {
			fmt.Println("Could not connect to LogEntries log endpoint...retrying")
			// Wait for 5 seconds before redialing
			time.Sleep(5000 * time.Millisecond)
		}
	}
}

func SendLogMessages(msg chan LogLine) {
	conn, err := DialLogEntries()
	if err != nil {
		fmt.Println("Could not connect to LogEntries log endpoint ", err.Error())
	}

	select {
	case logline, msg_ok := <-msg:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {
			t := logline.Line.Time
			line := fmt.Sprintf("%v %v %v %v\n", logline.Token, t.Format(time.RFC3339), logline.Line.Hostname, logline.Line.Content)
			_, err = conn.Write([]byte(line))
			if err != nil {
				log.Print("Send to Logentries endpoint failed.")
				break
			}
			// fmt.Printf("Sending line: %v", line)
		}
	}
}

func SyncLogEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(logentities)
	err = ioutil.WriteFile("logentries-logentities.gob", m.Bytes(), 0600)
	return (err)
}

func SyncHostEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(hostentities)
	err = ioutil.WriteFile("logentries-hostentities.gob", m.Bytes(), 0600)
	return (err)
}

func LoadLogEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-logentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&logentities)
	return (err)
}

func LoadHostEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-hostentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&hostentities)
	return (err)
}

func RegisterNewHost(h string) (he LogentriesHostEntity) {
	v := url.Values{}
	v.Set("request", "register")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("name", h)
	v.Set("hostname", h)
	v.Set("distver", "")
	v.Set("system", "")
	v.Set("distname", "")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &he)
	return (he)
}

func RegisterNewLog(e LogentriesHostEntity, n string) (logentity LogentriesLogEntity) {
	v := url.Values{}
	v.Set("request", "new_log")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("host_key", e.Host.Key)
	v.Set("name", n)
	v.Set("filename", "")
	v.Set("retention", "-1")
	v.Set("source", "token")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &logentity)
	return (logentity)
}

func main() {

	logconsumerPtr = flag.String("consumer", "api.logentries.com:10000", "Logentries log consumer endpoint <host:port> (Default: api.logentries.com:10000)")
	logentriesAPIKeyPtr = flag.String("apikey", "", "Logentries API key")
	listenAddrPtr = flag.String("listen", "0.0.0.0:1987", "Host/port to listen for syslog messages <host:port> (Default: 0.0.0.0:1987)")

	flag.Parse()

	if *logentriesAPIKeyPtr == "" {
		log.Fatal("Must pass a Logentries API key. Use -h for help.")
	}

	if _, err := os.Stat(logentities_filename); err == nil {
		err = LoadLogEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	if _, err := os.Stat(hostentities_filename); err == nil {
		err = LoadHostEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	// Create a server with one handler and run one listen gorutine
	s := syslog.NewServer()
	s.AddAllowedRunes("-._")
	s.AddHandler(newHandler())
	s.Listen(*listenAddrPtr)

	// Wait for terminating signal
	sc := make(chan os.Signal, 2)
	signal.Notify(sc, syscall.SIGTERM, syscall.SIGINT)
	<-sc

	// Shutdown the server
	fmt.Println("Shutdown the server...")
	s.Shutdown()
	fmt.Println("Server is down")
}
<commit_msg>Add buffering to channels; add debugging<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/chrissnell/syslog"
)

type handler struct {
	// To simplify implementation of our handler we embed helper
	// syslog.BaseHandler struct.
	*syslog.BaseHandler
}

type LogentriesHostEntity struct {
	Response  string `json:"response"`
	Host      Host   `json:"host"`
	Host_key  string `json:"host_key"`
	Worker    string `json:"worker"`
	Agent_key string `json:"agent_key"`
}

type Host struct {
	C        float64 `json:"c"`
	Name     string  `json:"name"`
	Distver  string  `json:"distver"`
	Hostname string  `json:"hostname"`
	Object   string  `json:"object"`
	Distname string  `json:"distname"`
	Key      string  `json:"key"`
}

type LogentriesLogEntity struct {
	Response string `json:"response"`
	Log_key  string `json:"log_key"`
	Log      Log    `json:"log"`
}

type Log struct {
	Token     string  `json:"token"`
	Created   float64 `json:"created"`
	Name      string  `json:"name`
	Retention float64 `json:"retention"`
	Filename  string  `json:"filename"`
	Object    string  `json:"object"`
	Type      string  `json:"type"`
	Key       string  `json:"key"`
	Follow    string  `json:"folow"`
}

type LogLine struct {
	Line  syslog.Message
	Token string
}

var (
	logconsumerPtr        *string
	logentriesAPIKeyPtr   *string
	listenAddrPtr         *string
	logentities           = make(map[string]LogentriesLogEntity)
	hostentities          = make(map[string]LogentriesHostEntity)
	tokenchan             = make(chan string)
	logentities_filename  = "logentries-logentities.gob"
	hostentities_filename = "logentries-hostentities.gob"
)

func newHandler() *handler {
	msg := make(chan syslog.Message, 100)
	// Filter function name set to nil to disable filtering
	h := handler{syslog.NewBaseHandler(5, nil, false)}
	go h.mainLoop(msg)
	go ProcessLogMessage(msg)
	return &h
}

func (h *handler) mainLoop(msg chan syslog.Message) {
	for {
		m := h.Get()
		if m == nil {
			break
		}
		msg <- *m
	}
	fmt.Println("Exit handler")
	h.End()
}

func ProcessLogMessage(msg chan syslog.Message) {
	tokenfetchdone := make(chan bool, 1)
	logentrieschan := make(chan LogLine, 100)
	lh := make(chan struct{ host, log string }, 100)

	var logline LogLine

	for m := range msg {
		if m.Hostname == "" {
			m.Hostname = "NONE"
		}

		// for debugging
		// if m.Hostname == "testpiper001" {
		//	log.Printf("recv from %v: %v\n", m.Hostname, m.Content)
		// }

		go GetTokenForLog(tokenfetchdone, lh)
		lh <- struct{ host, log string }{m.Hostname, m.Tag}
		token := <-tokenchan
		<-tokenfetchdone

		logline.Token = token
		logline.Line = m

		go SendLogMessages(logentrieschan)
		logentrieschan <- logline

	}
}

func GetTokenForLog(tokenfetchdone chan bool, lh chan struct{ host, log string }) {
	select {
	case lht, msg_ok := <-lh:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {

			var hostentity LogentriesHostEntity
			var logentity LogentriesLogEntity

			l := strings.Join([]string{lht.host, lht.log}, "::")

			hostentity = hostentities[lht.host]
			if hostentity.Host.Key == "" {
				log.Printf("Registering host entity: %v\n", lht.host)
				hostentity = RegisterNewHost(lht.host)

				// Store our new host token in our map and sync it to disk
				hostentities[lht.host] = hostentity
				err := SyncHostEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			}

			logentity = logentities[l]
			if logentity.Log.Token == "" {
				log.Printf("Registering new log entity [%v]: %v\n", lht.host, lht.log)
				logentity := RegisterNewLog(hostentity, lht.log)
				logentities[l] = logentity
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
				err := SyncLogEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			} else {
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
			}
		}
	}
}

func DialLogEntries() (conn net.Conn, err error) {
	for {
		conn, err = net.Dial("tcp", *logconsumerPtr)
		if err == nil {
			return conn, err
		} else {
			fmt.Println("Could not connect to LogEntries log endpoint...retrying")
			// Wait for 5 seconds before redialing
			time.Sleep(5000 * time.Millisecond)
		}
	}
}

func SendLogMessages(msg chan LogLine) {
	conn, err := DialLogEntries()
	if err != nil {
		fmt.Println("Could not connect to LogEntries log endpoint ", err.Error())
	}

	select {
	case logline, msg_ok := <-msg:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {
			// For debugging...
			// if logline.Line.Hostname == "testpiper001" {
			//	log.Printf("send to LE from %v: %v\n", logline.Line.Hostname, logline.Line.Content)
			// }
			t := logline.Line.Time
			line := fmt.Sprintf("%v %v %v %v\n", logline.Token, t.Format(time.RFC3339), logline.Line.Hostname, logline.Line.Content)
			_, err = conn.Write([]byte(line))
			if err != nil {
				log.Print("Send to Logentries endpoint failed.")
				break
			}
			// fmt.Printf("Sending line: %v", line)
		}
	}
}

func SyncLogEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(logentities)
	err = ioutil.WriteFile("logentries-logentities.gob", m.Bytes(), 0600)
	return (err)
}

func SyncHostEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(hostentities)
	err = ioutil.WriteFile("logentries-hostentities.gob", m.Bytes(), 0600)
	return (err)
}

func LoadLogEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-logentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&logentities)
	return (err)
}

func LoadHostEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-hostentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&hostentities)
	return (err)
}

func RegisterNewHost(h string) (he LogentriesHostEntity) {
	v := url.Values{}
	v.Set("request", "register")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("name", h)
	v.Set("hostname", h)
	v.Set("distver", "")
	v.Set("system", "")
	v.Set("distname", "")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &he)
	return (he)
}

func RegisterNewLog(e LogentriesHostEntity, n string) (logentity LogentriesLogEntity) {
	v := url.Values{}
	v.Set("request", "new_log")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("host_key", e.Host.Key)
	v.Set("name", n)
	v.Set("filename", "")
	v.Set("retention", "-1")
	v.Set("source", "token")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &logentity)
	return (logentity)
}

func main() {

	logconsumerPtr = flag.String("consumer", "api.logentries.com:10000", "Logentries log consumer endpoint <host:port> (Default: api.logentries.com:10000)")
	logentriesAPIKeyPtr = flag.String("apikey", "", "Logentries API key")
	listenAddrPtr = flag.String("listen", "0.0.0.0:1987", "Host/port to listen for syslog messages <host:port> (Default: 0.0.0.0:1987)")

	flag.Parse()

	if *logentriesAPIKeyPtr == "" {
		log.Fatal("Must pass a Logentries API key. Use -h for help.")
	}

	if _, err := os.Stat(logentities_filename); err == nil {
		err = LoadLogEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	if _, err := os.Stat(hostentities_filename); err == nil {
		err = LoadHostEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	// Create a server with one handler and run one listen gorutine
	s := syslog.NewServer()
	s.AddAllowedRunes("-._")
	s.AddHandler(newHandler())
	s.Listen(*listenAddrPtr)

	// Wait for terminating signal
	sc := make(chan os.Signal, 2)
	signal.Notify(sc, syscall.SIGTERM, syscall.SIGINT)
	<-sc

	// Shutdown the server
	fmt.Println("Shutdown the server...")
	s.Shutdown()
	fmt.Println("Server is down")
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

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

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

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

package ingress

import (
	"bytes"
	"io"
	nethttp "net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/cloudevents/sdk-go/v2/client"
	"github.com/cloudevents/sdk-go/v2/event"
	cehttp "github.com/cloudevents/sdk-go/v2/protocol/http"
	"github.com/google/go-cmp/cmp"
	"go.uber.org/zap"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"

	"knative.dev/eventing/pkg/apis/eventing"
	eventingv1 "knative.dev/eventing/pkg/apis/eventing/v1"
	broker "knative.dev/eventing/pkg/broker"
	"knative.dev/eventing/pkg/kncloudevents"
	reconcilertestingv1 "knative.dev/eventing/pkg/reconciler/testing/v1"
)

const (
	senderResponseStatusCode = nethttp.StatusAccepted
)

func TestHandler_ServeHTTP(t *testing.T) {
	t.Parallel()

	logger := zap.NewNop()

	tt := []struct {
		name            string
		method          string
		uri             string
		body            io.Reader
		headers         nethttp.Header
		expectedHeaders nethttp.Header
		statusCode      int
		handler         nethttp.Handler
		reporter        StatsReporter
		defaulter       client.EventDefaulter
		brokers         []*eventingv1.Broker
	}{
		{
			name:       "invalid method PATCH",
			method:     nethttp.MethodPatch,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method PUT",
			method:     nethttp.MethodPut,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method DELETE",
			method:     nethttp.MethodDelete,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method GET",
			method:     nethttp.MethodGet,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method OPTIONS",
			method:     nethttp.MethodOptions,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "valid (happy path)",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: senderResponseStatusCode,
			handler:    handler(),
			reporter:   &mockReporter{StatusCode: senderResponseStatusCode, EventDispatchTimeReported: true},
			defaulter:  broker.TTLDefaulter(logger, 100),
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "no TTL drop event",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusBadRequest,
			handler:    handler(),
			reporter:   &mockReporter{StatusCode: nethttp.StatusBadRequest},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "malformed request URI",
			method:     nethttp.MethodPost,
			uri:        "/knative/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusBadRequest,
			handler:    handler(),
			reporter:   &mockReporter{},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "root request URI",
			method:     nethttp.MethodPost,
			uri:        "/",
			body:       getValidEvent(),
			statusCode: nethttp.StatusNotFound,
			handler:    handler(),
			reporter:   &mockReporter{},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "pass headers to handler",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: senderResponseStatusCode,
			headers: nethttp.Header{
				"foo":              []string{"bar"},
				"Traceparent":      []string{"0"},
				"Knative-Foo":      []string{"123"},
				"X-Request-Id":     []string{"123"},
				cehttp.ContentType: []string{event.ApplicationCloudEventsJSON},
			},
			handler: &svc{},
			expectedHeaders: nethttp.Header{
				"Knative-Foo":  []string{"123"},
				"X-Request-Id": []string{"123"},
			},
			reporter:  &mockReporter{StatusCode: senderResponseStatusCode, EventDispatchTimeReported: true},
			defaulter: broker.TTLDefaulter(logger, 100),
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {

			s := httptest.NewServer(tc.handler)
			defer s.Close()

			recorder := httptest.NewRecorder()
			request := httptest.NewRequest(tc.method, tc.uri, tc.body)
			if tc.headers != nil {
				request.Header = tc.headers
			} else {
				tc.expectedHeaders = nethttp.Header{
					cehttp.ContentType: []string{event.ApplicationCloudEventsJSON},
				}
				request.Header.Add(cehttp.ContentType, event.ApplicationCloudEventsJSON)
			}

			annotatedBrokers := make([]runtime.Object, 0, len(tc.brokers))
			for _, b := range tc.brokers {
				// Write the channel address in the broker status annotation
				if b.Status.Annotations == nil {
					b.Status.Annotations = make(map[string]string, 1)
				}
				b.Status.Annotations[eventing.BrokerChannelAddressStatusAnnotationKey] = s.URL
				annotatedBrokers = append(annotatedBrokers, b)
			}
			listers := reconcilertestingv1.NewListers(annotatedBrokers)
			sender, _ := kncloudevents.NewHTTPMessageSenderWithTarget("")
			h := &Handler{
				Sender:       sender,
				Defaulter:    tc.defaulter,
				Reporter:     &mockReporter{},
				Logger:       logger,
				BrokerLister: listers.GetBrokerLister(),
			}

			h.ServeHTTP(recorder, request)

			result := recorder.Result()
			if result.StatusCode != tc.statusCode {
				t.Errorf("expected status code %d got %d", tc.statusCode, result.StatusCode)
			}

			if svc, ok := tc.handler.(*svc); ok {
				for k, expValue := range tc.expectedHeaders {
					if v, ok := svc.receivedHeaders[k]; !ok {
						t.Errorf("expected header %s - %v", k, svc.receivedHeaders)
					} else if diff := cmp.Diff(expValue, v); diff != "" {
						t.Error("(-want +got)", diff)
					}
				}
			}

			if diff := cmp.Diff(tc.reporter, h.Reporter); diff != "" {
				t.Errorf("expected reporter state %+v got %+v - diff %s", tc.reporter, h.Reporter, diff)
			}
		})
	}
}

type svc struct {
	receivedHeaders nethttp.Header
}

func (s *svc) ServeHTTP(w nethttp.ResponseWriter, req *nethttp.Request) {
	s.receivedHeaders = req.Header
	w.WriteHeader(senderResponseStatusCode)
}

func handler() nethttp.Handler {
	return nethttp.HandlerFunc(func(writer nethttp.ResponseWriter, request *nethttp.Request) {
		writer.WriteHeader(senderResponseStatusCode)
	})
}

type mockReporter struct {
	StatusCode                int
	EventDispatchTimeReported bool
}

func (r *mockReporter) ReportEventCount(_ *ReportArgs, responseCode int) error {
	r.StatusCode = responseCode
	return nil
}

func (r *mockReporter) ReportEventDispatchTime(_ *ReportArgs, _ int, _ time.Duration) error {
	r.EventDispatchTimeReported = true
	return nil
}

func getValidEvent() io.Reader {
	e := event.New()
	e.SetType("type")
	e.SetSource("source")
	e.SetID("1234")
	b, _ := e.MarshalJSON()
	return bytes.NewBuffer(b)
}

func makeBroker(name, namespace string) *eventingv1.Broker {
	return &eventingv1.Broker{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "eventing.knative.dev/v1",
			Kind:       "Broker",
		},
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      name,
		},
		Spec: eventingv1.BrokerSpec{},
	}
}
<commit_msg>Add tests for additional ingress error cases (#5281)<commit_after>/*
Copyright 2019 The Knative Authors

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

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

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

package ingress

import (
	"bytes"
	"io"
	nethttp "net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/cloudevents/sdk-go/v2/client"
	"github.com/cloudevents/sdk-go/v2/event"
	cehttp "github.com/cloudevents/sdk-go/v2/protocol/http"
	"github.com/google/go-cmp/cmp"
	"go.uber.org/zap"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"

	"knative.dev/eventing/pkg/apis/eventing"
	eventingv1 "knative.dev/eventing/pkg/apis/eventing/v1"
	broker "knative.dev/eventing/pkg/broker"
	"knative.dev/eventing/pkg/kncloudevents"
	reconcilertestingv1 "knative.dev/eventing/pkg/reconciler/testing/v1"
	duckv1 "knative.dev/pkg/apis/duck/v1"
)

const (
	senderResponseStatusCode = nethttp.StatusAccepted
)

func TestHandler_ServeHTTP(t *testing.T) {
	t.Parallel()

	logger := zap.NewNop()

	tt := []struct {
		name            string
		method          string
		uri             string
		body            io.Reader
		headers         nethttp.Header
		expectedHeaders nethttp.Header
		statusCode      int
		handler         nethttp.Handler
		reporter        StatsReporter
		defaulter       client.EventDefaulter
		brokers         []*eventingv1.Broker
	}{
		{
			name:       "invalid method PATCH",
			method:     nethttp.MethodPatch,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method PUT",
			method:     nethttp.MethodPut,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method DELETE",
			method:     nethttp.MethodDelete,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method GET",
			method:     nethttp.MethodGet,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "invalid method OPTIONS",
			method:     nethttp.MethodOptions,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusMethodNotAllowed,
			handler:    handler(),
			reporter:   &mockReporter{},
			defaulter:  broker.TTLDefaulter(logger, 100),
		},
		{
			name:       "valid (happy path)",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: senderResponseStatusCode,
			handler:    handler(),
			reporter:   &mockReporter{StatusCode: senderResponseStatusCode, EventDispatchTimeReported: true},
			defaulter:  broker.TTLDefaulter(logger, 100),
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "no TTL drop event",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusBadRequest,
			handler:    handler(),
			reporter:   &mockReporter{StatusCode: nethttp.StatusBadRequest},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "malformed request URI",
			method:     nethttp.MethodPost,
			uri:        "/knative/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusBadRequest,
			handler:    handler(),
			reporter:   &mockReporter{},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "malformed event",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       strings.NewReader("not an event"),
			statusCode: nethttp.StatusBadRequest,
			handler:    handler(),
			reporter:   &mockReporter{},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "no broker annotations",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: nethttp.StatusInternalServerError,
			handler:    handler(),
			reporter:   &mockReporter{StatusCode: nethttp.StatusInternalServerError, EventDispatchTimeReported: true},
			defaulter:  broker.TTLDefaulter(logger, 100),
			brokers: []*eventingv1.Broker{
				withUninitializedAnnotations(makeBroker("name", "ns")),
			},
		},
		{
			name:       "root request URI",
			method:     nethttp.MethodPost,
			uri:        "/",
			body:       getValidEvent(),
			statusCode: nethttp.StatusNotFound,
			handler:    handler(),
			reporter:   &mockReporter{},
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
		{
			name:       "pass headers to handler",
			method:     nethttp.MethodPost,
			uri:        "/ns/name",
			body:       getValidEvent(),
			statusCode: senderResponseStatusCode,
			headers: nethttp.Header{
				"foo":              []string{"bar"},
				"Traceparent":      []string{"0"},
				"Knative-Foo":      []string{"123"},
				"X-Request-Id":     []string{"123"},
				cehttp.ContentType: []string{event.ApplicationCloudEventsJSON},
			},
			handler: &svc{},
			expectedHeaders: nethttp.Header{
				"Knative-Foo":  []string{"123"},
				"X-Request-Id": []string{"123"},
			},
			reporter:  &mockReporter{StatusCode: senderResponseStatusCode, EventDispatchTimeReported: true},
			defaulter: broker.TTLDefaulter(logger, 100),
			brokers: []*eventingv1.Broker{
				makeBroker("name", "ns"),
			},
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {

			s := httptest.NewServer(tc.handler)
			defer s.Close()

			recorder := httptest.NewRecorder()
			request := httptest.NewRequest(tc.method, tc.uri, tc.body)
			if tc.headers != nil {
				request.Header = tc.headers
			} else {
				tc.expectedHeaders = nethttp.Header{
					cehttp.ContentType: []string{event.ApplicationCloudEventsJSON},
				}
				request.Header.Add(cehttp.ContentType, event.ApplicationCloudEventsJSON)
			}

			brokers := make([]runtime.Object, 0, len(tc.brokers))
			for _, b := range tc.brokers {
				// Write the channel address in the broker status annotation unless explicitly set to nil
				if b.Status.Annotations != nil {
					if _, set := b.Status.Annotations[eventing.BrokerChannelAddressStatusAnnotationKey]; !set {
						b.Status.Annotations = map[string]string{
							eventing.BrokerChannelAddressStatusAnnotationKey: s.URL,
						}
					}
				}
				brokers = append(brokers, b)
			}
			listers := reconcilertestingv1.NewListers(brokers)
			sender, _ := kncloudevents.NewHTTPMessageSenderWithTarget("")
			h := &Handler{
				Sender:       sender,
				Defaulter:    tc.defaulter,
				Reporter:     &mockReporter{},
				Logger:       logger,
				BrokerLister: listers.GetBrokerLister(),
			}

			h.ServeHTTP(recorder, request)

			result := recorder.Result()
			if result.StatusCode != tc.statusCode {
				t.Errorf("expected status code %d got %d", tc.statusCode, result.StatusCode)
			}

			if svc, ok := tc.handler.(*svc); ok {
				for k, expValue := range tc.expectedHeaders {
					if v, ok := svc.receivedHeaders[k]; !ok {
						t.Errorf("expected header %s - %v", k, svc.receivedHeaders)
					} else if diff := cmp.Diff(expValue, v); diff != "" {
						t.Error("(-want +got)", diff)
					}
				}
			}

			if diff := cmp.Diff(tc.reporter, h.Reporter); diff != "" {
				t.Errorf("expected reporter state %+v got %+v - diff %s", tc.reporter, h.Reporter, diff)
			}
		})
	}
}

type svc struct {
	receivedHeaders nethttp.Header
}

func (s *svc) ServeHTTP(w nethttp.ResponseWriter, req *nethttp.Request) {
	s.receivedHeaders = req.Header
	w.WriteHeader(senderResponseStatusCode)
}

func handler() nethttp.Handler {
	return nethttp.HandlerFunc(func(writer nethttp.ResponseWriter, request *nethttp.Request) {
		writer.WriteHeader(senderResponseStatusCode)
	})
}

type mockReporter struct {
	StatusCode                int
	EventDispatchTimeReported bool
}

func (r *mockReporter) ReportEventCount(_ *ReportArgs, responseCode int) error {
	r.StatusCode = responseCode
	return nil
}

func (r *mockReporter) ReportEventDispatchTime(_ *ReportArgs, _ int, _ time.Duration) error {
	r.EventDispatchTimeReported = true
	return nil
}

func getValidEvent() io.Reader {
	e := event.New()
	e.SetType("type")
	e.SetSource("source")
	e.SetID("1234")
	b, _ := e.MarshalJSON()
	return bytes.NewBuffer(b)
}

func makeBroker(name, namespace string) *eventingv1.Broker {
	return &eventingv1.Broker{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "eventing.knative.dev/v1",
			Kind:       "Broker",
		},
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      name,
		},
		Spec: eventingv1.BrokerSpec{},
		Status: eventingv1.BrokerStatus{
			Status: duckv1.Status{
				Annotations: map[string]string{},
			},
		},
	}
}

func withUninitializedAnnotations(b *eventingv1.Broker) *eventingv1.Broker {
	b.Status.Annotations = nil
	return b
}
<|endoftext|>


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

package portforward

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"reflect"
	"sync"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client"
	"k8s.io/kubernetes/pkg/util/httpstream"
)

func TestParsePortsAndNew(t *testing.T) {
	tests := []struct {
		input            []string
		expected         []ForwardedPort
		expectParseError bool
		expectNewError   bool
	}{
		{input: []string{}, expectNewError: true},
		{input: []string{"a"}, expectParseError: true, expectNewError: true},
		{input: []string{":a"}, expectParseError: true, expectNewError: true},
		{input: []string{"-1"}, expectParseError: true, expectNewError: true},
		{input: []string{"65536"}, expectParseError: true, expectNewError: true},
		{input: []string{"0"}, expectParseError: true, expectNewError: true},
		{input: []string{"0:0"}, expectParseError: true, expectNewError: true},
		{input: []string{"a:5000"}, expectParseError: true, expectNewError: true},
		{input: []string{"5000:a"}, expectParseError: true, expectNewError: true},
		{
			input: []string{"5000", "5000:5000", "8888:5000", "5000:8888", ":5000", "0:5000"},
			expected: []ForwardedPort{
				{5000, 5000},
				{5000, 5000},
				{8888, 5000},
				{5000, 8888},
				{0, 5000},
				{0, 5000},
			},
		},
	}

	for i, test := range tests {
		parsed, err := parsePorts(test.input)
		haveError := err != nil
		if e, a := test.expectParseError, haveError; e != a {
			t.Fatalf("%d: parsePorts: error expected=%t, got %t: %s", i, e, a, err)
		}

		expectedRequest := &client.Request{}
		expectedConfig := &client.Config{}
		expectedStopChan := make(chan struct{})
		pf, err := New(expectedRequest, expectedConfig, test.input, expectedStopChan)
		haveError = err != nil
		if e, a := test.expectNewError, haveError; e != a {
			t.Fatalf("%d: New: error expected=%t, got %t: %s", i, e, a, err)
		}

		if test.expectParseError || test.expectNewError {
			continue
		}

		for pi, expectedPort := range test.expected {
			if e, a := expectedPort.Local, parsed[pi].Local; e != a {
				t.Fatalf("%d: local expected: %d, got: %d", i, e, a)
			}
			if e, a := expectedPort.Remote, parsed[pi].Remote; e != a {
				t.Fatalf("%d: remote expected: %d, got: %d", i, e, a)
			}
		}

		if e, a := expectedRequest, pf.req; e != a {
			t.Fatalf("%d: req: expected %#v, got %#v", i, e, a)
		}
		if e, a := expectedConfig, pf.config; e != a {
			t.Fatalf("%d: config: expected %#v, got %#v", i, e, a)
		}
		if e, a := test.expected, pf.ports; !reflect.DeepEqual(e, a) {
			t.Fatalf("%d: ports: expected %#v, got %#v", i, e, a)
		}
		if e, a := expectedStopChan, pf.stopChan; e != a {
			t.Fatalf("%d: stopChan: expected %#v, got %#v", i, e, a)
		}
		if pf.Ready == nil {
			t.Fatalf("%d: Ready should be non-nil", i)
		}
	}
}

type fakeUpgrader struct {
	conn *fakeUpgradeConnection
	err  error
}

func (u *fakeUpgrader) upgrade(req *client.Request, config *client.Config) (httpstream.Connection, error) {
	return u.conn, u.err
}

type fakeUpgradeConnection struct {
	closeCalled bool
	lock        sync.Mutex
	streams     map[string]*fakeUpgradeStream
	portData    map[string]string
}

func newFakeUpgradeConnection() *fakeUpgradeConnection {
	return &fakeUpgradeConnection{
		streams:  make(map[string]*fakeUpgradeStream),
		portData: make(map[string]string),
	}
}

func (c *fakeUpgradeConnection) CreateStream(headers http.Header) (httpstream.Stream, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	stream := &fakeUpgradeStream{}
	c.streams[headers.Get(api.PortHeader)] = stream
	// only simulate data on the data stream for now, not the error stream
	if headers.Get(api.StreamType) == api.StreamTypeData {
		stream.data = c.portData[headers.Get(api.PortHeader)]
	}

	return stream, nil
}

func (c *fakeUpgradeConnection) Close() error {
	c.lock.Lock()
	defer c.lock.Unlock()

	c.closeCalled = true
	return nil
}

func (c *fakeUpgradeConnection) CloseChan() <-chan bool {
	return make(chan bool)
}

func (c *fakeUpgradeConnection) SetIdleTimeout(timeout time.Duration) {
}

type fakeUpgradeStream struct {
	readCalled  bool
	writeCalled bool
	dataWritten []byte
	closeCalled bool
	resetCalled bool
	data        string
	lock        sync.Mutex
}

func (s *fakeUpgradeStream) Read(p []byte) (int, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.readCalled = true
	b := []byte(s.data)
	n := copy(p, b)
	return n, io.EOF
}

func (s *fakeUpgradeStream) Write(p []byte) (int, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.writeCalled = true
	s.dataWritten = make([]byte, len(p))
	copy(s.dataWritten, p)
	return len(p), io.EOF
}

func (s *fakeUpgradeStream) Close() error {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.closeCalled = true
	return nil
}

func (s *fakeUpgradeStream) Reset() error {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.resetCalled = true
	return nil
}

func (s *fakeUpgradeStream) Headers() http.Header {
	s.lock.Lock()
	defer s.lock.Unlock()
	return http.Header{}
}

type GetListenerTestCase struct {
	Hostname                string
	Protocol                string
	ShouldRaiseError        bool
	ExpectedListenerAddress string
}

func TestGetListener(t *testing.T) {
	var pf PortForwarder
	testCases := []GetListenerTestCase{
		{
			Hostname:                "localhost",
			Protocol:                "tcp4",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "127.0.0.1",
		},
		{
			Hostname:                "127.0.0.1",
			Protocol:                "tcp4",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "127.0.0.1",
		},
		{
			Hostname:                "[::1]",
			Protocol:                "tcp6",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "::1",
		},
		{
			Hostname:         "[::1]",
			Protocol:         "tcp4",
			ShouldRaiseError: true,
		},
		{
			Hostname:         "127.0.0.1",
			Protocol:         "tcp6",
			ShouldRaiseError: true,
		},
		{
			// IPv6 address must be put into brackets. This test reveals this.
			Hostname:         "::1",
			Protocol:         "tcp6",
			ShouldRaiseError: true,
		},
	}

	for i, testCase := range testCases {
		expectedListenerPort := "12345"
		listener, err := pf.getListener(testCase.Protocol, testCase.Hostname, &ForwardedPort{12345, 12345})
		errorRaised := err != nil

		if testCase.ShouldRaiseError != errorRaised {
			t.Errorf("Test case #%d failed: Data %v an error has been raised(%t) where it should not (or reciprocally): %v", i, testCase, testCase.ShouldRaiseError, err)
			continue
		}
		if errorRaised {
			continue
		}

		if listener == nil {
			t.Errorf("Test case #%d did not raise an error but failed in initializing listener", i)
			continue
		}

		host, port, _ := net.SplitHostPort(listener.Addr().String())
		t.Logf("Asked a %s forward for: %s:%v, got listener %s:%s, expected: %s", testCase.Protocol, testCase.Hostname, 12345, host, port, expectedListenerPort)
		if host != testCase.ExpectedListenerAddress {
			t.Errorf("Test case #%d failed: Listener does not listen on exepected address: asked %v got %v", i, testCase.ExpectedListenerAddress, host)
		}
		if port != expectedListenerPort {
			t.Errorf("Test case #%d failed: Listener does not listen on exepected port: asked %v got %v", i, expectedListenerPort, port)

		}
		listener.Close()

	}
}

func TestForwardPorts(t *testing.T) {
	testCases := []struct {
		Upgrader *fakeUpgrader
		Ports    []string
		Send     map[uint16]string
		Receive  map[uint16]string
		Err      bool
	}{
		{
			Upgrader: &fakeUpgrader{err: errors.New("bail")},
			Err:      true,
		},
		{
			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
			Ports:    []string{"5000"},
		},
		{
			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
			Ports:    []string{"5000", "6000"},
			Send: map[uint16]string{
				5000: "abcd",
				6000: "ghij",
			},
			Receive: map[uint16]string{
				5000: "1234",
				6000: "5678",
			},
		},
	}

	for i, testCase := range testCases {
		stopChan := make(chan struct{}, 1)

		pf, err := New(&client.Request{}, &client.Config{}, testCase.Ports, stopChan)
		hasErr := err != nil
		if hasErr != testCase.Err {
			t.Fatalf("%d: New: expected %t, got %t: %v", i, testCase.Err, hasErr, err)
		}
		if pf == nil {
			continue
		}
		pf.upgrader = testCase.Upgrader
		if testCase.Upgrader.err != nil {
			err := pf.ForwardPorts()
			hasErr := err != nil
			if hasErr != testCase.Err {
				t.Fatalf("%d: ForwardPorts: expected %t, got %t: %v", i, testCase.Err, hasErr, err)
			}
			continue
		}

		doneChan := make(chan error)
		go func() {
			doneChan <- pf.ForwardPorts()
		}()
		<-pf.Ready

		conn := testCase.Upgrader.conn

		for port, data := range testCase.Send {
			conn.lock.Lock()
			conn.portData[fmt.Sprintf("%d", port)] = testCase.Receive[port]
			conn.lock.Unlock()

			clientConn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", port))
			if err != nil {
				t.Fatalf("%d: error dialing %d: %s", i, port, err)
			}
			defer clientConn.Close()

			n, err := clientConn.Write([]byte(data))
			if err != nil && err != io.EOF {
				t.Fatalf("%d: Error sending data '%s': %s", i, data, err)
			}
			if n == 0 {
				t.Fatalf("%d: unexpected write of 0 bytes", i)
			}
			b := make([]byte, 4)
			n, err = clientConn.Read(b)
			if err != nil && err != io.EOF {
				t.Fatalf("%d: Error reading data: %s", i, err)
			}
			if !bytes.Equal([]byte(testCase.Receive[port]), b) {
				t.Fatalf("%d: expected to read '%s', got '%s'", i, testCase.Receive[port], b)
			}
		}

		// tell r.ForwardPorts to stop
		close(stopChan)

		// wait for r.ForwardPorts to actually return
		err = <-doneChan
		if err != nil {
			t.Fatalf("%d: unexpected error: %s", i, err)
		}

		if e, a := len(testCase.Send), len(conn.streams); e != a {
			t.Fatalf("%d: expected %d streams to be created, got %d", i, e, a)
		}

		if !conn.closeCalled {
			t.Fatalf("%d: expected conn closure", i)
		}
	}

}
<commit_msg>Unit test fails when ipv6 is not available<commit_after>/*
Copyright 2015 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.
*/

package portforward

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client"
	"k8s.io/kubernetes/pkg/util/httpstream"
)

func TestParsePortsAndNew(t *testing.T) {
	tests := []struct {
		input            []string
		expected         []ForwardedPort
		expectParseError bool
		expectNewError   bool
	}{
		{input: []string{}, expectNewError: true},
		{input: []string{"a"}, expectParseError: true, expectNewError: true},
		{input: []string{":a"}, expectParseError: true, expectNewError: true},
		{input: []string{"-1"}, expectParseError: true, expectNewError: true},
		{input: []string{"65536"}, expectParseError: true, expectNewError: true},
		{input: []string{"0"}, expectParseError: true, expectNewError: true},
		{input: []string{"0:0"}, expectParseError: true, expectNewError: true},
		{input: []string{"a:5000"}, expectParseError: true, expectNewError: true},
		{input: []string{"5000:a"}, expectParseError: true, expectNewError: true},
		{
			input: []string{"5000", "5000:5000", "8888:5000", "5000:8888", ":5000", "0:5000"},
			expected: []ForwardedPort{
				{5000, 5000},
				{5000, 5000},
				{8888, 5000},
				{5000, 8888},
				{0, 5000},
				{0, 5000},
			},
		},
	}

	for i, test := range tests {
		parsed, err := parsePorts(test.input)
		haveError := err != nil
		if e, a := test.expectParseError, haveError; e != a {
			t.Fatalf("%d: parsePorts: error expected=%t, got %t: %s", i, e, a, err)
		}

		expectedRequest := &client.Request{}
		expectedConfig := &client.Config{}
		expectedStopChan := make(chan struct{})
		pf, err := New(expectedRequest, expectedConfig, test.input, expectedStopChan)
		haveError = err != nil
		if e, a := test.expectNewError, haveError; e != a {
			t.Fatalf("%d: New: error expected=%t, got %t: %s", i, e, a, err)
		}

		if test.expectParseError || test.expectNewError {
			continue
		}

		for pi, expectedPort := range test.expected {
			if e, a := expectedPort.Local, parsed[pi].Local; e != a {
				t.Fatalf("%d: local expected: %d, got: %d", i, e, a)
			}
			if e, a := expectedPort.Remote, parsed[pi].Remote; e != a {
				t.Fatalf("%d: remote expected: %d, got: %d", i, e, a)
			}
		}

		if e, a := expectedRequest, pf.req; e != a {
			t.Fatalf("%d: req: expected %#v, got %#v", i, e, a)
		}
		if e, a := expectedConfig, pf.config; e != a {
			t.Fatalf("%d: config: expected %#v, got %#v", i, e, a)
		}
		if e, a := test.expected, pf.ports; !reflect.DeepEqual(e, a) {
			t.Fatalf("%d: ports: expected %#v, got %#v", i, e, a)
		}
		if e, a := expectedStopChan, pf.stopChan; e != a {
			t.Fatalf("%d: stopChan: expected %#v, got %#v", i, e, a)
		}
		if pf.Ready == nil {
			t.Fatalf("%d: Ready should be non-nil", i)
		}
	}
}

type fakeUpgrader struct {
	conn *fakeUpgradeConnection
	err  error
}

func (u *fakeUpgrader) upgrade(req *client.Request, config *client.Config) (httpstream.Connection, error) {
	return u.conn, u.err
}

type fakeUpgradeConnection struct {
	closeCalled bool
	lock        sync.Mutex
	streams     map[string]*fakeUpgradeStream
	portData    map[string]string
}

func newFakeUpgradeConnection() *fakeUpgradeConnection {
	return &fakeUpgradeConnection{
		streams:  make(map[string]*fakeUpgradeStream),
		portData: make(map[string]string),
	}
}

func (c *fakeUpgradeConnection) CreateStream(headers http.Header) (httpstream.Stream, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	stream := &fakeUpgradeStream{}
	c.streams[headers.Get(api.PortHeader)] = stream
	// only simulate data on the data stream for now, not the error stream
	if headers.Get(api.StreamType) == api.StreamTypeData {
		stream.data = c.portData[headers.Get(api.PortHeader)]
	}

	return stream, nil
}

func (c *fakeUpgradeConnection) Close() error {
	c.lock.Lock()
	defer c.lock.Unlock()

	c.closeCalled = true
	return nil
}

func (c *fakeUpgradeConnection) CloseChan() <-chan bool {
	return make(chan bool)
}

func (c *fakeUpgradeConnection) SetIdleTimeout(timeout time.Duration) {
}

type fakeUpgradeStream struct {
	readCalled  bool
	writeCalled bool
	dataWritten []byte
	closeCalled bool
	resetCalled bool
	data        string
	lock        sync.Mutex
}

func (s *fakeUpgradeStream) Read(p []byte) (int, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.readCalled = true
	b := []byte(s.data)
	n := copy(p, b)
	return n, io.EOF
}

func (s *fakeUpgradeStream) Write(p []byte) (int, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.writeCalled = true
	s.dataWritten = make([]byte, len(p))
	copy(s.dataWritten, p)
	return len(p), io.EOF
}

func (s *fakeUpgradeStream) Close() error {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.closeCalled = true
	return nil
}

func (s *fakeUpgradeStream) Reset() error {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.resetCalled = true
	return nil
}

func (s *fakeUpgradeStream) Headers() http.Header {
	s.lock.Lock()
	defer s.lock.Unlock()
	return http.Header{}
}

type GetListenerTestCase struct {
	Hostname                string
	Protocol                string
	ShouldRaiseError        bool
	ExpectedListenerAddress string
}

func TestGetListener(t *testing.T) {
	var pf PortForwarder
	testCases := []GetListenerTestCase{
		{
			Hostname:                "localhost",
			Protocol:                "tcp4",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "127.0.0.1",
		},
		{
			Hostname:                "127.0.0.1",
			Protocol:                "tcp4",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "127.0.0.1",
		},
		{
			Hostname:                "[::1]",
			Protocol:                "tcp6",
			ShouldRaiseError:        false,
			ExpectedListenerAddress: "::1",
		},
		{
			Hostname:         "[::1]",
			Protocol:         "tcp4",
			ShouldRaiseError: true,
		},
		{
			Hostname:         "127.0.0.1",
			Protocol:         "tcp6",
			ShouldRaiseError: true,
		},
		{
			// IPv6 address must be put into brackets. This test reveals this.
			Hostname:         "::1",
			Protocol:         "tcp6",
			ShouldRaiseError: true,
		},
	}

	for i, testCase := range testCases {
		expectedListenerPort := "12345"
		listener, err := pf.getListener(testCase.Protocol, testCase.Hostname, &ForwardedPort{12345, 12345})
		if err != nil && strings.Contains(err.Error(), "cannot assign requested address") {
			t.Logf("Can't test #%d: %v", i, err)
			continue
		}
		errorRaised := err != nil

		if testCase.ShouldRaiseError != errorRaised {
			t.Errorf("Test case #%d failed: Data %v an error has been raised(%t) where it should not (or reciprocally): %v", i, testCase, testCase.ShouldRaiseError, err)
			continue
		}
		if errorRaised {
			continue
		}

		if listener == nil {
			t.Errorf("Test case #%d did not raise an error but failed in initializing listener", i)
			continue
		}

		host, port, _ := net.SplitHostPort(listener.Addr().String())
		t.Logf("Asked a %s forward for: %s:%v, got listener %s:%s, expected: %s", testCase.Protocol, testCase.Hostname, 12345, host, port, expectedListenerPort)
		if host != testCase.ExpectedListenerAddress {
			t.Errorf("Test case #%d failed: Listener does not listen on exepected address: asked %v got %v", i, testCase.ExpectedListenerAddress, host)
		}
		if port != expectedListenerPort {
			t.Errorf("Test case #%d failed: Listener does not listen on exepected port: asked %v got %v", i, expectedListenerPort, port)

		}
		listener.Close()

	}
}

func TestForwardPorts(t *testing.T) {
	testCases := []struct {
		Upgrader *fakeUpgrader
		Ports    []string
		Send     map[uint16]string
		Receive  map[uint16]string
		Err      bool
	}{
		{
			Upgrader: &fakeUpgrader{err: errors.New("bail")},
			Err:      true,
		},
		{
			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
			Ports:    []string{"5000"},
		},
		{
			Upgrader: &fakeUpgrader{conn: newFakeUpgradeConnection()},
			Ports:    []string{"5000", "6000"},
			Send: map[uint16]string{
				5000: "abcd",
				6000: "ghij",
			},
			Receive: map[uint16]string{
				5000: "1234",
				6000: "5678",
			},
		},
	}

	for i, testCase := range testCases {
		stopChan := make(chan struct{}, 1)

		pf, err := New(&client.Request{}, &client.Config{}, testCase.Ports, stopChan)
		hasErr := err != nil
		if hasErr != testCase.Err {
			t.Fatalf("%d: New: expected %t, got %t: %v", i, testCase.Err, hasErr, err)
		}
		if pf == nil {
			continue
		}
		pf.upgrader = testCase.Upgrader
		if testCase.Upgrader.err != nil {
			err := pf.ForwardPorts()
			hasErr := err != nil
			if hasErr != testCase.Err {
				t.Fatalf("%d: ForwardPorts: expected %t, got %t: %v", i, testCase.Err, hasErr, err)
			}
			continue
		}

		doneChan := make(chan error)
		go func() {
			doneChan <- pf.ForwardPorts()
		}()
		<-pf.Ready

		conn := testCase.Upgrader.conn

		for port, data := range testCase.Send {
			conn.lock.Lock()
			conn.portData[fmt.Sprintf("%d", port)] = testCase.Receive[port]
			conn.lock.Unlock()

			clientConn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", port))
			if err != nil {
				t.Fatalf("%d: error dialing %d: %s", i, port, err)
			}
			defer clientConn.Close()

			n, err := clientConn.Write([]byte(data))
			if err != nil && err != io.EOF {
				t.Fatalf("%d: Error sending data '%s': %s", i, data, err)
			}
			if n == 0 {
				t.Fatalf("%d: unexpected write of 0 bytes", i)
			}
			b := make([]byte, 4)
			n, err = clientConn.Read(b)
			if err != nil && err != io.EOF {
				t.Fatalf("%d: Error reading data: %s", i, err)
			}
			if !bytes.Equal([]byte(testCase.Receive[port]), b) {
				t.Fatalf("%d: expected to read '%s', got '%s'", i, testCase.Receive[port], b)
			}
		}

		// tell r.ForwardPorts to stop
		close(stopChan)

		// wait for r.ForwardPorts to actually return
		err = <-doneChan
		if err != nil {
			t.Fatalf("%d: unexpected error: %s", i, err)
		}

		if e, a := len(testCase.Send), len(conn.streams); e != a {
			t.Fatalf("%d: expected %d streams to be created, got %d", i, e, a)
		}

		if !conn.closeCalled {
			t.Fatalf("%d: expected conn closure", i)
		}
	}

}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"context"
	"sort"
	"time"

	"github.com/opentracing/opentracing-go"
	"github.com/pingcap/errors"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
)

// DirtyDB stores uncommitted write operations for a transaction.
// It is stored and retrieved by context.Value and context.SetValue method.
type DirtyDB struct {
	// tables is a map whose key is tableID.
	tables map[int64]*DirtyTable
}

// GetDirtyTable gets the DirtyTable by id from the DirtyDB.
func (udb *DirtyDB) GetDirtyTable(tid int64) *DirtyTable {
	dt, ok := udb.tables[tid]
	if !ok {
		dt = &DirtyTable{
			tid:         tid,
			addedRows:   make(map[int64]struct{}),
			deletedRows: make(map[int64]struct{}),
		}
		udb.tables[tid] = dt
	}
	return dt
}

// DirtyTable stores uncommitted write operation for a transaction.
type DirtyTable struct {
	tid int64
	// addedRows ...
	// the key is handle.
	addedRows   map[int64]struct{}
	deletedRows map[int64]struct{}
	truncated   bool
}

// AddRow adds a row to the DirtyDB.
func (dt *DirtyTable) AddRow(handle int64, row []types.Datum) {
	dt.addedRows[handle] = struct{}{}
}

// DeleteRow deletes a row from the DirtyDB.
func (dt *DirtyTable) DeleteRow(handle int64) {
	delete(dt.addedRows, handle)
	dt.deletedRows[handle] = struct{}{}
}

// TruncateTable truncates a table.
func (dt *DirtyTable) TruncateTable() {
	dt.addedRows = make(map[int64]struct{})
	dt.truncated = true
}

// GetDirtyDB returns the DirtyDB bind to the context.
func GetDirtyDB(ctx sessionctx.Context) *DirtyDB {
	var udb *DirtyDB
	x := ctx.GetSessionVars().TxnCtx.DirtyDB
	if x == nil {
		udb = &DirtyDB{tables: make(map[int64]*DirtyTable)}
		ctx.GetSessionVars().TxnCtx.DirtyDB = udb
	} else {
		udb = x.(*DirtyDB)
	}
	return udb
}

// UnionScanExec merges the rows from dirty table and the rows from distsql request.
type UnionScanExec struct {
	baseExecutor

	dirty *DirtyTable
	// usedIndex is the column offsets of the index which Src executor has used.
	usedIndex  []int
	desc       bool
	conditions []expression.Expression
	columns    []*model.ColumnInfo

	// belowHandleIndex is the handle's position of the below scan plan.
	belowHandleIndex int

	addedRows           [][]types.Datum
	cursor4AddRows      int
	sortErr             error
	snapshotRows        [][]types.Datum
	cursor4SnapshotRows int
	snapshotChunkBuffer *chunk.Chunk
}

// Open implements the Executor Open interface.
func (us *UnionScanExec) Open(ctx context.Context) error {
	if err := us.baseExecutor.Open(ctx); err != nil {
		return errors.Trace(err)
	}
	us.snapshotChunkBuffer = us.newFirstChunk()
	return nil
}

// Next implements the Executor Next interface.
func (us *UnionScanExec) Next(ctx context.Context, req *chunk.RecordBatch) error {
	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
		span1 := span.Tracer().StartSpan("unionScan.Next", opentracing.ChildOf(span.Context()))
		defer span1.Finish()
	}

	if us.runtimeStats != nil {
		start := time.Now()
		defer func() { us.runtimeStats.Record(time.Since(start), req.NumRows()) }()
	}
	req.GrowAndReset(us.maxChunkSize)
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	for i, batchSize := 0, req.Capacity(); i < batchSize; i++ {
		row, err := us.getOneRow(ctx)
		if err != nil {
			return errors.Trace(err)
		}
		// no more data.
		if row == nil {
			return nil
		}
		mutableRow.SetDatums(row...)
		req.AppendRow(mutableRow.ToRow())
	}
	return nil
}

// getOneRow gets one result row from dirty table or child.
func (us *UnionScanExec) getOneRow(ctx context.Context) ([]types.Datum, error) {
	for {
		snapshotRow, err := us.getSnapshotRow(ctx)
		if err != nil {
			return nil, errors.Trace(err)
		}
		addedRow := us.getAddedRow()
		var row []types.Datum
		var isSnapshotRow bool
		if addedRow == nil {
			row = snapshotRow
			isSnapshotRow = true
		} else if snapshotRow == nil {
			row = addedRow
		} else {
			isSnapshotRow, err = us.shouldPickFirstRow(snapshotRow, addedRow)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if isSnapshotRow {
				row = snapshotRow
			} else {
				row = addedRow
			}
		}
		if row == nil {
			return nil, nil
		}

		if isSnapshotRow {
			us.cursor4SnapshotRows++
		} else {
			us.cursor4AddRows++
		}
		return row, nil
	}
}

func (us *UnionScanExec) getSnapshotRow(ctx context.Context) ([]types.Datum, error) {
	if us.dirty.truncated {
		return nil, nil
	}
	if us.cursor4SnapshotRows < len(us.snapshotRows) {
		return us.snapshotRows[us.cursor4SnapshotRows], nil
	}
	var err error
	us.cursor4SnapshotRows = 0
	us.snapshotRows = us.snapshotRows[:0]
	for len(us.snapshotRows) == 0 {
		err = us.children[0].Next(ctx, chunk.NewRecordBatch(us.snapshotChunkBuffer))
		if err != nil || us.snapshotChunkBuffer.NumRows() == 0 {
			return nil, errors.Trace(err)
		}
		iter := chunk.NewIterator4Chunk(us.snapshotChunkBuffer)
		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
			snapshotHandle := row.GetInt64(us.belowHandleIndex)
			if _, ok := us.dirty.deletedRows[snapshotHandle]; ok {
				continue
			}
			if _, ok := us.dirty.addedRows[snapshotHandle]; ok {
				// If src handle appears in added rows, it means there is conflict and the transaction will fail to
				// commit, but for simplicity, we don't handle it here.
				continue
			}
			us.snapshotRows = append(us.snapshotRows, row.GetDatumRow(us.children[0].retTypes()))
		}
	}
	return us.snapshotRows[0], nil
}

func (us *UnionScanExec) getAddedRow() []types.Datum {
	var addedRow []types.Datum
	if us.cursor4AddRows < len(us.addedRows) {
		addedRow = us.addedRows[us.cursor4AddRows]
	}
	return addedRow
}

// shouldPickFirstRow picks the suitable row in order.
// The value returned is used to determine whether to pick the first input row.
func (us *UnionScanExec) shouldPickFirstRow(a, b []types.Datum) (bool, error) {
	var isFirstRow bool
	addedCmpSrc, err := us.compare(a, b)
	if err != nil {
		return isFirstRow, errors.Trace(err)
	}
	// Compare result will never be 0.
	if us.desc {
		if addedCmpSrc > 0 {
			isFirstRow = true
		}
	} else {
		if addedCmpSrc < 0 {
			isFirstRow = true
		}
	}
	return isFirstRow, nil
}

func (us *UnionScanExec) compare(a, b []types.Datum) (int, error) {
	sc := us.ctx.GetSessionVars().StmtCtx
	for _, colOff := range us.usedIndex {
		aColumn := a[colOff]
		bColumn := b[colOff]
		cmp, err := aColumn.CompareDatum(sc, &bColumn)
		if err != nil {
			return 0, errors.Trace(err)
		}
		if cmp != 0 {
			return cmp, nil
		}
	}
	aHandle := a[us.belowHandleIndex].GetInt64()
	bHandle := b[us.belowHandleIndex].GetInt64()
	var cmp int
	if aHandle == bHandle {
		cmp = 0
	} else if aHandle > bHandle {
		cmp = 1
	} else {
		cmp = -1
	}
	return cmp, nil
}

func (us *UnionScanExec) buildAndSortAddedRows(t table.Table) error {
	us.addedRows = make([][]types.Datum, 0, len(us.dirty.addedRows))
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	cols := t.WritableCols()
	for h := range us.dirty.addedRows {
		newData := make([]types.Datum, 0, us.schema.Len())
		data, err := t.RowWithCols(us.ctx, h, cols)
		if err != nil {
			return err
		}
		for _, col := range us.columns {
			if col.ID == model.ExtraHandleID {
				newData = append(newData, types.NewIntDatum(h))
			} else {
				newData = append(newData, data[col.Offset])
			}
		}
		mutableRow.SetDatums(newData...)
		matched, _, err := expression.EvalBool(us.ctx, us.conditions, mutableRow.ToRow())
		if err != nil {
			return errors.Trace(err)
		}
		if !matched {
			continue
		}
		us.addedRows = append(us.addedRows, newData)
	}
	if us.desc {
		sort.Sort(sort.Reverse(us))
	} else {
		sort.Sort(us)
	}
	if us.sortErr != nil {
		return errors.Trace(us.sortErr)
	}
	return nil
}

// Len implements sort.Interface interface.
func (us *UnionScanExec) Len() int {
	return len(us.addedRows)
}

// Less implements sort.Interface interface.
func (us *UnionScanExec) Less(i, j int) bool {
	cmp, err := us.compare(us.addedRows[i], us.addedRows[j])
	if err != nil {
		us.sortErr = errors.Trace(err)
		return true
	}
	return cmp < 0
}

// Swap implements sort.Interface interface.
func (us *UnionScanExec) Swap(i, j int) {
	us.addedRows[i], us.addedRows[j] = us.addedRows[j], us.addedRows[i]
}
<commit_msg>executor: only use the transaction buffer for union scan (#9428)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"context"
	"sort"
	"time"

	"github.com/opentracing/opentracing-go"
	"github.com/pingcap/errors"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/table/tables"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
)

// DirtyDB stores uncommitted write operations for a transaction.
// It is stored and retrieved by context.Value and context.SetValue method.
type DirtyDB struct {
	// tables is a map whose key is tableID.
	tables map[int64]*DirtyTable
}

// GetDirtyTable gets the DirtyTable by id from the DirtyDB.
func (udb *DirtyDB) GetDirtyTable(tid int64) *DirtyTable {
	dt, ok := udb.tables[tid]
	if !ok {
		dt = &DirtyTable{
			tid:         tid,
			addedRows:   make(map[int64]struct{}),
			deletedRows: make(map[int64]struct{}),
		}
		udb.tables[tid] = dt
	}
	return dt
}

// DirtyTable stores uncommitted write operation for a transaction.
type DirtyTable struct {
	tid int64
	// addedRows ...
	// the key is handle.
	addedRows   map[int64]struct{}
	deletedRows map[int64]struct{}
	truncated   bool
}

// AddRow adds a row to the DirtyDB.
func (dt *DirtyTable) AddRow(handle int64, row []types.Datum) {
	dt.addedRows[handle] = struct{}{}
}

// DeleteRow deletes a row from the DirtyDB.
func (dt *DirtyTable) DeleteRow(handle int64) {
	delete(dt.addedRows, handle)
	dt.deletedRows[handle] = struct{}{}
}

// TruncateTable truncates a table.
func (dt *DirtyTable) TruncateTable() {
	dt.addedRows = make(map[int64]struct{})
	dt.truncated = true
}

// GetDirtyDB returns the DirtyDB bind to the context.
func GetDirtyDB(ctx sessionctx.Context) *DirtyDB {
	var udb *DirtyDB
	x := ctx.GetSessionVars().TxnCtx.DirtyDB
	if x == nil {
		udb = &DirtyDB{tables: make(map[int64]*DirtyTable)}
		ctx.GetSessionVars().TxnCtx.DirtyDB = udb
	} else {
		udb = x.(*DirtyDB)
	}
	return udb
}

// UnionScanExec merges the rows from dirty table and the rows from distsql request.
type UnionScanExec struct {
	baseExecutor

	dirty *DirtyTable
	// usedIndex is the column offsets of the index which Src executor has used.
	usedIndex  []int
	desc       bool
	conditions []expression.Expression
	columns    []*model.ColumnInfo

	// belowHandleIndex is the handle's position of the below scan plan.
	belowHandleIndex int

	addedRows           [][]types.Datum
	cursor4AddRows      int
	sortErr             error
	snapshotRows        [][]types.Datum
	cursor4SnapshotRows int
	snapshotChunkBuffer *chunk.Chunk
}

// Open implements the Executor Open interface.
func (us *UnionScanExec) Open(ctx context.Context) error {
	if err := us.baseExecutor.Open(ctx); err != nil {
		return errors.Trace(err)
	}
	us.snapshotChunkBuffer = us.newFirstChunk()
	return nil
}

// Next implements the Executor Next interface.
func (us *UnionScanExec) Next(ctx context.Context, req *chunk.RecordBatch) error {
	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
		span1 := span.Tracer().StartSpan("unionScan.Next", opentracing.ChildOf(span.Context()))
		defer span1.Finish()
	}

	if us.runtimeStats != nil {
		start := time.Now()
		defer func() { us.runtimeStats.Record(time.Since(start), req.NumRows()) }()
	}
	req.GrowAndReset(us.maxChunkSize)
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	for i, batchSize := 0, req.Capacity(); i < batchSize; i++ {
		row, err := us.getOneRow(ctx)
		if err != nil {
			return errors.Trace(err)
		}
		// no more data.
		if row == nil {
			return nil
		}
		mutableRow.SetDatums(row...)
		req.AppendRow(mutableRow.ToRow())
	}
	return nil
}

// getOneRow gets one result row from dirty table or child.
func (us *UnionScanExec) getOneRow(ctx context.Context) ([]types.Datum, error) {
	for {
		snapshotRow, err := us.getSnapshotRow(ctx)
		if err != nil {
			return nil, errors.Trace(err)
		}
		addedRow := us.getAddedRow()
		var row []types.Datum
		var isSnapshotRow bool
		if addedRow == nil {
			row = snapshotRow
			isSnapshotRow = true
		} else if snapshotRow == nil {
			row = addedRow
		} else {
			isSnapshotRow, err = us.shouldPickFirstRow(snapshotRow, addedRow)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if isSnapshotRow {
				row = snapshotRow
			} else {
				row = addedRow
			}
		}
		if row == nil {
			return nil, nil
		}

		if isSnapshotRow {
			us.cursor4SnapshotRows++
		} else {
			us.cursor4AddRows++
		}
		return row, nil
	}
}

func (us *UnionScanExec) getSnapshotRow(ctx context.Context) ([]types.Datum, error) {
	if us.dirty.truncated {
		return nil, nil
	}
	if us.cursor4SnapshotRows < len(us.snapshotRows) {
		return us.snapshotRows[us.cursor4SnapshotRows], nil
	}
	var err error
	us.cursor4SnapshotRows = 0
	us.snapshotRows = us.snapshotRows[:0]
	for len(us.snapshotRows) == 0 {
		err = us.children[0].Next(ctx, chunk.NewRecordBatch(us.snapshotChunkBuffer))
		if err != nil || us.snapshotChunkBuffer.NumRows() == 0 {
			return nil, errors.Trace(err)
		}
		iter := chunk.NewIterator4Chunk(us.snapshotChunkBuffer)
		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
			snapshotHandle := row.GetInt64(us.belowHandleIndex)
			if _, ok := us.dirty.deletedRows[snapshotHandle]; ok {
				continue
			}
			if _, ok := us.dirty.addedRows[snapshotHandle]; ok {
				// If src handle appears in added rows, it means there is conflict and the transaction will fail to
				// commit, but for simplicity, we don't handle it here.
				continue
			}
			us.snapshotRows = append(us.snapshotRows, row.GetDatumRow(us.children[0].retTypes()))
		}
	}
	return us.snapshotRows[0], nil
}

func (us *UnionScanExec) getAddedRow() []types.Datum {
	var addedRow []types.Datum
	if us.cursor4AddRows < len(us.addedRows) {
		addedRow = us.addedRows[us.cursor4AddRows]
	}
	return addedRow
}

// shouldPickFirstRow picks the suitable row in order.
// The value returned is used to determine whether to pick the first input row.
func (us *UnionScanExec) shouldPickFirstRow(a, b []types.Datum) (bool, error) {
	var isFirstRow bool
	addedCmpSrc, err := us.compare(a, b)
	if err != nil {
		return isFirstRow, errors.Trace(err)
	}
	// Compare result will never be 0.
	if us.desc {
		if addedCmpSrc > 0 {
			isFirstRow = true
		}
	} else {
		if addedCmpSrc < 0 {
			isFirstRow = true
		}
	}
	return isFirstRow, nil
}

func (us *UnionScanExec) compare(a, b []types.Datum) (int, error) {
	sc := us.ctx.GetSessionVars().StmtCtx
	for _, colOff := range us.usedIndex {
		aColumn := a[colOff]
		bColumn := b[colOff]
		cmp, err := aColumn.CompareDatum(sc, &bColumn)
		if err != nil {
			return 0, errors.Trace(err)
		}
		if cmp != 0 {
			return cmp, nil
		}
	}
	aHandle := a[us.belowHandleIndex].GetInt64()
	bHandle := b[us.belowHandleIndex].GetInt64()
	var cmp int
	if aHandle == bHandle {
		cmp = 0
	} else if aHandle > bHandle {
		cmp = 1
	} else {
		cmp = -1
	}
	return cmp, nil
}

// rowWithColsInTxn gets the row from the transaction buffer.
func (us *UnionScanExec) rowWithColsInTxn(t table.Table, h int64, cols []*table.Column) ([]types.Datum, error) {
	key := t.RecordKey(h)
	txn, err := us.ctx.Txn(true)
	if err != nil {
		return nil, errors.Trace(err)
	}
	value, err := txn.GetMemBuffer().Get(key)
	if err != nil {
		return nil, errors.Trace(err)
	}
	v, _, err := tables.DecodeRawRowData(us.ctx, t.Meta(), h, cols, value)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return v, nil
}

func (us *UnionScanExec) buildAndSortAddedRows(t table.Table) error {
	us.addedRows = make([][]types.Datum, 0, len(us.dirty.addedRows))
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	cols := t.WritableCols()
	for h := range us.dirty.addedRows {
		newData := make([]types.Datum, 0, us.schema.Len())
		data, err := us.rowWithColsInTxn(t, h, cols)
		if err != nil {
			return err
		}
		for _, col := range us.columns {
			if col.ID == model.ExtraHandleID {
				newData = append(newData, types.NewIntDatum(h))
			} else {
				newData = append(newData, data[col.Offset])
			}
		}
		mutableRow.SetDatums(newData...)
		matched, _, err := expression.EvalBool(us.ctx, us.conditions, mutableRow.ToRow())
		if err != nil {
			return errors.Trace(err)
		}
		if !matched {
			continue
		}
		us.addedRows = append(us.addedRows, newData)
	}
	if us.desc {
		sort.Sort(sort.Reverse(us))
	} else {
		sort.Sort(us)
	}
	if us.sortErr != nil {
		return errors.Trace(us.sortErr)
	}
	return nil
}

// Len implements sort.Interface interface.
func (us *UnionScanExec) Len() int {
	return len(us.addedRows)
}

// Less implements sort.Interface interface.
func (us *UnionScanExec) Less(i, j int) bool {
	cmp, err := us.compare(us.addedRows[i], us.addedRows[j])
	if err != nil {
		us.sortErr = errors.Trace(err)
		return true
	}
	return cmp < 0
}

// Swap implements sort.Interface interface.
func (us *UnionScanExec) Swap(i, j int) {
	us.addedRows[i], us.addedRows[j] = us.addedRows[j], us.addedRows[i]
}
<|endoftext|>