<commit_before>package jsonrpc

import (
	"fmt"
)

func Example() {
	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	// create client
	rpcClient := NewRPCClient("http://my-rpc-service")

	// execute rpc to service
	response, _ := rpcClient.Call("getPersonByID", 12345)

	// parse result into struct
	var person Person
	response.GetObject(&person)

	// change result and send back using rpc
	person.Age = 35
	rpcClient.Call("setPersonByID", 12345, person)
}

func ExampleRPCClient_Call() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	// result processing omitted, see: RPCResponse methods
	rpcClient.Call("getTimestamp")

	rpcClient.Call("getPerson", 1234)

	rpcClient.Call("addNumbers", 5, 2, 3)

	rpcClient.Call("strangeFunction", 1, true, "alex", 3.4)

	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	person := Person{
		Name:    "alex",
		Age:     33,
		Country: "germany",
	}

	rpcClient.Call("setPersonByID", 123, person)
}

func ExampleRPCClient_CallNamed() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	// result processing omitted, see: RPCResponse methods
	rpcClient.CallNamed("createPerson", map[string]interface{}{
		"name":      "Bartholomew Allen",
		"nicknames": []string{"Barry", "Flash", },
		"male":      true,
		"age":       28,
		"address":   map[string]interface{}{"street": "Main Street", "city": "Central City"},
	})
}

func ExampleRPCResponse() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	response, _ := rpcClient.Call("addNumbers", 1, 2, 3)
	sum, _ := response.GetInt()
	fmt.Println(sum)

	response, _ = rpcClient.Call("isValidEmail", "my@ema.il")
	valid, _ := response.GetBool()
	fmt.Println(valid)

	response, _ = rpcClient.Call("getJoke")
	joke, _ := response.GetString()
	fmt.Println(joke)

	response, _ = rpcClient.Call("getPi", 10)
	piRounded, _ := response.GetFloat64()
	fmt.Println(piRounded)

	var rndNumbers []int
	response, _ = rpcClient.Call("getRndIntNumbers", 10)
	response.GetObject(&rndNumbers)
	fmt.Println(rndNumbers[0])

	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	var p Person
	response, _ = rpcClient.Call("getPersonByID", 1234)
	response.GetObject(&p)
	fmt.Println(p.Name)
}

func ExampleRPCClient_Batch() {
	rpcClient := NewRPCClient(httpServer.URL)

	req1 := rpcClient.NewRPCRequestObject("addNumbers", 1, 2, 3)
	req2 := rpcClient.NewRPCRequestObject("getTimestamp")
	responses, _ := rpcClient.Batch(req1, req2)

	response, _ := responses.GetResponseOf(req2)
	timestamp, _ := response.GetInt()

	fmt.Println(timestamp)
}
<commit_msg>fix example: removed unnecessary comma<commit_after>package jsonrpc

import (
	"fmt"
)

func Example() {
	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	// create client
	rpcClient := NewRPCClient("http://my-rpc-service")

	// execute rpc to service
	response, _ := rpcClient.Call("getPersonByID", 12345)

	// parse result into struct
	var person Person
	response.GetObject(&person)

	// change result and send back using rpc
	person.Age = 35
	rpcClient.Call("setPersonByID", 12345, person)
}

func ExampleRPCClient_Call() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	// result processing omitted, see: RPCResponse methods
	rpcClient.Call("getTimestamp")

	rpcClient.Call("getPerson", 1234)

	rpcClient.Call("addNumbers", 5, 2, 3)

	rpcClient.Call("strangeFunction", 1, true, "alex", 3.4)

	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	person := Person{
		Name:    "alex",
		Age:     33,
		Country: "germany",
	}

	rpcClient.Call("setPersonByID", 123, person)
}

func ExampleRPCClient_CallNamed() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	// result processing omitted, see: RPCResponse methods
	rpcClient.CallNamed("createPerson", map[string]interface{}{
		"name":      "Bartholomew Allen",
		"nicknames": []string{"Barry", "Flash"},
		"male":      true,
		"age":       28,
		"address":   map[string]interface{}{"street": "Main Street", "city": "Central City"},
	})
}

func ExampleRPCResponse() {
	rpcClient := NewRPCClient("http://my-rpc-service")

	response, _ := rpcClient.Call("addNumbers", 1, 2, 3)
	sum, _ := response.GetInt()
	fmt.Println(sum)

	response, _ = rpcClient.Call("isValidEmail", "my@ema.il")
	valid, _ := response.GetBool()
	fmt.Println(valid)

	response, _ = rpcClient.Call("getJoke")
	joke, _ := response.GetString()
	fmt.Println(joke)

	response, _ = rpcClient.Call("getPi", 10)
	piRounded, _ := response.GetFloat64()
	fmt.Println(piRounded)

	var rndNumbers []int
	response, _ = rpcClient.Call("getRndIntNumbers", 10)
	response.GetObject(&rndNumbers)
	fmt.Println(rndNumbers[0])

	type Person struct {
		Name    string `json:"name"`
		Age     int    `json:"age"`
		Country string `json:"country"`
	}

	var p Person
	response, _ = rpcClient.Call("getPersonByID", 1234)
	response.GetObject(&p)
	fmt.Println(p.Name)
}

func ExampleRPCClient_Batch() {
	rpcClient := NewRPCClient(httpServer.URL)

	req1 := rpcClient.NewRPCRequestObject("addNumbers", 1, 2, 3)
	req2 := rpcClient.NewRPCRequestObject("getTimestamp")
	responses, _ := rpcClient.Batch(req1, req2)

	response, _ := responses.GetResponseOf(req2)
	timestamp, _ := response.GetInt()

	fmt.Println(timestamp)
}
<|endoftext|>


<commit_before>package reflector_test

import (
	. "."
	"fmt"
)

func ExampleStructToMap() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := T{8, 3.14, "str", &pstr, 13}
	m := make(map[string]interface{})
	StructToMap(&s, m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m["Uint8"], m["f32"], m["String"], m["Pstring"], m["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleStructValueToMap() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := T{8, 3.14, "str", &pstr, 13}
	m := make(map[string]interface{})
	StructValueToMap(s, m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m["Uint8"], m["f32"], m["String"], m["Pstring"], m["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleStructsToMaps() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := []T{{8, 3.14, "str", &pstr, 13}}
	var m []map[string]interface{}
	StructsToMaps(s, &m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m[0]["Uint8"], m[0]["f32"], m[0]["String"], m[0]["Pstring"], m[0]["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleMapToStruct_NoConvert() {
	type T struct {
		Uint8   uint8   // no automatic type conversion
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s T
	m := map[string]interface{}{"Uint8": uint8(8), "f32": float32(3.14), "Pstring": "pstr", "foo": 13}
	MapToStruct(m, &s, NoConvert, "json")
	m["Pstring"] = "foo"
	fmt.Printf("%#v %#v %#v %#v %#v\n", s.Uint8, s.Float32, s.String, *s.Pstring, s.foo)
	// Output:
	// 0x8 3.14 "" "pstr" 0
}

func ExampleMapToStruct_Strconv() {
	type T struct {
		Uint8   uint8   // type conversion via strconv
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s T
	m := map[string]interface{}{"Uint8": 8, "f32": 3, "Pstring": "pstr", "foo": 13}
	MapToStruct(m, &s, Strconv, "json")
	m["Pstring"] = "foo"
	fmt.Printf("%#v %#v %#v %#v %#v\n", s.Uint8, s.Float32, s.String, *s.Pstring, s.foo)
	// Output:
	// 0x8 3 "" "pstr" 0
}

func ExampleMapsToStructs_NoConvert() {
	type T struct {
		Uint8   uint8   // no automatic type conversion
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in first map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s []T
	maps := []map[string]interface{}{
		{"Uint8": uint8(8), "Pstring": "pstr"},
		{"f32": float32(3.14), "String": "str", "foo": 13},
	}
	MapsToStructs(maps, &s, NoConvert, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[0].Uint8, s[0].Float32, s[0].String, *s[0].Pstring, s[0].foo)
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[1].Uint8, s[1].Float32, s[1].String, s[1].Pstring, s[1].foo)
	// Output:
	// 0x8 0 "" "pstr" 0
	// 0x0 3.14 "str" (*string)(nil) 0
}

func ExampleMapsToStructs_Strconv() {
	type T struct {
		Uint8   uint8   // type conversion via strconv
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in first map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s []T
	maps := []map[string]interface{}{
		{"Uint8": 8, "f32": 3, "foo": 13, "Pstring": "pstr"},
		{"Uint8": "9", "f32": "4", "String": "43", "foo": "13"},
	}
	MapsToStructs(maps, &s, Strconv, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[0].Uint8, s[0].Float32, s[0].String, *s[0].Pstring, s[0].foo)
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[1].Uint8, s[1].Float32, s[1].String, s[1].Pstring, s[1].foo)
	// Output:
	// 0x8 3 "" "pstr" 0
	// 0x9 4 "43" (*string)(nil) 0
}
<commit_msg>Fix examples.<commit_after>package reflector_test

import (
	. "."
	"fmt"
)

func ExampleStructToMap() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := T{8, 3.14, "str", &pstr, 13}
	m := make(map[string]interface{})
	StructToMap(&s, m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m["Uint8"], m["f32"], m["String"], m["Pstring"], m["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleStructValueToMap() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := T{8, 3.14, "str", &pstr, 13}
	m := make(map[string]interface{})
	StructValueToMap(s, m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m["Uint8"], m["f32"], m["String"], m["Pstring"], m["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleStructsToMaps() {
	type T struct {
		Uint8   uint8
		Float32 float32 `json:"f32"` // tag will be used
		String  string
		Pstring *string // pointer will be followed
		foo     int     // not exported
	}
	pstr := "pstr"
	s := []T{{8, 3.14, "str", &pstr, 13}}
	var m []map[string]interface{}
	StructsToMaps(s, &m, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", m[0]["Uint8"], m[0]["f32"], m[0]["String"], m[0]["Pstring"], m[0]["foo"])
	// Output:
	// 0x8 3.14 "str" "pstr" <nil>
}

func ExampleMapToStruct_noConvert() {
	type T struct {
		Uint8   uint8   // no automatic type conversion
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s T
	m := map[string]interface{}{"Uint8": uint8(8), "f32": float32(3.14), "Pstring": "pstr", "foo": 13}
	MapToStruct(m, &s, NoConvert, "json")
	m["Pstring"] = "foo"
	fmt.Printf("%#v %#v %#v %#v %#v\n", s.Uint8, s.Float32, s.String, *s.Pstring, s.foo)
	// Output:
	// 0x8 3.14 "" "pstr" 0
}

func ExampleMapToStruct_strconv() {
	type T struct {
		Uint8   uint8   // type conversion via strconv
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s T
	m := map[string]interface{}{"Uint8": 8, "f32": 3, "Pstring": "pstr", "foo": 13}
	MapToStruct(m, &s, Strconv, "json")
	m["Pstring"] = "foo"
	fmt.Printf("%#v %#v %#v %#v %#v\n", s.Uint8, s.Float32, s.String, *s.Pstring, s.foo)
	// Output:
	// 0x8 3 "" "pstr" 0
}

func ExampleMapsToStructs_noConvert() {
	type T struct {
		Uint8   uint8   // no automatic type conversion
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in first map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s []T
	maps := []map[string]interface{}{
		{"Uint8": uint8(8), "Pstring": "pstr"},
		{"f32": float32(3.14), "String": "str", "foo": 13},
	}
	MapsToStructs(maps, &s, NoConvert, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[0].Uint8, s[0].Float32, s[0].String, *s[0].Pstring, s[0].foo)
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[1].Uint8, s[1].Float32, s[1].String, s[1].Pstring, s[1].foo)
	// Output:
	// 0x8 0 "" "pstr" 0
	// 0x0 3.14 "str" (*string)(nil) 0
}

func ExampleMapsToStructs_strconv() {
	type T struct {
		Uint8   uint8   // type conversion via strconv
		Float32 float32 `json:"f32"` // tag will be used
		String  string  // not present in first map, will not be set
		Pstring *string // value will be deep-copied
		foo     int     // not exported, will not be set
	}
	var s []T
	maps := []map[string]interface{}{
		{"Uint8": 8, "f32": 3, "foo": 13, "Pstring": "pstr"},
		{"Uint8": "9", "f32": "4", "String": "43", "foo": "13"},
	}
	MapsToStructs(maps, &s, Strconv, "json")
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[0].Uint8, s[0].Float32, s[0].String, *s[0].Pstring, s[0].foo)
	fmt.Printf("%#v %#v %#v %#v %#v\n", s[1].Uint8, s[1].Float32, s[1].String, s[1].Pstring, s[1].foo)
	// Output:
	// 0x8 3 "" "pstr" 0
	// 0x9 4 "43" (*string)(nil) 0
}
<|endoftext|>


<commit_before>package amqp_test

import (
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net"
	"runtime"
	"time"

	"github.com/streadway/amqp"
)

//
// TLS setup instructions.
//
// These examples assume you have a RabbitMQ node running on localhost
// with TLS enabled. We explain how to generate the CA, certificates,
// keys and what node configuration you'd need below.
//
//
// The easiest way to re-create the CA, certificates and keys in these
// examples is by using tls-gen:
// https://github.com/michaelklishin/tls-gen.
//
// RabbitMQ TLS guide can be found at http://www.rabbitmq.com/ssl.html.
//
// Then in your rabbitmq.config, disable the plain AMQP port, verify clients
// and fail if no certificate is presented with the following:
//
//   [
//   {rabbit, [
//     {tcp_listeners, []},     % listens on 127.0.0.1:5672
//     {ssl_listeners, [5671]}, % listens on 0.0.0.0:5671
//     {ssl_options, [{cacertfile,"/path/to/your/testca/cacert.pem"},
//                    {certfile,"/path/to/your/server/cert.pem"},
//                    {keyfile,"/path/to/your/server/key.pem"},
//                    {verify,verify_peer},
//                    {fail_if_no_peer_cert,true}]}
//     ]}
//   ].
//
// The self-signing certificate authority's certificate must be included in
// the RootCAs to be trusted so that the server certificate can be verified.
//
// Alternatively to adding it to the tls.Config you can add the CA's cert to
// your system's root CAs.  The tls package will use the system roots
// specific to each support OS.  Under OS X, add (drag/drop) your cacert.pem
// file to the 'Certificates' section of KeyChain.app to add and always
// trust.
//
// Or with the command line add and trust the DER encoded certificate:
//
//   security add-certificate testca/cacert.cer
//   security add-trusted-cert testca/cacert.cer
//
// If you depend on the system root CAs, then use nil for the RootCAs field
// so the system roots will be loaded.
//
// Server names are validated by the crypto/tls package, so the server
// certificate must be made for the hostname in the URL.  Find the commonName
// (CN) and make sure the hostname in the URL matches this common name.  Per
// the RabbitMQ instructions (or tls-gen) for a self-signed cert, this defaults to the
// current hostname.
//
//   openssl x509 -noout -in /path/to/certificate.pem -subject
//
// If your server name in your certificate is different than the host you are
// connecting to, set the hostname used for verification in
// ServerName field of the tls.Config struct.

func ExampleConfig_timeout() {
	// Provide your own anonymous Dial function that delgates to net.DialTimout
	// for custom timeouts

	conn, err := amqp.DialConfig("amqp:///", amqp.Config{
		Dial: func(network, addr string) (net.Conn, error) {
			return net.DialTimeout(network, addr, 2*time.Second)
		},
	})

	log.Printf("conn: %v, err: %v", conn, err)
}

func ExampleDialTLS() {
	// see at the top
	cfg := new(tls.Config)

	// see at the top
	cfg.RootCAs = x509.NewCertPool()

	if ca, err := ioutil.ReadFile("testca/cacert.pem"); err == nil {
		cfg.RootCAs.AppendCertsFromPEM(ca)
	}

	// Move the client cert and key to a location specific to your application
	// and load them here.

	if cert, err := tls.LoadX509KeyPair("client/cert.pem", "client/key.pem"); err == nil {
		cfg.Certificates = append(cfg.Certificates, cert)
	}

	// see a note about Common Name (CN) at the top
	conn, err := amqp.DialTLS("amqps://server-name-from-certificate/", cfg)

	log.Printf("conn: %v, err: %v", conn, err)
}

func ExampleChannel_Confirm_bridge() {
	// This example acts as a bridge, shoveling all messages sent from the source
	// exchange "log" to destination exchange "log".

	// Confirming publishes can help from overproduction and ensure every message
	// is delivered.

	// Setup the source of the store and forward
	source, err := amqp.Dial("amqp://source/")
	if err != nil {
		log.Fatalf("connection.open source: %s", err)
	}
	defer source.Close()

	chs, err := source.Channel()
	if err != nil {
		log.Fatalf("channel.open source: %s", err)
	}

	if err := chs.ExchangeDeclare("log", "topic", true, false, false, false, nil); err != nil {
		log.Fatalf("exchange.declare destination: %s", err)
	}

	if _, err := chs.QueueDeclare("remote-tee", true, true, false, false, nil); err != nil {
		log.Fatalf("queue.declare source: %s", err)
	}

	if err := chs.QueueBind("remote-tee", "#", "logs", false, nil); err != nil {
		log.Fatalf("queue.bind source: %s", err)
	}

	shovel, err := chs.Consume("remote-tee", "shovel", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume source: %s", err)
	}

	// Setup the destination of the store and forward
	destination, err := amqp.Dial("amqp://destination/")
	if err != nil {
		log.Fatalf("connection.open destination: %s", err)
	}
	defer destination.Close()

	chd, err := destination.Channel()
	if err != nil {
		log.Fatalf("channel.open destination: %s", err)
	}

	if err := chd.ExchangeDeclare("log", "topic", true, false, false, false, nil); err != nil {
		log.Fatalf("exchange.declare destination: %s", err)
	}

	// Buffer of 1 for our single outstanding publishing
	confirms := chd.NotifyPublish(make(chan amqp.Confirmation, 1))

	if err := chd.Confirm(false); err != nil {
		log.Fatalf("confirm.select destination: %s", err)
	}

	// Now pump the messages, one by one, a smarter implementation
	// would batch the deliveries and use multiple ack/nacks
	for {
		msg, ok := <-shovel
		if !ok {
			log.Fatalf("source channel closed, see the reconnect example for handling this")
		}

		err = chd.Publish("logs", msg.RoutingKey, false, false, amqp.Publishing{
			// Copy all the properties
			ContentType:     msg.ContentType,
			ContentEncoding: msg.ContentEncoding,
			DeliveryMode:    msg.DeliveryMode,
			Priority:        msg.Priority,
			CorrelationId:   msg.CorrelationId,
			ReplyTo:         msg.ReplyTo,
			Expiration:      msg.Expiration,
			MessageId:       msg.MessageId,
			Timestamp:       msg.Timestamp,
			Type:            msg.Type,
			UserId:          msg.UserId,
			AppId:           msg.AppId,

			// Custom headers
			Headers: msg.Headers,

			// And the body
			Body: msg.Body,
		})

		if err != nil {
			msg.Nack(false, false)
			log.Fatalf("basic.publish destination: %+v", msg)
		}

		// only ack the source delivery when the destination acks the publishing
		if confirmed := <-confirms; confirmed.Ack {
			msg.Ack(false)
		} else {
			msg.Nack(false, false)
		}
	}
}

func ExampleChannel_Consume() {
	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}
	defer conn.Close()

	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("channel.open: %s", err)
	}

	// We declare our topology on both the publisher and consumer to ensure they
	// are the same.  This is part of AMQP being a programmable messaging model.
	//
	// See the Channel.Publish example for the complimentary declare.
	err = c.ExchangeDeclare("logs", "topic", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %s", err)
	}

	// Establish our queue topologies that we are responsible for
	type bind struct {
		queue string
		key   string
	}

	bindings := []bind{
		bind{"page", "alert"},
		bind{"email", "info"},
		bind{"firehose", "#"},
	}

	for _, b := range bindings {
		_, err = c.QueueDeclare(b.queue, true, false, false, false, nil)
		if err != nil {
			log.Fatalf("queue.declare: %v", err)
		}

		err = c.QueueBind(b.queue, b.key, "logs", false, nil)
		if err != nil {
			log.Fatalf("queue.bind: %v", err)
		}
	}

	// Set our quality of service.  Since we're sharing 3 consumers on the same
	// channel, we want at least 3 messages in flight.
	err = c.Qos(3, 0, false)
	if err != nil {
		log.Fatalf("basic.qos: %v", err)
	}

	// Establish our consumers that have different responsibilities.  Our first
	// two queues do not ack the messages on the server, so require to be acked
	// on the client.

	pages, err := c.Consume("page", "pager", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	go func() {
		for log := range pages {
			// ... this consumer is responsible for sending pages per log
			log.Ack(false)
		}
	}()

	// Notice how the concern for which messages arrive here are in the AMQP
	// topology and not in the queue.  We let the server pick a consumer tag this
	// time.

	emails, err := c.Consume("email", "", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	go func() {
		for log := range emails {
			// ... this consumer is responsible for sending emails per log
			log.Ack(false)
		}
	}()

	// This consumer requests that every message is acknowledged as soon as it's
	// delivered.

	firehose, err := c.Consume("firehose", "", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	// To show how to process the items in parallel, we'll use a work pool.
	for i := 0; i < runtime.NumCPU(); i++ {
		go func(work <-chan amqp.Delivery) {
			for _ = range work {
				// ... this consumer pulls from the firehose and doesn't need to acknowledge
			}
		}(firehose)
	}

	// Wait until you're ready to finish, could be a signal handler here.
	time.Sleep(10 * time.Second)

	// Cancelling a consumer by name will finish the range and gracefully end the
	// goroutine
	err = c.Cancel("pager", false)
	if err != nil {
		log.Fatalf("basic.cancel: %v", err)
	}

	// deferred closing the Connection will also finish the consumer's ranges of
	// their delivery chans.  If you need every delivery to be processed, make
	// sure to wait for all consumers goroutines to finish before exiting your
	// process.
}

func ExampleChannel_Publish() {
	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}

	// This waits for a server acknowledgment which means the sockets will have
	// flushed all outbound publishings prior to returning.  It's important to
	// block on Close to not lose any publishings.
	defer conn.Close()

	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("channel.open: %s", err)
	}

	// We declare our topology on both the publisher and consumer to ensure they
	// are the same.  This is part of AMQP being a programmable messaging model.
	//
	// See the Channel.Consume example for the complimentary declare.
	err = c.ExchangeDeclare("logs", "topic", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %v", err)
	}

	// Prepare this message to be persistent.  Your publishing requirements may
	// be different.
	msg := amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		ContentType:  "text/plain",
		Body:         []byte("Go Go AMQP!"),
	}

	// This is not a mandatory delivery, so it will be dropped if there are no
	// queues bound to the logs exchange.
	err = c.Publish("logs", "info", false, false, msg)
	if err != nil {
		// Since publish is asynchronous this can happen if the network connection
		// is reset or if the server has run out of resources.
		log.Fatalf("basic.publish: %v", err)
	}
}

func publishAllTheThings(conn *amqp.Connection) {
	// ... snarf snarf, barf barf
}

func ExampleConnection_NotifyBlocked() {
	// Simply logs when the server throttles the TCP connection for publishers

	// Test this by tuning your server to have a low memory watermark:
	// rabbitmqctl set_vm_memory_high_watermark 0.00000001

	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}
	defer conn.Close()

	blockings := conn.NotifyBlocked(make(chan amqp.Blocking))
	go func() {
		for b := range blockings {
			if b.Active {
				log.Printf("TCP blocked: %q", b.Reason)
			} else {
				log.Printf("TCP unblocked")
			}
		}
	}()

	// Your application domain channel setup publishings
	publishAllTheThings(conn)
}
<commit_msg>Fix godoc DialTLS Example<commit_after>package amqp_test

import (
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net"
	"runtime"
	"time"

	"github.com/streadway/amqp"
)

func ExampleConfig_timeout() {
	// Provide your own anonymous Dial function that delgates to net.DialTimout
	// for custom timeouts

	conn, err := amqp.DialConfig("amqp:///", amqp.Config{
		Dial: func(network, addr string) (net.Conn, error) {
			return net.DialTimeout(network, addr, 2*time.Second)
		},
	})

	log.Printf("conn: %v, err: %v", conn, err)
}

func ExampleDialTLS() {
	// This example assume you have a RabbitMQ node running on localhost
	// with TLS enabled.
	//
	// The easiest way to create the CA, certificates and keys required for these
	// examples is by using tls-gen: https://github.com/michaelklishin/tls-gen
	//
	// A comprehensive RabbitMQ TLS guide can be found at
	// http://www.rabbitmq.com/ssl.html
	//
	// Once you have the required TLS files in place, use the following
	// rabbitmq.config example for the RabbitMQ node that you will run on
	// localhost:
	//
	//   [
	//   {rabbit, [
	//     {tcp_listeners, []},     % listens on 127.0.0.1:5672
	//     {ssl_listeners, [5671]}, % listens on 0.0.0.0:5671
	//     {ssl_options, [{cacertfile,"/path/to/your/testca/cacert.pem"},
	//                    {certfile,"/path/to/your/server/cert.pem"},
	//                    {keyfile,"/path/to/your/server/key.pem"},
	//                    {verify,verify_peer},
	//                    {fail_if_no_peer_cert,true}]}
	//     ]}
	//   ].
	//
	//
	// In the above rabbitmq.config example, we are disabling the plain AMQP port
	// and verifying that clients and fail if no certificate is presented.
	//
	// The self-signing certificate authority's certificate (cacert.pem) must be
	// included in the RootCAs to be trusted, otherwise the server certificate
	// will fail certificate verification.
	//
	// Alternatively to adding it to the tls.Config. you can add the CA's cert to
	// your system's root CAs.  The tls package will use the system roots
	// specific to each support OS.  Under OS X, add (drag/drop) cacert.pem
	// file to the 'Certificates' section of KeyChain.app to add and always
	// trust.  You can also add it via the command line:
	//
	//   security add-certificate testca/cacert.pem
	//   security add-trusted-cert testca/cacert.pem
	//
	// If you depend on the system root CAs, then use nil for the RootCAs field
	// so the system roots will be loaded instead.
	//
	// Server names are validated by the crypto/tls package, so the server
	// certificate must be made for the hostname in the URL.  Find the commonName
	// (CN) and make sure the hostname in the URL matches this common name.  Per
	// the RabbitMQ instructions (or tls-gen) for a self-signed cert, this defaults to the
	// current hostname.
	//
	//   openssl x509 -noout -in /path/to/certificate.pem -subject
	//
	// If your server name in your certificate is different than the host you are
	// connecting to, set the hostname used for verification in
	// ServerName field of the tls.Config struct.
	cfg := new(tls.Config)

	// see at the top
	cfg.RootCAs = x509.NewCertPool()

	if ca, err := ioutil.ReadFile("testca/cacert.pem"); err == nil {
		cfg.RootCAs.AppendCertsFromPEM(ca)
	}

	// Move the client cert and key to a location specific to your application
	// and load them here.

	if cert, err := tls.LoadX509KeyPair("client/cert.pem", "client/key.pem"); err == nil {
		cfg.Certificates = append(cfg.Certificates, cert)
	}

	// see a note about Common Name (CN) at the top
	conn, err := amqp.DialTLS("amqps://server-name-from-certificate/", cfg)

	log.Printf("conn: %v, err: %v", conn, err)
}

func ExampleChannel_Confirm_bridge() {
	// This example acts as a bridge, shoveling all messages sent from the source
	// exchange "log" to destination exchange "log".

	// Confirming publishes can help from overproduction and ensure every message
	// is delivered.

	// Setup the source of the store and forward
	source, err := amqp.Dial("amqp://source/")
	if err != nil {
		log.Fatalf("connection.open source: %s", err)
	}
	defer source.Close()

	chs, err := source.Channel()
	if err != nil {
		log.Fatalf("channel.open source: %s", err)
	}

	if err := chs.ExchangeDeclare("log", "topic", true, false, false, false, nil); err != nil {
		log.Fatalf("exchange.declare destination: %s", err)
	}

	if _, err := chs.QueueDeclare("remote-tee", true, true, false, false, nil); err != nil {
		log.Fatalf("queue.declare source: %s", err)
	}

	if err := chs.QueueBind("remote-tee", "#", "logs", false, nil); err != nil {
		log.Fatalf("queue.bind source: %s", err)
	}

	shovel, err := chs.Consume("remote-tee", "shovel", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume source: %s", err)
	}

	// Setup the destination of the store and forward
	destination, err := amqp.Dial("amqp://destination/")
	if err != nil {
		log.Fatalf("connection.open destination: %s", err)
	}
	defer destination.Close()

	chd, err := destination.Channel()
	if err != nil {
		log.Fatalf("channel.open destination: %s", err)
	}

	if err := chd.ExchangeDeclare("log", "topic", true, false, false, false, nil); err != nil {
		log.Fatalf("exchange.declare destination: %s", err)
	}

	// Buffer of 1 for our single outstanding publishing
	confirms := chd.NotifyPublish(make(chan amqp.Confirmation, 1))

	if err := chd.Confirm(false); err != nil {
		log.Fatalf("confirm.select destination: %s", err)
	}

	// Now pump the messages, one by one, a smarter implementation
	// would batch the deliveries and use multiple ack/nacks
	for {
		msg, ok := <-shovel
		if !ok {
			log.Fatalf("source channel closed, see the reconnect example for handling this")
		}

		err = chd.Publish("logs", msg.RoutingKey, false, false, amqp.Publishing{
			// Copy all the properties
			ContentType:     msg.ContentType,
			ContentEncoding: msg.ContentEncoding,
			DeliveryMode:    msg.DeliveryMode,
			Priority:        msg.Priority,
			CorrelationId:   msg.CorrelationId,
			ReplyTo:         msg.ReplyTo,
			Expiration:      msg.Expiration,
			MessageId:       msg.MessageId,
			Timestamp:       msg.Timestamp,
			Type:            msg.Type,
			UserId:          msg.UserId,
			AppId:           msg.AppId,

			// Custom headers
			Headers: msg.Headers,

			// And the body
			Body: msg.Body,
		})

		if err != nil {
			msg.Nack(false, false)
			log.Fatalf("basic.publish destination: %+v", msg)
		}

		// only ack the source delivery when the destination acks the publishing
		if confirmed := <-confirms; confirmed.Ack {
			msg.Ack(false)
		} else {
			msg.Nack(false, false)
		}
	}
}

func ExampleChannel_Consume() {
	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}
	defer conn.Close()

	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("channel.open: %s", err)
	}

	// We declare our topology on both the publisher and consumer to ensure they
	// are the same.  This is part of AMQP being a programmable messaging model.
	//
	// See the Channel.Publish example for the complimentary declare.
	err = c.ExchangeDeclare("logs", "topic", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %s", err)
	}

	// Establish our queue topologies that we are responsible for
	type bind struct {
		queue string
		key   string
	}

	bindings := []bind{
		bind{"page", "alert"},
		bind{"email", "info"},
		bind{"firehose", "#"},
	}

	for _, b := range bindings {
		_, err = c.QueueDeclare(b.queue, true, false, false, false, nil)
		if err != nil {
			log.Fatalf("queue.declare: %v", err)
		}

		err = c.QueueBind(b.queue, b.key, "logs", false, nil)
		if err != nil {
			log.Fatalf("queue.bind: %v", err)
		}
	}

	// Set our quality of service.  Since we're sharing 3 consumers on the same
	// channel, we want at least 3 messages in flight.
	err = c.Qos(3, 0, false)
	if err != nil {
		log.Fatalf("basic.qos: %v", err)
	}

	// Establish our consumers that have different responsibilities.  Our first
	// two queues do not ack the messages on the server, so require to be acked
	// on the client.

	pages, err := c.Consume("page", "pager", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	go func() {
		for log := range pages {
			// ... this consumer is responsible for sending pages per log
			log.Ack(false)
		}
	}()

	// Notice how the concern for which messages arrive here are in the AMQP
	// topology and not in the queue.  We let the server pick a consumer tag this
	// time.

	emails, err := c.Consume("email", "", false, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	go func() {
		for log := range emails {
			// ... this consumer is responsible for sending emails per log
			log.Ack(false)
		}
	}()

	// This consumer requests that every message is acknowledged as soon as it's
	// delivered.

	firehose, err := c.Consume("firehose", "", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("basic.consume: %v", err)
	}

	// To show how to process the items in parallel, we'll use a work pool.
	for i := 0; i < runtime.NumCPU(); i++ {
		go func(work <-chan amqp.Delivery) {
			for _ = range work {
				// ... this consumer pulls from the firehose and doesn't need to acknowledge
			}
		}(firehose)
	}

	// Wait until you're ready to finish, could be a signal handler here.
	time.Sleep(10 * time.Second)

	// Cancelling a consumer by name will finish the range and gracefully end the
	// goroutine
	err = c.Cancel("pager", false)
	if err != nil {
		log.Fatalf("basic.cancel: %v", err)
	}

	// deferred closing the Connection will also finish the consumer's ranges of
	// their delivery chans.  If you need every delivery to be processed, make
	// sure to wait for all consumers goroutines to finish before exiting your
	// process.
}

func ExampleChannel_Publish() {
	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}

	// This waits for a server acknowledgment which means the sockets will have
	// flushed all outbound publishings prior to returning.  It's important to
	// block on Close to not lose any publishings.
	defer conn.Close()

	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("channel.open: %s", err)
	}

	// We declare our topology on both the publisher and consumer to ensure they
	// are the same.  This is part of AMQP being a programmable messaging model.
	//
	// See the Channel.Consume example for the complimentary declare.
	err = c.ExchangeDeclare("logs", "topic", true, false, false, false, nil)
	if err != nil {
		log.Fatalf("exchange.declare: %v", err)
	}

	// Prepare this message to be persistent.  Your publishing requirements may
	// be different.
	msg := amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		ContentType:  "text/plain",
		Body:         []byte("Go Go AMQP!"),
	}

	// This is not a mandatory delivery, so it will be dropped if there are no
	// queues bound to the logs exchange.
	err = c.Publish("logs", "info", false, false, msg)
	if err != nil {
		// Since publish is asynchronous this can happen if the network connection
		// is reset or if the server has run out of resources.
		log.Fatalf("basic.publish: %v", err)
	}
}

func publishAllTheThings(conn *amqp.Connection) {
	// ... snarf snarf, barf barf
}

func ExampleConnection_NotifyBlocked() {
	// Simply logs when the server throttles the TCP connection for publishers

	// Test this by tuning your server to have a low memory watermark:
	// rabbitmqctl set_vm_memory_high_watermark 0.00000001

	conn, err := amqp.Dial("amqp://guest:guest@localhost:5672/")
	if err != nil {
		log.Fatalf("connection.open: %s", err)
	}
	defer conn.Close()

	blockings := conn.NotifyBlocked(make(chan amqp.Blocking))
	go func() {
		for b := range blockings {
			if b.Active {
				log.Printf("TCP blocked: %q", b.Reason)
			} else {
				log.Printf("TCP unblocked")
			}
		}
	}()

	// Your application domain channel setup publishings
	publishAllTheThings(conn)
}
<|endoftext|>


<commit_before>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"fmt"
	glssh "github.com/gliderlabs/ssh"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/testdata"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"reflect"
	"testing"
)

type mockSSHKey struct {
	keyname string
	content []byte
	privkey agent.AddedKey
	pubkey  ssh.PublicKey
}

var (
	testPrivateKeys map[string]interface{}
	testSigners     map[string]ssh.Signer
	testPublicKeys  map[string]ssh.PublicKey
	sshAgentSocket  string
)

func init() {
	var err error

	n := len(testdata.PEMBytes)
	testSigners = make(map[string]ssh.Signer, n)
	testPrivateKeys = make(map[string]interface{}, n)
	testPublicKeys = make(map[string]ssh.PublicKey, n)

	for t, k := range testdata.PEMBytes {
		testPrivateKeys[t], err = ssh.ParseRawPrivateKey(k)
		if err != nil {
			panic(fmt.Sprintf("Unable to parse test key %s: %v", t, err))
		}
		testSigners[t], err = ssh.NewSignerFromKey(testPrivateKeys[t])
		if err != nil {
			panic(fmt.Sprintf("Unable to create signer for test key %s: %v", t, err))
		}
		testPublicKeys[t] = testSigners[t].PublicKey()
	}

	randomStr := fmt.Sprintf("%v", rand.Intn(5000))
	socketFile := "/tmp/gosocket" + randomStr + ".sock"
	setupSshAgent(socketFile)
	startSSHServer()
}

func setupSshAgent(socketFile string) {
	done := make(chan string, 1)
	a := agent.NewKeyring()
	go func(done chan<- string) {
		ln, err := net.Listen("unix", socketFile)
		if err != nil {
			panic(fmt.Sprintf("Couldn't create socket for tests %v", err))
		}
		// Need to wait until the socket is setup
		firstTime := true
		for {
			if firstTime == true {
				done <- socketFile
				firstTime = false
			}
			c, err := ln.Accept()
			defer c.Close()
			if err != nil {
				panic(fmt.Sprintf("Couldn't accept connection to agent tests %v", err))
			}
			go func(c io.ReadWriter) {
				err := agent.ServeAgent(a, c)
				if err != nil {
					panic(fmt.Sprintf("Couldn't serve ssh agent for tests %v", err))
				}

			}(c)
		}

	}(done)
	sshAgentSocket = <-done
}

func addKeytoSSHAgent(key agent.AddedKey) {
	aConn, _ := net.Dial("unix", sshAgentSocket)
	sshAgent := agent.NewClient(aConn)
	sshAgent.Add(key)
}

func removeKeyfromSSHAgent(key ssh.PublicKey) {
	aConn, _ := net.Dial("unix", sshAgentSocket)
	sshAgent := agent.NewClient(aConn)
	sshAgent.Remove(key)
}

func startSSHServer() {
	done := make(chan bool, 1)
	go func(done chan<- bool) {
		glssh.Handle(func(s glssh.Session) {
			authorizedKey := ssh.MarshalAuthorizedKey(s.PublicKey())
			io.WriteString(s, fmt.Sprintf("public key used by %s:\n", s.User()))
			s.Write(authorizedKey)
		})

		publicKeyOption := glssh.PublicKeyAuth(func(ctx glssh.Context, key glssh.PublicKey) bool {
			for _, pubk := range testPublicKeys {
				if glssh.KeysEqual(key, pubk) {
					return true
				}
			}
			return false // use glssh.KeysEqual() to compare against known keys
		})

		fmt.Println("starting ssh server on port 2222...")
		done <- true
		panic(glssh.ListenAndServe(":2222", nil, publicKeyOption))
	}(done)
	<-done
}

func TestMakeSigner(t *testing.T) {
	tests := []struct {
		name     string
		key      mockSSHKey
		expected ssh.Signer
	}{
		{name: "Basic key signer with valid rsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			expected: testSigners["rsa"],
		},
		{name: "Basic key signer with valid dsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			expected: testSigners["dsa"],
		},
		{name: "Basic key signer with valid ecdsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["ecdsa"],
			},
			expected: testSigners["ecdsa"],
		},
		{name: "Basic key signer with valid user key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["user"],
			},
			expected: testSigners["user"],
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Write content of the key to the keyname file
			ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			returned, _ := makeSigner(tt.key.keyname)
			if !reflect.DeepEqual(returned, tt.expected) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}
			os.Remove(tt.key.keyname)
		})
	}
}

func TestMakeKeyring(t *testing.T) {
	tests := []struct {
		name     string
		useagent bool
		key      mockSSHKey
		expected ssh.AuthMethod
	}{
		{name: "Basic key ring with valid rsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			expected: ssh.PublicKeys(testSigners["rsa"]),
		},
		{name: "Basic key ring with valid dsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			expected: ssh.PublicKeys(testSigners["dsa"]),
		},
		{name: "Basic key ring with valid ecdsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["ecdsa"],
			},
			expected: ssh.PublicKeys(testSigners["ecdsa"]),
		},
		{name: "Basic key ring with valid user key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["user"],
			},
			expected: ssh.PublicKeys(testSigners["user"]),
		},
		{name: "Basic key ring agent with valid rsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["rsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["rsa"]},
				pubkey:  testPublicKeys["rsa"],
			},
			expected: ssh.PublicKeys(testSigners["rsa"]),
		},
		{name: "Basic key ring agent with valid dsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["dsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["dsa"]},
				pubkey:  testPublicKeys["dsa"],
			},
			expected: ssh.PublicKeys(testSigners["dsa"]),
		},
		{name: "Basic key ring agent with valid ecdsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["ecdsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["ecdsa"]},
				pubkey:  testPublicKeys["ecdsa"],
			},
			expected: ssh.PublicKeys(testSigners["ecdsa"]),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.useagent == true {
				addKeytoSSHAgent(tt.key.privkey)
			}
			// Write content of the key to the keyname file
			if tt.key.keyname != "" {
				ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			}
			returned := makeKeyring(tt.key.keyname, tt.useagent)
			// DeepEqual always returns false for functions unless nil
			// hence converting to string to compare
			check1 := reflect.ValueOf(returned).String()
			check2 := reflect.ValueOf(tt.expected).String()
			if !reflect.DeepEqual(check1, check2) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}
			if tt.useagent == true {
				removeKeyfromSSHAgent(tt.key.pubkey)
			}
			if tt.key.keyname != "" {
				os.Remove(tt.key.keyname)
			}
		})
	}
}

//func Run(machines []string, cmd string, user string, key string, useAgent bool)
func TestRun(t *testing.T) {
	tests := []struct {
		name     string
		machines []string
		port     string
		user     string
		cmd      string
		key      mockSSHKey
		useagent bool
		expected bool
	}{
		{name: "Basic with valid rsa key",
			machines: []string{"localhost"},
			port:     "2222",
			cmd:      "ls",
			user:     "testuser",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			useagent: false,
			expected: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.useagent == true {
				addKeytoSSHAgent(tt.key.privkey)
			}
			// Write content of the key to the keyname file
			if tt.key.keyname != "" {
				ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			}
			returned := Run(tt.machines,
				tt.port,
				tt.cmd,
				tt.user,
				tt.key.keyname,
				tt.useagent)

			if !(returned == tt.expected) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}

			if tt.useagent == true {
				removeKeyfromSSHAgent(tt.key.pubkey)
			}
			if tt.key.keyname != "" {
				os.Remove(tt.key.keyname)
			}
		})
	}
}

func TestTearDown(t *testing.T) {
	tests := []struct {
		name string
		id   string
	}{
		{name: "Teardown SSH Agent",
			id: "sshAgentTdown"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.id == "sshAgentTdown" {
				os.Remove(sshAgentSocket)
			}

		})

	}
}
<commit_msg>Add more tests<commit_after>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"fmt"
	glssh "github.com/gliderlabs/ssh"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/testdata"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"reflect"
	"testing"
)

type mockSSHKey struct {
	keyname string
	content []byte
	privkey agent.AddedKey
	pubkey  ssh.PublicKey
}

var (
	testPrivateKeys map[string]interface{}
	testSigners     map[string]ssh.Signer
	testPublicKeys  map[string]ssh.PublicKey
	sshAgentSocket  string
)

func init() {
	var err error

	n := len(testdata.PEMBytes)
	testSigners = make(map[string]ssh.Signer, n)
	testPrivateKeys = make(map[string]interface{}, n)
	testPublicKeys = make(map[string]ssh.PublicKey, n)

	for t, k := range testdata.PEMBytes {
		testPrivateKeys[t], err = ssh.ParseRawPrivateKey(k)
		if err != nil {
			panic(fmt.Sprintf("Unable to parse test key %s: %v", t, err))
		}
		testSigners[t], err = ssh.NewSignerFromKey(testPrivateKeys[t])
		if err != nil {
			panic(fmt.Sprintf("Unable to create signer for test key %s: %v", t, err))
		}
		testPublicKeys[t] = testSigners[t].PublicKey()
	}

	randomStr := fmt.Sprintf("%v", rand.Intn(5000))
	socketFile := "/tmp/gosocket" + randomStr + ".sock"
	setupSshAgent(socketFile)
	startSSHServer()
}

func setupSshAgent(socketFile string) {
	done := make(chan string, 1)
	a := agent.NewKeyring()
	go func(done chan<- string) {
		ln, err := net.Listen("unix", socketFile)
		if err != nil {
			panic(fmt.Sprintf("Couldn't create socket for tests %v", err))
		}
		// Need to wait until the socket is setup
		firstTime := true
		for {
			if firstTime == true {
				done <- socketFile
				firstTime = false
			}
			c, err := ln.Accept()
			defer c.Close()
			if err != nil {
				panic(fmt.Sprintf("Couldn't accept connection to agent tests %v", err))
			}
			go func(c io.ReadWriter) {
				err := agent.ServeAgent(a, c)
				if err != nil {
					panic(fmt.Sprintf("Couldn't serve ssh agent for tests %v", err))
				}

			}(c)
		}

	}(done)
	sshAgentSocket = <-done
}

func addKeytoSSHAgent(key agent.AddedKey) {
	aConn, _ := net.Dial("unix", sshAgentSocket)
	sshAgent := agent.NewClient(aConn)
	sshAgent.Add(key)
}

func removeKeyfromSSHAgent(key ssh.PublicKey) {
	aConn, _ := net.Dial("unix", sshAgentSocket)
	sshAgent := agent.NewClient(aConn)
	sshAgent.Remove(key)
}

func startSSHServer() {
	done := make(chan bool, 1)
	go func(done chan<- bool) {
		glssh.Handle(func(s glssh.Session) {
			authorizedKey := ssh.MarshalAuthorizedKey(s.PublicKey())
			io.WriteString(s, fmt.Sprintf("public key used by %s:\n", s.User()))
			s.Write(authorizedKey)
		})

		publicKeyOption := glssh.PublicKeyAuth(func(ctx glssh.Context, key glssh.PublicKey) bool {
			for _, pubk := range testPublicKeys {
				if glssh.KeysEqual(key, pubk) {
					return true
				}
			}
			return false // use glssh.KeysEqual() to compare against known keys
		})

		fmt.Println("starting ssh server on port 2222...")
		done <- true
		panic(glssh.ListenAndServe(":2222", nil, publicKeyOption))
	}(done)
	<-done
}

func TestMakeSigner(t *testing.T) {
	tests := []struct {
		name     string
		key      mockSSHKey
		expected ssh.Signer
	}{
		{name: "Basic key signer with valid rsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			expected: testSigners["rsa"],
		},
		{name: "Basic key signer with valid dsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			expected: testSigners["dsa"],
		},
		{name: "Basic key signer with valid ecdsa key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["ecdsa"],
			},
			expected: testSigners["ecdsa"],
		},
		{name: "Basic key signer with valid user key",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["user"],
			},
			expected: testSigners["user"],
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Write content of the key to the keyname file
			ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			returned, _ := makeSigner(tt.key.keyname)
			if !reflect.DeepEqual(returned, tt.expected) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}
			os.Remove(tt.key.keyname)
		})
	}
}

func TestMakeKeyring(t *testing.T) {
	tests := []struct {
		name     string
		useagent bool
		key      mockSSHKey
		expected ssh.AuthMethod
	}{
		{name: "Basic key ring with valid rsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			expected: ssh.PublicKeys(testSigners["rsa"]),
		},
		{name: "Basic key ring with valid dsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			expected: ssh.PublicKeys(testSigners["dsa"]),
		},
		{name: "Basic key ring with valid ecdsa key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["ecdsa"],
			},
			expected: ssh.PublicKeys(testSigners["ecdsa"]),
		},
		{name: "Basic key ring with valid user key",
			useagent: false,
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["user"],
			},
			expected: ssh.PublicKeys(testSigners["user"]),
		},
		{name: "Basic key ring agent with valid rsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["rsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["rsa"]},
				pubkey:  testPublicKeys["rsa"],
			},
			expected: ssh.PublicKeys(testSigners["rsa"]),
		},
		{name: "Basic key ring agent with valid dsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["dsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["dsa"]},
				pubkey:  testPublicKeys["dsa"],
			},
			expected: ssh.PublicKeys(testSigners["dsa"]),
		},
		{name: "Basic key ring agent with valid ecdsa key",
			useagent: true,
			key: mockSSHKey{
				keyname: "",
				content: testdata.PEMBytes["ecdsa"],
				privkey: agent.AddedKey{PrivateKey: testPrivateKeys["ecdsa"]},
				pubkey:  testPublicKeys["ecdsa"],
			},
			expected: ssh.PublicKeys(testSigners["ecdsa"]),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.useagent == true {
				addKeytoSSHAgent(tt.key.privkey)
			}
			// Write content of the key to the keyname file
			if tt.key.keyname != "" {
				ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			}
			returned := makeKeyring(tt.key.keyname, tt.useagent)
			// DeepEqual always returns false for functions unless nil
			// hence converting to string to compare
			check1 := reflect.ValueOf(returned).String()
			check2 := reflect.ValueOf(tt.expected).String()
			if !reflect.DeepEqual(check1, check2) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}
			if tt.useagent == true {
				removeKeyfromSSHAgent(tt.key.pubkey)
			}
			if tt.key.keyname != "" {
				os.Remove(tt.key.keyname)
			}
		})
	}
}

//func Run(machines []string, cmd string, user string, key string, useAgent bool)
func TestRun(t *testing.T) {
	tests := []struct {
		name     string
		machines []string
		port     string
		user     string
		cmd      string
		key      mockSSHKey
		useagent bool
		expected bool
	}{
		{name: "Basic with valid rsa key",
			machines: []string{"localhost"},
			port:     "2222",
			cmd:      "ls",
			user:     "testuser",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			useagent: false,
			expected: true,
		},
		{name: "Basic with valid rsa key wrong hostname",
			machines: []string{"bogushost"},
			port:     "2222",
			cmd:      "ls",
			user:     "testuser",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["rsa"],
			},
			useagent: false,
			expected: false,
		},
		{name: "Basic with valid dsa key",
			machines: []string{"localhost"},
			port:     "2222",
			cmd:      "ls",
			user:     "testuser",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			useagent: false,
			expected: true,
		},
		{name: "Basic with valid dsa key wrong hostname",
			machines: []string{"bogushost"},
			port:     "2222",
			cmd:      "ls",
			user:     "testuser",
			key: mockSSHKey{
				keyname: "/tmp/mockkey",
				content: testdata.PEMBytes["dsa"],
			},
			useagent: false,
			expected: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.useagent == true {
				addKeytoSSHAgent(tt.key.privkey)
			}
			// Write content of the key to the keyname file
			if tt.key.keyname != "" {
				ioutil.WriteFile(tt.key.keyname, tt.key.content, 0644)
			}
			returned := Run(tt.machines,
				tt.port,
				tt.cmd,
				tt.user,
				tt.key.keyname,
				tt.useagent)

			if !(returned == tt.expected) {
				t.Errorf("Value received: %v expected %v", returned, tt.expected)
			}

			if tt.useagent == true {
				removeKeyfromSSHAgent(tt.key.pubkey)
			}
			if tt.key.keyname != "" {
				os.Remove(tt.key.keyname)
			}
		})
	}
}

func TestTearDown(t *testing.T) {
	tests := []struct {
		name string
		id   string
	}{
		{name: "Teardown SSH Agent",
			id: "sshAgentTdown"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.id == "sshAgentTdown" {
				os.Remove(sshAgentSocket)
			}

		})

	}
}
<|endoftext|>


<commit_before>package discover

//
//
// TODO:
//   should be able to guess article link format statistically
//   handle/allow subdomains (eg: www1.politicalbetting.com)
//   filter unwanted navlinks (eg "mirror.co.uk/all-about/fred bloggs")
//   HTTP error handling
//   multiple url formats (eg spectator has multiple cms's)
//   logging

import (
	"code.google.com/p/cascadia"
	"code.google.com/p/go.net/html"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
)

type Logger interface {
	Printf(format string, v ...interface{})
}

type NullLogger struct{}

func (l NullLogger) Printf(format string, v ...interface{}) {
}

type DiscovererDef struct {
	Name   string
	URL    string
	ArtPat []string
	NavSel string
	// BaseErrorThreshold is starting number of http errors to accept before
	// bailing out.
	// error threshold formula: base + 10% of successful request count
	BaseErrorThreshold int

	// Hostpat is a regex matching accepted domains
	// if empty, reject everything on a different domain
	HostPat string
}

type Discoverer struct {
	Name               string
	StartURL           url.URL
	ArtPats            []*regexp.Regexp
	NavLinkSel         cascadia.Selector
	BaseErrorThreshold int
	StripFragments     bool
	StripQuery         bool
	HostPat            *regexp.Regexp

	ErrorLog Logger
	InfoLog  Logger
	Stats    struct {
		ErrorCount int
		FetchCount int
	}
}

func NewDiscoverer(cfg DiscovererDef) (*Discoverer, error) {
	disc := &Discoverer{}
	u, err := url.Parse(cfg.URL)
	if err != nil {
		return nil, err
	}
	disc.Name = cfg.Name
	disc.StartURL = *u
	disc.ArtPats = make([]*regexp.Regexp, 0, len(cfg.ArtPat))
	for _, pat := range cfg.ArtPat {
		re, err := regexp.Compile(pat)
		if err != nil {
			return nil, err
		}
		disc.ArtPats = append(disc.ArtPats, re)
	}

	if cfg.NavSel == "" {
		disc.NavLinkSel = nil
	} else {
		sel, err := cascadia.Compile(cfg.NavSel)
		if err != nil {
			return nil, err
		}
		disc.NavLinkSel = sel
	}
	disc.BaseErrorThreshold = cfg.BaseErrorThreshold

	if cfg.HostPat != "" {
		re, err := regexp.Compile(cfg.HostPat)
		if err != nil {
			return nil, err
		}
		disc.HostPat = re
	}

	// defaults
	disc.StripFragments = true
	disc.StripQuery = true
	disc.ErrorLog = NullLogger{}
	disc.InfoLog = NullLogger{}
	return disc, nil
}

func (disc *Discoverer) Run(client *http.Client) (LinkSet, error) {

	queued := make(LinkSet) // nav pages to scan for article links
	seen := make(LinkSet)   // nav pages we've scanned
	arts := make(LinkSet)   // article links we've found so far

	queued.Add(disc.StartURL)

	for len(queued) > 0 {
		u := queued.Pop()
		seen.Add(u)
		//

		root, err := disc.fetchAndParse(client, &u)
		if err != nil {
			disc.ErrorLog.Printf("%s\n", err.Error())
			disc.Stats.ErrorCount++
			if disc.Stats.ErrorCount > disc.BaseErrorThreshold+(disc.Stats.FetchCount/10) {
				return nil, errors.New("Error threshold exceeded")
			} else {
				continue
			}
		}
		disc.Stats.FetchCount++

		navLinks, err := disc.findNavLinks(root)
		if err != nil {
			return nil, err
		}
		for navLink, _ := range navLinks {
			if _, got := seen[navLink]; !got {
				queued.Add(navLink)
			}
		}

		foo, err := disc.findArticles(root)
		if err != nil {
			return nil, err
		}
		arts.Merge(foo)

		disc.InfoLog.Printf("Visited %s, found %d articles\n", u.String(), len(foo))
	}

	return arts, nil
}

func (disc *Discoverer) fetchAndParse(client *http.Client, pageURL *url.URL) (*html.Node, error) {
	resp, err := client.Get(pageURL.String())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		err = errors.New(fmt.Sprintf("HTTP code %d (%s)", resp.StatusCode, pageURL.String()))

		return nil, err

	}

	root, err := html.Parse(resp.Body)
	if err != nil {
		return nil, err
	}

	return root, nil
}

var aSel cascadia.Selector = cascadia.MustCompile("a")

func (disc *Discoverer) findArticles(root *html.Node) (LinkSet, error) {
	arts := make(LinkSet)
	for _, a := range aSel.MatchAll(root) {
		// fetch url and extend to absolute
		link, err := disc.StartURL.Parse(GetAttr(a, "href"))
		if err != nil {
			continue
		}

		if !disc.isHostGood(link.Host) {
			continue
		}

		foo := link.RequestURI()
		accept := false
		for _, pat := range disc.ArtPats {
			if pat.MatchString(foo) {
				accept = true
				break
			}
		}
		if !accept {
			continue
		}

		if disc.StripFragments {
			link.Fragment = ""
		}
		if disc.StripQuery {
			link.RawQuery = ""
		}

		arts[*link] = true
	}
	return arts, nil
}

func (disc *Discoverer) findNavLinks(root *html.Node) (LinkSet, error) {
	navLinks := make(LinkSet)
	if disc.NavLinkSel == nil {
		return navLinks, nil
	}
	for _, a := range disc.NavLinkSel.MatchAll(root) {
		link, err := disc.StartURL.Parse(GetAttr(a, "href"))
		if err != nil {
			continue
		}

		if !disc.isHostGood(link.Host) {
			continue
		}

		link.Fragment = ""

		navLinks[*link] = true
	}
	return navLinks, nil
}

// is host domain one we'll accept?
func (disc *Discoverer) isHostGood(host string) bool {
	if disc.HostPat != nil {
		return disc.HostPat.MatchString(host)
	}
	return host == disc.StartURL.Host
}

// GetAttr retrieved the value of an attribute on a node.
// Returns empty string if attribute doesn't exist.
func GetAttr(n *html.Node, attr string) string {
	for _, a := range n.Attr {
		if a.Key == attr {
			return a.Val
		}
	}
	return ""
}

// GetTextContent recursively fetches the text for a node
func GetTextContent(n *html.Node) string {
	if n.Type == html.TextNode {
		return n.Data
	}
	txt := ""
	for child := n.FirstChild; child != nil; child = child.NextSibling {
		txt += GetTextContent(child)
	}

	return txt
}

// CompressSpace reduces all whitespace sequences (space, tabs, newlines etc) in a string to a single space.
// Leading/trailing space is trimmed.
// Has the effect of converting multiline strings to one line.
func CompressSpace(s string) string {
	multispacePat := regexp.MustCompile(`[\s]+`)
	s = strings.TrimSpace(multispacePat.ReplaceAllLiteralString(s, " "))
	return s
}
<commit_msg>make discover stats a public struct<commit_after>package discover

//
//
// TODO:
//   should be able to guess article link format statistically
//   handle/allow subdomains (eg: www1.politicalbetting.com)
//   filter unwanted navlinks (eg "mirror.co.uk/all-about/fred bloggs")
//   HTTP error handling
//   multiple url formats (eg spectator has multiple cms's)
//   logging

import (
	"code.google.com/p/cascadia"
	"code.google.com/p/go.net/html"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"strings"
)

type Logger interface {
	Printf(format string, v ...interface{})
}

type NullLogger struct{}

func (l NullLogger) Printf(format string, v ...interface{}) {
}

type DiscovererDef struct {
	Name   string
	URL    string
	ArtPat []string
	NavSel string
	// BaseErrorThreshold is starting number of http errors to accept before
	// bailing out.
	// error threshold formula: base + 10% of successful request count
	BaseErrorThreshold int

	// Hostpat is a regex matching accepted domains
	// if empty, reject everything on a different domain
	HostPat string
}

type DiscoverStats struct {
	ErrorCount int
	FetchCount int
}

type Discoverer struct {
	Name               string
	StartURL           url.URL
	ArtPats            []*regexp.Regexp
	NavLinkSel         cascadia.Selector
	BaseErrorThreshold int
	StripFragments     bool
	StripQuery         bool
	HostPat            *regexp.Regexp

	ErrorLog Logger
	InfoLog  Logger
	Stats    DiscoverStats
}

func NewDiscoverer(cfg DiscovererDef) (*Discoverer, error) {
	disc := &Discoverer{}
	u, err := url.Parse(cfg.URL)
	if err != nil {
		return nil, err
	}
	disc.Name = cfg.Name
	disc.StartURL = *u
	disc.ArtPats = make([]*regexp.Regexp, 0, len(cfg.ArtPat))
	for _, pat := range cfg.ArtPat {
		re, err := regexp.Compile(pat)
		if err != nil {
			return nil, err
		}
		disc.ArtPats = append(disc.ArtPats, re)
	}

	if cfg.NavSel == "" {
		disc.NavLinkSel = nil
	} else {
		sel, err := cascadia.Compile(cfg.NavSel)
		if err != nil {
			return nil, err
		}
		disc.NavLinkSel = sel
	}
	disc.BaseErrorThreshold = cfg.BaseErrorThreshold

	if cfg.HostPat != "" {
		re, err := regexp.Compile(cfg.HostPat)
		if err != nil {
			return nil, err
		}
		disc.HostPat = re
	}

	// defaults
	disc.StripFragments = true
	disc.StripQuery = true
	disc.ErrorLog = NullLogger{}
	disc.InfoLog = NullLogger{}
	return disc, nil
}

func (disc *Discoverer) Run(client *http.Client) (LinkSet, error) {
	// reset stats
	disc.Stats = DiscoverStats{}

	queued := make(LinkSet) // nav pages to scan for article links
	seen := make(LinkSet)   // nav pages we've scanned
	arts := make(LinkSet)   // article links we've found so far

	queued.Add(disc.StartURL)

	for len(queued) > 0 {
		u := queued.Pop()
		seen.Add(u)
		//

		root, err := disc.fetchAndParse(client, &u)
		if err != nil {
			disc.ErrorLog.Printf("%s\n", err.Error())
			disc.Stats.ErrorCount++
			if disc.Stats.ErrorCount > disc.BaseErrorThreshold+(disc.Stats.FetchCount/10) {
				return nil, errors.New("Error threshold exceeded")
			} else {
				continue
			}
		}
		disc.Stats.FetchCount++

		navLinks, err := disc.findNavLinks(root)
		if err != nil {
			return nil, err
		}
		for navLink, _ := range navLinks {
			if _, got := seen[navLink]; !got {
				queued.Add(navLink)
			}
		}

		foo, err := disc.findArticles(root)
		if err != nil {
			return nil, err
		}
		arts.Merge(foo)

		disc.InfoLog.Printf("Visited %s, found %d articles\n", u.String(), len(foo))
	}

	return arts, nil
}

func (disc *Discoverer) fetchAndParse(client *http.Client, pageURL *url.URL) (*html.Node, error) {
	resp, err := client.Get(pageURL.String())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		err = errors.New(fmt.Sprintf("HTTP code %d (%s)", resp.StatusCode, pageURL.String()))

		return nil, err

	}

	root, err := html.Parse(resp.Body)
	if err != nil {
		return nil, err
	}

	return root, nil
}

var aSel cascadia.Selector = cascadia.MustCompile("a")

func (disc *Discoverer) findArticles(root *html.Node) (LinkSet, error) {
	arts := make(LinkSet)
	for _, a := range aSel.MatchAll(root) {
		// fetch url and extend to absolute
		link, err := disc.StartURL.Parse(GetAttr(a, "href"))
		if err != nil {
			continue
		}

		if !disc.isHostGood(link.Host) {
			continue
		}

		foo := link.RequestURI()
		accept := false
		for _, pat := range disc.ArtPats {
			if pat.MatchString(foo) {
				accept = true
				break
			}
		}
		if !accept {
			continue
		}

		if disc.StripFragments {
			link.Fragment = ""
		}
		if disc.StripQuery {
			link.RawQuery = ""
		}

		arts[*link] = true
	}
	return arts, nil
}

func (disc *Discoverer) findNavLinks(root *html.Node) (LinkSet, error) {
	navLinks := make(LinkSet)
	if disc.NavLinkSel == nil {
		return navLinks, nil
	}
	for _, a := range disc.NavLinkSel.MatchAll(root) {
		link, err := disc.StartURL.Parse(GetAttr(a, "href"))
		if err != nil {
			continue
		}

		if !disc.isHostGood(link.Host) {
			continue
		}

		link.Fragment = ""

		navLinks[*link] = true
	}
	return navLinks, nil
}

// is host domain one we'll accept?
func (disc *Discoverer) isHostGood(host string) bool {
	if disc.HostPat != nil {
		return disc.HostPat.MatchString(host)
	}
	return host == disc.StartURL.Host
}

// GetAttr retrieved the value of an attribute on a node.
// Returns empty string if attribute doesn't exist.
func GetAttr(n *html.Node, attr string) string {
	for _, a := range n.Attr {
		if a.Key == attr {
			return a.Val
		}
	}
	return ""
}

// GetTextContent recursively fetches the text for a node
func GetTextContent(n *html.Node) string {
	if n.Type == html.TextNode {
		return n.Data
	}
	txt := ""
	for child := n.FirstChild; child != nil; child = child.NextSibling {
		txt += GetTextContent(child)
	}

	return txt
}

// CompressSpace reduces all whitespace sequences (space, tabs, newlines etc) in a string to a single space.
// Leading/trailing space is trimmed.
// Has the effect of converting multiline strings to one line.
func CompressSpace(s string) string {
	multispacePat := regexp.MustCompile(`[\s]+`)
	s = strings.TrimSpace(multispacePat.ReplaceAllLiteralString(s, " "))
	return s
}
<|endoftext|>


<commit_before>// Package extract provides a set of functions to extract parameters from an http:Request.
package extract

import (
	"fmt"
	"image/color"
	"net/http"
	"strconv"
	"strings"

	"github.com/taironas/tinygraphs/format"
)

// extract hexadecimal code background from HTTP request and return color.RGBA
func Background(req *http.Request) (color.RGBA, error) {
	bg := req.FormValue("bg")
	if len(bg) == 0 {
		return color.RGBA{}, fmt.Errorf("background: wrong input")
	}
	r, g, b, err := hexToRGB(bg)
	return color.RGBA{uint8(r), uint8(g), uint8(b), uint8(255)}, err
}

// extract hexadecimal code foreground from HTTP request and return color.RGBA
func Foreground(req *http.Request) (color.RGBA, error) {
	fg := req.FormValue("fg")
	if len(fg) == 0 {
		return color.RGBA{}, fmt.Errorf("background: wrong input")
	}
	r, g, b, err := hexToRGB(fg)
	return color.RGBA{uint8(r), uint8(g), uint8(b), uint8(255)}, err
}

// extract size from HTTP request and return it.
func Size(r *http.Request) int {
	strSize := r.FormValue("size")
	if len(strSize) > 0 {
		if size, errSize := strconv.ParseInt(strSize, 0, 64); errSize == nil {
			isize := int(size)
			if isize > 0 && isize < 1000 {
				return int(size)
			}
		}
	}
	return 210
}

func Format(r *http.Request) format.Format {
	strFmt := strings.ToLower(r.FormValue("fmt"))
	if len(strFmt) > 0 {
		if strFmt == "svg" {
			return format.SVG
		} else if strFmt == "jpeg" || strFmt == "jpg" {
			return format.JPEG
		}
	}
	return format.JPEG
}

// HexToRGB converts an Hex string to a RGB triple.
func hexToRGB(h string) (uint8, uint8, uint8, error) {
	if len(h) > 0 && h[0] == '#' {
		h = h[1:]
	}
	if len(h) == 3 {
		h = h[:1] + h[:1] + h[1:2] + h[1:2] + h[2:] + h[2:]
	}
	if len(h) == 6 {
		if rgb, err := strconv.ParseUint(string(h), 16, 32); err == nil {
			return uint8(rgb >> 16), uint8((rgb >> 8) & 0xFF), uint8(rgb & 0xFF), nil
		} else {
			return 0, 0, 0, err
		}
	}
	return 0, 0, 0, nil
}

func Theme(r *http.Request) string {
	strTheme := strings.ToLower(r.FormValue("theme"))
	if len(strTheme) > 0 {
		return strTheme
	}
	return "base"
}

// NumColors returns the number of colors in http request.
// Right now we support numbers between 2 and 4.
func NumColors(r *http.Request) int64 {
	s := strings.ToLower(r.FormValue("numcolors"))
	if len(s) > 0 {
		if n, err := strconv.ParseInt(s, 0, 64); err == nil {
			if n >= 2 && n <= 4 {
				return n
			}
		}
	}
	return 2
}

func Hexalines(r *http.Request) int64 {
	s := strings.ToLower(r.FormValue("hexalines"))
	if len(s) > 0 {
		if n, err := strconv.ParseInt(s, 0, 64); err == nil {
			if n%6 == 0 || n%4 == 0 {
				return n
			}
		}
	}
	return 6
}
<commit_msg>make size default value divisible by 3 and 2<commit_after>// Package extract provides a set of functions to extract parameters from an http:Request.
package extract

import (
	"fmt"
	"image/color"
	"net/http"
	"strconv"
	"strings"

	"github.com/taironas/tinygraphs/format"
)

// extract hexadecimal code background from HTTP request and return color.RGBA
func Background(req *http.Request) (color.RGBA, error) {
	bg := req.FormValue("bg")
	if len(bg) == 0 {
		return color.RGBA{}, fmt.Errorf("background: wrong input")
	}
	r, g, b, err := hexToRGB(bg)
	return color.RGBA{uint8(r), uint8(g), uint8(b), uint8(255)}, err
}

// extract hexadecimal code foreground from HTTP request and return color.RGBA
func Foreground(req *http.Request) (color.RGBA, error) {
	fg := req.FormValue("fg")
	if len(fg) == 0 {
		return color.RGBA{}, fmt.Errorf("background: wrong input")
	}
	r, g, b, err := hexToRGB(fg)
	return color.RGBA{uint8(r), uint8(g), uint8(b), uint8(255)}, err
}

// extract size from HTTP request and return it.
func Size(r *http.Request) int {
	strSize := r.FormValue("size")
	if len(strSize) > 0 {
		if size, errSize := strconv.ParseInt(strSize, 0, 64); errSize == nil {
			isize := int(size)
			if isize > 0 && isize < 1000 {
				return int(size)
			}
		}
	}
	return 240
}

func Format(r *http.Request) format.Format {
	strFmt := strings.ToLower(r.FormValue("fmt"))
	if len(strFmt) > 0 {
		if strFmt == "svg" {
			return format.SVG
		} else if strFmt == "jpeg" || strFmt == "jpg" {
			return format.JPEG
		}
	}
	return format.JPEG
}

// HexToRGB converts an Hex string to a RGB triple.
func hexToRGB(h string) (uint8, uint8, uint8, error) {
	if len(h) > 0 && h[0] == '#' {
		h = h[1:]
	}
	if len(h) == 3 {
		h = h[:1] + h[:1] + h[1:2] + h[1:2] + h[2:] + h[2:]
	}
	if len(h) == 6 {
		if rgb, err := strconv.ParseUint(string(h), 16, 32); err == nil {
			return uint8(rgb >> 16), uint8((rgb >> 8) & 0xFF), uint8(rgb & 0xFF), nil
		} else {
			return 0, 0, 0, err
		}
	}
	return 0, 0, 0, nil
}

func Theme(r *http.Request) string {
	strTheme := strings.ToLower(r.FormValue("theme"))
	if len(strTheme) > 0 {
		return strTheme
	}
	return "base"
}

// NumColors returns the number of colors in http request.
// Right now we support numbers between 2 and 4.
func NumColors(r *http.Request) int64 {
	s := strings.ToLower(r.FormValue("numcolors"))
	if len(s) > 0 {
		if n, err := strconv.ParseInt(s, 0, 64); err == nil {
			if n >= 2 && n <= 4 {
				return n
			}
		}
	}
	return 2
}

func Hexalines(r *http.Request) int64 {
	s := strings.ToLower(r.FormValue("hexalines"))
	if len(s) > 0 {
		if n, err := strconv.ParseInt(s, 0, 64); err == nil {
			if n%6 == 0 || n%4 == 0 {
				return n
			}
		}
	}
	return 6
}
<|endoftext|>


<commit_before>package extend
import (
	"reflect"
	"unsafe"
)

const initialSize = 4

type interfaceHeader struct {
	t uintptr
	data uintptr
}

// Pusher must be passed a pointer to a slice. It returns
// a function that pushes a new value onto the end of the
// slice, reallocating the slice if necessary.
func Pusher(ap interface{}) func(interface{}) {
	v := reflect.NewValue(ap).(*reflect.PtrValue).Elem().(*reflect.SliceValue)
	h := (*reflect.SliceHeader)(unsafe.Pointer(v.Addr()))
	t := v.Type().(*reflect.SliceType)
	elemType := t.Elem()
	esize := elemType.Size()
	_, isInterface := t.Elem().(*reflect.InterfaceType)
	unsafeCopy := !isInterface && esize <= uintptr(unsafe.Sizeof(uintptr(0)))
	if !unsafeCopy {
		return func(x interface{}) {
			len, cap := h.Len, h.Cap
			if len < cap {
				h.Len++
			}else{
				if cap == 0 {
					cap = initialSize
				}else{
					cap *= 2
				}
				b := reflect.MakeSlice(t, len+1, cap)
				reflect.ArrayCopy(b, v)
				v.SetValue(b)
			}
			v.Elem(len).SetValue(reflect.NewValue(x))
		}
	}
	// Nasty unsafe hackery:
	//
	// We know that the size of the type fits in a pointer,
	// so the value is held directly inside the interface value.
	// We copy each element to icopy, so that we can take
	// the address of it without triggering the allocator.
	// We set up e1 as a []byte alias to the data,
	// and do the actual copy by setting up e0 as a []byte
	// alias to the array element and invoking copy().
	// TODO: enable this code for type with size > sizeof(uintptr).
	var e0 []byte
	var e1 []byte
	he0 := (*reflect.SliceHeader)(unsafe.Pointer(&e0))
	he1 := (*reflect.SliceHeader)(unsafe.Pointer(&e1))
	var icopy interface{}
	he1.Data = uintptr(unsafe.Pointer(&icopy)) + uintptr(unsafe.Offsetof(interfaceHeader{}.data))
	he1.Len = int(esize)
	he1.Cap = int(esize)

	return func(x interface{}) {
		if reflect.Typeof(x) != elemType {
			panic("wrong type pushed")
		}
		len, cap := h.Len, h.Cap
		if len < cap {
			h.Len++
		}else{
			if cap == 0 {
				cap = initialSize
			}else{
				cap *= 2
			}
			b := reflect.MakeSlice(t, len+1, cap)
			reflect.ArrayCopy(b, v)
			v.SetValue(b)
		}
		icopy = x
		he0.Data = h.Data + esize * uintptr(len)
		he0.Len = int(esize)
		he0.Cap = int(esize)
		copy(e0, e1)
		e0 = nil
	}
}
<commit_msg>make unsafe hackery work for larger types.<commit_after>package extend
import (
	"reflect"
	"unsafe"
)

const initialSize = 4

// This type mirrors the actual runtime interface type.
type interfaceHeader struct {
	t uintptr
	data uintptr
}

// Pusher must be passed a pointer to a slice. It returns
// a function that pushes a new value onto the end of the
// slice, reallocating the slice if necessary.
func Pusher(ap interface{}) func(interface{}) {
	v := reflect.NewValue(ap).(*reflect.PtrValue).Elem().(*reflect.SliceValue)
	h := (*reflect.SliceHeader)(unsafe.Pointer(v.Addr()))
	t := v.Type().(*reflect.SliceType)
	elemType := t.Elem()

	if _, ok := elemType.(*reflect.InterfaceType); ok {
		// If the array's element type is an interface,
		// then we must go the slow route, because
		// the type being pushed will not be identical
		// the type in the array, so we rely on reflect
		// to do the value->interface conversion.
		return func(x interface{}) {
			len, cap := h.Len, h.Cap
			if len < cap {
				h.Len++
			}else{
				if cap == 0 {
					cap = initialSize
				}else{
					cap *= 2
				}
				b := reflect.MakeSlice(t, len+1, cap)
				reflect.ArrayCopy(b, v)
				v.SetValue(b)
			}
			v.Elem(len).SetValue(reflect.NewValue(x))
		}
	}
	// Nasty unsafe hackery:
	//
	// We know that the size of the type fits in a pointer,
	// so the value is held directly inside the interface value.
	// We copy each element to icopy, so that we can take
	// the address of it without triggering the allocator.
	// We set up e1 as a []byte alias to the data,
	// and do the actual copy by setting up e0 as a []byte
	// alias to the array element and invoking copy().
	// TODO: enable this code for type with size > sizeof(uintptr).
	esize := elemType.Size()
	valueInsideInterface := esize <= uintptr(unsafe.Sizeof(uintptr(0)))

	var e0 []byte
	var e1 []byte
	he0 := (*reflect.SliceHeader)(unsafe.Pointer(&e0))
	he1 := (*reflect.SliceHeader)(unsafe.Pointer(&e1))

	var icopy interface{}
	ih := (*interfaceHeader)(unsafe.Pointer(&icopy))
	if valueInsideInterface {
		he1.Data = uintptr(unsafe.Pointer(&ih.data))
		he1.Len = int(esize)
		he1.Cap = int(esize)
	}

	return func(x interface{}) {
		if reflect.Typeof(x) != elemType {
			panic("wrong type pushed")
		}
		len, cap := h.Len, h.Cap
		if len < cap {
			h.Len++
		}else{
			if cap == 0 {
				cap = initialSize
			}else{
				cap *= 2
			}
			b := reflect.MakeSlice(t, len+1, cap)
			reflect.ArrayCopy(b, v)
			v.SetValue(b)
		}
		icopy = x
		he0.Data = h.Data + esize * uintptr(len)
		he0.Len = int(esize)
		he0.Cap = int(esize)
		if valueInsideInterface {
			copy(e0, e1)
		}else{
			he1.Data = ih.data
			he1.Len = int(esize)
			he1.Cap = int(esize)
			copy(e0, e1)
			e1 = nil
		}
		e0 = nil
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package fileutil

import (
	"os"
	"regexp"
	"strings"
)

var (
	shebangRe = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)\s`)
	extRe     = regexp.MustCompile(`\.(sh|bash)$`)
)

func HasShebang(bs []byte) bool {
	return shebangRe.Match(bs)
}

type ScriptConfidence int

const (
	ConfNotScript ScriptConfidence = iota
	ConfIfShebang
	ConfIsScript
)

func CouldBeScript(info os.FileInfo) ScriptConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.', !info.Mode().IsRegular():
		return ConfNotScript
	case extRe.MatchString(name):
		return ConfIsScript
	case strings.Contains(name, "."):
		return ConfNotScript // different extension
	case info.Size() < int64(len("#/bin/sh\n")):
		return ConfNotScript // cannot possibly hold valid shebang
	default:
		return ConfIfShebang
	}
}
<commit_msg>fileutil: include package doc<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

// Package fileutil contains code to work with shell files, also known
// as shell scripts.
package fileutil

import (
	"os"
	"regexp"
	"strings"
)

var (
	shebangRe = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)\s`)
	extRe     = regexp.MustCompile(`\.(sh|bash)$`)
)

func HasShebang(bs []byte) bool {
	return shebangRe.Match(bs)
}

type ScriptConfidence int

const (
	ConfNotScript ScriptConfidence = iota
	ConfIfShebang
	ConfIsScript
)

func CouldBeScript(info os.FileInfo) ScriptConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.', !info.Mode().IsRegular():
		return ConfNotScript
	case extRe.MatchString(name):
		return ConfIsScript
	case strings.Contains(name, "."):
		return ConfNotScript // different extension
	case info.Size() < int64(len("#/bin/sh\n")):
		return ConfNotScript // cannot possibly hold valid shebang
	default:
		return ConfIfShebang
	}
}
<|endoftext|>


<commit_before>package finder

import (
	"path/filepath"
	"os"
	"regexp"
)

type Finder interface {
	Find(basedir string) <-chan string
}

type Filter interface {
	Accept(path string, info os.FileInfo) bool
}

type minSizeFilter struct {
	size int64
}

func (filter minSizeFilter) Accept(path string, info os.FileInfo) bool {
	return info.Size() >= filter.size
}

type regexFilter struct {
	regex    *regexp.Regexp
	negative bool
}

func newRegexFilter(regex string, negative bool) regexFilter {
	return regexFilter{regexp.MustCompile(regex), negative}
}

func (filter regexFilter) Accept(path string, info os.FileInfo) bool {
	return filter.negative != filter.regex.MatchString(path)
}

var Filters = struct {
	MinSize      func(size int64) Filter
	IncludeRegex func(pattern string) Filter
	ExcludeRegex func(pattern string) Filter
}{
	MinSize:      func(size int64) Filter { return minSizeFilter{size} },
	IncludeRegex: func(regex string) Filter { return newRegexFilter(regex, false) },
	ExcludeRegex: func(regex string) Filter { return newRegexFilter(regex, true) },
}

type defaultFinder struct {
	filters []Filter
}

func (finder defaultFinder) Find(basedir string) <-chan string {
	paths := make(chan string)
	walkfn := func(path string, info os.FileInfo, err error) error {
		for _, filter := range finder.filters {
			if !filter.Accept(path, info) {
				return nil
			}
		}
		if info.Mode().IsRegular() {
			paths <- path
		}
		return nil
	}
	go func() {
		filepath.Walk(basedir, walkfn)
		close(paths)
	}()
	return paths
}

func NewFinder(filters ... Filter) Finder {
	return defaultFinder{filters: filters}
}
<commit_msg>Simplify filter declaration using type aliases<commit_after>package finder

import (
	"path/filepath"
	"os"
	"regexp"
)

type Finder interface {
	Find(basedir string) <-chan string
}

type Filter interface {
	Accept(path string, info os.FileInfo) bool
}

type minSizeFilter struct {
	size int64
}

func (filter minSizeFilter) Accept(path string, info os.FileInfo) bool {
	return info.Size() >= filter.size
}

type regexFilter struct {
	regex    *regexp.Regexp
	negative bool
}

func newRegexFilter(regex string, negative bool) regexFilter {
	return regexFilter{regexp.MustCompile(regex), negative}
}

func (filter regexFilter) Accept(path string, info os.FileInfo) bool {
	return filter.negative != filter.regex.MatchString(path)
}

type filterByInt64 func(n int64) Filter
type filterByString func(s string) Filter

var Filters = struct {
	MinSize      filterByInt64
	IncludeRegex filterByString
	ExcludeRegex filterByString
}{
	MinSize:      func(size int64) Filter { return minSizeFilter{size} },
	IncludeRegex: func(regex string) Filter { return newRegexFilter(regex, false) },
	ExcludeRegex: func(regex string) Filter { return newRegexFilter(regex, true) },
}

type defaultFinder struct {
	filters []Filter
}

func (finder defaultFinder) Find(basedir string) <-chan string {
	paths := make(chan string)
	walkfn := func(path string, info os.FileInfo, err error) error {
		for _, filter := range finder.filters {
			if !filter.Accept(path, info) {
				return nil
			}
		}
		if info.Mode().IsRegular() {
			paths <- path
		}
		return nil
	}
	go func() {
		filepath.Walk(basedir, walkfn)
		close(paths)
	}()
	return paths
}

func NewFinder(filters ... Filter) Finder {
	return defaultFinder{filters: filters}
}
<|endoftext|>


<commit_before>package finder

import (
	"bytes"
	"context"
	"fmt"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/go-graphite/carbonapi/pkg/parser"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/pkg/where"
)

type TaggedTermOp int

const (
	TaggedTermEq       TaggedTermOp = 1
	TaggedTermMatch    TaggedTermOp = 2
	TaggedTermNe       TaggedTermOp = 3
	TaggedTermNotMatch TaggedTermOp = 4
)

type TaggedTerm struct {
	Key   string
	Op    TaggedTermOp
	Value string
}

type TaggedTermList []TaggedTerm

func (s TaggedTermList) Len() int {
	return len(s)
}
func (s TaggedTermList) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s TaggedTermList) Less(i, j int) bool {
	if s[i].Op < s[j].Op {
		return true
	}
	if s[i].Op > s[j].Op {
		return false
	}
	if s[i].Key == "__name__" && s[j].Key != "__name__" {
		return true
	}
	return false
}

type TaggedFinder struct {
	url            string             // clickhouse dsn
	table          string             // graphite_tag table
	absKeepEncoded bool               // Abs returns url encoded value. For queries from prometheus
	opts           clickhouse.Options // clickhouse query timeout
	body           []byte             // clickhouse response
}

func NewTagged(url string, table string, absKeepEncoded bool, opts clickhouse.Options) *TaggedFinder {
	return &TaggedFinder{
		url:            url,
		table:          table,
		absKeepEncoded: absKeepEncoded,
		opts:           opts,
	}
}

func (term *TaggedTerm) concat() string {
	return fmt.Sprintf("%s=%s", term.Key, term.Value)
}

func TaggedTermWhere1(term *TaggedTerm) string {
	// positive expression check only in Tag1
	// negative check in all Tags
	switch term.Op {
	case TaggedTermEq:
		return where.Eq("Tag1", term.concat())
	case TaggedTermNe:
		if term.Value == "" {
			// special case
			// container_name!=""  ==> container_name exists and it is not empty
			return where.HasPrefixAndNotEq("Tag1", term.Key+"=")
		}
		return fmt.Sprintf("NOT arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermMatch:
		return where.Match("Tag1", term.concat())
	case TaggedTermNotMatch:
		return fmt.Sprintf("NOT arrayExists((x) -> %s), Tags)", where.Match("x", term.concat()))
	default:
		return ""
	}
}

func TaggedTermWhereN(term *TaggedTerm) string {
	// arrayExists((x) -> %s, Tags)
	switch term.Op {
	case TaggedTermEq:
		return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermNe:
		if term.Value == "" {
			// special case
			// container_name!=""  ==> container_name exists and it is not empty
			return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.HasPrefixAndNotEq("x", term.Key+"="))
		}
		return fmt.Sprintf("NOT arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermMatch:
		return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.Match("x", term.concat()))
	case TaggedTermNotMatch:
		return fmt.Sprintf("NOT arrayExists((x) -> %s), Tags)", where.Match("x", term.concat()))
	default:
		return ""
	}
}

func ParseTaggedConditions(conditions []string) ([]TaggedTerm, error) {
	terms := make([]TaggedTerm, len(conditions))

	for i := 0; i < len(conditions); i++ {
		s := conditions[i]

		a := strings.SplitN(s, "=", 2)
		if len(a) != 2 {
			return nil, fmt.Errorf("wrong seriesByTag expr: %#v", s)
		}

		a[0] = strings.TrimSpace(a[0])
		a[1] = strings.TrimSpace(a[1])

		op := "="

		if len(a[0]) > 0 && a[0][len(a[0])-1] == '!' {
			op = "!" + op
			a[0] = strings.TrimSpace(a[0][:len(a[0])-1])
		}

		if len(a[1]) > 0 && a[1][0] == '~' {
			op = op + "~"
			a[1] = strings.TrimSpace(a[1][1:])
		}

		terms[i].Key = a[0]
		terms[i].Value = a[1]

		if terms[i].Key == "name" {
			terms[i].Key = "__name__"
		}

		switch op {
		case "=":
			terms[i].Op = TaggedTermEq
		case "!=":
			terms[i].Op = TaggedTermNe
		case "=~":
			terms[i].Op = TaggedTermMatch
		case "!=~":
			terms[i].Op = TaggedTermNotMatch
		default:
			return nil, fmt.Errorf("wrong seriesByTag expr: %#v", s)
		}
	}

	sort.Sort(TaggedTermList(terms))

	return terms, nil
}

func ParseSeriesByTag(query string) ([]TaggedTerm, error) {
	expr, _, err := parser.ParseExpr(query)
	if err != nil {
		return nil, err
	}

	validationError := fmt.Errorf("wrong seriesByTag call: %#v", query)

	// check
	if !expr.IsFunc() {
		return nil, validationError
	}
	if expr.Target() != "seriesByTag" {
		return nil, validationError
	}

	args := expr.Args()
	if len(args) < 1 {
		return nil, validationError
	}

	for i := 0; i < len(args); i++ {
		if !args[i].IsString() {
			return nil, validationError
		}
	}

	conditions := make([]string, 0, len(args))
	for i := 0; i < len(args); i++ {
		s := args[i].StringValue()
		if s == "" {
			continue
		}
		conditions = append(conditions, s)
	}

	return ParseTaggedConditions(conditions)
}

func TaggedWhere(terms []TaggedTerm) (*where.Where, *where.Where) {
	w := where.New()
	pw := where.New()
	x := TaggedTermWhere1(&terms[0])
	if terms[0].Op == TaggedTermMatch {
		pw.And(x)
	}
	w.And(x)

	for i := 1; i < len(terms); i++ {
		w.And(TaggedTermWhereN(&terms[i]))
	}

	return w, pw
}

func (t *TaggedFinder) Execute(ctx context.Context, query string, from int64, until int64) error {
	terms, err := ParseSeriesByTag(query)
	if err != nil {
		return err
	}

	return t.ExecutePrepared(ctx, terms, from, until)
}

func (t *TaggedFinder) ExecutePrepared(ctx context.Context, terms []TaggedTerm, from int64, until int64) error {
	var err error
	w, pw := TaggedWhere(terms)

	w.Andf(
		"Date >='%s' AND Date <= '%s'",
		time.Unix(from, 0).Format("2006-01-02"),
		time.Unix(until, 0).Format("2006-01-02"),
	)

	sql := fmt.Sprintf("SELECT Path FROM %s %s %s GROUP BY Path", t.table, pw.PreWhereSQL(), w.SQL())
	t.body, err = clickhouse.Query(scope.WithTable(ctx, t.table), t.url, sql, t.opts)
	return err
}

func (t *TaggedFinder) List() [][]byte {
	if t.body == nil {
		return [][]byte{}
	}

	rows := bytes.Split(t.body, []byte{'\n'})

	skip := 0
	for i := 0; i < len(rows); i++ {
		if len(rows[i]) == 0 {
			skip++
			continue
		}
		if skip > 0 {
			rows[i-skip] = rows[i]
		}
	}

	rows = rows[:len(rows)-skip]

	return rows
}

func (t *TaggedFinder) Series() [][]byte {
	return t.List()
}

func (t *TaggedFinder) Abs(v []byte) []byte {
	if t.absKeepEncoded {
		return v
	}

	u, err := url.Parse(string(v))
	if err != nil {
		return v
	}

	tags := make([]string, 0, len(u.Query()))
	for k, v := range u.Query() {
		tags = append(tags, fmt.Sprintf("%s=%s", k, v[0]))
	}

	sort.Strings(tags)
	if len(tags) == 0 {
		return []byte(u.Path)
	}

	return []byte(fmt.Sprintf("%s;%s", u.Path, strings.Join(tags, ";")))
}
<commit_msg>fix not match expression<commit_after>package finder

import (
	"bytes"
	"context"
	"fmt"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/go-graphite/carbonapi/pkg/parser"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/pkg/where"
)

type TaggedTermOp int

const (
	TaggedTermEq       TaggedTermOp = 1
	TaggedTermMatch    TaggedTermOp = 2
	TaggedTermNe       TaggedTermOp = 3
	TaggedTermNotMatch TaggedTermOp = 4
)

type TaggedTerm struct {
	Key   string
	Op    TaggedTermOp
	Value string
}

type TaggedTermList []TaggedTerm

func (s TaggedTermList) Len() int {
	return len(s)
}
func (s TaggedTermList) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s TaggedTermList) Less(i, j int) bool {
	if s[i].Op < s[j].Op {
		return true
	}
	if s[i].Op > s[j].Op {
		return false
	}
	if s[i].Key == "__name__" && s[j].Key != "__name__" {
		return true
	}
	return false
}

type TaggedFinder struct {
	url            string             // clickhouse dsn
	table          string             // graphite_tag table
	absKeepEncoded bool               // Abs returns url encoded value. For queries from prometheus
	opts           clickhouse.Options // clickhouse query timeout
	body           []byte             // clickhouse response
}

func NewTagged(url string, table string, absKeepEncoded bool, opts clickhouse.Options) *TaggedFinder {
	return &TaggedFinder{
		url:            url,
		table:          table,
		absKeepEncoded: absKeepEncoded,
		opts:           opts,
	}
}

func (term *TaggedTerm) concat() string {
	return fmt.Sprintf("%s=%s", term.Key, term.Value)
}

func TaggedTermWhere1(term *TaggedTerm) string {
	// positive expression check only in Tag1
	// negative check in all Tags
	switch term.Op {
	case TaggedTermEq:
		return where.Eq("Tag1", term.concat())
	case TaggedTermNe:
		if term.Value == "" {
			// special case
			// container_name!=""  ==> container_name exists and it is not empty
			return where.HasPrefixAndNotEq("Tag1", term.Key+"=")
		}
		return fmt.Sprintf("NOT arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermMatch:
		return where.Match("Tag1", term.concat())
	case TaggedTermNotMatch:
		return fmt.Sprintf("NOT arrayExists((x) -> %s), Tags)", where.Match("x", term.concat()))
	default:
		return ""
	}
}

func TaggedTermWhereN(term *TaggedTerm) string {
	// arrayExists((x) -> %s, Tags)
	switch term.Op {
	case TaggedTermEq:
		return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermNe:
		if term.Value == "" {
			// special case
			// container_name!=""  ==> container_name exists and it is not empty
			return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.HasPrefixAndNotEq("x", term.Key+"="))
		}
		return fmt.Sprintf("NOT arrayExists((x) -> %s, Tags)", where.Eq("x", term.concat()))
	case TaggedTermMatch:
		return fmt.Sprintf("arrayExists((x) -> %s, Tags)", where.Match("x", term.concat()))
	case TaggedTermNotMatch:
		return fmt.Sprintf("NOT arrayExists((x) -> %s, Tags)", where.Match("x", term.concat()))
	default:
		return ""
	}
}

func ParseTaggedConditions(conditions []string) ([]TaggedTerm, error) {
	terms := make([]TaggedTerm, len(conditions))

	for i := 0; i < len(conditions); i++ {
		s := conditions[i]

		a := strings.SplitN(s, "=", 2)
		if len(a) != 2 {
			return nil, fmt.Errorf("wrong seriesByTag expr: %#v", s)
		}

		a[0] = strings.TrimSpace(a[0])
		a[1] = strings.TrimSpace(a[1])

		op := "="

		if len(a[0]) > 0 && a[0][len(a[0])-1] == '!' {
			op = "!" + op
			a[0] = strings.TrimSpace(a[0][:len(a[0])-1])
		}

		if len(a[1]) > 0 && a[1][0] == '~' {
			op = op + "~"
			a[1] = strings.TrimSpace(a[1][1:])
		}

		terms[i].Key = a[0]
		terms[i].Value = a[1]

		if terms[i].Key == "name" {
			terms[i].Key = "__name__"
		}

		switch op {
		case "=":
			terms[i].Op = TaggedTermEq
		case "!=":
			terms[i].Op = TaggedTermNe
		case "=~":
			terms[i].Op = TaggedTermMatch
		case "!=~":
			terms[i].Op = TaggedTermNotMatch
		default:
			return nil, fmt.Errorf("wrong seriesByTag expr: %#v", s)
		}
	}

	sort.Sort(TaggedTermList(terms))

	return terms, nil
}

func ParseSeriesByTag(query string) ([]TaggedTerm, error) {
	expr, _, err := parser.ParseExpr(query)
	if err != nil {
		return nil, err
	}

	validationError := fmt.Errorf("wrong seriesByTag call: %#v", query)

	// check
	if !expr.IsFunc() {
		return nil, validationError
	}
	if expr.Target() != "seriesByTag" {
		return nil, validationError
	}

	args := expr.Args()
	if len(args) < 1 {
		return nil, validationError
	}

	for i := 0; i < len(args); i++ {
		if !args[i].IsString() {
			return nil, validationError
		}
	}

	conditions := make([]string, 0, len(args))
	for i := 0; i < len(args); i++ {
		s := args[i].StringValue()
		if s == "" {
			continue
		}
		conditions = append(conditions, s)
	}

	return ParseTaggedConditions(conditions)
}

func TaggedWhere(terms []TaggedTerm) (*where.Where, *where.Where) {
	w := where.New()
	pw := where.New()
	x := TaggedTermWhere1(&terms[0])
	if terms[0].Op == TaggedTermMatch {
		pw.And(x)
	}
	w.And(x)

	for i := 1; i < len(terms); i++ {
		w.And(TaggedTermWhereN(&terms[i]))
	}

	return w, pw
}

func (t *TaggedFinder) Execute(ctx context.Context, query string, from int64, until int64) error {
	terms, err := ParseSeriesByTag(query)
	if err != nil {
		return err
	}

	return t.ExecutePrepared(ctx, terms, from, until)
}

func (t *TaggedFinder) ExecutePrepared(ctx context.Context, terms []TaggedTerm, from int64, until int64) error {
	var err error
	w, pw := TaggedWhere(terms)

	w.Andf(
		"Date >='%s' AND Date <= '%s'",
		time.Unix(from, 0).Format("2006-01-02"),
		time.Unix(until, 0).Format("2006-01-02"),
	)

	sql := fmt.Sprintf("SELECT Path FROM %s %s %s GROUP BY Path", t.table, pw.PreWhereSQL(), w.SQL())
	t.body, err = clickhouse.Query(scope.WithTable(ctx, t.table), t.url, sql, t.opts)
	return err
}

func (t *TaggedFinder) List() [][]byte {
	if t.body == nil {
		return [][]byte{}
	}

	rows := bytes.Split(t.body, []byte{'\n'})

	skip := 0
	for i := 0; i < len(rows); i++ {
		if len(rows[i]) == 0 {
			skip++
			continue
		}
		if skip > 0 {
			rows[i-skip] = rows[i]
		}
	}

	rows = rows[:len(rows)-skip]

	return rows
}

func (t *TaggedFinder) Series() [][]byte {
	return t.List()
}

func (t *TaggedFinder) Abs(v []byte) []byte {
	if t.absKeepEncoded {
		return v
	}

	u, err := url.Parse(string(v))
	if err != nil {
		return v
	}

	tags := make([]string, 0, len(u.Query()))
	for k, v := range u.Query() {
		tags = append(tags, fmt.Sprintf("%s=%s", k, v[0]))
	}

	sort.Strings(tags)
	if len(tags) == 0 {
		return []byte(u.Path)
	}

	return []byte(fmt.Sprintf("%s;%s", u.Path, strings.Join(tags, ";")))
}
<|endoftext|>