<commit_before>//
// Copyright 2014, Sander van Harmelen
//
// 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 main

import (
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

func (cg *ChefGuard) executeChecks() (int, error) {
	if cfg.Tests.Foodcritic != "" {
		if errCode, err := runFoodcritic(cg.ChefOrg, cg.CookbookPath); err != nil {
			if errCode == http.StatusInternalServerError || !cg.continueAfterFailedCheck("foodcritic") {
				return errCode, err
			}
		}
	}
	if cfg.Tests.Rubocop != "" {
		if errCode, err := runRubocop(cg.CookbookPath); err != nil {
			if errCode == http.StatusInternalServerError || !cg.continueAfterFailedCheck("rubocop") {
				return errCode, err
			}
		}
	}
	return 0, nil
}

func (cg *ChefGuard) continueAfterFailedCheck(check string) bool {
	WARNING.Printf("%s errors when uploading cookbook '%s' for '%s'\n", strings.Title(check), cg.Cookbook.Name, cg.User)
	if getEffectiveConfig("Mode", cg.ChefOrg).(string) == "permissive" && cg.ForcedUpload {
		return true
	}
	return false
}

func runFoodcritic(org, cookbookPath string) (int, error) {
	args := getFoodcriticArgs(org, cookbookPath)
	cmd := exec.Command(cfg.Tests.Foodcritic, args...)

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, "RUBY_THREAD_VM_STACK_SIZE=2097152")

	output, err := cmd.CombinedOutput()
	if err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			if exitError.Sys().(syscall.WaitStatus).ExitStatus() == 3 {
				errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
				return http.StatusPreconditionFailed, fmt.Errorf("\n=== Foodcritic errors found ===\n%s\n===============================\n", errText)
			}
		}
		return http.StatusInternalServerError, fmt.Errorf("Failed to execute foodcritic tests: %s - %s", output, err)
	}

	// This is still needed for Foodcritic > v9.x.x
	if strings.TrimSpace(string(output)) != "" {
		errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
		return http.StatusPreconditionFailed, fmt.Errorf("\n=== Foodcritic errors found ===\n%s\n===============================\n", errText)
	}

	return 0, nil
}

func getFoodcriticArgs(org, cookbookPath string) []string {
	excludes := cfg.Default.ExcludeFCs
	custExcludes := getEffectiveConfig("ExcludeFCs", org)
	if excludes != custExcludes {
		excludes = fmt.Sprintf("%s,%s", excludes, custExcludes)
	}
	excls := strings.Split(excludes, ",")
	args := []string{}
	for _, excl := range excls {
		args = append(args, fmt.Sprintf("-t ~%s", strings.TrimSpace(excl)))
	}
	if cfg.Default.IncludeFCs != "" {
		args = append(args, "-I", cfg.Default.IncludeFCs)
	}
	return append(args, "-B", cookbookPath)
}

func runRubocop(cookbookPath string) (int, error) {
	cmd := exec.Command(cfg.Tests.Rubocop, cookbookPath)
	cmd.Env = []string{"HOME=" + cfg.Default.Tempdir}
	output, err := cmd.CombinedOutput()
	if err != nil {
		if strings.Contains(string(output), "offense") {
			errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
			return http.StatusPreconditionFailed, fmt.Errorf("\n=== Rubocop errors found ===\n%s\n============================\n", errText)
		}
		return http.StatusInternalServerError, fmt.Errorf("Failed to execute rubocop tests: %s - %s", output, err)
	}
	return 0, nil
}
<commit_msg>Maintain backwards compatibility (#120)<commit_after>//
// Copyright 2014, Sander van Harmelen
//
// 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 main

import (
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

func (cg *ChefGuard) executeChecks() (int, error) {
	if cfg.Tests.Foodcritic != "" {
		if errCode, err := runFoodcritic(cg.ChefOrg, cg.CookbookPath); err != nil {
			if errCode == http.StatusInternalServerError || !cg.continueAfterFailedCheck("foodcritic") {
				return errCode, err
			}
		}
	}
	if cfg.Tests.Rubocop != "" {
		if errCode, err := runRubocop(cg.CookbookPath); err != nil {
			if errCode == http.StatusInternalServerError || !cg.continueAfterFailedCheck("rubocop") {
				return errCode, err
			}
		}
	}
	return 0, nil
}

func (cg *ChefGuard) continueAfterFailedCheck(check string) bool {
	WARNING.Printf("%s errors when uploading cookbook '%s' for '%s'\n", strings.Title(check), cg.Cookbook.Name, cg.User)
	if getEffectiveConfig("Mode", cg.ChefOrg).(string) == "permissive" && cg.ForcedUpload {
		return true
	}
	return false
}

func runFoodcritic(org, cookbookPath string) (int, error) {
	args := getFoodcriticArgs(org, cookbookPath)
	cmd := exec.Command(cfg.Tests.Foodcritic, args...)

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, "RUBY_THREAD_VM_STACK_SIZE=2097152")

	output, err := cmd.CombinedOutput()
	if err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			if exitError.Sys().(syscall.WaitStatus).ExitStatus() == 3 {
				errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
				return http.StatusPreconditionFailed, fmt.Errorf("\n=== Foodcritic errors found ===\n%s\n===============================\n", errText)
			}
		}

		return http.StatusInternalServerError, fmt.Errorf("Failed to execute \"foodcritic %s\": %s - %s", strings.Join(cmd.Args, " "), output, err)
	}

	// This is still needed for Foodcritic > v9.x.x
	if strings.TrimSpace(string(output)) != "" {
		errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
		return http.StatusPreconditionFailed, fmt.Errorf("\n=== Foodcritic errors found ===\n%s\n===============================\n", errText)
	}

	return 0, nil
}

func getFoodcriticArgs(org, cookbookPath string) []string {
	excludes := cfg.Default.ExcludeFCs
	custExcludes := getEffectiveConfig("ExcludeFCs", org)
	if excludes != custExcludes {
		excludes = fmt.Sprintf("%s,%s", excludes, custExcludes)
	}
	args := []string{}
	if excludes != "" {
		args = append(args, "--tags", "~"+strings.Replace(excludes, ",", ",~", -1))
	}
	if cfg.Default.IncludeFCs != "" {
		args = append(args, "--include", cfg.Default.IncludeFCs)
	}
	return append(args, "--no-progress", "--cookbook-path", cookbookPath)
}

func runRubocop(cookbookPath string) (int, error) {
	cmd := exec.Command(cfg.Tests.Rubocop, cookbookPath)
	cmd.Env = []string{"HOME=" + cfg.Default.Tempdir}
	output, err := cmd.CombinedOutput()
	if err != nil {
		if strings.Contains(string(output), "offense") {
			errText := strings.TrimSpace(strings.Replace(string(output), fmt.Sprintf("%s/", cookbookPath), "", -1))
			return http.StatusPreconditionFailed, fmt.Errorf("\n=== Rubocop errors found ===\n%s\n============================\n", errText)
		}
		return http.StatusInternalServerError, fmt.Errorf("Failed to execute \"rubocop %s\": %s - %s", cookbookPath, output, err)
	}
	return 0, nil
}
<|endoftext|>


<commit_before>package article

import "time"

const (
	English  Language = "ENG"
	Japanese Language = "JAP"
)

type Language string

type Article struct {
	Id         string              `json:"id"`
	Title      map[Language]string `json:"title"`
	Body       map[Language]string `json:"body,omitempty"`
	DatePosted time.Time           `json:"posted"`
	DateEdited time.Time           `json:"edited,omitempty"`
}

type ArticleList []*Article

// Make slices of article pointers sortable by posting time.
func (l ArticleList) Len() int {
	return len(l)
}

func (l ArticleList) Less(i, j int) bool {
	return l[i].DatePosted.Unix() < l[j].DatePosted.Unix()
}

func (l ArticleList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}
<commit_msg>Serve japanese language name in japanese.<commit_after>package article

import "time"

const (
	English  Language = "ENG"
	Japanese Language = "日本語"
)

type Language string

type Article struct {
	Id         string              `json:"id"`
	Title      map[Language]string `json:"title"`
	Body       map[Language]string `json:"body,omitempty"`
	DatePosted time.Time           `json:"posted"`
	DateEdited time.Time           `json:"edited,omitempty"`
}

type ArticleList []*Article

// Make slices of article pointers sortable by posting time.
func (l ArticleList) Len() int {
	return len(l)
}

func (l ArticleList) Less(i, j int) bool {
	return l[i].DatePosted.Unix() < l[j].DatePosted.Unix()
}

func (l ArticleList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}
<|endoftext|>


<commit_before>// Copyright 2012 Gary Burd
//
// 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 redis_test

import (
	"fmt"
	"reflect"
	"sync"
	"testing"

	"github.com/garyburd/redigo/redis"
)

func publish(channel, value interface{}) {
	c, err := dial()
	if err != nil {
		fmt.Println(err)
		return
	}
	defer c.Close()
	c.Do("PUBLISH", channel, value)
}

// Applications can receive pushed messages from one goroutine and manage subscriptions from another goroutine.
func ExamplePubSubConn() {
	c, err := dial()
	if err != nil {
		fmt.Println(err)
		return
	}
	defer c.Close()
	var wg sync.WaitGroup
	wg.Add(2)

	psc := redis.PubSubConn{Conn: c}

	// This goroutine receives and prints pushed notifications from the server.
	// The goroutine exits when the connection is unsubscribed from all
	// channels or there is an error.
	go func() {
		defer wg.Done()
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
			case redis.PMessage:
				fmt.Printf("PMessage: %s %s %s\n", n.Pattern, n.Channel, n.Data)
			case redis.Subscription:
				fmt.Printf("Subscription: %s %s %d\n", n.Kind, n.Channel, n.Count)
				if n.Count == 0 {
					return
				}
			case error:
				fmt.Printf("error: %v\n", n)
				return
			}
		}
	}()

	// This goroutine manages subscriptions for the connection.
	go func() {
		defer wg.Done()

		psc.Subscribe("example")
		psc.PSubscribe("p*")

		// The following function calls publish a message using another
		// connection to the Redis server.
		publish("example", "hello")
		publish("example", "world")
		publish("pexample", "foo")
		publish("pexample", "bar")

		// Unsubscribe from all connections. This will cause the receiving
		// goroutine to exit.
		psc.Unsubscribe()
		psc.PUnsubscribe()
	}()

	wg.Wait()

	// Output:
	// Subscription: subscribe example 1
	// Subscription: psubscribe p* 2
	// Message: example hello
	// Message: example world
	// PMessage: p* pexample foo
	// PMessage: p* pexample bar
	// Subscription: unsubscribe example 1
	// Subscription: punsubscribe p* 0
}

func expectPushed(t *testing.T, c redis.PubSubConn, message string, expected interface{}) {
	actual := c.Receive()
	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("%s = %v, want %v", message, actual, expected)
	}
}

func TestPushed(t *testing.T) {
	pc, err := redis.DialDefaultServer()
	if err != nil {
		t.Fatalf("error connection to database, %v", err)
	}
	defer pc.Close()

	sc, err := redis.DialDefaultServer()
	if err != nil {
		t.Fatalf("error connection to database, %v", err)
	}
	defer sc.Close()

	c := redis.PubSubConn{Conn: sc}

	c.Subscribe("c1")
	expectPushed(t, c, "Subscribe(c1)", redis.Subscription{Kind: "subscribe", Channel: "c1", Count: 1})
	c.Subscribe("c2")
	expectPushed(t, c, "Subscribe(c2)", redis.Subscription{Kind: "subscribe", Channel: "c2", Count: 2})
	c.PSubscribe("p1")
	expectPushed(t, c, "PSubscribe(p1)", redis.Subscription{Kind: "psubscribe", Channel: "p1", Count: 3})
	c.PSubscribe("p2")
	expectPushed(t, c, "PSubscribe(p2)", redis.Subscription{Kind: "psubscribe", Channel: "p2", Count: 4})
	c.PUnsubscribe()
	expectPushed(t, c, "Punsubscribe(p1)", redis.Subscription{Kind: "punsubscribe", Channel: "p1", Count: 3})
	expectPushed(t, c, "Punsubscribe()", redis.Subscription{Kind: "punsubscribe", Channel: "p2", Count: 2})

	pc.Do("PUBLISH", "c1", "hello")
	expectPushed(t, c, "PUBLISH c1 hello", redis.Message{Channel: "c1", Data: []byte("hello")})

	c.Ping("hello")
	expectPushed(t, c, `Ping("hello")`, redis.Pong{"hello"})

	c.Conn.Send("PING")
	c.Conn.Flush()
	expectPushed(t, c, `Send("PING")`, redis.Pong{})
}
<commit_msg>Fix go vet warnings<commit_after>// Copyright 2012 Gary Burd
//
// 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 redis_test

import (
	"fmt"
	"reflect"
	"sync"
	"testing"

	"github.com/garyburd/redigo/redis"
)

func publish(channel, value interface{}) {
	c, err := dial()
	if err != nil {
		fmt.Println(err)
		return
	}
	defer c.Close()
	c.Do("PUBLISH", channel, value)
}

// Applications can receive pushed messages from one goroutine and manage subscriptions from another goroutine.
func ExamplePubSubConn() {
	c, err := dial()
	if err != nil {
		fmt.Println(err)
		return
	}
	defer c.Close()
	var wg sync.WaitGroup
	wg.Add(2)

	psc := redis.PubSubConn{Conn: c}

	// This goroutine receives and prints pushed notifications from the server.
	// The goroutine exits when the connection is unsubscribed from all
	// channels or there is an error.
	go func() {
		defer wg.Done()
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
			case redis.PMessage:
				fmt.Printf("PMessage: %s %s %s\n", n.Pattern, n.Channel, n.Data)
			case redis.Subscription:
				fmt.Printf("Subscription: %s %s %d\n", n.Kind, n.Channel, n.Count)
				if n.Count == 0 {
					return
				}
			case error:
				fmt.Printf("error: %v\n", n)
				return
			}
		}
	}()

	// This goroutine manages subscriptions for the connection.
	go func() {
		defer wg.Done()

		psc.Subscribe("example")
		psc.PSubscribe("p*")

		// The following function calls publish a message using another
		// connection to the Redis server.
		publish("example", "hello")
		publish("example", "world")
		publish("pexample", "foo")
		publish("pexample", "bar")

		// Unsubscribe from all connections. This will cause the receiving
		// goroutine to exit.
		psc.Unsubscribe()
		psc.PUnsubscribe()
	}()

	wg.Wait()

	// Output:
	// Subscription: subscribe example 1
	// Subscription: psubscribe p* 2
	// Message: example hello
	// Message: example world
	// PMessage: p* pexample foo
	// PMessage: p* pexample bar
	// Subscription: unsubscribe example 1
	// Subscription: punsubscribe p* 0
}

func expectPushed(t *testing.T, c redis.PubSubConn, message string, expected interface{}) {
	actual := c.Receive()
	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("%s = %v, want %v", message, actual, expected)
	}
}

func TestPushed(t *testing.T) {
	pc, err := redis.DialDefaultServer()
	if err != nil {
		t.Fatalf("error connection to database, %v", err)
	}
	defer pc.Close()

	sc, err := redis.DialDefaultServer()
	if err != nil {
		t.Fatalf("error connection to database, %v", err)
	}
	defer sc.Close()

	c := redis.PubSubConn{Conn: sc}

	c.Subscribe("c1")
	expectPushed(t, c, "Subscribe(c1)", redis.Subscription{Kind: "subscribe", Channel: "c1", Count: 1})
	c.Subscribe("c2")
	expectPushed(t, c, "Subscribe(c2)", redis.Subscription{Kind: "subscribe", Channel: "c2", Count: 2})
	c.PSubscribe("p1")
	expectPushed(t, c, "PSubscribe(p1)", redis.Subscription{Kind: "psubscribe", Channel: "p1", Count: 3})
	c.PSubscribe("p2")
	expectPushed(t, c, "PSubscribe(p2)", redis.Subscription{Kind: "psubscribe", Channel: "p2", Count: 4})
	c.PUnsubscribe()
	expectPushed(t, c, "Punsubscribe(p1)", redis.Subscription{Kind: "punsubscribe", Channel: "p1", Count: 3})
	expectPushed(t, c, "Punsubscribe()", redis.Subscription{Kind: "punsubscribe", Channel: "p2", Count: 2})

	pc.Do("PUBLISH", "c1", "hello")
	expectPushed(t, c, "PUBLISH c1 hello", redis.Message{Channel: "c1", Data: []byte("hello")})

	c.Ping("hello")
	expectPushed(t, c, `Ping("hello")`, redis.Pong{Data: "hello"})

	c.Conn.Send("PING")
	c.Conn.Flush()
	expectPushed(t, c, `Send("PING")`, redis.Pong{})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
	"golang.org/x/net/websocket"
)
// IMPORTS, BITCH!

func onInit() string {
	return "[INIT] Golang 10 min lib for Shitty mongo/Chatango (@Coil) initiated. " + time.Now().String()
}

func main() {
	init := 1
	fmt.Println(onInit()) // Friendly little init message.
	origin := "http://st.chatango.com" // Origin required, thank Chatango for that.
	url := "ws://s23.chatango.com:1800" // Websocket server for the chat.
	ws, err := websocket.Dial(url, "", origin) // Create a websocket connection
	if err != nil { // IS THIS THING ON?
		fmt.Println("[FATAL] ERR NIGNOG")
	}
	ws.Write([]byte("bauth:hazerdklan::Coil:password\x00")) // tested at the klan
	time.Sleep(10000000000) // Needs a delay. HAVE YOUR FUCKING DELAY, CHATANGO
	ws.Write([]byte("bm:t12j:256:<nED1C24/><f x11555555=\"0\">get rekt mofucka\r\n\x00")) // Send a message
	for init == 1 { // Golang's while loop equivalent
		time.Sleep(100) // MORE SLEEPING FUUUUUUUUUUUUUUUUUUUUU
		buf := make([]byte,1024) // Receive from the websocket
		var n int // cursor
		if n, err = ws.Read(buf); err != nil { // Make sure this shit isn't null, nignog
			fmt.Println("[FATAL] ERR NIGNOG") // but in the event that it is null...
		}
		fmt.Printf("%s\n", buf[:n]) // Print the data!
	}
}
<commit_msg>Update cholib.go<commit_after>package main

import (
	"fmt"
	"time"
	"golang.org/x/net/websocket"
)
// IMPORTS, YAY!

func onInit() string {
	return "[INIT] Golang 10 min lib for Chatango (@Coil) initiated. " + time.Now().String()
}

func main() {
	init := 1
	fmt.Println(onInit()) // Friendly little init message.
	origin := "http://st.chatango.com" // Origin required, thank Chatango for that.
	url := "ws://s23.chatango.com:1800" // Websocket server for the chat.
	ws, err := websocket.Dial(url, "", origin) // Create a websocket connection
	if err != nil { // IS THIS THING ON?
		fmt.Println("[FATAL] ERR NIGNOG")
	}
	ws.Write([]byte("bauth:hazerdklan::Coil:password\x00")) // tested at the klan
	time.Sleep(10000000000) // Needs a delay. HAVE YOUR DELAY, CHATANGO
	ws.Write([]byte("bm:t12j:256:<nED1C24/><f x11555555=\"0\">get rekt mofucka\r\n\x00")) // Send a message
	for init == 1 { // Golang's while loop equivalent
		time.Sleep(100) // MORE SLEEPING FUUUUUUUUUUUUUUUUUUUUU
		buf := make([]byte,1024) // Receive from the websocket
		var n int // cursor
		if n, err = ws.Read(buf); err != nil { // Make sure the data isn't null, nignog
			fmt.Println("[FATAL] ERR NIGNOG") // but in the event that it is null...
		}
		fmt.Printf("%s\n", buf[:n]) // Print the data!
	}
}
<|endoftext|>


<commit_before>// The chroot package is able to create an Amazon AMI without requiring
// the launch of a new instance for every build. It does this by attaching
// and mounting the root volume of another AMI and chrooting into that
// directory. It then creates an AMI from that attached drive.
package chroot

import (
	"errors"
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/multistep"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"log"
	"runtime"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazon.chroot"

// Config is the configuration that is chained through the steps and
// settable from the template.
type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`

	ChrootMounts   [][]string `mapstructure:"chroot_mounts"`
	CommandWrapper string     `mapstructure:"command_wrapper"`
	CopyFiles      []string   `mapstructure:"copy_files"`
	DevicePath     string     `mapstructure:"device_path"`
	MountPath      string     `mapstructure:"mount_path"`
	SourceAmi      string     `mapstructure:"source_ami"`

	tpl *packer.ConfigTemplate
}

type wrappedCommandTemplate struct {
	Command string
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	md, err := common.DecodeConfig(&b.config, raws...)
	if err != nil {
		return nil, err
	}

	b.config.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return nil, err
	}
	b.config.tpl.UserVars = b.config.PackerUserVars
	b.config.tpl.Funcs(awscommon.TemplateFuncs)

	// Defaults
	if b.config.ChrootMounts == nil {
		b.config.ChrootMounts = make([][]string, 0)
	}

	if b.config.CopyFiles == nil {
		b.config.CopyFiles = make([]string, 0)
	}

	if len(b.config.ChrootMounts) == 0 {
		b.config.ChrootMounts = [][]string{
			[]string{"proc", "proc", "/proc"},
			[]string{"sysfs", "sysfs", "/sys"},
			[]string{"bind", "/dev", "/dev"},
			[]string{"devpts", "devpts", "/dev/pts"},
			[]string{"binfmt_misc", "binfmt_misc", "/proc/sys/fs/binfmt_misc"},
		}
	}

	if len(b.config.CopyFiles) == 0 {
		b.config.CopyFiles = []string{"/etc/resolv.conf"}
	}

	if b.config.CommandWrapper == "" {
		b.config.CommandWrapper = "{{.Command}}"
	}

	if b.config.MountPath == "" {
		b.config.MountPath = "packer-amazon-chroot-volumes/{{.Device}}"
	}

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)
	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(b.config.tpl)...)
	errs = packer.MultiErrorAppend(errs, b.config.AMIConfig.Prepare(b.config.tpl)...)

	for i, mounts := range b.config.ChrootMounts {
		if len(mounts) != 3 {
			errs = packer.MultiErrorAppend(
				errs, errors.New("Each chroot_mounts entry should be three elements."))
			break
		}

		for j, entry := range mounts {
			b.config.ChrootMounts[i][j], err = b.config.tpl.Process(entry, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(errs,
					fmt.Errorf("Error processing chroot_mounts[%d][%d]: %s",
						i, j, err))
			}
		}
	}

	for i, file := range b.config.CopyFiles {
		var err error
		b.config.CopyFiles[i], err = b.config.tpl.Process(file, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing copy_files[%d]: %s",
					i, err))
		}
	}

	if b.config.SourceAmi == "" {
		errs = packer.MultiErrorAppend(errs, errors.New("source_ami is required."))
	}

	templates := map[string]*string{
		"device_path": &b.config.DevicePath,
		"source_ami":  &b.config.SourceAmi,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = b.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return nil, errs
	}

	log.Println(common.ScrubConfig(b.config, b.config.AccessKey, b.config.SecretKey))
	return nil, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	if runtime.GOOS != "linux" {
		return nil, errors.New("The amazon-chroot builder only works on Linux environments.")
	}

	region, err := b.config.Region()
	if err != nil {
		return nil, err
	}

	auth, err := b.config.AccessConfig.Auth()
	if err != nil {
		return nil, err
	}

	ec2conn := ec2.New(auth, region)

	wrappedCommand := func(command string) (string, error) {
		return b.config.tpl.Process(
			b.config.CommandWrapper, &wrappedCommandTemplate{
				Command: command,
			})
	}

	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("ec2", ec2conn)
	state.Put("hook", hook)
	state.Put("ui", ui)
	state.Put("wrappedCommand", CommandWrapper(wrappedCommand))

	// Build the steps
	steps := []multistep.Step{
		&StepInstanceInfo{},
		&StepSourceAMIInfo{},
		&StepFlock{},
		&StepPrepareDevice{},
		&StepCreateVolume{},
		&StepAttachVolume{},
		&StepEarlyUnflock{},
		&StepMountDevice{},
		&StepMountExtra{},
		&StepCopyFiles{},
		&StepChrootProvision{},
		&StepEarlyCleanup{},
		&StepSnapshot{},
		&StepRegisterAMI{},
		&awscommon.StepAMIRegionCopy{
			Regions: b.config.AMIRegions,
		},
		&awscommon.StepModifyAMIAttributes{
			Description: b.config.AMIDescription,
			Users:       b.config.AMIUsers,
			Groups:      b.config.AMIGroups,
		},
		&awscommon.StepCreateTags{
			Tags: b.config.AMITags,
		},
	}

	// Run!
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If there are no AMIs, then just return
	if _, ok := state.GetOk("amis"); !ok {
		return nil, nil
	}

	// Build the artifact and return it
	artifact := &awscommon.Artifact{
		Amis:           state.Get("amis").(map[string]string),
		BuilderIdValue: BuilderId,
		Conn:           ec2conn,
	}

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<commit_msg>builder/amazon/chroot: mount in device-specific places<commit_after>// The chroot package is able to create an Amazon AMI without requiring
// the launch of a new instance for every build. It does this by attaching
// and mounting the root volume of another AMI and chrooting into that
// directory. It then creates an AMI from that attached drive.
package chroot

import (
	"errors"
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/multistep"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"log"
	"runtime"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazon.chroot"

// Config is the configuration that is chained through the steps and
// settable from the template.
type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`

	ChrootMounts   [][]string `mapstructure:"chroot_mounts"`
	CommandWrapper string     `mapstructure:"command_wrapper"`
	CopyFiles      []string   `mapstructure:"copy_files"`
	DevicePath     string     `mapstructure:"device_path"`
	MountPath      string     `mapstructure:"mount_path"`
	SourceAmi      string     `mapstructure:"source_ami"`

	tpl *packer.ConfigTemplate
}

type wrappedCommandTemplate struct {
	Command string
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	md, err := common.DecodeConfig(&b.config, raws...)
	if err != nil {
		return nil, err
	}

	b.config.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return nil, err
	}
	b.config.tpl.UserVars = b.config.PackerUserVars
	b.config.tpl.Funcs(awscommon.TemplateFuncs)

	// Defaults
	if b.config.ChrootMounts == nil {
		b.config.ChrootMounts = make([][]string, 0)
	}

	if b.config.CopyFiles == nil {
		b.config.CopyFiles = make([]string, 0)
	}

	if len(b.config.ChrootMounts) == 0 {
		b.config.ChrootMounts = [][]string{
			[]string{"proc", "proc", "/proc"},
			[]string{"sysfs", "sysfs", "/sys"},
			[]string{"bind", "/dev", "/dev"},
			[]string{"devpts", "devpts", "/dev/pts"},
			[]string{"binfmt_misc", "binfmt_misc", "/proc/sys/fs/binfmt_misc"},
		}
	}

	if len(b.config.CopyFiles) == 0 {
		b.config.CopyFiles = []string{"/etc/resolv.conf"}
	}

	if b.config.CommandWrapper == "" {
		b.config.CommandWrapper = "{{.Command}}"
	}

	if b.config.MountPath == "" {
		b.config.MountPath = "/mnt/packer-amazon-chroot-volumes/{{.Device}}"
	}

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)
	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(b.config.tpl)...)
	errs = packer.MultiErrorAppend(errs, b.config.AMIConfig.Prepare(b.config.tpl)...)

	for i, mounts := range b.config.ChrootMounts {
		if len(mounts) != 3 {
			errs = packer.MultiErrorAppend(
				errs, errors.New("Each chroot_mounts entry should be three elements."))
			break
		}

		for j, entry := range mounts {
			b.config.ChrootMounts[i][j], err = b.config.tpl.Process(entry, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(errs,
					fmt.Errorf("Error processing chroot_mounts[%d][%d]: %s",
						i, j, err))
			}
		}
	}

	for i, file := range b.config.CopyFiles {
		var err error
		b.config.CopyFiles[i], err = b.config.tpl.Process(file, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Error processing copy_files[%d]: %s",
					i, err))
		}
	}

	if b.config.SourceAmi == "" {
		errs = packer.MultiErrorAppend(errs, errors.New("source_ami is required."))
	}

	templates := map[string]*string{
		"device_path": &b.config.DevicePath,
		"source_ami":  &b.config.SourceAmi,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = b.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return nil, errs
	}

	log.Println(common.ScrubConfig(b.config, b.config.AccessKey, b.config.SecretKey))
	return nil, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	if runtime.GOOS != "linux" {
		return nil, errors.New("The amazon-chroot builder only works on Linux environments.")
	}

	region, err := b.config.Region()
	if err != nil {
		return nil, err
	}

	auth, err := b.config.AccessConfig.Auth()
	if err != nil {
		return nil, err
	}

	ec2conn := ec2.New(auth, region)

	wrappedCommand := func(command string) (string, error) {
		return b.config.tpl.Process(
			b.config.CommandWrapper, &wrappedCommandTemplate{
				Command: command,
			})
	}

	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("ec2", ec2conn)
	state.Put("hook", hook)
	state.Put("ui", ui)
	state.Put("wrappedCommand", CommandWrapper(wrappedCommand))

	// Build the steps
	steps := []multistep.Step{
		&StepInstanceInfo{},
		&StepSourceAMIInfo{},
		&StepFlock{},
		&StepPrepareDevice{},
		&StepCreateVolume{},
		&StepAttachVolume{},
		&StepEarlyUnflock{},
		&StepMountDevice{},
		&StepMountExtra{},
		&StepCopyFiles{},
		&StepChrootProvision{},
		&StepEarlyCleanup{},
		&StepSnapshot{},
		&StepRegisterAMI{},
		&awscommon.StepAMIRegionCopy{
			Regions: b.config.AMIRegions,
		},
		&awscommon.StepModifyAMIAttributes{
			Description: b.config.AMIDescription,
			Users:       b.config.AMIUsers,
			Groups:      b.config.AMIGroups,
		},
		&awscommon.StepCreateTags{
			Tags: b.config.AMITags,
		},
	}

	// Run!
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If there are no AMIs, then just return
	if _, ok := state.GetOk("amis"); !ok {
		return nil, nil
	}

	// Build the artifact and return it
	artifact := &awscommon.Artifact{
		Amis:           state.Get("amis").(map[string]string),
		BuilderIdValue: BuilderId,
		Conn:           ec2conn,
	}

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<|endoftext|>


<commit_before>// The googlecompute package contains a packer.Builder implementation that
// builds images for Google Compute Engine.
package googlecompute

import (
	"log"
	"time"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

// The unique ID for this builder.
const BuilderId = "packer.googlecompute"

// Builder represents a Packer Builder.
type Builder struct {
	config *Config
	runner multistep.Runner
}

// Prepare processes the build configuration parameters.
func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	c, warnings, errs := NewConfig(raws...)
	if errs != nil {
		return warnings, errs
	}
	b.config = c

	return warnings, nil
}

// Run executes a googlecompute Packer build and returns a packer.Artifact
// representing a GCE machine image.
func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Initialize the Google Compute Engine API.
	client, err := New(b.config.ProjectId, b.config.Zone, b.config.clientSecrets, b.config.privateKeyBytes)
	if err != nil {
		log.Println("Failed to create the Google Compute Engine client.")
		return nil, err
	}
	// Set up the state.
	state := new(multistep.BasicStateBag)
	state.Put("config", b.config)
	state.Put("client", client)
	state.Put("hook", hook)
	state.Put("ui", ui)
	// Build the steps.
	steps := []multistep.Step{
		new(stepCreateSSHKey),
		new(stepCreateInstance),
		new(stepInstanceInfo),
		&common.StepConnectSSH{
			SSHAddress:     sshAddress,
			SSHConfig:      sshConfig,
			SSHWaitTimeout: 5 * time.Minute,
		},
		new(common.StepProvision),
		new(stepUpdateGsutil),
		new(stepCreateImage),
		new(stepUploadImage),
		new(stepRegisterImage),
	}
	// Run the steps.
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}
	b.runner.Run(state)
	// Report any errors.
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}
	if _, ok := state.GetOk("image_name"); !ok {
		log.Println("Failed to find image_name in state. Bug?")
		return nil, nil
	}
	artifact := &Artifact{
		imageName: state.Get("image_name").(string),
		client:    client,
	}
	return artifact, nil
}

// Cancel.
func (b *Builder) Cancel() {}
<commit_msg>builder/googlecompute: cancel works<commit_after>// The googlecompute package contains a packer.Builder implementation that
// builds images for Google Compute Engine.
package googlecompute

import (
	"log"
	"time"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

// The unique ID for this builder.
const BuilderId = "packer.googlecompute"

// Builder represents a Packer Builder.
type Builder struct {
	config *Config
	runner multistep.Runner
}

// Prepare processes the build configuration parameters.
func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	c, warnings, errs := NewConfig(raws...)
	if errs != nil {
		return warnings, errs
	}
	b.config = c

	return warnings, nil
}

// Run executes a googlecompute Packer build and returns a packer.Artifact
// representing a GCE machine image.
func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Initialize the Google Compute Engine API.
	client, err := New(b.config.ProjectId, b.config.Zone, b.config.clientSecrets, b.config.privateKeyBytes)
	if err != nil {
		log.Println("Failed to create the Google Compute Engine client.")
		return nil, err
	}

	// Set up the state.
	state := new(multistep.BasicStateBag)
	state.Put("config", b.config)
	state.Put("client", client)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Build the steps.
	steps := []multistep.Step{
		new(stepCreateSSHKey),
		new(stepCreateInstance),
		new(stepInstanceInfo),
		&common.StepConnectSSH{
			SSHAddress:     sshAddress,
			SSHConfig:      sshConfig,
			SSHWaitTimeout: 5 * time.Minute,
		},
		new(common.StepProvision),
		new(stepUpdateGsutil),
		new(stepCreateImage),
		new(stepUploadImage),
		new(stepRegisterImage),
	}

	// Run the steps.
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}
	b.runner.Run(state)

	// Report any errors.
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}
	if _, ok := state.GetOk("image_name"); !ok {
		log.Println("Failed to find image_name in state. Bug?")
		return nil, nil
	}

	artifact := &Artifact{
		imageName: state.Get("image_name").(string),
		client:    client,
	}
	return artifact, nil
}

// Cancel.
func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<|endoftext|>


<commit_before>package receiver

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"os"
)

// Sync message from sbc
// echo -ne '\x00\x0A\x00\x30\x59\x41\x37\x38\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x20\x00\x01\x00\x02\x00\xFC\x77\x31\x00\x00\x00\x1E\x00\x00\x00\x00\x43\x5A\x07\x03\x00\x06\x65\x63\x7A\x37\x33\x30' | nc localhost 4739

// Iheader is the Oracle IPFIX Header
//
// Wire format:
//
// Bytes:  0                   1                   2                   3
// Bits:   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |             Version (uint16)          |            Length (uint16)            |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                   ExportTime (unit32)                         |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                 SequenceNumber (uint32)                       |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                 ObservationID (uint32)                        |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//
//
type IpfixHeader struct {
	Version       uint16
	Length        uint16
	ExportTime    uint32
	SeqNum        uint32
	ObservationID uint32
}

// SetHeader represents set header fields
type SetHeader struct {
	ID     uint16
	Length uint16
}

type DataSet struct {
	HandShake Hs
	SIP       Su
}

type Hs struct {
	MaVer       uint16
	MiVer       uint16
	CFlags1     uint16
	CFlags2     uint16
	SFlags      uint16
	Timeout     uint16
	SystemID    uint32
	Product     uint16
	SMaVer      uint8
	SMiVer      uint8
	Revision    uint8
	HostnameLen uint8
	Hostname    []byte
}

// IPFIX is the Oracle IPFIX Handshake Message
type IPFIX struct {
	Header IpfixHeader
	Set    SetHeader
	Data   DataSet
}

type Su struct {
	TimeSec   uint32
	TimeMic   uint32
	IntSlot   uint8
	IntPort   uint8
	IntVlan   uint16
	CallIDLen uint8
	CallID    []byte
	CallIDEnd uint8
	IPlen     uint16
	VL        uint8
	TOS       uint8
	TLen      uint16
	TID       uint16
	TFlags    uint16
	TTL       uint8
	TProto    uint8
	TPos      uint16
	SrcIP     uint32
	DstIP     uint32
	DstPort   uint16
	SrcPort   uint16
	UDPlen    uint16
	MsgLen    uint16
	SipMsg    []byte
}

// BufToInt8 casts []byte to []int8
func BufToInt8(b []byte) []int8 {
	bi := make([]int8, len(b))
	for i, v := range b {
		bi[i] = int8(v)
	}
	return bi
}

// NewIPFIX fills the IPFIX struct with structured binary data from r
func NewIPFIX(header []byte) *IPFIX {
	var ni IPFIX
	rni := bytes.NewReader(header)
	binary.Read(rni, binary.BigEndian, &ni.Header.Version)
	binary.Read(rni, binary.BigEndian, &ni.Header.Length)
	binary.Read(rni, binary.BigEndian, &ni.Header.ExportTime)
	binary.Read(rni, binary.BigEndian, &ni.Header.SeqNum)
	binary.Read(rni, binary.BigEndian, &ni.Header.ObservationID)
	binary.Read(rni, binary.BigEndian, &ni.Set.ID)
	binary.Read(rni, binary.BigEndian, &ni.Set.Length)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.MaVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.MiVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.CFlags1)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.CFlags2)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SFlags)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Timeout)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SystemID)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Product)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SMaVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SMiVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Revision)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.HostnameLen)
	ni.Data.HandShake.Hostname = make([]byte, ni.Data.HandShake.HostnameLen)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Hostname)
	return &ni
}

func NewSipRecUdp(header []byte) *IPFIX {
	var niu IPFIX
	rniu := bytes.NewReader(header)

	binary.Read(rniu, binary.BigEndian, &niu.Header.Version)
	binary.Read(rniu, binary.BigEndian, &niu.Header.Length)
	binary.Read(rniu, binary.BigEndian, &niu.Header.ExportTime)
	binary.Read(rniu, binary.BigEndian, &niu.Header.SeqNum)
	binary.Read(rniu, binary.BigEndian, &niu.Header.ObservationID)
	binary.Read(rniu, binary.BigEndian, &niu.Set.ID)
	binary.Read(rniu, binary.BigEndian, &niu.Set.Length)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TimeSec)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TimeMic)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntSlot)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntVlan)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallIDLen)
	niu.Data.SIP.CallID = make([]byte, niu.Data.SIP.CallIDLen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallID)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallIDEnd)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IPlen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.VL)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TOS)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TLen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TID)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TFlags)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TTL)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TProto)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TPos)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SrcIP)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.DstIP)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.DstPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SrcPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.UDPlen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.MsgLen)
	/*
		fmt.Println(rniu.Len())
		fmt.Println("headerlen:", len(header))
	*/
	niu.Data.SIP.SipMsg = make([]byte, rniu.Len())
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SipMsg)

	return &niu
}

// HandShake writes the binary IPFIX representation into the buffer
func (fi *IPFIX) HandShake() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, &fi.Header.Version)
	binary.Write(buf, binary.BigEndian, &fi.Header.Length)
	binary.Write(buf, binary.BigEndian, &fi.Header.ExportTime)
	binary.Write(buf, binary.BigEndian, &fi.Header.SeqNum)
	binary.Write(buf, binary.BigEndian, &fi.Header.ObservationID)
	binary.Write(buf, binary.BigEndian, &fi.Set.ID)
	binary.Write(buf, binary.BigEndian, &fi.Set.Length)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.MaVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.MiVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.CFlags1)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.CFlags2)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SFlags)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Timeout)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SystemID)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Product)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SMaVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SMiVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Revision)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.HostnameLen)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Hostname)

	return buf.Bytes()
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}

func SyncClient(conn net.Conn) {
	fmt.Println("Handling new connection...")

	// Close connection when this function ends
	defer func() {
		fmt.Println("Closing connection...")
		conn.Close()
	}()

	hs := make([]byte, 512)
	n, err := conn.Read(hs)
	checkError(err)

	packet := NewIPFIX(hs[:n])
	if packet.Set.ID == 256 {
		packet.Set.ID++
	}

	fmt.Printf("Write: %#v\n", BufToInt8(packet.HandShake()))

	if packet.Set.ID == 257 && packet.Header.Length > 20 {
		binary.Write(conn, binary.BigEndian, BufToInt8(packet.HandShake()))
	}

	//bufReader := bufio.NewReader(conn)

	for {
		b := make([]byte, 4096)
		n, err := conn.Read(b)
		checkError(err)

		data := NewSipRecUdp(b[:n])

		/*		// Read IPFIX Messages delimited by newline
				bytes, err := bufReader.ReadBytes('\n')
				checkError(err)*/

		//fmt.Printf("%s", (binary.BigEndian.Uint16(bytes)))
		fmt.Printf("%s", data.Data.SIP)

	}
}
<commit_msg>Rename dataset 259<commit_after>package receiver

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"os"
)

// Sync message from sbc
// echo -ne '\x00\x0A\x00\x30\x59\x41\x37\x38\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x20\x00\x01\x00\x02\x00\xFC\x77\x31\x00\x00\x00\x1E\x00\x00\x00\x00\x43\x5A\x07\x03\x00\x06\x65\x63\x7A\x37\x33\x30' | nc localhost 4739

// Iheader is the Oracle IPFIX Header
//
// Wire format:
//
// Bytes:  0                   1                   2                   3
// Bits:   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |             Version (uint16)          |            Length (uint16)            |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                   ExportTime (unit32)                         |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                 SequenceNumber (uint32)                       |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//         |                                 ObservationID (uint32)                        |
//         +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//
//
type IpfixHeader struct {
	Version       uint16
	Length        uint16
	ExportTime    uint32
	SeqNum        uint32
	ObservationID uint32
}

// SetHeader represents set header fields
type SetHeader struct {
	ID     uint16
	Length uint16
}

type DataSet struct {
	HandShake Hs
	SIP       Su
}

type Hs struct {
	MaVer       uint16
	MiVer       uint16
	CFlags1     uint16
	CFlags2     uint16
	SFlags      uint16
	Timeout     uint16
	SystemID    uint32
	Product     uint16
	SMaVer      uint8
	SMiVer      uint8
	Revision    uint8
	HostnameLen uint8
	Hostname    []byte
}

// IPFIX is the Oracle IPFIX Handshake Message
type IPFIX struct {
	Header IpfixHeader
	Set    SetHeader
	Data   DataSet
}

type Su struct {
	TimeSec   uint32
	TimeMic   uint32
	IntSlot   uint8
	IntPort   uint8
	IntVlan   uint16
	CallIDLen uint8
	CallID    []byte
	CallIDEnd uint8
	IPlen     uint16
	VL        uint8
	TOS       uint8
	TLen      uint16
	TID       uint16
	TFlags    uint16
	TTL       uint8
	TProto    uint8
	TPos      uint16
	SrcIP     uint32
	DstIP     uint32
	DstPort   uint16
	SrcPort   uint16
	UDPlen    uint16
	MsgLen    uint16
	SipMsg    []byte
}

// BufToInt8 casts []byte to []int8
func BufToInt8(b []byte) []int8 {
	bi := make([]int8, len(b))
	for i, v := range b {
		bi[i] = int8(v)
	}
	return bi
}

// NewIPFIX fills the IPFIX struct with structured binary data from r
func NewIPFIX(header []byte) *IPFIX {
	var ni IPFIX
	rni := bytes.NewReader(header)
	binary.Read(rni, binary.BigEndian, &ni.Header.Version)
	binary.Read(rni, binary.BigEndian, &ni.Header.Length)
	binary.Read(rni, binary.BigEndian, &ni.Header.ExportTime)
	binary.Read(rni, binary.BigEndian, &ni.Header.SeqNum)
	binary.Read(rni, binary.BigEndian, &ni.Header.ObservationID)
	binary.Read(rni, binary.BigEndian, &ni.Set.ID)
	binary.Read(rni, binary.BigEndian, &ni.Set.Length)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.MaVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.MiVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.CFlags1)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.CFlags2)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SFlags)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Timeout)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SystemID)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Product)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SMaVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.SMiVer)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Revision)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.HostnameLen)
	ni.Data.HandShake.Hostname = make([]byte, ni.Data.HandShake.HostnameLen)
	binary.Read(rni, binary.BigEndian, &ni.Data.HandShake.Hostname)
	return &ni
}

func NewSendSipUDP(header []byte) *IPFIX {
	var niu IPFIX
	rniu := bytes.NewReader(header)

	binary.Read(rniu, binary.BigEndian, &niu.Header.Version)
	binary.Read(rniu, binary.BigEndian, &niu.Header.Length)
	binary.Read(rniu, binary.BigEndian, &niu.Header.ExportTime)
	binary.Read(rniu, binary.BigEndian, &niu.Header.SeqNum)
	binary.Read(rniu, binary.BigEndian, &niu.Header.ObservationID)
	binary.Read(rniu, binary.BigEndian, &niu.Set.ID)
	binary.Read(rniu, binary.BigEndian, &niu.Set.Length)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TimeSec)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TimeMic)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntSlot)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IntVlan)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallIDLen)
	niu.Data.SIP.CallID = make([]byte, niu.Data.SIP.CallIDLen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallID)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.CallIDEnd)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.IPlen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.VL)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TOS)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TLen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TID)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TFlags)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TTL)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TProto)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.TPos)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SrcIP)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.DstIP)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.DstPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SrcPort)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.UDPlen)
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.MsgLen)
	/*
		fmt.Println(rniu.Len())
		fmt.Println("headerlen:", len(header))
	*/
	niu.Data.SIP.SipMsg = make([]byte, rniu.Len())
	binary.Read(rniu, binary.BigEndian, &niu.Data.SIP.SipMsg)

	return &niu
}

// HandShake writes the binary IPFIX representation into the buffer
func (fi *IPFIX) HandShake() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, &fi.Header.Version)
	binary.Write(buf, binary.BigEndian, &fi.Header.Length)
	binary.Write(buf, binary.BigEndian, &fi.Header.ExportTime)
	binary.Write(buf, binary.BigEndian, &fi.Header.SeqNum)
	binary.Write(buf, binary.BigEndian, &fi.Header.ObservationID)
	binary.Write(buf, binary.BigEndian, &fi.Set.ID)
	binary.Write(buf, binary.BigEndian, &fi.Set.Length)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.MaVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.MiVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.CFlags1)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.CFlags2)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SFlags)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Timeout)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SystemID)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Product)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SMaVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.SMiVer)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Revision)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.HostnameLen)
	binary.Write(buf, binary.BigEndian, &fi.Data.HandShake.Hostname)

	return buf.Bytes()
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}

func SyncClient(conn net.Conn) {
	fmt.Println("Handling new connection...")

	// Close connection when this function ends
	defer func() {
		fmt.Println("Closing connection...")
		conn.Close()
	}()

	hs := make([]byte, 512)
	n, err := conn.Read(hs)
	checkError(err)

	packet := NewIPFIX(hs[:n])
	if packet.Set.ID == 256 {
		packet.Set.ID++
	}

	fmt.Printf("Write: %#v\n", BufToInt8(packet.HandShake()))

	if packet.Set.ID == 257 && packet.Header.Length > 20 {
		binary.Write(conn, binary.BigEndian, BufToInt8(packet.HandShake()))
	}

	//bufReader := bufio.NewReader(conn)

	for {
		b := make([]byte, 4096)
		n, err := conn.Read(b)
		checkError(err)

		data := NewSendSipUDP(b[:n])

		/*		// Read IPFIX Messages delimited by newline
				bytes, err := bufReader.ReadBytes('\n')
				checkError(err)*/

		//fmt.Printf("%s", (binary.BigEndian.Uint16(bytes)))
		fmt.Printf("%s", data.Data.SIP)

	}
}
<|endoftext|>


<commit_before>package carbonserver

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"github.com/go-graphite/carbonzipper/zipper/httpHeaders"
	protov2 "github.com/go-graphite/protocol/carbonapi_v2_pb"
	protov3 "github.com/go-graphite/protocol/carbonapi_v3_pb"
	pickle "github.com/lomik/og-rek"
)

type findResponse struct {
	data        []byte
	contentType string
	files       int
}

func (listener *CarbonserverListener) findHandler(wr http.ResponseWriter, req *http.Request) {
	// URL: /metrics/find/?local=1&format=pickle&query=the.metric.path.with.glob

	t0 := time.Now()
	ctx := req.Context()

	atomic.AddUint64(&listener.metrics.FindRequests, 1)

	format := req.FormValue("format")
	query := req.Form["query"]

	var response *findResponse

	logger := TraceContextToZap(ctx, listener.logger.With(
		zap.String("handler", "find"),
		zap.String("url", req.URL.RequestURI()),
		zap.String("peer", req.RemoteAddr),
	))

	accessLogger := TraceContextToZap(ctx, listener.accessLogger.With(
		zap.String("handler", "find"),
		zap.String("url", req.URL.RequestURI()),
		zap.String("peer", req.RemoteAddr),
	))

	accepts := req.Header["Accept"]
	for _, accept := range accepts {
		if accept == httpHeaders.ContentTypeCarbonAPIv3PB {
			format = "carbonapi_v3_pb"
			break
		}
	}

	if format == "" {
		format = "json"
	}

	formatCode, ok := knownFormats[format]
	if !ok {
		atomic.AddUint64(&listener.metrics.FindErrors, 1)
		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", "unsupported format"),
			zap.Int("http_code", http.StatusBadRequest),
		)
		http.Error(wr, "Bad request (unsupported format)",
			http.StatusBadRequest)
		return
	}

	if formatCode == protoV3Format {
		body, err := ioutil.ReadAll(req.Body)
		if err != nil {
			accessLogger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", err.Error()),
				zap.Int("http_code", http.StatusBadRequest),
			)
			http.Error(wr, "Bad request (unsupported format)",
				http.StatusBadRequest)
		}

		var pv3Request protov3.MultiGlobRequest
		pv3Request.Unmarshal(body)

		fmt.Printf("\n\n%+v\n\n", pv3Request)

		query = pv3Request.Metrics
	}

	logger = logger.With(
		zap.Strings("query", query),
		zap.String("format", format),
	)

	accessLogger = accessLogger.With(
		zap.Strings("query", query),
		zap.String("format", format),
	)

	if len(query) == 0 {
		atomic.AddUint64(&listener.metrics.FindErrors, 1)
		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", "empty query"),
			zap.Int("http_code", http.StatusBadRequest),
		)
		http.Error(wr, "Bad request (no query)", http.StatusBadRequest)
		return
	}

	var err error
	fromCache := false
	//rch := make(chan *findResponse)
	//response := make(chan findResponse)
	//var resp findResponse
	if listener.findCacheEnabled {
		key := strings.Join(query, ",") + "&" + format
		size := uint64(100 * 1024 * 1024)
		item := listener.findCache.getQueryItem(key, size, 300)
		res, ok := item.FetchOrLock()
		listener.prometheus.cacheRequest("find", ok)
		if !ok {
			logger.Debug("find cache miss")
			atomic.AddUint64(&listener.metrics.FindCacheMiss, 1)
			response, err = listener.findMetrics(ctx, logger, t0, formatCode, query)
			if err != nil {
				item.StoreAbort()
			} else {
				item.StoreAndUnlock(response)
			}
		} else if res != nil {
			logger.Debug("query cache hit")
			atomic.AddUint64(&listener.metrics.FindCacheHit, 1)
			response = res.(*findResponse)
			fromCache = true
		}
	} else {
		response, err = listener.findMetrics(ctx, logger, t0, formatCode, query)
	}

	if err != nil || response == nil {
		var code int
		var reason string
		if _, ok := err.(errorNotFound); ok {
			reason = "Not Found"
			code = http.StatusNotFound
		} else {
			reason = "Internal error while processing request"
			code = http.StatusInternalServerError
		}

		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", reason),
			zap.Error(err),
			zap.Int("http_code", code),
		)
		http.Error(wr, fmt.Sprintf("%s (%v)", reason, err), code)

		return
	}

	wr.Header().Set("Content-Type", response.contentType)
	wr.Write(response.data)

	if response.files == 0 {
		// to get an idea how often we search for nothing
		atomic.AddUint64(&listener.metrics.FindZero, 1)
	}

	accessLogger.Info("find success",
		zap.Duration("runtime_seconds", time.Since(t0)),
		zap.Int("Files", response.files),
		zap.Bool("find_cache_enabled", listener.findCacheEnabled),
		zap.Bool("from_cache", fromCache),
		zap.Int("http_code", http.StatusOK),
	)
	return
}

type errorNotFound struct{}

func (err errorNotFound) Error() string {
	return "Not Found"
}

type findError struct {
	name string
	err  error
}

type globs struct {
	Name  string
	Files []string
	Leafs []bool
}

func (listener *CarbonserverListener) findMetrics(ctx context.Context, logger *zap.Logger, t0 time.Time, format responseFormat, names []string) (*findResponse, error) {
	var result findResponse
	metricsCount := uint64(0)
	expandedGlobs, err := listener.getExpandedGlobs(ctx, logger, t0, names)
	if expandedGlobs == nil {
		return nil, err
	}
	atomic.AddUint64(&listener.metrics.MetricsFound, metricsCount)
	switch format {
	case protoV3Format, jsonFormat:
		var err error
		multiResponse := protov3.MultiGlobResponse{}
		for _, glob := range expandedGlobs {
			result.files += len(glob.Files)
			response := protov3.GlobResponse{
				Name:    glob.Name,
				Matches: make([]protov3.GlobMatch, 0),
			}

			for i, p := range glob.Files {
				if glob.Leafs[i] {
					metricsCount++
				}
				response.Matches = append(response.Matches, protov3.GlobMatch{Path: p, IsLeaf: glob.Leafs[i]})
			}
			multiResponse.Metrics = append(multiResponse.Metrics, response)
		}

		logger.Debug("will send out response",
			zap.Any("response", multiResponse),
		)

		switch format {
		case jsonFormat:
			result.contentType = httpHeaders.ContentTypeJSON
			result.data, err = json.Marshal(multiResponse)
		case protoV3Format:
			result.contentType = httpHeaders.ContentTypeCarbonAPIv3PB
			result.data, err = multiResponse.Marshal()
		}

		if err != nil {
			atomic.AddUint64(&listener.metrics.FindErrors, 1)

			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "response encode failed"),
				zap.Error(err),
			)
			return nil, err
		}

		if len(multiResponse.Metrics) == 0 {
			return nil, errorNotFound{}
		}

		return &result, err

	case protoV2Format:
		result.contentType = httpHeaders.ContentTypeProtobuf
		var err error
		response := protov2.GlobResponse{
			Name:    names[0],
			Matches: make([]protov2.GlobMatch, 0),
		}

		result.files += len(expandedGlobs[0].Files)
		for i, p := range expandedGlobs[0].Files {
			if expandedGlobs[0].Leafs[i] {
				metricsCount++
			}
			response.Matches = append(response.Matches, protov2.GlobMatch{Path: p, IsLeaf: expandedGlobs[0].Leafs[i]})
		}
		result.data, err = response.Marshal()
		result.contentType = httpHeaders.ContentTypeProtobuf

		if err != nil {
			atomic.AddUint64(&listener.metrics.FindErrors, 1)

			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "response encode failed"),
				zap.Error(err),
			)
			return nil, err
		}

		if len(response.Matches) == 0 {
			return nil, errorNotFound{}
		}

		return &result, err

	case pickleFormat:
		// [{'metric_path': 'metric', 'intervals': [(x,y)], 'isLeaf': True},]
		var metrics []map[string]interface{}
		var m map[string]interface{}
		files := 0

		glob := expandedGlobs[0]
		files += len(glob.Files)
		for i, p := range glob.Files {
			if glob.Leafs[i] {
				metricsCount++
			}
			m = make(map[string]interface{})
			m["metric_path"] = p
			m["isLeaf"] = glob.Leafs[i]

			metrics = append(metrics, m)
		}
		var buf bytes.Buffer
		pEnc := pickle.NewEncoder(&buf)
		pEnc.Encode(metrics)
		return &findResponse{buf.Bytes(), httpHeaders.ContentTypePickle, files}, nil
	}
	return nil, nil
}

func (listener *CarbonserverListener) getExpandedGlobs(ctx context.Context, logger *zap.Logger, t0 time.Time, names []string) ([]globs, error) {
	var expandedGlobs []globs
	var errors []findError
	var err error
	expGlobResultChan := make(chan *ExpandedGlobResponse, len(names))

	for _, name := range names {
		go listener.expandGlobs(name, expGlobResultChan)
	}
	responseCount := 0
GATHER:
	for {
		select {
		case expandedResult := <-expGlobResultChan:
			responseCount++
			glob := globs{
				Name: expandedResult.Name,
			}
			glob.Files, glob.Leafs, err = expandedResult.Files, expandedResult.Leafs, expandedResult.Err
			if err != nil {
				errors = append(errors, findError{name: expandedResult.Name, err: err})
			}
			expandedGlobs = append(expandedGlobs, glob)
			if responseCount == len(names) {
				break GATHER
			}

		case <-ctx.Done():
			done <- struct{}{}
			return nil, fmt.Errorf("find failed, timeout")
		}
	}

	if len(errors) > 0 {
		atomic.AddUint64(&listener.metrics.FindErrors, uint64(len(errors)))

		if len(errors) == len(names) {
			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "can't expand globs"),
				zap.Any("errors", errors),
			)
			return nil, fmt.Errorf("find failed, can't expand globs")
		} else {
			logger.Warn("find partly failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "can't expand globs for some metrics"),
				zap.Any("errors", errors),
			)
		}
	}

	logger.Debug("expandGlobs result",
		zap.String("action", "expandGlobs"),
		zap.Strings("metrics", names),
		//zap.String("format", format.String()),
		zap.Any("result", expandedGlobs),
	)
	return expandedGlobs, nil
}
<commit_msg>Remove done channel.<commit_after>package carbonserver

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"github.com/go-graphite/carbonzipper/zipper/httpHeaders"
	protov2 "github.com/go-graphite/protocol/carbonapi_v2_pb"
	protov3 "github.com/go-graphite/protocol/carbonapi_v3_pb"
	pickle "github.com/lomik/og-rek"
)

type findResponse struct {
	data        []byte
	contentType string
	files       int
}

func (listener *CarbonserverListener) findHandler(wr http.ResponseWriter, req *http.Request) {
	// URL: /metrics/find/?local=1&format=pickle&query=the.metric.path.with.glob

	t0 := time.Now()
	ctx := req.Context()

	atomic.AddUint64(&listener.metrics.FindRequests, 1)

	format := req.FormValue("format")
	query := req.Form["query"]

	var response *findResponse

	logger := TraceContextToZap(ctx, listener.logger.With(
		zap.String("handler", "find"),
		zap.String("url", req.URL.RequestURI()),
		zap.String("peer", req.RemoteAddr),
	))

	accessLogger := TraceContextToZap(ctx, listener.accessLogger.With(
		zap.String("handler", "find"),
		zap.String("url", req.URL.RequestURI()),
		zap.String("peer", req.RemoteAddr),
	))

	accepts := req.Header["Accept"]
	for _, accept := range accepts {
		if accept == httpHeaders.ContentTypeCarbonAPIv3PB {
			format = "carbonapi_v3_pb"
			break
		}
	}

	if format == "" {
		format = "json"
	}

	formatCode, ok := knownFormats[format]
	if !ok {
		atomic.AddUint64(&listener.metrics.FindErrors, 1)
		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", "unsupported format"),
			zap.Int("http_code", http.StatusBadRequest),
		)
		http.Error(wr, "Bad request (unsupported format)",
			http.StatusBadRequest)
		return
	}

	if formatCode == protoV3Format {
		body, err := ioutil.ReadAll(req.Body)
		if err != nil {
			accessLogger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", err.Error()),
				zap.Int("http_code", http.StatusBadRequest),
			)
			http.Error(wr, "Bad request (unsupported format)",
				http.StatusBadRequest)
		}

		var pv3Request protov3.MultiGlobRequest
		pv3Request.Unmarshal(body)

		fmt.Printf("\n\n%+v\n\n", pv3Request)

		query = pv3Request.Metrics
	}

	logger = logger.With(
		zap.Strings("query", query),
		zap.String("format", format),
	)

	accessLogger = accessLogger.With(
		zap.Strings("query", query),
		zap.String("format", format),
	)

	if len(query) == 0 {
		atomic.AddUint64(&listener.metrics.FindErrors, 1)
		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", "empty query"),
			zap.Int("http_code", http.StatusBadRequest),
		)
		http.Error(wr, "Bad request (no query)", http.StatusBadRequest)
		return
	}

	var err error
	fromCache := false
	//rch := make(chan *findResponse)
	//response := make(chan findResponse)
	//var resp findResponse
	if listener.findCacheEnabled {
		key := strings.Join(query, ",") + "&" + format
		size := uint64(100 * 1024 * 1024)
		item := listener.findCache.getQueryItem(key, size, 300)
		res, ok := item.FetchOrLock()
		listener.prometheus.cacheRequest("find", ok)
		if !ok {
			logger.Debug("find cache miss")
			atomic.AddUint64(&listener.metrics.FindCacheMiss, 1)
			response, err = listener.findMetrics(ctx, logger, t0, formatCode, query)
			if err != nil {
				item.StoreAbort()
			} else {
				item.StoreAndUnlock(response)
			}
		} else if res != nil {
			logger.Debug("query cache hit")
			atomic.AddUint64(&listener.metrics.FindCacheHit, 1)
			response = res.(*findResponse)
			fromCache = true
		}
	} else {
		response, err = listener.findMetrics(ctx, logger, t0, formatCode, query)
	}

	if err != nil || response == nil {
		var code int
		var reason string
		if _, ok := err.(errorNotFound); ok {
			reason = "Not Found"
			code = http.StatusNotFound
		} else {
			reason = "Internal error while processing request"
			code = http.StatusInternalServerError
		}

		accessLogger.Error("find failed",
			zap.Duration("runtime_seconds", time.Since(t0)),
			zap.String("reason", reason),
			zap.Error(err),
			zap.Int("http_code", code),
		)
		http.Error(wr, fmt.Sprintf("%s (%v)", reason, err), code)

		return
	}

	wr.Header().Set("Content-Type", response.contentType)
	wr.Write(response.data)

	if response.files == 0 {
		// to get an idea how often we search for nothing
		atomic.AddUint64(&listener.metrics.FindZero, 1)
	}

	accessLogger.Info("find success",
		zap.Duration("runtime_seconds", time.Since(t0)),
		zap.Int("Files", response.files),
		zap.Bool("find_cache_enabled", listener.findCacheEnabled),
		zap.Bool("from_cache", fromCache),
		zap.Int("http_code", http.StatusOK),
	)
	return
}

type errorNotFound struct{}

func (err errorNotFound) Error() string {
	return "Not Found"
}

type findError struct {
	name string
	err  error
}

type globs struct {
	Name  string
	Files []string
	Leafs []bool
}

func (listener *CarbonserverListener) findMetrics(ctx context.Context, logger *zap.Logger, t0 time.Time, format responseFormat, names []string) (*findResponse, error) {
	var result findResponse
	metricsCount := uint64(0)
	expandedGlobs, err := listener.getExpandedGlobs(ctx, logger, t0, names)
	if expandedGlobs == nil {
		return nil, err
	}
	atomic.AddUint64(&listener.metrics.MetricsFound, metricsCount)
	switch format {
	case protoV3Format, jsonFormat:
		var err error
		multiResponse := protov3.MultiGlobResponse{}
		for _, glob := range expandedGlobs {
			result.files += len(glob.Files)
			response := protov3.GlobResponse{
				Name:    glob.Name,
				Matches: make([]protov3.GlobMatch, 0),
			}

			for i, p := range glob.Files {
				if glob.Leafs[i] {
					metricsCount++
				}
				response.Matches = append(response.Matches, protov3.GlobMatch{Path: p, IsLeaf: glob.Leafs[i]})
			}
			multiResponse.Metrics = append(multiResponse.Metrics, response)
		}

		logger.Debug("will send out response",
			zap.Any("response", multiResponse),
		)

		switch format {
		case jsonFormat:
			result.contentType = httpHeaders.ContentTypeJSON
			result.data, err = json.Marshal(multiResponse)
		case protoV3Format:
			result.contentType = httpHeaders.ContentTypeCarbonAPIv3PB
			result.data, err = multiResponse.Marshal()
		}

		if err != nil {
			atomic.AddUint64(&listener.metrics.FindErrors, 1)

			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "response encode failed"),
				zap.Error(err),
			)
			return nil, err
		}

		if len(multiResponse.Metrics) == 0 {
			return nil, errorNotFound{}
		}

		return &result, err

	case protoV2Format:
		result.contentType = httpHeaders.ContentTypeProtobuf
		var err error
		response := protov2.GlobResponse{
			Name:    names[0],
			Matches: make([]protov2.GlobMatch, 0),
		}

		result.files += len(expandedGlobs[0].Files)
		for i, p := range expandedGlobs[0].Files {
			if expandedGlobs[0].Leafs[i] {
				metricsCount++
			}
			response.Matches = append(response.Matches, protov2.GlobMatch{Path: p, IsLeaf: expandedGlobs[0].Leafs[i]})
		}
		result.data, err = response.Marshal()
		result.contentType = httpHeaders.ContentTypeProtobuf

		if err != nil {
			atomic.AddUint64(&listener.metrics.FindErrors, 1)

			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "response encode failed"),
				zap.Error(err),
			)
			return nil, err
		}

		if len(response.Matches) == 0 {
			return nil, errorNotFound{}
		}

		return &result, err

	case pickleFormat:
		// [{'metric_path': 'metric', 'intervals': [(x,y)], 'isLeaf': True},]
		var metrics []map[string]interface{}
		var m map[string]interface{}
		files := 0

		glob := expandedGlobs[0]
		files += len(glob.Files)
		for i, p := range glob.Files {
			if glob.Leafs[i] {
				metricsCount++
			}
			m = make(map[string]interface{})
			m["metric_path"] = p
			m["isLeaf"] = glob.Leafs[i]

			metrics = append(metrics, m)
		}
		var buf bytes.Buffer
		pEnc := pickle.NewEncoder(&buf)
		pEnc.Encode(metrics)
		return &findResponse{buf.Bytes(), httpHeaders.ContentTypePickle, files}, nil
	}
	return nil, nil
}

func (listener *CarbonserverListener) getExpandedGlobs(ctx context.Context, logger *zap.Logger, t0 time.Time, names []string) ([]globs, error) {
	var expandedGlobs []globs
	var errors []findError
	var err error
	expGlobResultChan := make(chan *ExpandedGlobResponse, len(names))

	for _, name := range names {
		go listener.expandGlobs(name, expGlobResultChan)
	}
	responseCount := 0
GATHER:
	for {
		select {
		case expandedResult := <-expGlobResultChan:
			responseCount++
			glob := globs{
				Name: expandedResult.Name,
			}
			glob.Files, glob.Leafs, err = expandedResult.Files, expandedResult.Leafs, expandedResult.Err
			if err != nil {
				errors = append(errors, findError{name: expandedResult.Name, err: err})
			}
			expandedGlobs = append(expandedGlobs, glob)
			if responseCount == len(names) {
				break GATHER
			}

		case <-ctx.Done():
			return nil, fmt.Errorf("find failed, timeout")
		}
	}

	if len(errors) > 0 {
		atomic.AddUint64(&listener.metrics.FindErrors, uint64(len(errors)))

		if len(errors) == len(names) {
			logger.Error("find failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "can't expand globs"),
				zap.Any("errors", errors),
			)
			return nil, fmt.Errorf("find failed, can't expand globs")
		} else {
			logger.Warn("find partly failed",
				zap.Duration("runtime_seconds", time.Since(t0)),
				zap.String("reason", "can't expand globs for some metrics"),
				zap.Any("errors", errors),
			)
		}
	}

	logger.Debug("expandGlobs result",
		zap.String("action", "expandGlobs"),
		zap.Strings("metrics", names),
		//zap.String("format", format.String()),
		zap.Any("result", expandedGlobs),
	)
	return expandedGlobs, nil
}
<|endoftext|>


<commit_before>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dns

// A client implementation.

import (
	"bytes"
	"io"
	"net"
	"time"
)

const dnsTimeout time.Duration = 2 * 1e9
const tcpIdleTimeout time.Duration = 8 * time.Second

// A Conn represents a connection to a DNS server.
type Conn struct {
	net.Conn                         // a net.Conn holding the connection
	UDPSize        uint16            // Minimum receive buffer for UDP messages
	TsigSecret     map[string]string // Secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	rtt            time.Duration
	t              time.Time
	tsigRequestMAC string
}

// A Client defines parameters for a DNS client.
type Client struct {
	Net            string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	DialTimeout    time.Duration     // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	SingleInflight bool              // if true suppress multiple outstanding queries for the same Qname, Qtype and Qclass
	group          singleflight
}

// Exchange performs a synchronous UDP query. It sends the message m to the address
// contained in a and waits for an reply. Exchange does not retry a failed query, nor
// will it fall back to TCP in case of truncation.
func Exchange(m *Msg, a string) (r *Msg, err error) {
	var co *Conn
	co, err = DialTimeout("udp", a, dnsTimeout)
	if err != nil {
		return nil, err
	}

	defer co.Close()
	co.SetReadDeadline(time.Now().Add(dnsTimeout))
	co.SetWriteDeadline(time.Now().Add(dnsTimeout))
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// ExchangeConn performs a synchronous query. It sends the message m via the connection
// c and waits for a reply. The connection c is not closed by ExchangeConn.
// This function is going away, but can easily be mimicked:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, _  := co.ReadMsg()
//	co.Close()
//
func ExchangeConn(c net.Conn, m *Msg) (r *Msg, err error) {
	println("dns: this function is deprecated")
	co := new(Conn)
	co.Conn = c
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *dns.Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
// Exchange does not retry a failed query, nor will it fall back to TCP in
// case of truncation.
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	if !c.SingleInflight {
		return c.exchange(m, a)
	}
	// This adds a bunch of garbage, TODO(miek).
	t := "nop"
	if t1, ok := TypeToString[m.Question[0].Qtype]; ok {
		t = t1
	}
	cl := "nop"
	if cl1, ok := ClassToString[m.Question[0].Qclass]; ok {
		cl = cl1
	}
	r, rtt, err, shared := c.group.Do(m.Question[0].Name+t+cl, func() (*Msg, time.Duration, error) {
		return c.exchange(m, a)
	})
	if err != nil {
		return r, rtt, err
	}
	if shared {
		r1 := r.copy()
		r = r1
	}
	return r, rtt, nil
}

func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	timeout := dnsTimeout
	var co *Conn
	if c.DialTimeout != 0 {
		timeout = c.DialTimeout
	}
	if c.Net == "" {
		co, err = DialTimeout("udp", a, timeout)
	} else {
		co, err = DialTimeout(c.Net, a, timeout)
	}
	if err != nil {
		return nil, 0, err
	}
	timeout = dnsTimeout
	if c.ReadTimeout != 0 {
		timeout = c.ReadTimeout
	}
	co.SetReadDeadline(time.Now().Add(timeout))
	timeout = dnsTimeout
	if c.WriteTimeout != 0 {
		timeout = c.WriteTimeout
	}
	co.SetWriteDeadline(time.Now().Add(timeout))
	defer co.Close()
	opt := m.IsEdns0()
	if opt != nil && opt.UDPSize() >= MinMsgSize {
		co.UDPSize = opt.UDPSize()
	}
	co.TsigSecret = c.TsigSecret
	if err = co.WriteMsg(m); err != nil {
		return nil, 0, err
	}
	r, err = co.ReadMsg()
	return r, co.rtt, err
}

// ReadMsg reads a message from the connection co.
// If the received message contains a TSIG record the transaction
// signature is verified.
func (co *Conn) ReadMsg() (*Msg, error) {
	var p []byte
	m := new(Msg)
	if _, ok := co.Conn.(*net.TCPConn); ok {
		p = make([]byte, MaxMsgSize)
	} else {
		if co.UDPSize >= 512 {
			p = make([]byte, co.UDPSize)
		} else {
			p = make([]byte, MinMsgSize)
		}
	}
	n, err := co.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	co.rtt = time.Since(co.t)
	if t := m.IsTsig(); t != nil {
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
	}
	return m, err
}

// Read implements the net.Conn read method.
func (co *Conn) Read(p []byte) (n int, err error) {
	if co.Conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	if t, ok := co.Conn.(*net.TCPConn); ok {
		n, err = t.Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = t.Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := t.Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
		return n, err
	}
	// assume udp connection
	n, err = co.Conn.Read(p)
	if err != nil {
		return n, err
	}
	return n, err
}

// WriteMsg sends a message throught the connection co.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (co *Conn) WriteMsg(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
		// Set for the next read, allthough only used in zone transfers
		co.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	co.t = time.Now()
	if _, err = co.Write(out); err != nil {
		return err
	}
	return nil
}

// Write implements the net.Conn Write method.
func (co *Conn) Write(p []byte) (n int, err error) {
	if t, ok := co.Conn.(*net.TCPConn); ok {
		lp := len(p)
		if lp < 2 {
			return 0, io.ErrShortBuffer
		}
		if lp > MaxMsgSize {
			return 0, &Error{err: "message too large"}
		}
		l := make([]byte, 2, lp+2)
		l[0], l[1] = packUint16(uint16(lp))
		p = append(l, p...)
		n, err := io.Copy(t, bytes.NewReader(p))
		return int(n), err
	}
	n, err = co.Conn.(*net.UDPConn).Write(p)
	return n, err
}

// Dial connects to the address on the named network.
func Dial(network, address string) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Dialtimeout acts like Dial but takes a timeout.
func DialTimeout(network, address string, timeout time.Duration) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.DialTimeout(network, address, timeout)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Close implements the net.Conn Close method.
func (co *Conn) Close() error { return co.Conn.Close() }

// LocalAddr implements the net.Conn LocalAddr method.
func (co *Conn) LocalAddr() net.Addr { return co.Conn.LocalAddr() }

// RemoteAddr implements the net.Conn RemoteAddr method.
func (co *Conn) RemoteAddr() net.Addr { return co.Conn.RemoteAddr() }

// SetDeadline implements the net.Conn SetDeadline method.
func (co *Conn) SetDeadline(t time.Time) error { return co.Conn.SetDeadline(t) }

// SetReadDeadline implements the net.Conn SetReadDeadline method.
func (co *Conn) SetReadDeadline(t time.Time) error { return co.Conn.SetReadDeadline(t) }

// SetWriteDeadline implements the net.Conn SetWriteDeadline method.
func (co *Conn) SetWriteDeadline(t time.Time) error { return co.Conn.SetWriteDeadline(t) }
<commit_msg>Cleanup singleFlight a bit<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dns

// A client implementation.

import (
	"bytes"
	"io"
	"net"
	"time"
)

const dnsTimeout time.Duration = 2 * 1e9
const tcpIdleTimeout time.Duration = 8 * time.Second

// A Conn represents a connection to a DNS server.
type Conn struct {
	net.Conn                         // a net.Conn holding the connection
	UDPSize        uint16            // Minimum receive buffer for UDP messages
	TsigSecret     map[string]string // Secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	rtt            time.Duration
	t              time.Time
	tsigRequestMAC string
}

// A Client defines parameters for a DNS client.
type Client struct {
	Net            string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	DialTimeout    time.Duration     // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	SingleInflight bool              // if true suppress multiple outstanding queries for the same Qname, Qtype and Qclass
	group          singleflight
}

// Exchange performs a synchronous UDP query. It sends the message m to the address
// contained in a and waits for an reply. Exchange does not retry a failed query, nor
// will it fall back to TCP in case of truncation.
func Exchange(m *Msg, a string) (r *Msg, err error) {
	var co *Conn
	co, err = DialTimeout("udp", a, dnsTimeout)
	if err != nil {
		return nil, err
	}

	defer co.Close()
	co.SetReadDeadline(time.Now().Add(dnsTimeout))
	co.SetWriteDeadline(time.Now().Add(dnsTimeout))
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// ExchangeConn performs a synchronous query. It sends the message m via the connection
// c and waits for a reply. The connection c is not closed by ExchangeConn.
// This function is going away, but can easily be mimicked:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, _  := co.ReadMsg()
//	co.Close()
//
func ExchangeConn(c net.Conn, m *Msg) (r *Msg, err error) {
	println("dns: this function is deprecated")
	co := new(Conn)
	co.Conn = c
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *dns.Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
// Exchange does not retry a failed query, nor will it fall back to TCP in
// case of truncation.
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	if !c.SingleInflight {
		return c.exchange(m, a)
	}
	// This adds a bunch of garbage, TODO(miek).
	t := "nop"
	if t1, ok := TypeToString[m.Question[0].Qtype]; ok {
		t = t1
	}
	cl := "nop"
	if cl1, ok := ClassToString[m.Question[0].Qclass]; ok {
		cl = cl1
	}
	r, rtt, err, shared := c.group.Do(m.Question[0].Name+t+cl, func() (*Msg, time.Duration, error) {
		return c.exchange(m, a)
	})
	if err != nil {
		return r, rtt, err
	}
	if shared {
		return r.copy(), rtt, nil
	}
	return r, rtt, nil
}

func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	timeout := dnsTimeout
	var co *Conn
	if c.DialTimeout != 0 {
		timeout = c.DialTimeout
	}
	if c.Net == "" {
		co, err = DialTimeout("udp", a, timeout)
	} else {
		co, err = DialTimeout(c.Net, a, timeout)
	}
	if err != nil {
		return nil, 0, err
	}
	timeout = dnsTimeout
	if c.ReadTimeout != 0 {
		timeout = c.ReadTimeout
	}
	co.SetReadDeadline(time.Now().Add(timeout))
	timeout = dnsTimeout
	if c.WriteTimeout != 0 {
		timeout = c.WriteTimeout
	}
	co.SetWriteDeadline(time.Now().Add(timeout))
	defer co.Close()
	opt := m.IsEdns0()
	if opt != nil && opt.UDPSize() >= MinMsgSize {
		co.UDPSize = opt.UDPSize()
	}
	co.TsigSecret = c.TsigSecret
	if err = co.WriteMsg(m); err != nil {
		return nil, 0, err
	}
	r, err = co.ReadMsg()
	return r, co.rtt, err
}

// ReadMsg reads a message from the connection co.
// If the received message contains a TSIG record the transaction
// signature is verified.
func (co *Conn) ReadMsg() (*Msg, error) {
	var p []byte
	m := new(Msg)
	if _, ok := co.Conn.(*net.TCPConn); ok {
		p = make([]byte, MaxMsgSize)
	} else {
		if co.UDPSize >= 512 {
			p = make([]byte, co.UDPSize)
		} else {
			p = make([]byte, MinMsgSize)
		}
	}
	n, err := co.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	co.rtt = time.Since(co.t)
	if t := m.IsTsig(); t != nil {
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
	}
	return m, err
}

// Read implements the net.Conn read method.
func (co *Conn) Read(p []byte) (n int, err error) {
	if co.Conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	if t, ok := co.Conn.(*net.TCPConn); ok {
		n, err = t.Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = t.Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := t.Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
		return n, err
	}
	// assume udp connection
	n, err = co.Conn.Read(p)
	if err != nil {
		return n, err
	}
	return n, err
}

// WriteMsg sends a message throught the connection co.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (co *Conn) WriteMsg(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
		// Set for the next read, allthough only used in zone transfers
		co.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	co.t = time.Now()
	if _, err = co.Write(out); err != nil {
		return err
	}
	return nil
}

// Write implements the net.Conn Write method.
func (co *Conn) Write(p []byte) (n int, err error) {
	if t, ok := co.Conn.(*net.TCPConn); ok {
		lp := len(p)
		if lp < 2 {
			return 0, io.ErrShortBuffer
		}
		if lp > MaxMsgSize {
			return 0, &Error{err: "message too large"}
		}
		l := make([]byte, 2, lp+2)
		l[0], l[1] = packUint16(uint16(lp))
		p = append(l, p...)
		n, err := io.Copy(t, bytes.NewReader(p))
		return int(n), err
	}
	n, err = co.Conn.(*net.UDPConn).Write(p)
	return n, err
}

// Dial connects to the address on the named network.
func Dial(network, address string) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Dialtimeout acts like Dial but takes a timeout.
func DialTimeout(network, address string, timeout time.Duration) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.DialTimeout(network, address, timeout)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Close implements the net.Conn Close method.
func (co *Conn) Close() error { return co.Conn.Close() }

// LocalAddr implements the net.Conn LocalAddr method.
func (co *Conn) LocalAddr() net.Addr { return co.Conn.LocalAddr() }

// RemoteAddr implements the net.Conn RemoteAddr method.
func (co *Conn) RemoteAddr() net.Addr { return co.Conn.RemoteAddr() }

// SetDeadline implements the net.Conn SetDeadline method.
func (co *Conn) SetDeadline(t time.Time) error { return co.Conn.SetDeadline(t) }

// SetReadDeadline implements the net.Conn SetReadDeadline method.
func (co *Conn) SetReadDeadline(t time.Time) error { return co.Conn.SetReadDeadline(t) }

// SetWriteDeadline implements the net.Conn SetWriteDeadline method.
func (co *Conn) SetWriteDeadline(t time.Time) error { return co.Conn.SetWriteDeadline(t) }
<|endoftext|>