<commit_before>package flowgraph

import (
	"github.com/vectaport/fgbase"

	"fmt"
)

// Hub interface for flowgraph hubs that are connected by flowgraph streams
type Hub interface {

	/* TRACING */

	// Tracef for debug trace printing.
	Tracef(format string, v ...interface{})

	// LogError for logging of error messages.
	LogError(format string, v ...interface{})

	/* PROBING */

	// Name returns the hub name
	Name() string

	// Source returns source stream by index
	Source(n int) Stream

	// Destination returns destination stream by index
	Destination(n int) Stream

	// FindSource returns source stream by port name
	FindSource(name string) Stream

	// FindDestination returns destination stream by port name
	FindDestination(name string) Stream

	// NumSource returns the number of source ports
	NumSource() int

	// NumDestination returns the number of destination ports
	NumDestination() int

	// SourceNames returns the names of the source ports
	SourceNames() []string

	// DestinationNames returns the names of the destination ports
	DestinationNames() []string

	/* CONFIGURING */

	// AddSource adds a source port for each stream
	AddSource(s ...Stream)

	// AddDestination adds a destination port for each stream
	AddDestination(s ...Stream)

	// SetSourceNames names the source ports
	SetSourceNames(nm ...string)

	// SetDestinationNames names the destination ports
	SetDestinationNames(nm ...string)

	// SetSource sets the stream for a named source port
	SetSource(port string, s Stream) error

	// SetDestination sets the stream for a named destination port
	SetDestination(port string, s Stream) error

	/* IMPLEMENTATION */

	// Base returns the value that implements this hub
	// The type of this value identifies the implementation.
	Base() interface{}
}

// implementation of Hub
type hub struct {
	base *fgbase.Node
}

// Tracef for debug trace printing.  Uses atomic log mechanism.
func (h hub) Tracef(format string, v ...interface{}) {
	h.base.Tracef(format, v)
}

// LogError for logging of error messages.  Uses atomic log mechanism.
func (h hub) LogError(format string, v ...interface{}) {
	h.base.LogError(format, v)
}

// Name returns the hub name
func (h hub) Name() string {
	return h.base.Name
}

// Source returns source stream by index
func (h hub) Source(i int) Stream {
	return stream{h.base.SrcGet(i)}
}

// Destination returns destination stream by index
func (h hub) Destination(i int) Stream {
	return stream{h.base.DstGet(i)}
}

// FindSource returns source stream by port name
func (h hub) FindSource(name string) Stream {
	e, ok := h.base.FindSrc(name)
	if !ok {
		return nil
	}
	if e == nil {
		return stream{nil}
	}
	return stream{e}
}

// FindDestination returns destination stream by port name
func (h hub) FindDestination(name string) Stream {
	e, ok := h.base.FindDst(name)
	if !ok {
		return nil
	}
	if e == nil {
		return stream{nil}
	}
	return stream{e}
}

// AddSource adds a source port for each stream
func (h hub) AddSource(s ...Stream) {
	for _, sv := range s {
		h.base.SrcAppend(sv.Base().(*fgbase.Edge))
	}
}

// AddDestination adds a destination port for each stream
func (h hub) AddDestination(s ...Stream) {
	for _, sv := range s {
		h.base.DstAppend(sv.Base().(*fgbase.Edge))
	}
}

// NumSource returns the number of source ports
func (h hub) NumSource() int {
	return h.base.SrcCnt()
}

// NumDestination returns the number of destination ports
func (h hub) NumDestination() int {
	return h.base.DstCnt()
}

// SetSourceNames names the source ports
func (h hub) SetSourceNames(nm ...string) {
	h.base.SetSrcNames(nm...)
}

// SetDestinationNames names the destination ports
func (h hub) SetDestinationNames(nm ...string) {
	h.base.SetDstNames(nm...)
}

// SourceNames returns the names of the source ports
func (h hub) SourceNames() []string {
	return h.base.SrcNames()
}

// DestinationNames returns the names of the destination ports
func (h hub) DestinationNames() []string {
	return h.base.DstNames()
}

// SetSource sets the stream for a named source port
func (h hub) SetSource(port string, s Stream) error {
	i, ok := h.base.FindSrcIndex(port)
	if !ok {
		return fmt.Errorf("source port %s not found on hub %s\n", port, h.Name())
	}
	h.base.SrcSet(i, s.Base().(*fgbase.Edge))
	return nil
}

// SetDestination sets the stream for a named destination port
func (h hub) SetDestination(port string, s Stream) error {
	i, ok := h.base.FindDstIndex(port)
	if !ok {
		return fmt.Errorf("destination port %s not found on hub %s\n", port, h.Name())
	}
	h.base.DstSet(i, s.Base().(*fgbase.Edge))
	return nil
}

// Base returns the value that implements this hub
// The type of this value identifies the implementation.
func (h hub) Base() interface{} {
	return h.base
}
<commit_msg>check number of ports before naming<commit_after>package flowgraph

import (
	"github.com/vectaport/fgbase"

	"fmt"
)

// Hub interface for flowgraph hubs that are connected by flowgraph streams
type Hub interface {

	/* TRACING */

	// Tracef for debug trace printing.
	Tracef(format string, v ...interface{})

	// LogError for logging of error messages.
	LogError(format string, v ...interface{})

	/* PROBING */

	// Name returns the hub name
	Name() string

	// Source returns source stream by index
	Source(n int) Stream

	// Destination returns destination stream by index
	Destination(n int) Stream

	// FindSource returns source stream by port name
	FindSource(name string) Stream

	// FindDestination returns destination stream by port name
	FindDestination(name string) Stream

	// NumSource returns the number of source ports
	NumSource() int

	// NumDestination returns the number of destination ports
	NumDestination() int

	// SourceNames returns the names of the source ports
	SourceNames() []string

	// DestinationNames returns the names of the destination ports
	DestinationNames() []string

	/* CONFIGURING */

	// AddSource adds a source port for each stream
	AddSource(s ...Stream)

	// AddDestination adds a destination port for each stream
	AddDestination(s ...Stream)

	// SetSourceNames names the source ports
	SetSourceNames(nm ...string)

	// SetDestinationNames names the destination ports
	SetDestinationNames(nm ...string)

	// SetSource sets the stream for a named source port
	SetSource(port string, s Stream) error

	// SetDestination sets the stream for a named destination port
	SetDestination(port string, s Stream) error

	/* IMPLEMENTATION */

	// Base returns the value that implements this hub
	// The type of this value identifies the implementation.
	Base() interface{}
}

// implementation of Hub
type hub struct {
	base *fgbase.Node
}

// Tracef for debug trace printing.  Uses atomic log mechanism.
func (h hub) Tracef(format string, v ...interface{}) {
	h.base.Tracef(format, v)
}

// LogError for logging of error messages.  Uses atomic log mechanism.
func (h hub) LogError(format string, v ...interface{}) {
	h.base.LogError(format, v)
}

// Name returns the hub name
func (h hub) Name() string {
	return h.base.Name
}

// Source returns source stream by index
func (h hub) Source(i int) Stream {
	return stream{h.base.SrcGet(i)}
}

// Destination returns destination stream by index
func (h hub) Destination(i int) Stream {
	return stream{h.base.DstGet(i)}
}

// FindSource returns source stream by port name
func (h hub) FindSource(name string) Stream {
	e, ok := h.base.FindSrc(name)
	if !ok {
		return nil
	}
	if e == nil {
		return stream{nil}
	}
	return stream{e}
}

// FindDestination returns destination stream by port name
func (h hub) FindDestination(name string) Stream {
	e, ok := h.base.FindDst(name)
	if !ok {
		return nil
	}
	if e == nil {
		return stream{nil}
	}
	return stream{e}
}

// AddSource adds a source port for each stream
func (h hub) AddSource(s ...Stream) {
	for _, sv := range s {
		h.base.SrcAppend(sv.Base().(*fgbase.Edge))
	}
}

// AddDestination adds a destination port for each stream
func (h hub) AddDestination(s ...Stream) {
	for _, sv := range s {
		h.base.DstAppend(sv.Base().(*fgbase.Edge))
	}
}

// NumSource returns the number of source ports
func (h hub) NumSource() int {
	return h.base.SrcCnt()
}

// NumDestination returns the number of destination ports
func (h hub) NumDestination() int {
	return h.base.DstCnt()
}

// SetSourceNames names the source ports
func (h hub) SetSourceNames(nm ...string) {
        if len(nm)!=h.base.SrcCnt() {
	   h.LogError("Number of source ports does not match number of names (%d!=%d)\n", h.base.SrcCnt(), len(nm))
	}
	h.base.SetSrcNames(nm...)
}

// SetDestinationNames names the destination ports
func (h hub) SetDestinationNames(nm ...string) {
        if len(nm)!=h.base.DstCnt() {
	   h.LogError("Number of destination ports does not match number of names (%d!=%d)\n", h.base.DstCnt(), len(nm))
	}
	h.base.SetDstNames(nm...)
}

// SourceNames returns the names of the source ports
func (h hub) SourceNames() []string {
	return h.base.SrcNames()
}

// DestinationNames returns the names of the destination ports
func (h hub) DestinationNames() []string {
	return h.base.DstNames()
}

// SetSource sets the stream for a named source port
func (h hub) SetSource(port string, s Stream) error {
	i, ok := h.base.FindSrcIndex(port)
	if !ok {
		return fmt.Errorf("source port %s not found on hub %s\n", port, h.Name())
	}
	h.base.SrcSet(i, s.Base().(*fgbase.Edge))
	return nil
}

// SetDestination sets the stream for a named destination port
func (h hub) SetDestination(port string, s Stream) error {
	i, ok := h.base.FindDstIndex(port)
	if !ok {
		return fmt.Errorf("destination port %s not found on hub %s\n", port, h.Name())
	}
	h.base.DstSet(i, s.Base().(*fgbase.Edge))
	return nil
}

// Base returns the value that implements this hub
// The type of this value identifies the implementation.
func (h hub) Base() interface{} {
	return h.base
}
<|endoftext|>


<commit_before>package main

import (
    "net"
    "bufio"
    "fmt"
    "os"
    "time"
    "strings"
)

const (
    CONNECTION_TYPE = "tcp"
    CONNECTION_HOST = "localhost"
    CONNECTION_PORT = "7010"
)

type Client struct {
    in chan string
    out chan string
    user_id uint64
    reader *bufio.Reader
    writer *bufio.Writer
}

type Hub struct {
    clients []*Client
    connections chan net.Conn
    in chan string
    out chan string
    writer *bufio.Writer
}

func (client *Client) Listen() {
    go client.Read()
    go client.Write()
}

func (client *Client) Read() {
    for {
        line, _ := client.reader.ReadString('\n')
        client.in <- line
    }
}

func (client *Client) Write() {
    for data := range client.out {
        client.writer.WriteString(data)
        client.writer.Flush()
    }
}

func (hub *Hub) Join(connection net.Conn) {
    client := CreateClient(connection)
    hub.clients = append(hub.clients, client)

    go hub.ListenClient(client)
}

func (hub *Hub) ListenClient(client *Client) {
    for {
        in := <-client.in
        if (strings.HasPrefix(in, "/whoami")) {
            client.out <- "hub> " + fmt.Sprintf("%d", client.user_id) + "\n"
        } else if (strings.HasPrefix(in, "/list")) {
            hub.ListClients(client)
        } else if (strings.HasPrefix(in, "/msg")) {
            hub.SendMessage(client, in)
        } else if (strings.HasPrefix(in, "/quit")) {
            hub.QuitClient(client)
        }
    }
}

func (hub *Hub) Write(data string) {
    hub.writer.WriteString(data)
    hub.writer.Flush()
}

func (hub *Hub) Listen() {
    for {
        select {
        case data := <-hub.in:
            hub.Write(data)
        case conn := <-hub.connections:
            hub.Join(conn)
        }
    }
}

func (hub *Hub) SendMessage(forClient *Client, message string) {
    s := strings.Split(message, " ");
    receivers, body := s[1], s[2]
    r := strings.Split(receivers, ",")
    for _, client := range hub.clients {
        for _, receiver := range r {
            if (fmt.Sprintf("%d", client.user_id) == receiver) {
                client.out <- fmt.Sprintf("%d", forClient.user_id) + "> " + body
            }
        }
    }
}

func (hub *Hub) ListClients(forClient *Client) {
    onlyMe := true
    for _, client := range hub.clients {
        if (forClient.user_id != client.user_id) {
            forClient.out <- "hub> " + fmt.Sprintf("%d", client.user_id) + "\n"
            onlyMe = false
        }
    }
    if (onlyMe == true) {
        forClient.out <- "hub> No one else here :(\n"
    }
}

func (hub *Hub) QuitClient(client *Client) {
    var tmpClients = make([]*Client, 0)
    for _, c := range hub.clients {
        if (c.user_id != client.user_id) {
            tmpClients = append(tmpClients, c)
        }
    }
    hub.clients = tmpClients
}

func CreateHub() *Hub {
    writer := bufio.NewWriter(os.Stdout)

    hub := &Hub{
        clients: make([]*Client, 0),
        connections: make(chan net.Conn),
        in: make(chan string),
        out: make(chan string),
        writer: writer,
    }

    go hub.Listen()

    return hub
}

func CreateClient(connection net.Conn) *Client {
    writer := bufio.NewWriter(connection)
    reader := bufio.NewReader(connection)

    user_id := uint64(time.Now().UnixNano())

    client := &Client{
        in: make(chan string),
        out: make(chan string),
        reader: reader,
        writer: writer,
        user_id: user_id,
    }

    client.Listen()

    return client
}

func HandleError(err error, message string) {
    if (err != nil) {
        fmt.Println("ERROR (" + message + "): ", err.Error())
        os.Exit(1)
    }
}

func main() {
    fmt.Println("Server initializing...")
    hub := CreateHub()

    listener, err := net.Listen(CONNECTION_TYPE, CONNECTION_HOST + ":" + CONNECTION_PORT)
    HandleError(err, "LISTEN")

    for {
        conn, err := listener.Accept()
        HandleError(err, "ACCEPT")
        hub.connections <- conn
    }
}

<commit_msg>bugfix: invalid msg command crashes hub<commit_after>package main

import (
    "net"
    "bufio"
    "fmt"
    "os"
    "time"
    "strings"
)

const (
    CONNECTION_TYPE = "tcp"
    CONNECTION_HOST = "localhost"
    CONNECTION_PORT = "7010"
)

type Client struct {
    in chan string
    out chan string
    user_id uint64
    reader *bufio.Reader
    writer *bufio.Writer
}

type Hub struct {
    clients []*Client
    connections chan net.Conn
    in chan string
    out chan string
    writer *bufio.Writer
}

func (client *Client) Listen() {
    go client.Read()
    go client.Write()
}

func (client *Client) Read() {
    for {
        line, _ := client.reader.ReadString('\n')
        client.in <- line
    }
}

func (client *Client) Write() {
    for data := range client.out {
        client.writer.WriteString(data)
        client.writer.Flush()
    }
}

func (hub *Hub) Join(connection net.Conn) {
    client := CreateClient(connection)
    hub.clients = append(hub.clients, client)

    go hub.ListenClient(client)
}

func (hub *Hub) ListenClient(client *Client) {
    for {
        in := <-client.in
        if (strings.HasPrefix(in, "/whoami")) {
            client.out <- "hub> " + fmt.Sprintf("%d", client.user_id) + "\n"
        } else if (strings.HasPrefix(in, "/list")) {
            hub.ListClients(client)
        } else if (strings.HasPrefix(in, "/msg")) {
            hub.SendMessage(client, in)
        } else if (strings.HasPrefix(in, "/quit")) {
            hub.QuitClient(client)
        }
    }
}

func (hub *Hub) Write(data string) {
    hub.writer.WriteString(data)
    hub.writer.Flush()
}

func (hub *Hub) Listen() {
    for {
        select {
        case data := <-hub.in:
            hub.Write(data)
        case conn := <-hub.connections:
            hub.Join(conn)
        }
    }
}

func (hub *Hub) SendMessage(fromClient *Client, message string) {
    if (strings.Count(message, " ") <= 1) {
        fromClient.out <- "hub> Invalid /msg command parameters. Use /msg [user_id1,user_id2,...] [msg]\n"
        return
    }
    s := strings.Split(message, " ");
    receivers, body := s[1], s[2]
    r := strings.Split(receivers, ",")
    for _, client := range hub.clients {
        for _, receiver := range r {
            if (fmt.Sprintf("%d", client.user_id) == receiver) {
                client.out <- fmt.Sprintf("%d", fromClient.user_id) + "> " + body
            }
        }
    }
}

func (hub *Hub) ListClients(forClient *Client) {
    onlyMe := true
    for _, client := range hub.clients {
        if (forClient.user_id != client.user_id) {
            forClient.out <- "hub> " + fmt.Sprintf("%d", client.user_id) + "\n"
            onlyMe = false
        }
    }
    if (onlyMe == true) {
        forClient.out <- "hub> No one else here :(\n"
    }
}

func (hub *Hub) QuitClient(client *Client) {
    var tmpClients = make([]*Client, 0)
    for _, c := range hub.clients {
        if (c.user_id != client.user_id) {
            tmpClients = append(tmpClients, c)
        }
    }
    hub.clients = tmpClients
}

func CreateHub() *Hub {
    writer := bufio.NewWriter(os.Stdout)

    hub := &Hub{
        clients: make([]*Client, 0),
        connections: make(chan net.Conn),
        in: make(chan string),
        out: make(chan string),
        writer: writer,
    }

    go hub.Listen()

    return hub
}

func CreateClient(connection net.Conn) *Client {
    writer := bufio.NewWriter(connection)
    reader := bufio.NewReader(connection)

    user_id := uint64(time.Now().UnixNano())

    client := &Client{
        in: make(chan string),
        out: make(chan string),
        reader: reader,
        writer: writer,
        user_id: user_id,
    }

    client.Listen()

    return client
}

func HandleError(err error, message string) {
    if (err != nil) {
        fmt.Println("ERROR (" + message + "): ", err.Error())
        os.Exit(1)
    }
}

func main() {
    fmt.Println("Server initializing...")
    hub := CreateHub()

    listener, err := net.Listen(CONNECTION_TYPE, CONNECTION_HOST + ":" + CONNECTION_PORT)
    HandleError(err, "LISTEN")

    for {
        conn, err := listener.Accept()
        HandleError(err, "ACCEPT")
        hub.connections <- conn
    }
}

<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
)

// pow computes z to the (nonnegative integer) power i.
func pow(z complex128, i int) complex128 {
	if i < 0 {
		panic("negative power")
	}

	var res complex128 = 1

	for i > 0 {
		res *= z
		i--
	}

	return res
}

// A Function can be evaluated at a complex point.
type Function interface {
	// evaluate computes the value of the function at z.
	Evaluate(z complex128) complex128
}

// A poly is a (complex) polynomial.
type poly []complex128

// Degree returns the degree of the polynomial.
//
// The degree of the zero polynomial is 0.
func (f poly) Degree() int {
	if len(f) == 0 {
		return 0
	}

	var deg int

	for i, v := range f {
		if v != 0 {
			deg = i
		}
	}

	return deg
}

// Normalise returns a new poly with any excess zeroes culled from the end.
//
// The normalised form of zero is {0}, not {}.
func (f poly) Normalise() poly {
	deg := f.Degree()

	if len(f) == 0 {
		return poly([]complex128{0})
	}

	if deg < len(f) {
		return f[:deg+1]
	}

	return f
}

func (f poly) IsConstant() bool {
	return f.Degree() == 0
}

func (f poly) Add(g poly) poly {
	deg1 := f.Degree()
	deg2 := g.Degree()
	var max_deg int

	if deg1 >= deg2 {
		max_deg = deg1
	} else {
		max_deg = deg2
	}

	res := make([]complex128, max_deg+1)

	for i := range res {
		var fi, gi complex128

		if i >= len(f) {
			fi = 0
		} else {
			fi = f[i]
		}

		if i >= len(g) {
			gi = 0
		} else {
			gi = g[i]
		}

		res[i] = fi + gi
	}

	return poly(res)
}

func (f poly) Negative() poly {
	res := make([]complex128, len(f))

	for i := range res {
		res[i] = -f[i]
	}

	return res
}

func (f poly) Subtract(g poly) poly {
	return f.Add(g.Negative())
}

func (f poly) Multiply(g poly) poly {
	deg := f.Degree() + g.Degree()

	res := make([]complex128, deg+1)

	for i := range res {
		for j := 0; j <= i; j++ {
			var fj, gi_j complex128

			if j >= len(f) {
				fj = 0
			} else {
				fj = f[j]
			}

			if i-j >= len(g) {
				gi_j = 0
			} else {
				gi_j = g[i-j]
			}

			res[i] += fj * gi_j
		}
	}

	return res
}

func (f poly) Equal(g poly) bool {
	deg1 := f.Degree()
	deg2 := g.Degree()

	if deg1 != deg2 {
		return false
	}

	f = f.Normalise()
	g = g.Normalise()

	for i := range f {
		if f[i] != g[i] {
			return false
		}
	}

	return true
}

func (f poly) Derivative() poly {
	if f.IsConstant() {
		return poly([]complex128{0})
	}

	res := make([]complex128, f.Degree())

	for i := range res {
		res[i] = complex(float64(i+1), 0) * f[i+1]
	}

	return poly(res)
}

func (f poly) String() string {
	b := new(bytes.Buffer)

	f = f.Normalise()
	fmt.Fprint(b, f[0])

	if len(f) >= 2 {
		if f[1] != 0 {
			fmt.Fprintf(b, " + %vt", f[1])
		}
	} else {
		return b.String()
	}

	for i, z := range f[2:] {
		if z != 0 {
			fmt.Fprintf(b, " + %vt^%v", z, i+2)
		}
	}

	return b.String()
}

func (f poly) Evaluate(z complex128) complex128 {
	var res complex128

	for i, coeff := range f {
		res += coeff * pow(z, i)
	}

	return res
}

// NewPoly returns a poly with coefficients coeffs, starting with the constant term.
func NewPoly(coeffs ...complex128) poly {
	return poly(coeffs).Normalise()
}

// A rat is a rational function.
type rat struct {
	numerator   poly
	denominator poly
}

func (f rat) Evaluate(z complex128) complex128 {
	num := f.numerator.Evaluate(z)
	denom := f.denominator.Evaluate(z)
	return num / denom
}
<commit_msg>Add Orbit function<commit_after>package main

import (
	"bytes"
	"fmt"
	"math/cmplx"
)

// pow computes z to the (nonnegative integer) power i.
func pow(z complex128, i int) complex128 {
	if i < 0 {
		panic("negative power")
	}

	var res complex128 = 1

	for i > 0 {
		res *= z
		i--
	}

	return res
}

// Orbit computes the result of iterating f starting at z.
//
// -1 indicates the orbit is bounded by bound after n_iter, and
// a nonnegative return indicates that the bound was exceeded after
// that many iterations.
func Orbit(f Function, z complex128, bound float64, n int) int {
	for i := 0; i < n; i++ {
		if cmplx.Abs(z) >= bound {
			return i
		}
		z = f.Evaluate(z)
	}

	return -1
}

// A Function can be evaluated at a complex point.
type Function interface {
	// evaluate computes the value of the function at z.
	Evaluate(z complex128) complex128
}

// A poly is a (complex) polynomial.
type poly []complex128

// Degree returns the degree of the polynomial.
//
// The degree of the zero polynomial is 0.
func (f poly) Degree() int {
	if len(f) == 0 {
		return 0
	}

	var deg int

	for i, v := range f {
		if v != 0 {
			deg = i
		}
	}

	return deg
}

// Normalise returns a new poly with any excess zeroes culled from the end.
//
// The normalised form of zero is {0}, not {}.
func (f poly) Normalise() poly {
	deg := f.Degree()

	if len(f) == 0 {
		return poly([]complex128{0})
	}

	if deg < len(f) {
		return f[:deg+1]
	}

	return f
}

func (f poly) IsConstant() bool {
	return f.Degree() == 0
}

func (f poly) Add(g poly) poly {
	deg1 := f.Degree()
	deg2 := g.Degree()
	var max_deg int

	if deg1 >= deg2 {
		max_deg = deg1
	} else {
		max_deg = deg2
	}

	res := make([]complex128, max_deg+1)

	for i := range res {
		var fi, gi complex128

		if i >= len(f) {
			fi = 0
		} else {
			fi = f[i]
		}

		if i >= len(g) {
			gi = 0
		} else {
			gi = g[i]
		}

		res[i] = fi + gi
	}

	return poly(res)
}

func (f poly) Negative() poly {
	res := make([]complex128, len(f))

	for i := range res {
		res[i] = -f[i]
	}

	return res
}

func (f poly) Subtract(g poly) poly {
	return f.Add(g.Negative())
}

func (f poly) Multiply(g poly) poly {
	deg := f.Degree() + g.Degree()

	res := make([]complex128, deg+1)

	for i := range res {
		for j := 0; j <= i; j++ {
			var fj, gi_j complex128

			if j >= len(f) {
				fj = 0
			} else {
				fj = f[j]
			}

			if i-j >= len(g) {
				gi_j = 0
			} else {
				gi_j = g[i-j]
			}

			res[i] += fj * gi_j
		}
	}

	return res
}

func (f poly) Equal(g poly) bool {
	deg1 := f.Degree()
	deg2 := g.Degree()

	if deg1 != deg2 {
		return false
	}

	f = f.Normalise()
	g = g.Normalise()

	for i := range f {
		if f[i] != g[i] {
			return false
		}
	}

	return true
}

func (f poly) Derivative() poly {
	if f.IsConstant() {
		return poly([]complex128{0})
	}

	res := make([]complex128, f.Degree())

	for i := range res {
		res[i] = complex(float64(i+1), 0) * f[i+1]
	}

	return poly(res)
}

func (f poly) String() string {
	b := new(bytes.Buffer)

	f = f.Normalise()
	fmt.Fprint(b, f[0])

	if len(f) >= 2 {
		if f[1] != 0 {
			fmt.Fprintf(b, " + %vt", f[1])
		}
	} else {
		return b.String()
	}

	for i, z := range f[2:] {
		if z != 0 {
			fmt.Fprintf(b, " + %vt^%v", z, i+2)
		}
	}

	return b.String()
}

func (f poly) Evaluate(z complex128) complex128 {
	var res complex128

	for i, coeff := range f {
		res += coeff * pow(z, i)
	}

	return res
}

// NewPoly returns a poly with coefficients coeffs, starting with the constant term.
func NewPoly(coeffs ...complex128) poly {
	return poly(coeffs).Normalise()
}

// A rat is a rational function.
type rat struct {
	numerator   poly
	denominator poly
}

func (f rat) Evaluate(z complex128) complex128 {
	num := f.numerator.Evaluate(z)
	denom := f.denominator.Evaluate(z)
	return num / denom
}
<|endoftext|>


<commit_before>/*
Copyright 2013 Google 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,
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 encrypt

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"time"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/jsonconfig"
)

// TODO:
// http://godoc.org/code.google.com/p/go.crypto/scrypt
// crypto/aes
// index.Storage (initially: memindex) for all metadata.

/*
TODO: decide meta format. One argument is to stick with JSON, like
option (a) below. But that means we can't easily read it incrementally
during enumerate, which argues for a line-based format, with a magic
header.

Option a)
{"camliVersion": 1,
"camliType": "encryptedMeta",
"encryptedBlobs": [
  ["sha1-plainplainplain", "sha1-encencencenc", "iviviviviviviv" (%xx)]
],
}

*/

type storage struct {
	*blobserver.SimpleBlobHubPartitionMap

	// Encryption key.
	key []byte

	// blobs holds encrypted versions of all plaintext blobs.
	blobs blobserver.Storage

	// meta holds metadata mapping between the names of plaintext
	// blobs and their after-encryption name, as well as their
	// IV. Each blob in meta contains 1 or more blob
	// description. All new insertions generate both a new
	// encrypted blob in 'blobs' and one single-meta blob in
	// 'meta'. The small metadata blobs are occasionally rolled up
	// into bigger blobs with multiple blob descriptions.
	//
	// TODO: which IV are these written with? safe to just use key
	// with a zero IV? ask experts.
	meta blobserver.Storage
}

func (s *storage) RemoveBlobs(blobs []*blobref.BlobRef) error {
	panic("TODO: implement")
}

func (s *storage) StatBlobs(dest chan<- blobref.SizedBlobRef, blobs []*blobref.BlobRef, wait time.Duration) error {
	panic("TODO: implement")
}

func (s *storage) ReceiveBlob(b *blobref.BlobRef, source io.Reader) (sb blobref.SizedBlobRef, err error) {
	panic("TODO: implement")
}

func (s *storage) FetchStreaming(b *blobref.BlobRef) (file io.ReadCloser, size int64, err error) {
	panic("TODO: implement")
}

func (s *storage) EnumerateBlobs(dest chan<- blobref.SizedBlobRef, after string, limit int, wait time.Duration) error {
	panic("TODO: implement")
}

func init() {
	blobserver.RegisterStorageConstructor("encrypt", blobserver.StorageConstructor(newFromConfig))
}

func newFromConfig(ld blobserver.Loader, config jsonconfig.Obj) (bs blobserver.Storage, err error) {
	sto := &storage{
		SimpleBlobHubPartitionMap: &blobserver.SimpleBlobHubPartitionMap{},
	}

	key := config.OptionalString("key", "")
	keyFile := config.OptionalString("keyFile", "")
	switch {
	case key != "":
		sto.key, err = hex.DecodeString(key)
		if err != nil || len(sto.key) != 16 {
			return nil, fmt.Errorf("The 'key' parameter must be 16 bytes of 32 hex digits. (currently fixed at AES-128")
		}
	case keyFile != "":
		// TODO: check that keyFile's unix permissions aren't too permissive.
		sto.key, err = ioutil.ReadFile(keyFile)
		if err != nil {
			return nil, fmt.Errorf("Reading key file %v: %v", keyFile, err)
		}
	}
	blobStorage := config.RequiredString("blobs")
	metaStorage := config.RequiredString("meta")
	if err := config.Validate(); err != nil {
		return nil, err
	}

	sto.blobs, err = ld.GetStorage(blobStorage)
	if err != nil {
		return
	}
	sto.meta, err = ld.GetStorage(metaStorage)
	if err != nil {
		return
	}
	if sto.key == nil {
		// TODO: add a way to prompt from stdin on start? or keychain support?
		return nil, errors.New("no encryption key set with 'key' or 'keyFile'")
	}
	return sto, nil
}
<commit_msg>encrypt: write encrypted blobs.<commit_after>/*
Copyright 2013 Google 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,
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 encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"time"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/jsonconfig"
)

// TODO:
// http://godoc.org/code.google.com/p/go.crypto/scrypt
// crypto/aes
// index.Storage (initially: memindex) for all metadata.

/*
TODO: decide meta format. One argument is to stick with JSON, like
option (a) below. But that means we can't easily read it incrementally
during enumerate, which argues for a line-based format, with a magic
header.

Option a)
{"camliVersion": 1,
"camliType": "encryptedMeta",
"encryptedBlobs": [
  ["sha1-plainplainplain", "sha1-encencencenc", "iviviviviviviv" (%xx)]
],
}

*/

type storage struct {
	*blobserver.SimpleBlobHubPartitionMap

	block cipher.Block

	// Encryption key.
	key []byte

	// blobs holds encrypted versions of all plaintext blobs.
	blobs blobserver.Storage

	// meta holds metadata mapping between the names of plaintext
	// blobs and their after-encryption name, as well as their
	// IV. Each blob in meta contains 1 or more blob
	// description. All new insertions generate both a new
	// encrypted blob in 'blobs' and one single-meta blob in
	// 'meta'. The small metadata blobs are occasionally rolled up
	// into bigger blobs with multiple blob descriptions.
	//
	// TODO: which IV are these written with? safe to just use key
	// with a zero IV? ask experts.
	meta blobserver.Storage
}

func (s *storage) randIV() []byte {
	iv := make([]byte, s.block.BlockSize())
	n, err := rand.Read(iv)
	if err != nil {
		panic(err)
	}
	if n != len(iv) {
		panic("short read from crypto/rand")
	}
	return iv
}

func (s *storage) RemoveBlobs(blobs []*blobref.BlobRef) error {
	panic("TODO: implement")
}

func (s *storage) StatBlobs(dest chan<- blobref.SizedBlobRef, blobs []*blobref.BlobRef, wait time.Duration) error {
	panic("TODO: implement")
}

func (s *storage) ReceiveBlob(b *blobref.BlobRef, source io.Reader) (sb blobref.SizedBlobRef, err error) {
	iv := s.randIV()
	stream := cipher.NewCTR(s.block, iv)

	var buf bytes.Buffer
	buf.Write(iv) // TODO: write more structured header w/ version & IV length? or does that weaken it?
	sw := cipher.StreamWriter{S: stream, W: &buf}
	n, err := io.Copy(sw, source)
	if err != nil {
		return sb, err
	}
	if err := sw.Close(); err != nil {
		return sb, err
	}

	// TODO: upload buf.Bytes() to s.blobs
	// TODO: upload meta blob with two blobrefs & IV to s.meta
	// TODO: update index with mapping

	return blobref.SizedBlobRef{b, n}, nil
}

func (s *storage) FetchStreaming(b *blobref.BlobRef) (file io.ReadCloser, size int64, err error) {
	panic("TODO: implement")
}

func (s *storage) EnumerateBlobs(dest chan<- blobref.SizedBlobRef, after string, limit int, wait time.Duration) error {
	panic("TODO: implement")
}

func init() {
	blobserver.RegisterStorageConstructor("encrypt", blobserver.StorageConstructor(newFromConfig))
}

func newFromConfig(ld blobserver.Loader, config jsonconfig.Obj) (bs blobserver.Storage, err error) {
	sto := &storage{
		SimpleBlobHubPartitionMap: &blobserver.SimpleBlobHubPartitionMap{},
	}

	key := config.OptionalString("key", "")
	keyFile := config.OptionalString("keyFile", "")
	switch {
	case key != "":
		sto.key, err = hex.DecodeString(key)
		if err != nil || len(sto.key) != 16 {
			return nil, fmt.Errorf("The 'key' parameter must be 16 bytes of 32 hex digits. (currently fixed at AES-128)")
		}
	case keyFile != "":
		// TODO: check that keyFile's unix permissions aren't too permissive.
		sto.key, err = ioutil.ReadFile(keyFile)
		if err != nil {
			return nil, fmt.Errorf("Reading key file %v: %v", keyFile, err)
		}
	}
	blobStorage := config.RequiredString("blobs")
	metaStorage := config.RequiredString("meta")
	if err := config.Validate(); err != nil {
		return nil, err
	}

	sto.blobs, err = ld.GetStorage(blobStorage)
	if err != nil {
		return
	}
	sto.meta, err = ld.GetStorage(metaStorage)
	if err != nil {
		return
	}
	if sto.key == nil {
		// TODO: add a way to prompt from stdin on start? or keychain support?
		return nil, errors.New("no encryption key set with 'key' or 'keyFile'")
	}
	sto.block, err = aes.NewCipher(sto.key)
	if err != nil {
		return nil, fmt.Errorf("The key must be exactly 16 bytes (currently only AES-128 is supported): %v", err)
	}
	return sto, nil
}
<|endoftext|>


<commit_before>package logging

import (
	"context"
	"os"
	"path/filepath"

	"github.com/telepresenceio/telepresence/v2/pkg/client"

	"github.com/sirupsen/logrus"
	"golang.org/x/term"

	"github.com/datawire/dlib/dlog"
	"github.com/telepresenceio/telepresence/v2/pkg/filelocation"
)

// IsTerminal returns whether the given file descriptor is a terminal
var IsTerminal = term.IsTerminal

// loggerForTest exposes internals to initcontext_test.go
var loggerForTest *logrus.Logger

// InitContext sets up standard Telepresence logging for a background process
func InitContext(ctx context.Context, name string) (context.Context, error) {
	logger := logrus.New()
	loggerForTest = logger
	logLevels := client.GetConfig(ctx).LogLevels
	if name == "daemon" {
		logger.SetLevel(logLevels.RootDaemon)
	} else if name == "connector" {
		logger.SetLevel(logLevels.UserDaemon)
	}
	logger.ReportCaller = true

	if IsTerminal(int(os.Stdout.Fd())) {
		logger.Formatter = NewFormatter("15:04:05.0000")
	} else {
		logger.Formatter = NewFormatter("2006/01/02 15:04:05.0000")
		dir, err := filelocation.AppUserLogDir(ctx)
		if err != nil {
			return ctx, err
		}
		rf, err := OpenRotatingFile(filepath.Join(dir, name+".log"), "20060102T150405", true, true, 0600, NewRotateOnce(), 5)
		if err != nil {
			return ctx, err
		}
		logger.SetOutput(rf)
	}
	return dlog.WithLogger(ctx, dlog.WrapLogrus(logger)), nil
}
<commit_msg>Have logging enabled when reading config<commit_after>package logging

import (
	"context"
	"os"
	"path/filepath"

	"github.com/telepresenceio/telepresence/v2/pkg/client"

	"github.com/sirupsen/logrus"
	"golang.org/x/term"

	"github.com/datawire/dlib/dlog"
	"github.com/telepresenceio/telepresence/v2/pkg/filelocation"
)

// IsTerminal returns whether the given file descriptor is a terminal
var IsTerminal = term.IsTerminal

// loggerForTest exposes internals to initcontext_test.go
var loggerForTest *logrus.Logger

// InitContext sets up standard Telepresence logging for a background process
func InitContext(ctx context.Context, name string) (context.Context, error) {
	logger := logrus.New()
	loggerForTest = logger

	// Start with DebugLevel so that the config is read using that level
	logger.SetLevel(logrus.DebugLevel)
	logger.ReportCaller = true

	if IsTerminal(int(os.Stdout.Fd())) {
		logger.Formatter = NewFormatter("15:04:05.0000")
	} else {
		logger.Formatter = NewFormatter("2006/01/02 15:04:05.0000")
		dir, err := filelocation.AppUserLogDir(ctx)
		if err != nil {
			return ctx, err
		}
		rf, err := OpenRotatingFile(filepath.Join(dir, name+".log"), "20060102T150405", true, true, 0600, NewRotateOnce(), 5)
		if err != nil {
			return ctx, err
		}
		logger.SetOutput(rf)
	}
	ctx = dlog.WithLogger(ctx, dlog.WrapLogrus(logger))

	// Read the config and set the configured level.
	logLevels := client.GetConfig(ctx).LogLevels
	if name == "daemon" {
		logger.SetLevel(logLevels.RootDaemon)
	} else if name == "connector" {
		logger.SetLevel(logLevels.UserDaemon)
	}
	return ctx, nil
}
<|endoftext|>


<commit_before>package internal

import (
	"net/http"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/pkg/errors"
)

// MaybeRetryRequest is an internal implementation detail of this module. It
// shouldn't be used by users of the geoipupdate Go library. You can use the
// RetryFor field of geoipupdate.Config if you'd like to retry failed requests
// when using the library directly.
func MaybeRetryRequest(c *http.Client, retryFor time.Duration, req *http.Request) (*http.Response, error) {
	if retryFor < 0 {
		return nil, errors.New("negative retry duration")
	}
	if req.Body != nil {
		return nil, errors.New("can't retry requests with bodies")
	}
	exp := backoff.NewExponentialBackOff()
	exp.MaxElapsedTime = retryFor
	var resp *http.Response
	err := backoff.Retry(
		func() error {
			var err error
			resp, err = c.Do(req)
			return errors.Wrap(err, "error performing http request")
		},
		exp,
	)
	return resp, err
}
<commit_msg>Don't wrap nil errors<commit_after>package internal

import (
	"net/http"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/pkg/errors"
)

// MaybeRetryRequest is an internal implementation detail of this module. It
// shouldn't be used by users of the geoipupdate Go library. You can use the
// RetryFor field of geoipupdate.Config if you'd like to retry failed requests
// when using the library directly.
func MaybeRetryRequest(c *http.Client, retryFor time.Duration, req *http.Request) (*http.Response, error) {
	if retryFor < 0 {
		return nil, errors.New("negative retry duration")
	}
	if req.Body != nil {
		return nil, errors.New("can't retry requests with bodies")
	}
	exp := backoff.NewExponentialBackOff()
	exp.MaxElapsedTime = retryFor
	var resp *http.Response
	err := backoff.Retry(
		func() error {
			var err error
			resp, err = c.Do(req)
			if err != nil {
				return errors.Wrap(err, "error performing http request")
			}
			return err
		},
		exp,
	)
	return resp, err
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 generator

import (
	"fmt"
	"strings"

	"k8s.io/kube-state-metrics/v2/pkg/metric"
)

// FamilyGenerator provides everything needed to generate a metric family with a
// Kubernetes object.
// DeprecatedVersion is defined only if the metric for which this options applies is,
// in fact, deprecated.
type FamilyGenerator struct {
	Name              string
	Help              string
	Type              metric.Type
	DeprecatedVersion string
	GenerateFunc      func(obj interface{}) *metric.Family
}

// NewFamilyGenerator creates new FamilyGenerator instances.
func NewFamilyGenerator(name string, help string, metricType metric.Type, deprecatedVersion string, generateFunc func(obj interface{}) *metric.Family) *FamilyGenerator {
	f := &FamilyGenerator{
		Name:              name,
		Type:              metricType,
		Help:              help,
		DeprecatedVersion: deprecatedVersion,
		GenerateFunc:      generateFunc,
	}
	if deprecatedVersion != "" {
		f.Help = fmt.Sprintf("(Deprecated since %s) %s", deprecatedVersion, help)
	}
	return f
}

// Generate calls the FamilyGenerator.GenerateFunc and gives the family its
// name. The reasoning behind injecting the name at such a late point in time is
// deduplication in the code, preventing typos made by developers as
// well as saving memory.
func (g *FamilyGenerator) Generate(obj interface{}) *metric.Family {
	family := g.GenerateFunc(obj)
	family.Name = g.Name
	family.Type = g.Type
	return family
}

func (g *FamilyGenerator) generateHeader() string {
	header := strings.Builder{}
	header.WriteString("# HELP ")
	header.WriteString(g.Name)
	header.WriteByte(' ')
	header.WriteString(g.Help)
	header.WriteByte('\n')
	header.WriteString("# TYPE ")
	header.WriteString(g.Name)
	header.WriteByte(' ')
	header.WriteString(string(g.Type))

	return header.String()
}

// ExtractMetricFamilyHeaders takes in a slice of FamilyGenerator metrics and
// returns the extracted headers.
func ExtractMetricFamilyHeaders(families []FamilyGenerator) []string {
	headers := make([]string, len(families))

	for i, f := range families {
		headers[i] = f.generateHeader()
	}

	return headers
}

// ComposeMetricGenFuncs takes a slice of metric families and returns a function
// that composes their metric generation functions into a single one.
func ComposeMetricGenFuncs(familyGens []FamilyGenerator) func(obj interface{}) []metric.FamilyInterface {
	return func(obj interface{}) []metric.FamilyInterface {
		families := make([]metric.FamilyInterface, len(familyGens))

		for i, gen := range familyGens {
			families[i] = gen.Generate(obj)
		}

		return families
	}
}
<commit_msg>Add an OptIn property to the FamilyGenerator struct to identify an opt-in metric family (#1574)<commit_after>/*
Copyright 2019 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 generator

import (
	"fmt"
	"strings"

	"k8s.io/kube-state-metrics/v2/pkg/metric"
)

// FamilyGenerator provides everything needed to generate a metric family with a
// Kubernetes object.
// DeprecatedVersion is defined only if the metric for which this options applies is,
// in fact, deprecated.
type FamilyGenerator struct {
	Name              string
	Help              string
	Type              metric.Type
	OptIn             bool
	DeprecatedVersion string
	GenerateFunc      func(obj interface{}) *metric.Family
}

// NewFamilyGenerator creates new FamilyGenerator instances.
func NewFamilyGenerator(name string, help string, metricType metric.Type, deprecatedVersion string, generateFunc func(obj interface{}) *metric.Family) *FamilyGenerator {
	f := &FamilyGenerator{
		Name:              name,
		Type:              metricType,
		Help:              help,
		OptIn:             false,
		DeprecatedVersion: deprecatedVersion,
		GenerateFunc:      generateFunc,
	}
	if deprecatedVersion != "" {
		f.Help = fmt.Sprintf("(Deprecated since %s) %s", deprecatedVersion, help)
	}
	return f
}

// NewOptInFamilyGenerator creates new FamilyGenerator instances for opt-in metric families.
func NewOptInFamilyGenerator(name string, help string, metricType metric.Type, deprecatedVersion string, generateFunc func(obj interface{}) *metric.Family) *FamilyGenerator {
	f := NewFamilyGenerator(name, help, metricType, deprecatedVersion, generateFunc)
	f.OptIn = true
	return f
}

// Generate calls the FamilyGenerator.GenerateFunc and gives the family its
// name. The reasoning behind injecting the name at such a late point in time is
// deduplication in the code, preventing typos made by developers as
// well as saving memory.
func (g *FamilyGenerator) Generate(obj interface{}) *metric.Family {
	family := g.GenerateFunc(obj)
	family.Name = g.Name
	family.Type = g.Type
	return family
}

func (g *FamilyGenerator) generateHeader() string {
	header := strings.Builder{}
	header.WriteString("# HELP ")
	header.WriteString(g.Name)
	header.WriteByte(' ')
	header.WriteString(g.Help)
	header.WriteByte('\n')
	header.WriteString("# TYPE ")
	header.WriteString(g.Name)
	header.WriteByte(' ')
	header.WriteString(string(g.Type))

	return header.String()
}

// ExtractMetricFamilyHeaders takes in a slice of FamilyGenerator metrics and
// returns the extracted headers.
func ExtractMetricFamilyHeaders(families []FamilyGenerator) []string {
	headers := make([]string, len(families))

	for i, f := range families {
		headers[i] = f.generateHeader()
	}

	return headers
}

// ComposeMetricGenFuncs takes a slice of metric families and returns a function
// that composes their metric generation functions into a single one.
func ComposeMetricGenFuncs(familyGens []FamilyGenerator) func(obj interface{}) []metric.FamilyInterface {
	return func(obj interface{}) []metric.FamilyInterface {
		families := make([]metric.FamilyInterface, len(familyGens))

		for i, gen := range familyGens {
			families[i] = gen.Generate(obj)
		}

		return families
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes 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 components

import (
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// OpenStackOptionsBulder adds options for OpenStack to the model
type OpenStackOptionsBulder struct {
	Context *OptionsContext
}

var _ loader.OptionsBuilder = &OpenStackOptionsBulder{}

func (b *OpenStackOptionsBulder) BuildOptions(o interface{}) error {

	clusterSpec := o.(*kops.ClusterSpec)

	if kops.CloudProviderID(clusterSpec.CloudProvider) != kops.CloudProviderOpenstack {
		return nil
	}

	if clusterSpec.CloudConfig == nil {
		clusterSpec.CloudConfig = &kops.CloudConfiguration{}
	}

	if clusterSpec.CloudConfig.Openstack == nil {
		clusterSpec.CloudConfig.Openstack = &kops.OpenstackConfiguration{}
	}
	if clusterSpec.CloudConfig.Openstack.BlockStorage == nil {
		clusterSpec.CloudConfig.Openstack.BlockStorage = &kops.OpenstackBlockStorageConfig{}
	}

	if clusterSpec.CloudConfig.Openstack.BlockStorage.CreateStorageClass == nil {
		clusterSpec.CloudConfig.Openstack.BlockStorage.CreateStorageClass = fi.Bool(true)
	}
	return nil

}
<commit_msg>Force external cloud controller manager on OS<commit_after>/*
Copyright 2020 The Kubernetes 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 components

import (
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// OpenStackOptionsBulder adds options for OpenStack to the model
type OpenStackOptionsBulder struct {
	Context *OptionsContext
}

var _ loader.OptionsBuilder = &OpenStackOptionsBulder{}

func (b *OpenStackOptionsBulder) BuildOptions(o interface{}) error {

	clusterSpec := o.(*kops.ClusterSpec)

	if kops.CloudProviderID(clusterSpec.CloudProvider) != kops.CloudProviderOpenstack {
		return nil
	}

	if clusterSpec.CloudConfig == nil {
		clusterSpec.CloudConfig = &kops.CloudConfiguration{}
	}

	if clusterSpec.CloudConfig.Openstack == nil {
		clusterSpec.CloudConfig.Openstack = &kops.OpenstackConfiguration{}
	}
	if clusterSpec.CloudConfig.Openstack.BlockStorage == nil {
		clusterSpec.CloudConfig.Openstack.BlockStorage = &kops.OpenstackBlockStorageConfig{}
	}

	if clusterSpec.CloudConfig.Openstack.BlockStorage.CreateStorageClass == nil {
		clusterSpec.CloudConfig.Openstack.BlockStorage.CreateStorageClass = fi.Bool(true)
	}

	if b.Context.IsKubernetesGTE("1.13.0") {
		if clusterSpec.ExternalCloudControllerManager == nil {
			clusterSpec.ExternalCloudControllerManager = &kops.CloudControllerManagerConfig{}
		}
	}
	return nil

}
<|endoftext|>


<commit_before>package staticlib

import (
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"mime"
	"path/filepath"
)

type pipelineAction func(file File) File
type pipelineChan <-chan File

func pipelineStage(in pipelineChan, f pipelineAction) pipelineChan {
	out := make(chan File)
	go func() {
		for inFile := range in {
			out <- f(inFile)
		}
		close(out)
	}()
	return out
}

func setContentType(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		inFile.contentType = mime.TypeByExtension(filepath.Ext(inFile.path))
		return inFile
	})
}

func setSize(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		body := inFile.Body()
		defer body.Close()
		if fi, err := body.Stat(); err != nil {
			panic(err)
		} else {
			inFile.size = fi.Size()
		}
		return inFile
	})
}

func setCacheControl(cfg Config, in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		maxAge := cfg.MaxAge(inFile.Key())
		inFile.cacheControl = fmt.Sprintf("public; max-age=%d", maxAge)
		return inFile
	})
}

func gzipProcessor(workingDir workingDir, cfg Config, in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		if cfg.ShouldGzip(inFile.Key()) {
			return gzipFile(workingDir, inFile)
		}
		return inFile
	})
}

func gzipFile(workingDir workingDir, in File) File {
	reader := in.Body()
	defer reader.Close()
	writer := workingDir.tempFile()
	defer writer.Close()
	compressor := gzip.NewWriter(writer)
	_, err := io.Copy(compressor, reader)
	if err != nil {
		panic(err)
	}
	compressor.Close()
	out := in
	out.path = writer.Name()
	out.contentEncoding = "gzip"
	// fmt.Printf("compression saved %%%f (%d, %d)\n", float64(out.Size()) / float64(in.Size()), out.Size(), in.Size())
	if in.Size() <= out.Size() {
		return in
	}
	return out
}

func digestProcessor(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		hash := md5.New()
		inBody := inFile.Body()
		defer inBody.Close()
		var result []byte
		if _, err := io.Copy(hash, inBody); err != nil {
			panic(err)
		}
		inFile.digest = hash.Sum(result)
		return inFile
	})
}
<commit_msg>Auto detect contentType by looking at file contents<commit_after>package staticlib

import (
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"net/http"
)

type pipelineAction func(file File) File
type pipelineChan <-chan File

func pipelineStage(in pipelineChan, f pipelineAction) pipelineChan {
	out := make(chan File)
	go func() {
		for inFile := range in {
			out <- f(inFile)
		}
		close(out)
	}()
	return out
}

func setContentType(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		body := inFile.Body()
		defer body.Close()
		buff := make([]byte, 512)
		if _, err := body.Read(buff); err != nil {
			panic(err)
		}
		inFile.contentType = http.DetectContentType(buff)
		return inFile
	})
}

func setSize(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		body := inFile.Body()
		defer body.Close()
		if fi, err := body.Stat(); err != nil {
			panic(err)
		} else {
			inFile.size = fi.Size()
		}
		return inFile
	})
}

func setCacheControl(cfg Config, in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		maxAge := cfg.MaxAge(inFile.Key())
		inFile.cacheControl = fmt.Sprintf("public; max-age=%d", maxAge)
		return inFile
	})
}

func gzipProcessor(workingDir workingDir, cfg Config, in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		if cfg.ShouldGzip(inFile.Key()) {
			return gzipFile(workingDir, inFile)
		}
		return inFile
	})
}

func gzipFile(workingDir workingDir, in File) File {
	reader := in.Body()
	defer reader.Close()
	writer := workingDir.tempFile()
	defer writer.Close()
	compressor := gzip.NewWriter(writer)
	_, err := io.Copy(compressor, reader)
	if err != nil {
		panic(err)
	}
	compressor.Close()
	out := in
	out.path = writer.Name()
	out.contentEncoding = "gzip"
	// fmt.Printf("compression saved %%%f (%d, %d)\n", float64(out.Size()) / float64(in.Size()), out.Size(), in.Size())
	if in.Size() <= out.Size() {
		return in
	}
	return out
}

func digestProcessor(in pipelineChan) pipelineChan {
	return pipelineStage(in, func(inFile File) File {
		hash := md5.New()
		inBody := inFile.Body()
		defer inBody.Close()
		var result []byte
		if _, err := io.Copy(hash, inBody); err != nil {
			panic(err)
		}
		inFile.digest = hash.Sum(result)
		return inFile
	})
}
<|endoftext|>


<commit_before>package reqlimit

import (
	"net"
	"net/http"
	"sync"
	"time"
)

// limiter is a http.Handler that rate-limits a http.Handler on a per-IP basis.
type limiter struct {
	mu sync.Mutex

	requests     map[string](chan struct{})
	limitTimeout time.Duration
	limit        uint64

	nextHandler http.Handler
}

// New creates a new Limiter, limiting each remote host served by the supplied
// handler to `limit` requests over duration of the supplied `timeout`.
func New(h http.Handler, limit uint64, timeout time.Duration) *limiter {
	return &limiter{
		requests:     make(map[string]chan struct{}),
		limitTimeout: timeout,
		limit:        limit,

		nextHandler: h,
	}
}

// Listener's ServeHTTP implements the http.Handler interface and checks if the
// remote host has exceeded the request limit. If it has, it returns a
// http.Error with http.StatusTooManyRequests. Otherwise, the protected handler
// will be called.
func (l *limiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteIP, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// grab the requests channel for this remote host, or create it if it doesn't
	// exist
	l.mu.Lock()
	requests, exists := l.requests[remoteIP]
	if !exists {
		requests = make(chan struct{}, l.limit)
		l.requests[remoteIP] = requests
	}
	l.mu.Unlock()

	// add to the request channel, throw an error if it is currently full.
	select {
	case requests <- struct{}{}:
		// drain the request channel after the limit timeout
		go func() {
			time.Sleep(l.limitTimeout)
			<-requests
		}()

	default:
		http.Error(w, "request limit exceeded", http.StatusTooManyRequests)
		return
	}

	l.nextHandler.ServeHTTP(w, r)
}
<commit_msg>Realistic get ip method - faster locks<commit_after>package reqlimit

import (
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// limiter is a http.Handler that rate-limits a http.Handler on a per-IP basis.
type limiter struct {
	mu       sync.RWMutex // locks for requests
	requests map[string](chan struct{})

	limitTimeout time.Duration
	limit        uint64

	nextHandler http.Handler
}

// New creates a new Limiter, limiting each remote host served by the supplied
// handler to `limit` requests over duration of the supplied `timeout`.
func New(h http.Handler, limit uint64, timeout time.Duration) *limiter {
	return &limiter{
		requests:     make(map[string]chan struct{}),
		limitTimeout: timeout,
		limit:        limit,

		nextHandler: h,
	}
}

// getIP tries to find the real ip address of a client.
func getIP(r *http.Request) (string, error) {
	header := r.Header.Get("X-Real-Ip")
	realIP := strings.TrimSpace(header)
	if realIP != "" {
		return realIP, nil
	}

	realIP = r.Header.Get("X-Forwarded-For")
	idx := strings.IndexByte(realIP, ',')
	if idx >= 0 {
		realIP = realIP[0:idx]
	}
	realIP = strings.TrimSpace(realIP)
	if realIP != "" {
		return realIP, nil
	}

	addr := strings.TrimSpace(r.RemoteAddr)

	// if addr has port use the net.SplitHostPort otherwise(error occurs) take as it is
	ip, _, err := net.SplitHostPort(addr)
	return ip, err
}

// Listener's ServeHTTP implements the http.Handler interface and checks if the
// remote host has exceeded the request limit. If it has, it returns a
// http.Error with http.StatusTooManyRequests. Otherwise, the protected handler
// will be called.
func (l *limiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteIP, err := getIP(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// grab the requests channel for this remote host
	l.mu.RLock() // read lock
	requests, exists := l.requests[remoteIP]
	l.mu.RUnlock()

	// or create it if it doesn't exist
	if !exists {
		requests = make(chan struct{}, l.limit)

		l.mu.Lock() // write & read lock
		l.requests[remoteIP] = requests
		l.mu.Unlock()
	}

	// add to the request channel, throw an error if it is currently full.
	select {
	case requests <- struct{}{}:
		// drain the request channel after the limit timeout
		go func() {
			time.Sleep(l.limitTimeout)
			<-requests
		}()

	default:
		http.Error(w, "request limit exceeded", http.StatusTooManyRequests)
		return
	}

	l.nextHandler.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2017 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cloud

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

const (
	// RequestTimeout is the default cloud request timeout
	RequestTimeout = 20 * time.Second
	// RetryInterval is the default cloud request retry interval
	RetryInterval = 500 * time.Millisecond
	// MaxRetries specifies max retry attempts
	MaxRetries = 3

	k6IdempotencyKeyHeader = "k6-Idempotency-Key"
)

// Client handles communication with Load Impact cloud API.
type Client struct {
	client  *http.Client
	token   string
	baseURL string
	version string

	retries       int
	retryInterval time.Duration
}

func NewClient(token, host, version string) *Client {
	c := &Client{
		client:        &http.Client{Timeout: RequestTimeout},
		token:         token,
		baseURL:       fmt.Sprintf("%s/v1", host),
		version:       version,
		retries:       MaxRetries,
		retryInterval: RetryInterval,
	}
	return c
}

func (c *Client) NewRequest(method, url string, data interface{}) (*http.Request, error) {
	var buf io.Reader

	if data != nil {
		b, err := json.Marshal(&data)
		if err != nil {
			return nil, err
		}

		buf = bytes.NewBuffer(b)
	}

	req, err := http.NewRequest(method, url, buf)
	if err != nil {
		return nil, err
	}

	if shouldAddIdempotencyKey(req) {
		req.Header.Set(k6IdempotencyKeyHeader, randomStrHex())
	}

	return req, nil
}

func (c *Client) Do(req *http.Request, v interface{}) error {
	var originalBody []byte
	var err error

	if req.Body != nil {
		originalBody, err = ioutil.ReadAll(req.Body)
		if err != nil {
			return err
		}

		if cerr := req.Body.Close(); cerr != nil {
			err = cerr
		}
	}

	c.prepareHeader(req)

	for i := 1; i <= c.retries; i++ {
		if len(originalBody) > 0 {
			req.Body = ioutil.NopCloser(bytes.NewBuffer(originalBody))
		}

		retry, err := c.do(req, v, i)

		if retry {
			time.Sleep(c.retryInterval)
			continue
		}

		return err
	}

	return err
}

func (c *Client) prepareHeader(req *http.Request) {
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.token != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.token))
	}
	req.Header.Set("User-Agent", "k6cloud/"+c.version)
}

func (c *Client) do(req *http.Request, v interface{}, attempt int) (retry bool, err error) {
	resp, err := c.client.Do(req)

	defer func() {
		if resp != nil {
			if cerr := resp.Body.Close(); cerr != nil && err == nil {
				err = cerr
			}
		}
	}()

	if shouldRetry(resp, err, attempt, c.retries) {
		return true, err
	}

	if err != nil {
		return false, err
	}

	if err = checkResponse(resp); err != nil {
		return false, err
	}

	if v != nil {
		if err = json.NewDecoder(resp.Body).Decode(v); err == io.EOF {
			err = nil // Ignore EOF from empty body
		}
	}

	return false, err
}

func checkResponse(r *http.Response) error {
	if r == nil {
		return ErrUnknown
	}

	if c := r.StatusCode; c >= 200 && c <= 299 {
		return nil
	}

	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	var payload struct {
		Error ErrorResponse `json:"error"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		if r.StatusCode == http.StatusUnauthorized {
			return ErrNotAuthenticated
		}
		if r.StatusCode == http.StatusForbidden {
			return ErrNotAuthorized
		}
		return errors.Errorf(
			"Unexpected HTTP error from %s: %d %s",
			r.Request.URL,
			r.StatusCode,
			http.StatusText(r.StatusCode),
		)
	}
	payload.Error.Response = r
	return payload.Error
}

func shouldRetry(resp *http.Response, err error, attempt, maxAttempts int) bool {
	if attempt >= maxAttempts {
		return false
	}

	if resp == nil || err != nil {
		return true
	}

	if resp.StatusCode >= 500 || resp.StatusCode == 429 {
		return true
	}

	return false
}

func shouldAddIdempotencyKey(req *http.Request) bool {
	switch req.Method {
	case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace:
		return false
	default:
		return req.Header.Get(k6IdempotencyKeyHeader) == ""
	}
}

// randomStrHex returns a hex string which can be used
// for session token id or idempotency key.
//nolint:gosec
func randomStrHex() string {
	// 16 hex characters
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}
<commit_msg>stats/cloud: prepareHeader -> prepareHeaders<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2017 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cloud

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

const (
	// RequestTimeout is the default cloud request timeout
	RequestTimeout = 20 * time.Second
	// RetryInterval is the default cloud request retry interval
	RetryInterval = 500 * time.Millisecond
	// MaxRetries specifies max retry attempts
	MaxRetries = 3

	k6IdempotencyKeyHeader = "k6-Idempotency-Key"
)

// Client handles communication with Load Impact cloud API.
type Client struct {
	client  *http.Client
	token   string
	baseURL string
	version string

	retries       int
	retryInterval time.Duration
}

func NewClient(token, host, version string) *Client {
	c := &Client{
		client:        &http.Client{Timeout: RequestTimeout},
		token:         token,
		baseURL:       fmt.Sprintf("%s/v1", host),
		version:       version,
		retries:       MaxRetries,
		retryInterval: RetryInterval,
	}
	return c
}

func (c *Client) NewRequest(method, url string, data interface{}) (*http.Request, error) {
	var buf io.Reader

	if data != nil {
		b, err := json.Marshal(&data)
		if err != nil {
			return nil, err
		}

		buf = bytes.NewBuffer(b)
	}

	req, err := http.NewRequest(method, url, buf)
	if err != nil {
		return nil, err
	}

	if shouldAddIdempotencyKey(req) {
		req.Header.Set(k6IdempotencyKeyHeader, randomStrHex())
	}

	return req, nil
}

func (c *Client) Do(req *http.Request, v interface{}) error {
	var originalBody []byte
	var err error

	if req.Body != nil {
		originalBody, err = ioutil.ReadAll(req.Body)
		if err != nil {
			return err
		}

		if cerr := req.Body.Close(); cerr != nil {
			err = cerr
		}
	}

	c.prepareHeaders(req)

	for i := 1; i <= c.retries; i++ {
		if len(originalBody) > 0 {
			req.Body = ioutil.NopCloser(bytes.NewBuffer(originalBody))
		}

		retry, err := c.do(req, v, i)

		if retry {
			time.Sleep(c.retryInterval)
			continue
		}

		return err
	}

	return err
}

func (c *Client) prepareHeaders(req *http.Request) {
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.token != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Token %s", c.token))
	}
	req.Header.Set("User-Agent", "k6cloud/"+c.version)
}

func (c *Client) do(req *http.Request, v interface{}, attempt int) (retry bool, err error) {
	resp, err := c.client.Do(req)

	defer func() {
		if resp != nil {
			if cerr := resp.Body.Close(); cerr != nil && err == nil {
				err = cerr
			}
		}
	}()

	if shouldRetry(resp, err, attempt, c.retries) {
		return true, err
	}

	if err != nil {
		return false, err
	}

	if err = checkResponse(resp); err != nil {
		return false, err
	}

	if v != nil {
		if err = json.NewDecoder(resp.Body).Decode(v); err == io.EOF {
			err = nil // Ignore EOF from empty body
		}
	}

	return false, err
}

func checkResponse(r *http.Response) error {
	if r == nil {
		return ErrUnknown
	}

	if c := r.StatusCode; c >= 200 && c <= 299 {
		return nil
	}

	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	var payload struct {
		Error ErrorResponse `json:"error"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		if r.StatusCode == http.StatusUnauthorized {
			return ErrNotAuthenticated
		}
		if r.StatusCode == http.StatusForbidden {
			return ErrNotAuthorized
		}
		return errors.Errorf(
			"Unexpected HTTP error from %s: %d %s",
			r.Request.URL,
			r.StatusCode,
			http.StatusText(r.StatusCode),
		)
	}
	payload.Error.Response = r
	return payload.Error
}

func shouldRetry(resp *http.Response, err error, attempt, maxAttempts int) bool {
	if attempt >= maxAttempts {
		return false
	}

	if resp == nil || err != nil {
		return true
	}

	if resp.StatusCode >= 500 || resp.StatusCode == 429 {
		return true
	}

	return false
}

func shouldAddIdempotencyKey(req *http.Request) bool {
	switch req.Method {
	case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace:
		return false
	default:
		return req.Header.Get(k6IdempotencyKeyHeader) == ""
	}
}

// randomStrHex returns a hex string which can be used
// for session token id or idempotency key.
//nolint:gosec
func randomStrHex() string {
	// 16 hex characters
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}
<|endoftext|>


<commit_before>package kepub

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/html"

	"github.com/PuerkitoBio/goquery"
	"github.com/beevik/etree"
)

func cleanFiles(basepath string) error {
	_ = os.Remove(filepath.Join(basepath, "META-INF", "calibre_bookmarks.txt"))
	_ = os.Remove(filepath.Join(basepath, "META-INF", "iTunesMetadata.plist"))
	_ = os.Remove(filepath.Join(basepath, "iTunesMetadata.plist"))
	_ = os.Remove(filepath.Join(basepath, "META-INF", "iTunesArtwork.plist"))
	_ = os.Remove(filepath.Join(basepath, "iTunesArtwork.plist"))
	_ = os.Remove(filepath.Join(basepath, "META-INF", ".DS_STORE"))
	_ = os.Remove(filepath.Join(basepath, ".DS_STORE"))
	_ = os.Remove(filepath.Join(basepath, "META-INF", "thumbs.db"))
	_ = os.Remove(filepath.Join(basepath, "thumbs.db"))
	return nil
}

// processOPF cleans up extra calibre metadata from the content.opf file, and adds a reference to the cover image.
func processOPF(opftext *string) error {
	opf := etree.NewDocument()
	err := opf.ReadFromString(*opftext)
	if err != nil {
		return err
	}

	for _, e := range opf.FindElements("//meta[@name='cover']") {
		coverid := e.SelectAttrValue("content", "")
		if coverid == "" {
			coverid = "cover"
		}
		for _, f := range opf.FindElements("//[@id='" + coverid + "']") {
			f.CreateAttr("properties", "cover-image")
		}
	}

	for _, e := range opf.FindElements("//meta[@name='calibre:timestamp']") {
		e.Parent().RemoveChild(e)
	}

	for _, e := range opf.FindElements("//dc:contributor[@role='bkp']") {
		e.Parent().RemoveChild(e)
	}

	opf.Indent(4)

	*opftext, err = opf.WriteToString()
	if err != nil {
		return err
	}

	return nil
}

// addDivs adds kobo divs.
func addDivs(doc *goquery.Document) error {
	if len(doc.Find("div").Nodes) > len(doc.Find("p").Nodes) {
		// If there are more divs than ps, divs are probably being used as paragraphs, and adding the kobo divs will most likely break the book.
		return nil
	}
	doc.Find("body>*").WrapAllHtml(`<div class="book-inner"></div>`)
	doc.Find("body>*").WrapAllHtml(`<div class="book-columns"></div>`)
	return nil
}

// createSpan creates a Kobo span
func createSpan(paragraph, segment int, text string) *html.Node {
	span := &html.Node{
		Type: html.ElementNode,
		Data: "span",
		Attr: []html.Attribute{
			html.Attribute{
				Key: "class",
				Val: "koboSpan",
			},
			html.Attribute{
				Key: "id",
				Val: fmt.Sprintf("kobo.%v.%v", paragraph, segment),
			},
		},
	}

	span.AppendChild(&html.Node{
		Type: html.TextNode,
		Data: text,
	})

	return span
}

// addSpansToNode is a recursive helper function for addSpans.
func addSpansToNode(node *html.Node, paragraph *int, segment *int) {
	sentencere := regexp.MustCompile(`((?ms).*?[\.\!\?\:]['"”’“…]?\s*)`)

	nextNodes := []*html.Node{}
	for c := node.FirstChild; c != nil; c = c.NextSibling {
		nextNodes = append(nextNodes, c)
	}

	if node.Type == html.TextNode {
		if node.Parent.Data == "pre" {
			// Do not add spans to pre elements
			return
		}
		*segment++

		sentencesindexes := sentencere.FindAllStringIndex(node.Data, -1)
		sentences := []string{}
		lasti := []int{0, 0}
		for _, i := range sentencesindexes {
			if lasti[1] != i[0] {
				// If gap in regex matches, add the gap to the sentence list to avoid losing text
				sentences = append(sentences, node.Data[lasti[1]:i[0]])
			}
			sentences = append(sentences, node.Data[i[0]:i[1]])
			lasti = i
		}
		if lasti[1] != len(node.Data) {
			// If gap in regex matches, add the gap to the sentence list to avoid losing text
			sentences = append(sentences, node.Data[lasti[1]:len(node.Data)])
		}

		for _, sentence := range sentences {
			if strings.TrimSpace(sentence) != "" {
				node.Parent.InsertBefore(createSpan(*paragraph, *segment, sentence), node)
				*segment++
			}
		}

		node.Parent.RemoveChild(node)

		return
	}

	if node.Type != html.ElementNode {
		return
	}
	if node.Data == "img" {
		return
	}
	if node.Data == "p" || node.Data == "ol" || node.Data == "ul" {
		*segment = 0
		*paragraph++
	}

	for _, c := range nextNodes {
		addSpansToNode(c, paragraph, segment)
	}
}

// addSpans adds kobo spans.
func addSpans(doc *goquery.Document) error {
	alreadyHasSpans := false
	doc.Find("span").Each(func(i int, s *goquery.Selection) {
		if val, _ := s.Attr("class"); strings.Contains(val, "koboSpan") {
			alreadyHasSpans = true
		}
	})
	if alreadyHasSpans {
		return nil
	}

	paragraph := 0
	segment := 0

	for _, n := range doc.Find("body").Nodes {
		addSpansToNode(n, &paragraph, &segment)
	}

	return nil
}

// addKoboStyles adds kobo styles.
func addKoboStyles(doc *goquery.Document) error {
	s := doc.Find("head").First().AppendHtml(`<style type="text/css">div#book-inner{margin-top: 0;margin-bottom: 0;}</style>`)
	if s.Length() != 1 {
		return fmt.Errorf("could not append kobo styles")
	}
	return nil
}

// smartenPunctuation smartens punctuation in html code. It must be run last.
func smartenPunctuation(html *string) error {
	// em and en dashes
	*html = strings.Replace(*html, "---", " – ", -1)
	*html = strings.Replace(*html, "--", " — ", -1)

	// TODO: smart quotes

	// Fix comments
	*html = strings.Replace(*html, "<! — ", "<!-- ", -1)
	*html = strings.Replace(*html, " — >", " -->", -1)
	return nil
}

// cleanHTML cleans up html for a kobo epub.
func cleanHTML(doc *goquery.Document) error {
	// Remove Adobe DRM tags
	doc.Find(`meta[name="Adept.expected.resource"]`).Remove()

	// Remove empty MS <o:p> tags
	doc.Find(`o\:p`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.Trim(s.Text(), "\t \n") == ""
	}).Remove()

	// Remove empty headings
	doc.Find(`h1,h2,h3,h4,h5,h6`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.Trim(s.Text(), "\t \n") == ""
	}).Remove()

	// Remove MS <st1:whatever> tags
	doc.Find(`*`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.HasPrefix(goquery.NodeName(s), "st1:")
	}).Remove()

	// Open self closing p tags
	doc.Find(`p`).Each(func(_ int, s *goquery.Selection) {
		if s.Children().Length() == 0 && strings.Trim(s.Text(), "\n \t") == "" {
			s.SetHtml("")
		}
	})

	// Add type to style tags
	doc.Find(`style`).SetAttr("type", "text/css")

	return nil
}

// process processes the html of a content file in an ordinary epub and converts it into a kobo epub by adding kobo divs, kobo spans, smartening punctuation, and cleaning html.
func process(content *string) error {
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(*content))
	if err != nil {
		return err
	}

	if err := addDivs(doc); err != nil {
		return err
	}

	if err := addSpans(doc); err != nil {
		return err
	}

	if err := addKoboStyles(doc); err != nil {
		return err
	}

	if err := cleanHTML(doc); err != nil {
		return err
	}

	h, err := doc.Html()
	if err != nil {
		return err
	}

	if err := smartenPunctuation(&h); err != nil {
		return err
	}

	// Remove unicode replacement chars
	h = strings.Replace(h, "�", "", -1)

	// Fix commented xml tag
	h = strings.Replace(h, `<!-- ?xml version="1.0" encoding="utf-8"? -->`, `<?xml version="1.0" encoding="utf-8"?>`, 1)
	h = strings.Replace(h, `<!--?xml version="1.0" encoding="utf-8"?-->`, `<?xml version="1.0" encoding="utf-8"?>`, 1)

	*content = h

	return nil
}
<commit_msg>Refactored epub file cleanup<commit_after>package kepub

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/html"

	"github.com/PuerkitoBio/goquery"
	"github.com/beevik/etree"
)

func cleanFiles(basepath string) error {
	toRemove := []string{
		"META-INF/calibre_bookmarks.txt",
		"META-INF/iTunesMetadata.plist",
		"META-INF/iTunesArtwork.plist",
		"META-INF/.DS_STORE",
		"META-INF/thumbs.db",
		".DS_STORE",
		"thumbs.db",
		"iTunesMetadata.plist",
		"iTunesArtwork.plist",
	}

	for _, file := range toRemove {
		os.Remove(filepath.Join(basepath, file))
	}

	return nil
}

// processOPF cleans up extra calibre metadata from the content.opf file, and adds a reference to the cover image.
func processOPF(opftext *string) error {
	opf := etree.NewDocument()
	err := opf.ReadFromString(*opftext)
	if err != nil {
		return err
	}

	for _, e := range opf.FindElements("//meta[@name='cover']") {
		coverid := e.SelectAttrValue("content", "")
		if coverid == "" {
			coverid = "cover"
		}
		for _, f := range opf.FindElements("//[@id='" + coverid + "']") {
			f.CreateAttr("properties", "cover-image")
		}
	}

	for _, e := range opf.FindElements("//meta[@name='calibre:timestamp']") {
		e.Parent().RemoveChild(e)
	}

	for _, e := range opf.FindElements("//dc:contributor[@role='bkp']") {
		e.Parent().RemoveChild(e)
	}

	opf.Indent(4)

	*opftext, err = opf.WriteToString()
	if err != nil {
		return err
	}

	return nil
}

// addDivs adds kobo divs.
func addDivs(doc *goquery.Document) error {
	if len(doc.Find("div").Nodes) > len(doc.Find("p").Nodes) {
		// If there are more divs than ps, divs are probably being used as paragraphs, and adding the kobo divs will most likely break the book.
		return nil
	}
	doc.Find("body>*").WrapAllHtml(`<div class="book-inner"></div>`)
	doc.Find("body>*").WrapAllHtml(`<div class="book-columns"></div>`)
	return nil
}

// createSpan creates a Kobo span
func createSpan(paragraph, segment int, text string) *html.Node {
	span := &html.Node{
		Type: html.ElementNode,
		Data: "span",
		Attr: []html.Attribute{
			html.Attribute{
				Key: "class",
				Val: "koboSpan",
			},
			html.Attribute{
				Key: "id",
				Val: fmt.Sprintf("kobo.%v.%v", paragraph, segment),
			},
		},
	}

	span.AppendChild(&html.Node{
		Type: html.TextNode,
		Data: text,
	})

	return span
}

// addSpansToNode is a recursive helper function for addSpans.
func addSpansToNode(node *html.Node, paragraph *int, segment *int) {
	sentencere := regexp.MustCompile(`((?ms).*?[\.\!\?\:]['"”’“…]?\s*)`)

	nextNodes := []*html.Node{}
	for c := node.FirstChild; c != nil; c = c.NextSibling {
		nextNodes = append(nextNodes, c)
	}

	if node.Type == html.TextNode {
		if node.Parent.Data == "pre" {
			// Do not add spans to pre elements
			return
		}
		*segment++

		sentencesindexes := sentencere.FindAllStringIndex(node.Data, -1)
		sentences := []string{}
		lasti := []int{0, 0}
		for _, i := range sentencesindexes {
			if lasti[1] != i[0] {
				// If gap in regex matches, add the gap to the sentence list to avoid losing text
				sentences = append(sentences, node.Data[lasti[1]:i[0]])
			}
			sentences = append(sentences, node.Data[i[0]:i[1]])
			lasti = i
		}
		if lasti[1] != len(node.Data) {
			// If gap in regex matches, add the gap to the sentence list to avoid losing text
			sentences = append(sentences, node.Data[lasti[1]:len(node.Data)])
		}

		for _, sentence := range sentences {
			if strings.TrimSpace(sentence) != "" {
				node.Parent.InsertBefore(createSpan(*paragraph, *segment, sentence), node)
				*segment++
			}
		}

		node.Parent.RemoveChild(node)

		return
	}

	if node.Type != html.ElementNode {
		return
	}
	if node.Data == "img" {
		return
	}
	if node.Data == "p" || node.Data == "ol" || node.Data == "ul" {
		*segment = 0
		*paragraph++
	}

	for _, c := range nextNodes {
		addSpansToNode(c, paragraph, segment)
	}
}

// addSpans adds kobo spans.
func addSpans(doc *goquery.Document) error {
	alreadyHasSpans := false
	doc.Find("span").Each(func(i int, s *goquery.Selection) {
		if val, _ := s.Attr("class"); strings.Contains(val, "koboSpan") {
			alreadyHasSpans = true
		}
	})
	if alreadyHasSpans {
		return nil
	}

	paragraph := 0
	segment := 0

	for _, n := range doc.Find("body").Nodes {
		addSpansToNode(n, &paragraph, &segment)
	}

	return nil
}

// addKoboStyles adds kobo styles.
func addKoboStyles(doc *goquery.Document) error {
	s := doc.Find("head").First().AppendHtml(`<style type="text/css">div#book-inner{margin-top: 0;margin-bottom: 0;}</style>`)
	if s.Length() != 1 {
		return fmt.Errorf("could not append kobo styles")
	}
	return nil
}

// smartenPunctuation smartens punctuation in html code. It must be run last.
func smartenPunctuation(html *string) error {
	// em and en dashes
	*html = strings.Replace(*html, "---", " – ", -1)
	*html = strings.Replace(*html, "--", " — ", -1)

	// TODO: smart quotes

	// Fix comments
	*html = strings.Replace(*html, "<! — ", "<!-- ", -1)
	*html = strings.Replace(*html, " — >", " -->", -1)
	return nil
}

// cleanHTML cleans up html for a kobo epub.
func cleanHTML(doc *goquery.Document) error {
	// Remove Adobe DRM tags
	doc.Find(`meta[name="Adept.expected.resource"]`).Remove()

	// Remove empty MS <o:p> tags
	doc.Find(`o\:p`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.Trim(s.Text(), "\t \n") == ""
	}).Remove()

	// Remove empty headings
	doc.Find(`h1,h2,h3,h4,h5,h6`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.Trim(s.Text(), "\t \n") == ""
	}).Remove()

	// Remove MS <st1:whatever> tags
	doc.Find(`*`).FilterFunction(func(_ int, s *goquery.Selection) bool {
		return strings.HasPrefix(goquery.NodeName(s), "st1:")
	}).Remove()

	// Open self closing p tags
	doc.Find(`p`).Each(func(_ int, s *goquery.Selection) {
		if s.Children().Length() == 0 && strings.Trim(s.Text(), "\n \t") == "" {
			s.SetHtml("")
		}
	})

	// Add type to style tags
	doc.Find(`style`).SetAttr("type", "text/css")

	return nil
}

// process processes the html of a content file in an ordinary epub and converts it into a kobo epub by adding kobo divs, kobo spans, smartening punctuation, and cleaning html.
func process(content *string) error {
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(*content))
	if err != nil {
		return err
	}

	if err := addDivs(doc); err != nil {
		return err
	}

	if err := addSpans(doc); err != nil {
		return err
	}

	if err := addKoboStyles(doc); err != nil {
		return err
	}

	if err := cleanHTML(doc); err != nil {
		return err
	}

	h, err := doc.Html()
	if err != nil {
		return err
	}

	if err := smartenPunctuation(&h); err != nil {
		return err
	}

	// Remove unicode replacement chars
	h = strings.Replace(h, "�", "", -1)

	// Fix commented xml tag
	h = strings.Replace(h, `<!-- ?xml version="1.0" encoding="utf-8"? -->`, `<?xml version="1.0" encoding="utf-8"?>`, 1)
	h = strings.Replace(h, `<!--?xml version="1.0" encoding="utf-8"?-->`, `<?xml version="1.0" encoding="utf-8"?>`, 1)

	*content = h

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. 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 cachingfs

import (
	"time"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
)

// Constants that define the relative offsets of the inodes exported by the
// file system. See notes on the RenumberInodes method.
const (
	FooInodeOffset = iota
	DirInodeOffset
	BarInodeOffset
)

// A file system with a fixed structure that looks like this:
//
//     foo
//     dir/
//         bar
//
// The file system is configured with durations that specify how long to allow
// inode entries and attributes to be cached, used when responding to fuse
// requests. It also exposes methods for renumbering inodes and updating mtimes
// that are useful in testing that these durations are honored.
type CachingFS struct {
	fuseutil.NotImplementedFileSystem
}

var _ fuse.FileSystem = &CachingFS{}

// Create a file system that issues cacheable responses according to the
// following rules:
//
//  *  LookUpInodeResponse.Entry.EntryExpiration is set according to
//     lookupEntryTimeout.
//
//  *  GetInodeAttributesResponse.AttributesExpiration is set according to
//     getattrTimeout.
//
//  *  Nothing else is marked cacheable. (In particular, the attributes
//     returned by LookUpInode are not cacheable.)
//
func NewCachingFS(
	lookupEntryTimeout time.Duration,
	getattrTimeout time.Duration) (fs *CachingFS, err error)

// Cause inodes to receive IDs according to the following rules in further
// responses to fuse:
//
//  *  The ID of "foo" is base + FooInodeOffset.
//  *  The ID of "dir" is base + DirInodeOffset.
//  *  The ID of "dir/bar" is base + BarInodeOffset.
//
// If this method has never been called, the file system behaves as if it were
// called with base set to fuse.RootInodeID + 1.
//
// REQUIRES: base > fuse.RootInodeID
func (fs *CachingFS) RenumberInodes(base fuse.InodeID)

// Cause further queries for the attributes of inodes to use the supplied time
// as the inode's mtime.
func (fs *CachingFS) SetMtime(mtime time.Time)
<commit_msg>Implemented public functions.<commit_after>// Copyright 2015 Google Inc. 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 cachingfs

import (
	"sync"
	"time"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
)

// Constants that define the relative offsets of the inodes exported by the
// file system. See notes on the RenumberInodes method.
const (
	FooInodeOffset = iota
	DirInodeOffset
	BarInodeOffset
)

// A file system with a fixed structure that looks like this:
//
//     foo
//     dir/
//         bar
//
// The file system is configured with durations that specify how long to allow
// inode entries and attributes to be cached, used when responding to fuse
// requests. It also exposes methods for renumbering inodes and updating mtimes
// that are useful in testing that these durations are honored.
type CachingFS struct {
	fuseutil.NotImplementedFileSystem
	mu sync.Mutex

	// GUARDED_BY(mu)
	inodeIDBase fuse.InodeID

	// GUARDED_BY(mu)
	mtime time.Time
}

var _ fuse.FileSystem = &CachingFS{}

// Create a file system that issues cacheable responses according to the
// following rules:
//
//  *  LookUpInodeResponse.Entry.EntryExpiration is set according to
//     lookupEntryTimeout.
//
//  *  GetInodeAttributesResponse.AttributesExpiration is set according to
//     getattrTimeout.
//
//  *  Nothing else is marked cacheable. (In particular, the attributes
//     returned by LookUpInode are not cacheable.)
//
func NewCachingFS(
	lookupEntryTimeout time.Duration,
	getattrTimeout time.Duration) (fs *CachingFS, err error) {
	fs = &CachingFS{
		inodeIDBase: fuse.RootInodeID + 1,
		mtime:       time.Now(),
	}

	return
}

// Cause inodes to receive IDs according to the following rules in further
// responses to fuse:
//
//  *  The ID of "foo" is base + FooInodeOffset.
//  *  The ID of "dir" is base + DirInodeOffset.
//  *  The ID of "dir/bar" is base + BarInodeOffset.
//
// If this method has never been called, the file system behaves as if it were
// called with base set to fuse.RootInodeID + 1.
//
// REQUIRES: base > fuse.RootInodeID
func (fs *CachingFS) RenumberInodes(base fuse.InodeID) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	fs.inodeIDBase = base
}

// Cause further queries for the attributes of inodes to use the supplied time
// as the inode's mtime.
func (fs *CachingFS) SetMtime(mtime time.Time) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	fs.mtime = mtime
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess 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 unsharded

import (
	"context"
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	cell            = "zone1"
	hostname        = "localhost"
	KeyspaceName    = "customer"
	SchemaSQL       = `
CREATE TABLE t1 (
    c1 BIGINT NOT NULL,
    c2 BIGINT NOT NULL,
    c3 BIGINT,
    c4 varchar(100),
    PRIMARY KEY (c1),
    UNIQUE KEY (c2),
    UNIQUE KEY (c3),
    UNIQUE KEY (c4)
) ENGINE=Innodb;
`
	VSchema = `
{
    "sharded": false,
    "tables": {
        "t1": {
            "columns": [
                {
                    "name": "c1",
                    "type": "INT64"
                },
                {
                    "name": "c2",
                    "type": "INT64"
                },
                {
                    "name": "c3",
                    "type": "INT64"
                },
                {
                    "name": "c4",
                    "type": "VARCHAR"
                }
            ]
        }
    }
}
`
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		if err := clusterInstance.StartTopo(); err != nil {
			return 1
		}

		// Start keyspace
		Keyspace := &cluster.Keyspace{
			Name:      KeyspaceName,
			SchemaSQL: SchemaSQL,
			VSchema:   VSchema,
		}
		if err := clusterInstance.StartUnshardedKeyspace(*Keyspace, 0, false); err != nil {
			return 1
		}

		// Start vtgate
		if err := clusterInstance.StartVtgate(); err != nil {
			return 1
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestSelectIntoAndLoadFrom(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()

	defer exec(t, conn, `delete from t1`)
	exec(t, conn, `insert into t1(c1, c2, c3, c4) values (300,100,300,'abc')`)

	exec(t, conn, `select * from t1 into outfile 'x.txt'`)
	execAssertError(t, conn, `load data infile 'x.txt' into table t1`, "ERROR 1062 (23000): Duplicate entry '300' for key 'PRIMARY'")
	exec(t, conn, `delete from t1`)
	exec(t, conn, `load data infile 'x.txt' into table t1`)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
	exec(t, conn, `select * from t1 into dumpfile 'x1.txt'`)
	exec(t, conn, `select * from t1 into outfile 'x2.txt' Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`)
	exec(t, conn, `load data infile 'x.txt' into replace table t1 Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
}

func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err)
	return qr
}

func execAssertError(t *testing.T, conn *mysql.Conn, query string, errorString string) {
	t.Helper()
	_, err := conn.ExecuteFetch(query, 1000, true)
	require.Error(t, err)
	assert.Contains(t, err.Error(), errorString)
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := exec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}
<commit_msg>Modified test to prevent secure-file-priv error<commit_after>/*
Copyright 2020 The Vitess 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 unsharded

import (
	"context"
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	cell            = "zone1"
	hostname        = "localhost"
	KeyspaceName    = "customer"
	SchemaSQL       = `
CREATE TABLE t1 (
    c1 BIGINT NOT NULL,
    c2 BIGINT NOT NULL,
    c3 BIGINT,
    c4 varchar(100),
    PRIMARY KEY (c1),
    UNIQUE KEY (c2),
    UNIQUE KEY (c3),
    UNIQUE KEY (c4)
) ENGINE=Innodb;
`
	VSchema = `
{
    "sharded": false,
    "tables": {
        "t1": {
            "columns": [
                {
                    "name": "c1",
                    "type": "INT64"
                },
                {
                    "name": "c2",
                    "type": "INT64"
                },
                {
                    "name": "c3",
                    "type": "INT64"
                },
                {
                    "name": "c4",
                    "type": "VARCHAR"
                }
            ]
        }
    }
}
`
)

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		if err := clusterInstance.StartTopo(); err != nil {
			return 1
		}

		// Start keyspace
		Keyspace := &cluster.Keyspace{
			Name:      KeyspaceName,
			SchemaSQL: SchemaSQL,
			VSchema:   VSchema,
		}
		if err := clusterInstance.StartUnshardedKeyspace(*Keyspace, 0, false); err != nil {
			return 1
		}

		// Start vtgate
		if err := clusterInstance.StartVtgate(); err != nil {
			return 1
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestSelectIntoAndLoadFrom(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()

	defer exec(t, conn, `delete from t1`)
	exec(t, conn, `insert into t1(c1, c2, c3, c4) values (300,100,300,'abc')`)
	res := exec(t, conn, `select @@secure_file_priv;`)
	directory := res.Rows[0][0].ToString()
	query := `select * from t1 into outfile '` + directory + `x.txt'`
	exec(t, conn, query)
	query = `load data infile '` + directory + `x.txt' into table t1`
	execAssertError(t, conn, query, "ERROR 1062 (23000): Duplicate entry '300' for key 'PRIMARY'")
	exec(t, conn, `delete from t1`)
	exec(t, conn, query)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
	query = `select * from t1 into dumpfile '` + directory + `x1.txt'`
	exec(t, conn, query)
	query = `select * from t1 into outfile '` + directory + `x2.txt' Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`
	exec(t, conn, query)
	query = `load data infile '` + directory + `x.txt' into replace table t1 Fields terminated by ';' optionally enclosed by '"' escaped by '\t' lines terminated by '\n'`
	exec(t, conn, query)
	assertMatches(t, conn, `select c1,c2,c3 from t1`, `[[INT64(300) INT64(100) INT64(300)]]`)
}

func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err)
	return qr
}

func execAssertError(t *testing.T, conn *mysql.Conn, query string, errorString string) {
	t.Helper()
	_, err := conn.ExecuteFetch(query, 1000, true)
	require.Error(t, err)
	assert.Contains(t, err.Error(), errorString)
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := exec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}
<|endoftext|>