<commit_before>// Copyright 2015 xeipuuv ( https://github.com/xeipuuv )
//
// 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.

// author           xeipuuv
// author-github    https://github.com/xeipuuv
// author-mail      xeipuuv@gmail.com
//
// repository-name  gojsonschema
// repository-desc  An implementation of JSON Schema, based on IETF's draft v4 - Go language.
//
// description		Different strategies to load JSON files.
// 					Includes References (file and HTTP), JSON strings and Go types.
//
// created          01-02-2015

package gojsonschema

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/xeipuuv/gojsonreference"
)

var osFS = osFileSystem(os.Open)

// JSON loader interface

type JSONLoader interface {
	JsonSource() interface{}
	LoadJSON() (interface{}, error)
	JsonReference() (gojsonreference.JsonReference, error)
	LoaderFactory() JSONLoaderFactory
}

type JSONLoaderFactory interface {
	New(source string) JSONLoader
}

type DefaultJSONLoaderFactory struct {
}

type FileSystemJSONLoaderFactory struct {
	fs http.FileSystem
}

func (d DefaultJSONLoaderFactory) New(source string) JSONLoader {
	return &jsonReferenceLoader{
		fs:     osFS,
		source: source,
	}
}

func (f FileSystemJSONLoaderFactory) New(source string) JSONLoader {
	return &jsonReferenceLoader{
		fs:     f.fs,
		source: source,
	}
}

// osFileSystem is a functional wrapper for os.Open that implements http.FileSystem.
type osFileSystem func(string) (*os.File, error)

func (o osFileSystem) Open(name string) (http.File, error) {
	return o(name)
}

// JSON Reference loader
// references are used to load JSONs from files and HTTP

type jsonReferenceLoader struct {
	fs     http.FileSystem
	source string
}

func (l *jsonReferenceLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonReferenceLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference(l.JsonSource().(string))
}

func (l *jsonReferenceLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

// NewReferenceLoader returns a JSON reference loader using the given source and the local OS file system.
func NewReferenceLoader(source string) *jsonReferenceLoader {
	return &jsonReferenceLoader{
		fs:     osFS,
		source: source,
	}
}

// NewReferenceLoaderFileSystem returns a JSON reference loader using the given source and file system.
func NewReferenceLoaderFileSystem(source string, fs http.FileSystem) *jsonReferenceLoader {
	return &jsonReferenceLoader{
		fs:     fs,
		source: source,
	}
}

func (l *jsonReferenceLoader) LoadJSON() (interface{}, error) {

	var err error

	reference, err := gojsonreference.NewJsonReference(l.JsonSource().(string))
	if err != nil {
		return nil, err
	}

	refToUrl := reference
	refToUrl.GetUrl().Fragment = ""

	var document interface{}

	if reference.HasFileScheme {

		filename := strings.Replace(refToUrl.String(), "file://", "", -1)
		if runtime.GOOS == "windows" {
			// on Windows, a file URL may have an extra leading slash, use slashes
			// instead of backslashes, and have spaces escaped
			if strings.HasPrefix(filename, "/") {
				filename = filename[1:]
			}
			filename = filepath.FromSlash(filename)
			filename = strings.Replace(filename, "%20", " ", -1)
		}

		document, err = l.loadFromFile(filename)
		if err != nil {
			return nil, err
		}

	} else {

		document, err = l.loadFromHTTP(refToUrl.String())
		if err != nil {
			return nil, err
		}

	}

	return document, nil

}

func (l *jsonReferenceLoader) loadFromHTTP(address string) (interface{}, error) {

	resp, err := http.Get(address)
	if err != nil {
		return nil, err
	}

	// must return HTTP Status 200 OK
	if resp.StatusCode != http.StatusOK {
		return nil, errors.New(formatErrorDescription(Locale.httpBadStatus(), ErrorDetails{"status": resp.Status}))
	}

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

	return decodeJsonUsingNumber(bytes.NewReader(bodyBuff))

}

func (l *jsonReferenceLoader) loadFromFile(path string) (interface{}, error) {
	f, err := l.fs.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	bodyBuff, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	return decodeJsonUsingNumber(bytes.NewReader(bodyBuff))

}

// JSON string loader

type jsonStringLoader struct {
	source string
}

func (l *jsonStringLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonStringLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference("#")
}

func (l *jsonStringLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

func NewStringLoader(source string) *jsonStringLoader {
	return &jsonStringLoader{source: source}
}

func (l *jsonStringLoader) LoadJSON() (interface{}, error) {

	return decodeJsonUsingNumber(strings.NewReader(l.JsonSource().(string)))

}

// JSON Go (types) loader
// used to load JSONs from the code as maps, interface{}, structs ...

type jsonGoLoader struct {
	source interface{}
}

func (l *jsonGoLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonGoLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference("#")
}

func (l *jsonGoLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

func NewGoLoader(source interface{}) *jsonGoLoader {
	return &jsonGoLoader{source: source}
}

func (l *jsonGoLoader) LoadJSON() (interface{}, error) {

	// convert it to a compliant JSON first to avoid types "mismatches"

	jsonBytes, err := json.Marshal(l.JsonSource())
	if err != nil {
		return nil, err
	}

	return decodeJsonUsingNumber(bytes.NewReader(jsonBytes))

}

func decodeJsonUsingNumber(r io.Reader) (interface{}, error) {

	var document interface{}

	decoder := json.NewDecoder(r)
	decoder.UseNumber()

	err := decoder.Decode(&document)
	if err != nil {
		return nil, err
	}

	return document, nil

}
<commit_msg>avoid url-encoded characters in file name<commit_after>// Copyright 2015 xeipuuv ( https://github.com/xeipuuv )
//
// 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.

// author           xeipuuv
// author-github    https://github.com/xeipuuv
// author-mail      xeipuuv@gmail.com
//
// repository-name  gojsonschema
// repository-desc  An implementation of JSON Schema, based on IETF's draft v4 - Go language.
//
// description		Different strategies to load JSON files.
// 					Includes References (file and HTTP), JSON strings and Go types.
//
// created          01-02-2015

package gojsonschema

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/xeipuuv/gojsonreference"
)

var osFS = osFileSystem(os.Open)

// JSON loader interface

type JSONLoader interface {
	JsonSource() interface{}
	LoadJSON() (interface{}, error)
	JsonReference() (gojsonreference.JsonReference, error)
	LoaderFactory() JSONLoaderFactory
}

type JSONLoaderFactory interface {
	New(source string) JSONLoader
}

type DefaultJSONLoaderFactory struct {
}

type FileSystemJSONLoaderFactory struct {
	fs http.FileSystem
}

func (d DefaultJSONLoaderFactory) New(source string) JSONLoader {
	return &jsonReferenceLoader{
		fs:     osFS,
		source: source,
	}
}

func (f FileSystemJSONLoaderFactory) New(source string) JSONLoader {
	return &jsonReferenceLoader{
		fs:     f.fs,
		source: source,
	}
}

// osFileSystem is a functional wrapper for os.Open that implements http.FileSystem.
type osFileSystem func(string) (*os.File, error)

func (o osFileSystem) Open(name string) (http.File, error) {
	return o(name)
}

// JSON Reference loader
// references are used to load JSONs from files and HTTP

type jsonReferenceLoader struct {
	fs     http.FileSystem
	source string
}

func (l *jsonReferenceLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonReferenceLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference(l.JsonSource().(string))
}

func (l *jsonReferenceLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

// NewReferenceLoader returns a JSON reference loader using the given source and the local OS file system.
func NewReferenceLoader(source string) *jsonReferenceLoader {
	return &jsonReferenceLoader{
		fs:     osFS,
		source: source,
	}
}

// NewReferenceLoaderFileSystem returns a JSON reference loader using the given source and file system.
func NewReferenceLoaderFileSystem(source string, fs http.FileSystem) *jsonReferenceLoader {
	return &jsonReferenceLoader{
		fs:     fs,
		source: source,
	}
}

func (l *jsonReferenceLoader) LoadJSON() (interface{}, error) {

	var err error

	reference, err := gojsonreference.NewJsonReference(l.JsonSource().(string))
	if err != nil {
		return nil, err
	}

	refToUrl := reference
	refToUrl.GetUrl().Fragment = ""

	var document interface{}

	if reference.HasFileScheme {

		filename := strings.Replace(refToUrl.GetUrl().Path, "file://", "", -1)
		if runtime.GOOS == "windows" {
			// on Windows, a file URL may have an extra leading slash, use slashes
			// instead of backslashes, and have spaces escaped
			if strings.HasPrefix(filename, "/") {
				filename = filename[1:]
			}
			filename = filepath.FromSlash(filename)
		}

		document, err = l.loadFromFile(filename)
		if err != nil {
			return nil, err
		}

	} else {

		document, err = l.loadFromHTTP(refToUrl.String())
		if err != nil {
			return nil, err
		}

	}

	return document, nil

}

func (l *jsonReferenceLoader) loadFromHTTP(address string) (interface{}, error) {

	resp, err := http.Get(address)
	if err != nil {
		return nil, err
	}

	// must return HTTP Status 200 OK
	if resp.StatusCode != http.StatusOK {
		return nil, errors.New(formatErrorDescription(Locale.httpBadStatus(), ErrorDetails{"status": resp.Status}))
	}

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

	return decodeJsonUsingNumber(bytes.NewReader(bodyBuff))

}

func (l *jsonReferenceLoader) loadFromFile(path string) (interface{}, error) {
	f, err := l.fs.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	bodyBuff, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	return decodeJsonUsingNumber(bytes.NewReader(bodyBuff))

}

// JSON string loader

type jsonStringLoader struct {
	source string
}

func (l *jsonStringLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonStringLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference("#")
}

func (l *jsonStringLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

func NewStringLoader(source string) *jsonStringLoader {
	return &jsonStringLoader{source: source}
}

func (l *jsonStringLoader) LoadJSON() (interface{}, error) {

	return decodeJsonUsingNumber(strings.NewReader(l.JsonSource().(string)))

}

// JSON Go (types) loader
// used to load JSONs from the code as maps, interface{}, structs ...

type jsonGoLoader struct {
	source interface{}
}

func (l *jsonGoLoader) JsonSource() interface{} {
	return l.source
}

func (l *jsonGoLoader) JsonReference() (gojsonreference.JsonReference, error) {
	return gojsonreference.NewJsonReference("#")
}

func (l *jsonGoLoader) LoaderFactory() JSONLoaderFactory {
	return &DefaultJSONLoaderFactory{}
}

func NewGoLoader(source interface{}) *jsonGoLoader {
	return &jsonGoLoader{source: source}
}

func (l *jsonGoLoader) LoadJSON() (interface{}, error) {

	// convert it to a compliant JSON first to avoid types "mismatches"

	jsonBytes, err := json.Marshal(l.JsonSource())
	if err != nil {
		return nil, err
	}

	return decodeJsonUsingNumber(bytes.NewReader(jsonBytes))

}

func decodeJsonUsingNumber(r io.Reader) (interface{}, error) {

	var document interface{}

	decoder := json.NewDecoder(r)
	decoder.UseNumber()

	err := decoder.Decode(&document)
	if err != nil {
		return nil, err
	}

	return document, nil

}
<|endoftext|>


<commit_before>package jsonwriter

import (
	"encoding/json"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"time"
	"unicode/utf8"
)

var (
	quote        = []byte(`"`)
	keyStart     = quote
	null         = []byte("null")
	_true        = []byte("true")
	_false       = []byte("false")
	comma        = []byte(",")
	keyEnd       = []byte(`":`)
	startObject  = []byte("{")
	endObject    = []byte("}")
	startArray   = []byte("[")
	endArray     = []byte("]")
	escapedQuote = []byte(`\"`)
	escapedSlash = []byte(`\\`)
	escapedBS    = []byte(`\b`)
	escapedFF    = []byte(`\f`)
	escapedNL    = []byte(`\n`)
	escapedLF    = []byte(`\r`)
	escapedTab   = []byte(`\t`)
)

type Writer struct {
	depth int
	first bool
	array bool
	W     io.Writer
}

// Creates a JsonWriter that writes to the provided io.Writer
func New(w io.Writer) *Writer {
	return &Writer{
		W:     w,
		first: true,
	}
}

// Starts the writing process by creating an object.
// Should only be called once
func (w *Writer) RootObject(f func()) {
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
}

// Starts the writing process by creating an array.
// Should only be called once
func (w *Writer) RootArray(f func()) {
	w.array = true
	w.W.Write(startArray)
	f()
	w.W.Write(endArray)
}

// Star an object with the given key
func (w *Writer) Object(key string, f func()) {
	w.Key(key)
	w.first = true
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
}

// Star an array with the given key
func (w *Writer) Array(key string, f func()) {
	w.Key(key)
	w.first, w.array = true, true
	w.W.Write(startArray)
	f()
	w.array = false
	w.W.Write(endArray)
}

// Star an object within an array (a keyless object)
func (w *Writer) ArrayObject(f func()) {
	w.Separator()
	w.first = true
	w.array = false
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
	w.array = true
}

// Writes an array with the given array of values
// (values must be an array or a slice)
func (w *Writer) ArrayValues(key string, values interface{}) {
	w.Key(key)
	w.first, w.array = true, true
	w.W.Write(startArray)

	vo := reflect.ValueOf(values)
	kind := vo.Kind()
	if kind == reflect.Array || kind == reflect.Slice {
		for i, l := 0, vo.Len(); i < l; i++ {
			w.Value(vo.Index(i).Interface())
		}
	}
	w.array = false
	w.W.Write(endArray)
}

// Writes a key. The key is placed within quotes and ends
// with a colon
func (w *Writer) Key(key string) {
	w.Separator()
	w.W.Write(keyStart)
	w.writeString(key)
	w.W.Write(keyEnd)
}

// value can be a string, byte, u?int(8|16|32|64)?, float(32|64)?,
// time.Time, bool, nil or encoding/json.Marshaller
func (w *Writer) Value(value interface{}) {
	if w.array {
		w.Separator()
	}

	if value == nil {
		w.W.Write(null)
		return
	}

	switch t := value.(type) {
	case bool:
		if t == true {
			w.W.Write(_true)
		} else {
			w.W.Write(_false)
		}
	case uint8:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint16:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint32:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint64:
		w.W.Write([]byte(strconv.FormatUint(t, 10)))
	case int8:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int16:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int32:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int64:
		w.W.Write([]byte(strconv.FormatInt(t, 10)))
	case float32:
		w.W.Write([]byte(strconv.FormatFloat(float64(t), 'g', -1, 32)))
	case float64:
		w.W.Write([]byte(strconv.FormatFloat(t, 'g', -1, 64)))
	case json.Marshaler:
		b, _ := t.MarshalJSON()
		w.W.Write(b)
	case string:
		w.W.Write(quote)
		w.writeString(t)
		w.W.Write(quote)
	case time.Time:
		w.W.Write(quote)
		w.writeString(t.Format(time.RFC3339Nano))
		w.W.Write(quote)
	default:
		panic(fmt.Sprintf("unsuported valued type %v", value))
	}
}

// Writes the value as-is, leaving delimiters and encoding up to the caller
func (w *Writer) Raw(data []byte) {
	w.W.Write(data)
}

// writes a key: value
// This is the same as calling WriteKey(key) followe by WriteValue(value)
func (w *Writer) KeyValue(key string, value interface{}) {
	w.Key(key)
	w.Value(value)
}

// writes a key: value
// write a key'd raw value
func (w *Writer) KeyRaw(key string, data []byte) {
	w.Key(key)
	w.W.Write(data)
}

// writes a key: value where value is a string
// This is an optimized Write method
func (w *Writer) KeyString(key string, value string) {
	w.Key(key)
	w.W.Write(quote)
	w.writeString(value)
	w.W.Write(quote)
}

// writes a key: value where value is a int
// This is an optimized Write method
func (w *Writer) KeyInt(key string, value int) {
	w.Key(key)
	w.W.Write([]byte(strconv.Itoa(value)))
}

// writes a key: value where value is a bool
// This is an optimized Write method
func (w *Writer) KeyBool(key string, value bool) {
	w.Key(key)
	if value {
		w.W.Write(_true)
	} else {
		w.W.Write(_false)
	}
}

func (w *Writer) Separator() {
	if w.first == false {
		w.W.Write(comma)
	} else {
		w.first = false
	}
}

func (w *Writer) writeString(s string) {
	start, end := 0, 0
	var special []byte
L:
	for i, r := range s {
		switch r {
		case '"':
			special = escapedQuote
		case '\\':
			special = escapedSlash
		case '\b':
			special = escapedBS
		case '\f':
			special = escapedFF
		case '\n':
			special = escapedNL
		case '\r':
			special = escapedLF
		case '\t':
			special = escapedTab
		default:
			end += utf8.RuneLen(r)
			continue L
		}

		if end > start {
			w.W.Write([]byte(s[start:end]))
		}
		w.W.Write(special)
		start = i + 1
		end = start
	}
	if end > start {
		w.W.Write([]byte(s[start:end]))
	}
}
<commit_msg>ability to reset the writer to a starting-state, useful if you're creating a re-usable pool of these<commit_after>package jsonwriter

import (
	"encoding/json"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"time"
	"unicode/utf8"
)

var (
	quote        = []byte(`"`)
	keyStart     = quote
	null         = []byte("null")
	_true        = []byte("true")
	_false       = []byte("false")
	comma        = []byte(",")
	keyEnd       = []byte(`":`)
	startObject  = []byte("{")
	endObject    = []byte("}")
	startArray   = []byte("[")
	endArray     = []byte("]")
	escapedQuote = []byte(`\"`)
	escapedSlash = []byte(`\\`)
	escapedBS    = []byte(`\b`)
	escapedFF    = []byte(`\f`)
	escapedNL    = []byte(`\n`)
	escapedLF    = []byte(`\r`)
	escapedTab   = []byte(`\t`)
)

type Writer struct {
	depth int
	first bool
	array bool
	W     io.Writer
}

// Creates a JsonWriter that writes to the provided io.Writer
func New(w io.Writer) *Writer {
	return &Writer{
		W:     w,
		first: true,
	}
}

// Starts the writing process by creating an object.
// Should only be called once
func (w *Writer) RootObject(f func()) {
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
}

// Starts the writing process by creating an array.
// Should only be called once
func (w *Writer) RootArray(f func()) {
	w.array = true
	w.W.Write(startArray)
	f()
	w.W.Write(endArray)
}

// Star an object with the given key
func (w *Writer) Object(key string, f func()) {
	w.Key(key)
	w.first = true
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
}

// Star an array with the given key
func (w *Writer) Array(key string, f func()) {
	w.Key(key)
	w.first, w.array = true, true
	w.W.Write(startArray)
	f()
	w.array = false
	w.W.Write(endArray)
}

// Star an object within an array (a keyless object)
func (w *Writer) ArrayObject(f func()) {
	w.Separator()
	w.first = true
	w.array = false
	w.W.Write(startObject)
	f()
	w.W.Write(endObject)
	w.array = true
}

// Writes an array with the given array of values
// (values must be an array or a slice)
func (w *Writer) ArrayValues(key string, values interface{}) {
	w.Key(key)
	w.first, w.array = true, true
	w.W.Write(startArray)

	vo := reflect.ValueOf(values)
	kind := vo.Kind()
	if kind == reflect.Array || kind == reflect.Slice {
		for i, l := 0, vo.Len(); i < l; i++ {
			w.Value(vo.Index(i).Interface())
		}
	}
	w.array = false
	w.W.Write(endArray)
}

// Writes a key. The key is placed within quotes and ends
// with a colon
func (w *Writer) Key(key string) {
	w.Separator()
	w.W.Write(keyStart)
	w.writeString(key)
	w.W.Write(keyEnd)
}

// value can be a string, byte, u?int(8|16|32|64)?, float(32|64)?,
// time.Time, bool, nil or encoding/json.Marshaller
func (w *Writer) Value(value interface{}) {
	if w.array {
		w.Separator()
	}

	if value == nil {
		w.W.Write(null)
		return
	}

	switch t := value.(type) {
	case bool:
		if t == true {
			w.W.Write(_true)
		} else {
			w.W.Write(_false)
		}
	case uint8:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint16:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint32:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint:
		w.W.Write([]byte(strconv.FormatUint(uint64(t), 10)))
	case uint64:
		w.W.Write([]byte(strconv.FormatUint(t, 10)))
	case int8:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int16:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int32:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int:
		w.W.Write([]byte(strconv.FormatInt(int64(t), 10)))
	case int64:
		w.W.Write([]byte(strconv.FormatInt(t, 10)))
	case float32:
		w.W.Write([]byte(strconv.FormatFloat(float64(t), 'g', -1, 32)))
	case float64:
		w.W.Write([]byte(strconv.FormatFloat(t, 'g', -1, 64)))
	case json.Marshaler:
		b, _ := t.MarshalJSON()
		w.W.Write(b)
	case string:
		w.W.Write(quote)
		w.writeString(t)
		w.W.Write(quote)
	case time.Time:
		w.W.Write(quote)
		w.writeString(t.Format(time.RFC3339Nano))
		w.W.Write(quote)
	default:
		panic(fmt.Sprintf("unsuported valued type %v", value))
	}
}

// Writes the value as-is, leaving delimiters and encoding up to the caller
func (w *Writer) Raw(data []byte) {
	w.W.Write(data)
}

// writes a key: value
// This is the same as calling WriteKey(key) followe by WriteValue(value)
func (w *Writer) KeyValue(key string, value interface{}) {
	w.Key(key)
	w.Value(value)
}

// writes a key: value
// write a key'd raw value
func (w *Writer) KeyRaw(key string, data []byte) {
	w.Key(key)
	w.W.Write(data)
}

// writes a key: value where value is a string
// This is an optimized Write method
func (w *Writer) KeyString(key string, value string) {
	w.Key(key)
	w.W.Write(quote)
	w.writeString(value)
	w.W.Write(quote)
}

// writes a key: value where value is a int
// This is an optimized Write method
func (w *Writer) KeyInt(key string, value int) {
	w.Key(key)
	w.W.Write([]byte(strconv.Itoa(value)))
}

// writes a key: value where value is a bool
// This is an optimized Write method
func (w *Writer) KeyBool(key string, value bool) {
	w.Key(key)
	if value {
		w.W.Write(_true)
	} else {
		w.W.Write(_false)
	}
}

func (w *Writer) Separator() {
	if w.first == false {
		w.W.Write(comma)
	} else {
		w.first = false
	}
}

func (w *Writer) writeString(s string) {
	start, end := 0, 0
	var special []byte
L:
	for i, r := range s {
		switch r {
		case '"':
			special = escapedQuote
		case '\\':
			special = escapedSlash
		case '\b':
			special = escapedBS
		case '\f':
			special = escapedFF
		case '\n':
			special = escapedNL
		case '\r':
			special = escapedLF
		case '\t':
			special = escapedTab
		default:
			end += utf8.RuneLen(r)
			continue L
		}

		if end > start {
			w.W.Write([]byte(s[start:end]))
		}
		w.W.Write(special)
		start = i + 1
		end = start
	}
	if end > start {
		w.W.Write([]byte(s[start:end]))
	}
}

func (w *Writer) Reset() {
	w.depth = 0
	w.first = true
	w.array = false
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/kataras/iris"
)

func main() {
	app := iris.New()

	// this works as expected now,
	// will handle *all* expect DELETE requests, even if there is no routes
	app.Get("/action/{p}", h)

	app.Run(iris.Addr(":8080"), iris.WithoutServerError(iris.ErrServerClosed))
}

func h(ctx iris.Context) {
	ctx.Writef("[%s] %s : Parameter = `%s`", ctx.Method(), ctx.Path(), ctx.Params().Get("p"))
}

func fallbackHandler(ctx iris.Context) {
	if ctx.Method() == "DELETE" {
		ctx.Next()

		return
	}

	ctx.Writef("[%s] %s : From fallback handler", ctx.Method(), ctx.Path())
}
<commit_msg>Add a missing call for `fallbackHandler`<commit_after>package main

import (
	"github.com/kataras/iris"
)

func main() {
	app := iris.New()

	// add a fallback handler to process requests that would not be declared in the router.
	app.Fallback(fallbackHandler)

	// this works as expected now,
	// will handle *all* expect DELETE requests, even if there is no routes
	app.Get("/action/{p}", h)

	app.Run(iris.Addr(":8080"), iris.WithoutServerError(iris.ErrServerClosed))
}

func h(ctx iris.Context) {
	ctx.Writef("[%s] %s : Parameter = `%s`", ctx.Method(), ctx.Path(), ctx.Params().Get("p"))
}

func fallbackHandler(ctx iris.Context) {
	if ctx.Method() == "DELETE" {
		ctx.Next()

		return
	}

	ctx.Writef("[%s] %s : From fallback handler", ctx.Method(), ctx.Path())
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"compress/gzip"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	runewidth "github.com/mattn/go-runewidth"
)

const delimiter string = "::::"

func fail(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func failOn(b bool, message string) {
	if b {
		fail(errors.New(message))
	}
}

type Stack []string

func (s *Stack) Push(dirname string) {
	*s = append(*s, dirname)
}

func (s *Stack) DiscardTop() {
	stack := *s
	failOn(len(stack) <= 0, "Invalid directory state. Corrupted database?")
	i := len(stack) - 1
	*s = stack[:i]
}

func keyval(line string) (string, string) {
	i := strings.Index(line, ":")
	if i == -1 || i == len(line)-1 {
		return line, ""
	}
	return line[:i], line[i+2:]
}

type Track struct {
	Album    string
	Artist   string
	Date     string
	Filename string
	Genre    string
	Path     string
	Time     string
	Title    string
}

func (t *Track) Set(key, value string) {
	switch key {
	case "Album":
		t.Album = value
	case "Artist":
		t.Artist = value
	case "Date":
		t.Date = value
	case "Genre":
		t.Genre = value
	case "Time":
		t.Time = value
	case "Title":
		t.Title = value
	}
}

func formatDurationString(str string) string {
	duration, err := time.ParseDuration(str + "s")
	if err != nil {
		return "(-:--)"
	}
	t := time.Time{}.Add(duration)
	if duration < time.Hour {
		return t.Format("(4:05)")
	}
	return fmt.Sprintf("(%d:%s)", int(duration.Hours()), t.Format("04:05"))
}

func withoutExt(path string) string {
	basename := filepath.Base(path)
	return strings.TrimSuffix(basename, filepath.Ext(basename))
}

func alignLeftRight(maxlen int, description, duration string) string {
	stop := maxlen - len(duration)
	s := runewidth.Truncate(description, stop, "... ")
	return runewidth.FillRight(s, stop) + duration
}

func termWidth() int {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	fail(err)
	var height, width int
	_, err = fmt.Sscanf(string(out), "%d %d\n", &height, &width)
	fail(err)
	return width
}

func trackFormatter() func(*Track) string {
	// Remove 5 from screen width for correct fzf display at right edge.
	width := termWidth() - 5
	return func(t *Track) string {
		info := t.Title
		if info == "" {
			info = withoutExt(t.Filename)
		}
		if t.Artist != "" {
			info = t.Artist + " - " + info
		}
		if t.Album != "" {
			info += " {" + t.Album + "}"
		}
		info = strings.Replace(info, delimiter, "", -1)
		duration := formatDurationString(t.Time)
		// Right align duration
		info = alignLeftRight(width, info, duration)
		return info + delimiter + t.Path
	}
}

func groupByArtist(tracks []*Track) []*Track {
	// group by artist, then shuffle to stop same order, but keep artist together
	artists := map[string][]*Track{}
	for _, t := range tracks {
		artists[t.Artist] = append(artists[t.Artist], t)
	}
	shuffled := make([]*Track, len(tracks))
	i := 0
	for _, tracks := range artists {
		for _, t := range tracks {
			shuffled[i] = t
			i += 1
		}
	}
	return shuffled
}

func parse(scan *bufio.Scanner) []*Track {

	tracks, track := []*Track{}, new(Track)
	dirstack := make(Stack, 0, 64)

	for scan.Scan() {
		key, value := keyval(scan.Text())
		switch key {
		case "directory":
			dirstack.Push(value)
		case "end":
			dirstack.DiscardTop()
		case "Artist", "Album", "Date", "Genre", "Time", "Title":
			track.Set(key, value)
		case "song_begin":
			track.Filename = value
			track.Path = filepath.Join(append(dirstack, track.Filename)...)
		case "song_end":
			tracks = append(tracks, track)
			track = new(Track)
		}
	}
	fail(scan.Err())
	return tracks
}

func expandUser(path, home string) string {
	if path[:2] == "~/" {
		path = strings.Replace(path, "~", home, 1)
	}
	return path
}

func findDbFile() string {
	usr, err := user.Current()
	fail(err)
	home := usr.HomeDir
	paths := []string{
		filepath.Join(os.Getenv("XDG_CONFIG_HOME"), "/mpd/mpd.conf"),
		filepath.Join(home, ".config", "/mpd/mpd.conf"),
		filepath.Join(home, ".mpdconf"),
		"/etc/mpd.conf",
	}
	var file *os.File
	var confpath string
	for _, path := range paths {
		file, err = os.Open(path)
		if err == nil {
			confpath = path
			break
		}
	}
	failOn(file == nil, "No config file found")

	expDb := regexp.MustCompile(`^\s*db_file\s*"([^"]+)"`)
	scan := bufio.NewScanner(file)
	var dbFile string
	for scan.Scan() {
		m := expDb.FindStringSubmatch(scan.Text())
		if m != nil {
			dbFile = expandUser(m[1], home)
			break
		}
	}
	fail(scan.Err())
	fail(file.Close())
	failOn(dbFile == "", fmt.Sprintf("Could not find 'db_file' in configuration file '%s'", confpath))
	return dbFile
}

func failNotify(message string) {
	err := exec.Command("tmux", "display", message).Run()
	if err != nil {
		exec.Command("notify-send", message).Run()
	}
	fail(errors.New(message))
}

func mpcRun(args ...string) string {
	out, err := exec.Command("mpc", args...).CombinedOutput()
	if err != nil {
		failNotify(string(out))
	}
	return string(out)
}

func mpcPlay(path string) {
	pos, found := mpcFindOnPlaylist(path)
	if !found {
		mpcRun("add", path)
	}
	mpcRun("play", strconv.Itoa(pos))
}

func mpcFindOnPlaylist(path string) (int, bool) {
	playlist := mpcRun("playlist", "-f", "%file%")
	lines := strings.Split(playlist, "\n")
	for i, line := range lines {
		if line == path {
			return i + 1, true
		}
	}
	return len(lines), false
}

func cmdPlay(fzfline string) {
	fields := strings.SplitN(fzfline, delimiter, 2)
	if len(fields) != 2 {
		failNotify("mpd-fzf: assert split failure")
	}
	path := fields[1]
	mpcPlay(path)
}

func fzfcmd() *exec.Cmd {
	bind := "--bind=enter:execute-silent(mpd-fzf _play {})"
	fzf := exec.Command("fzf", "--no-hscroll", bind)
	fzf.Stderr = os.Stderr
	return fzf
}

func ignoreExitInterrupt(err error) error {
	if strings.HasSuffix(err.Error(), "130") {
		return nil
	}
	return err
}

func main() {
	if len(os.Args) > 2 && os.Args[1] == "_play" {
		cmdPlay(os.Args[2])
		return
	}
	dbFile := findDbFile()
	format := trackFormatter()

	file, err := os.Open(dbFile)
	fail(err)
	gz, err := gzip.NewReader(file)
	fail(err)

	scan := bufio.NewScanner(gz)
	tracks := groupByArtist(parse(scan))

	fail(gz.Close())
	fail(file.Close())

	fzf := fzfcmd()
	in, err := fzf.StdinPipe()
	fail(err)
	fail(fzf.Start())
	for _, t := range tracks {
		fmt.Fprintln(in, format(t))
	}
	fail(in.Close())
	fail(ignoreExitInterrupt(fzf.Wait()))
}
<commit_msg>Handle a couple of indexing edge cases<commit_after>package main

import (
	"bufio"
	"compress/gzip"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	runewidth "github.com/mattn/go-runewidth"
)

const delimiter string = "::::"

func fail(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func failOn(b bool, message string) {
	if b {
		fail(errors.New(message))
	}
}

type Stack []string

func (s *Stack) Push(dirname string) {
	*s = append(*s, dirname)
}

func (s *Stack) DiscardTop() {
	stack := *s
	failOn(len(stack) <= 0, "Invalid directory state. Corrupted database?")
	i := len(stack) - 1
	*s = stack[:i]
}

func keyValue(line string) (string, string) {
	i := strings.Index(line, ":")
	if i < 0 {
		return line, ""
	}
	if i == len(line)-1 {
		return line[:i], ""
	}
	// The value is always preceded by a space, so +2
	n := 2
	// But check anyway...
	if line[i+1] != ' ' {
		n = 1
	}
	return line[:i], line[i+n:]
}

type Track struct {
	Album    string
	Artist   string
	Date     string
	Filename string
	Genre    string
	Path     string
	Time     string
	Title    string
}

func (t *Track) Set(key, value string) {
	switch key {
	case "Album":
		t.Album = value
	case "Artist":
		t.Artist = value
	case "Date":
		t.Date = value
	case "Genre":
		t.Genre = value
	case "Time":
		t.Time = value
	case "Title":
		t.Title = value
	}
}

func formatDurationString(str string) string {
	duration, err := time.ParseDuration(str + "s")
	if err != nil {
		return "(-:--)"
	}
	t := time.Time{}.Add(duration)
	if duration < time.Hour {
		return t.Format("(4:05)")
	}
	return fmt.Sprintf("(%d:%s)", int(duration.Hours()), t.Format("04:05"))
}

func withoutExt(path string) string {
	basename := filepath.Base(path)
	return strings.TrimSuffix(basename, filepath.Ext(basename))
}

func alignLeftRight(maxlen int, description, duration string) string {
	stop := maxlen - len(duration)
	s := runewidth.Truncate(description, stop, "... ")
	return runewidth.FillRight(s, stop) + duration
}

func termWidth() int {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	fail(err)
	var height, width int
	_, err = fmt.Sscanf(string(out), "%d %d\n", &height, &width)
	fail(err)
	return width
}

func trackFormatter() func(*Track) string {
	// Remove 5 from screen width for correct fzf display at right edge.
	width := termWidth() - 5
	return func(t *Track) string {
		info := t.Title
		if info == "" {
			info = withoutExt(t.Filename)
		}
		if t.Artist != "" {
			info = t.Artist + " - " + info
		}
		if t.Album != "" {
			info += " {" + t.Album + "}"
		}
		info = strings.Replace(info, delimiter, "", -1)
		duration := formatDurationString(t.Time)
		// Right align duration
		info = alignLeftRight(width, info, duration)
		return info + delimiter + t.Path
	}
}

func groupByArtist(tracks []*Track) []*Track {
	// group by artist, then shuffle to stop same order, but keep artist together
	artists := map[string][]*Track{}
	for _, t := range tracks {
		artists[t.Artist] = append(artists[t.Artist], t)
	}
	shuffled := make([]*Track, len(tracks))
	i := 0
	for _, tracks := range artists {
		for _, t := range tracks {
			shuffled[i] = t
			i += 1
		}
	}
	return shuffled
}

func parse(scan *bufio.Scanner) []*Track {

	tracks, track := []*Track{}, new(Track)
	dirstack := make(Stack, 0, 64)

	for scan.Scan() {
		key, value := keyValue(scan.Text())
		switch key {
		case "directory":
			dirstack.Push(value)
		case "end":
			dirstack.DiscardTop()
		case "Artist", "Album", "Date", "Genre", "Time", "Title":
			track.Set(key, value)
		case "song_begin":
			track.Filename = value
			track.Path = filepath.Join(append(dirstack, track.Filename)...)
		case "song_end":
			tracks = append(tracks, track)
			track = new(Track)
		}
	}
	fail(scan.Err())
	return tracks
}

func expandUser(path, home string) string {
	if strings.HasPrefix(path, "~/") {
		path = strings.Replace(path, "~", home, 1)
	}
	return path
}

func findDbFile() string {
	usr, err := user.Current()
	fail(err)
	home := usr.HomeDir
	paths := []string{
		filepath.Join(os.Getenv("XDG_CONFIG_HOME"), "/mpd/mpd.conf"),
		filepath.Join(home, ".config", "/mpd/mpd.conf"),
		filepath.Join(home, ".mpdconf"),
		"/etc/mpd.conf",
	}
	var file *os.File
	var confpath string
	for _, path := range paths {
		file, err = os.Open(path)
		if err == nil {
			confpath = path
			break
		}
	}
	failOn(file == nil, "No config file found")

	expDb := regexp.MustCompile(`^\s*db_file\s*"([^"]+)"`)
	scan := bufio.NewScanner(file)
	var dbFile string
	for scan.Scan() {
		m := expDb.FindStringSubmatch(scan.Text())
		if m != nil {
			dbFile = expandUser(m[1], home)
			break
		}
	}
	fail(scan.Err())
	fail(file.Close())
	failOn(dbFile == "", fmt.Sprintf("Could not find 'db_file' in configuration file '%s'", confpath))
	return dbFile
}

func failNotify(message string) {
	err := exec.Command("tmux", "display", message).Run()
	if err != nil {
		exec.Command("notify-send", message).Run()
	}
	fail(errors.New(message))
}

func mpcRun(args ...string) string {
	out, err := exec.Command("mpc", args...).CombinedOutput()
	if err != nil {
		failNotify(string(out))
	}
	return string(out)
}

func mpcPlay(path string) {
	pos, found := mpcFindOnPlaylist(path)
	if !found {
		mpcRun("add", path)
	}
	mpcRun("play", strconv.Itoa(pos))
}

func mpcFindOnPlaylist(path string) (int, bool) {
	playlist := mpcRun("playlist", "-f", "%file%")
	lines := strings.Split(playlist, "\n")
	for i, line := range lines {
		if line == path {
			return i + 1, true
		}
	}
	return len(lines), false
}

func cmdPlay(fzfline string) {
	fields := strings.SplitN(fzfline, delimiter, 2)
	if len(fields) != 2 {
		failNotify("mpd-fzf: assert split failure")
	}
	path := fields[1]
	mpcPlay(path)
}

func fzfcmd() *exec.Cmd {
	bind := "--bind=enter:execute-silent(mpd-fzf _play {})"
	fzf := exec.Command("fzf", "--no-hscroll", bind)
	fzf.Stderr = os.Stderr
	return fzf
}

func ignoreExitInterrupt(err error) error {
	if strings.HasSuffix(err.Error(), "130") {
		return nil
	}
	return err
}

func main() {
	if len(os.Args) > 2 && os.Args[1] == "_play" {
		cmdPlay(os.Args[2])
		return
	}
	dbFile := findDbFile()
	format := trackFormatter()

	file, err := os.Open(dbFile)
	fail(err)
	gz, err := gzip.NewReader(file)
	fail(err)

	scan := bufio.NewScanner(gz)
	tracks := groupByArtist(parse(scan))

	fail(gz.Close())
	fail(file.Close())

	fzf := fzfcmd()
	in, err := fzf.StdinPipe()
	fail(err)
	fail(fzf.Start())
	for _, t := range tracks {
		fmt.Fprintln(in, format(t))
	}
	fail(in.Close())
	fail(ignoreExitInterrupt(fzf.Wait()))
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Joseph D Poirier
// Distributable under the terms of The New BSD License
// that can be found in the LICENSE file.

// Package mpu wraps libmpu, an AHRS.
//
// Build example
//
// libmpu.so:
//   $ make -f Makefile-native-shared
//
// mpu go wrapper:
//   $ go build -o mpu.a mpu.go
//

package mpu

/*
#cgo linux LDFLAGS: -limu -lm
#cgo darwin LDFLAGS: -limu -lm
#cgo windows CFLAGS: -IC:/WINDOWS/system32
#cgo windows LDFLAGS: -L. -limu -LC:/WINDOWS/system32

#include <stdlib.h>
#include <stdint.h>
#include "../libimu.h"
*/
import "C"
import (
	"errors"
	"unsafe"
)

// Current version.
var PackageVersion = "v0.2b"

// InitMPU
func InitMPU(sample_rate, yaw_mix_factor int) int {
	return int(C.init_mpu(C.int(sample_rate), C.int(yaw_mix_factor)))
}

func EnableFusion() int {
	return int(C.enableFusion())
}

func DisableFusion() int {
	return int(C.disableFusion())
}

// CloseMPU
func CloseMPU() {
	C.close_mpu()
}

// ReadMPU
func ReadMPU() (pitch, roll, heading float32, err error) {
	i := int(C.read_mpu((*C.float)(unsafe.Pointer(&pitch)),
		(*C.float)(unsafe.Pointer(&roll)),
		(*C.float)(unsafe.Pointer(&heading))))
	if i == -1 {
		err = errors.New("error reading MPU")
	}
	return
}

// ReadMPURaw
func ReadMPURaw() (gx, gy, gz, ax, ay, az, mx, my, mz float32, err error) {
	i := int(C.read_mpu_raw(
		(*C.float)(unsafe.Pointer(&gx)),
		(*C.float)(unsafe.Pointer(&gy)),
		(*C.float)(unsafe.Pointer(&gz)),
		(*C.float)(unsafe.Pointer(&ax)),
		(*C.float)(unsafe.Pointer(&ay)),
		(*C.float)(unsafe.Pointer(&az)),
		(*C.float)(unsafe.Pointer(&mx)),
		(*C.float)(unsafe.Pointer(&my)),
		(*C.float)(unsafe.Pointer(&mz))))
	if i == -1 {
		err = errors.New("error reading MPU")
	}
	return
}
<commit_msg>Added RawData struct and converted client return value<commit_after>// Copyright (c) 2015 Joseph D Poirier
// Distributable under the terms of The New BSD License
// that can be found in the LICENSE file.

// Package mpu wraps libmpu, an AHRS.
//
// Build example
//
// libmpu.so:
//   $ make -f Makefile-native-shared
//
// mpu go wrapper:
//   $ go build -o mpu.a mpu.go
//

package mpu

/*
#cgo linux LDFLAGS: -limu -lm
#cgo darwin LDFLAGS: -limu -lm
#cgo windows CFLAGS: -IC:/WINDOWS/system32
#cgo windows LDFLAGS: -L. -limu -LC:/WINDOWS/system32

#include <stdlib.h>
#include <stdint.h>
#include "../libimu.h"
*/
import "C"
import (
	"errors"
	"unsafe"
)

// Raw data struct
type RawData struct {
	gx, gy, gz, ax, ay, az, mx, my, mz float32
}

// Current version.
var PackageVersion = "v0.2b"

// InitMPU
func InitMPU(sample_rate, yaw_mix_factor int) int {
	return int(C.init_mpu(C.int(sample_rate), C.int(yaw_mix_factor)))
}

func EnableFusion() int {
	return int(C.enableFusion())
}

func DisableFusion() int {
	return int(C.disableFusion())
}

// CloseMPU
func CloseMPU() {
	C.close_mpu()
}

// ReadMPU
func ReadMPU() (pitch, roll, heading float32, err error) {
	i := int(C.read_mpu((*C.float)(unsafe.Pointer(&pitch)),
		(*C.float)(unsafe.Pointer(&roll)),
		(*C.float)(unsafe.Pointer(&heading))))
	if i == -1 {
		err = errors.New("error reading MPU")
	}
	return
}

// ReadMPURaw
func ReadMPURaw() (d RawData, err error) {
	i := int(C.read_mpu_raw(
		(*C.float)(unsafe.Pointer(&d.gx)),
		(*C.float)(unsafe.Pointer(&d.gy)),
		(*C.float)(unsafe.Pointer(&d.gz)),
		(*C.float)(unsafe.Pointer(&d.ax)),
		(*C.float)(unsafe.Pointer(&d.ay)),
		(*C.float)(unsafe.Pointer(&d.az)),
		(*C.float)(unsafe.Pointer(&d.mx)),
		(*C.float)(unsafe.Pointer(&d.my)),
		(*C.float)(unsafe.Pointer(&d.mz))))
	if i == -1 {
		err = errors.New("error reading MPU")
	}
	return
}
<|endoftext|>


<commit_before>// This file is part of Monsti, a web content management system.
// Copyright 2012-2013 Christian Neumann
//
// Monsti is free software: you can redistribute it and/or modify it under the
// terms of the GNU Affero General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option) any
// later version.
//
// Monsti is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
// A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more
// details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Monsti.  If not, see <http://www.gnu.org/licenses/>.

/*
 Monsti is a simple and resource efficient CMS.

 This package implements the data service.
*/
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"

	"pkg.monsti.org/monsti/api/service"
	"pkg.monsti.org/monsti/api/util"
)

// DataService implements RPC methods for the Data service.
type DataService struct {
	Settings settings
}

// getNode looks up the given node.
// If no such node exists, return nil.
// It adds a path attribute with the given path.
func getNode(root, path string) (node []byte, err error) {
	node_path := filepath.Join(root, path[1:], "node.json")
	node, err = ioutil.ReadFile(node_path)
	if err != nil {
		return
	}
	pathJSON := fmt.Sprintf(`{"path":%q,`, path)
	node = bytes.Replace(node, []byte("{"), []byte(pathJSON), 1)
	return
}

// getChildren looks up child nodes of the given node.
func getChildren(root, path string) (nodes [][]byte, err error) {
	files, err := ioutil.ReadDir(filepath.Join(root, path))
	if err != nil {
		return
	}
	for _, file := range files {
		node, _ := getNode(root, filepath.Join(path, file.Name()))
		if node != nil {
			nodes = append(nodes, node)
		}
	}
	return
}

type GetChildrenArgs struct {
	Site, Path string
}

func (i *DataService) GetChildren(args GetChildrenArgs,
	reply *[][]byte) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	ret, err := getChildren(site, args.Path)
	*reply = ret
	return err
}

type GetNodeDataArgs struct{ Site, Path, File string }

func (i *DataService) GetNodeData(args *GetNodeDataArgs,
	reply *[]byte) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	path := filepath.Join(site, args.Path[1:], args.File)
	ret, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		*reply = nil
		return nil
	}
	*reply = ret
	return err
}

type WriteNodeDataArgs struct {
	Site, Path, File string
	Content          []byte
}

func (i *DataService) WriteNodeData(args *WriteNodeDataArgs,
	reply *int) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	path := filepath.Join(site, args.Path[1:], args.File)
	err := os.MkdirAll(filepath.Dir(path), 0700)
	if err != nil {
		return fmt.Errorf("Could not create node directory: %v", err)
	}
	err = ioutil.WriteFile(path, []byte(args.Content), 0600)
	if err != nil {
		return fmt.Errorf("Could not write node data: %v", err)
	}
	return nil
}

type RemoveNodeArgs struct {
	Site, Node string
}

func (i *DataService) RemoveNode(args *RemoveNodeArgs, reply *int) error {
	root := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	nodePath := filepath.Join(root, args.Node[1:])
	if err := os.RemoveAll(nodePath); err != nil {
		return fmt.Errorf("Can't remove node: %v", err)
	}
	return nil
}

// getConfig returns the configuration value or section for the given name.
func getConfig(path, name string) ([]byte, error) {
	content, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("Could not read configuration: %v", err)
	}
	var target interface{}
	err = json.Unmarshal(content, &target)
	if err != nil {
		return nil, fmt.Errorf("Could not parse configuration: %v", err)
	}
	subs := strings.Split(name, ".")
	for _, sub := range subs {
		if sub == "" {
			break
		}
		targetT := reflect.TypeOf(target)
		if targetT != reflect.TypeOf(map[string]interface{}{}) {
			target = nil
			break
		}
		var ok bool
		if target, ok = (target.(map[string]interface{}))[sub]; !ok {
			target = nil
			break
		}
	}
	target = map[string]interface{}{"Value": target}
	ret, err := json.Marshal(target)
	if err != nil {
		return nil, fmt.Errorf("Could not encode configuration: %v", err)
	}
	return ret, nil
}

type GetConfigArgs struct{ Site, Module, Name string }

func (i *DataService) GetConfig(args *GetConfigArgs,
	reply *[]byte) error {
	configPath := i.Settings.Monsti.GetSiteConfigPath(args.Site)
	config, err := getConfig(filepath.Join(configPath, args.Module+".json"),
		args.Name)
	if err != nil {
		reply = nil
		return err
	}
	*reply = config
	return nil
}

type settings struct {
	Monsti util.MonstiSettings
}

func main() {
	logger := log.New(os.Stderr, "", 0)

	// Load configuration
	flag.Parse()
	cfgPath := util.GetConfigPath(flag.Arg(0))
	var settings settings
	if err := util.LoadModuleSettings("data", cfgPath, &settings); err != nil {
		logger.Fatal("Could not load settings: ", err)
	}

	// Connect to Info service
	info, err := service.NewInfoConnection(settings.Monsti.GetServicePath(
		service.Info.String()))
	if err != nil {
		logger.Fatalf("Could not connect to Info service: %v", err)
	}

	// Start own Data service
	var waitGroup sync.WaitGroup
	logger.Println("Setting up Data service")
	waitGroup.Add(1)
	dataPath := settings.Monsti.GetServicePath(service.Data.String())
	go func() {
		defer waitGroup.Done()
		var data_ DataService
		data_.Settings = settings
		provider := service.NewProvider("Data", &data_)
		provider.Logger = logger
		if err := provider.Listen(dataPath); err != nil {
			logger.Fatalf("Could not start Data service: %v", err)
		}
		if err := provider.Accept(); err != nil {
			logger.Fatalf("Could not accept at Data service: %v", err)
		}
	}()

	if err := info.PublishService("Data", dataPath); err != nil {
		logger.Fatalf("Could not publish Data service: %v", err)
	}

	waitGroup.Wait()
}
<commit_msg>Fix path attribute name in node JSON.<commit_after>// This file is part of Monsti, a web content management system.
// Copyright 2012-2013 Christian Neumann
//
// Monsti is free software: you can redistribute it and/or modify it under the
// terms of the GNU Affero General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option) any
// later version.
//
// Monsti is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
// A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more
// details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Monsti.  If not, see <http://www.gnu.org/licenses/>.

/*
 Monsti is a simple and resource efficient CMS.

 This package implements the data service.
*/
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"

	"pkg.monsti.org/monsti/api/service"
	"pkg.monsti.org/monsti/api/util"
)

// DataService implements RPC methods for the Data service.
type DataService struct {
	Settings settings
}

// getNode looks up the given node.
// If no such node exists, return nil.
// It adds a path attribute with the given path.
func getNode(root, path string) (node []byte, err error) {
	node_path := filepath.Join(root, path[1:], "node.json")
	node, err = ioutil.ReadFile(node_path)
	if err != nil {
		return
	}
	pathJSON := fmt.Sprintf(`{"Path":%q,`, path)
	node = bytes.Replace(node, []byte("{"), []byte(pathJSON), 1)
	return
}

// getChildren looks up child nodes of the given node.
func getChildren(root, path string) (nodes [][]byte, err error) {
	files, err := ioutil.ReadDir(filepath.Join(root, path))
	if err != nil {
		return
	}
	for _, file := range files {
		node, _ := getNode(root, filepath.Join(path, file.Name()))
		if node != nil {
			nodes = append(nodes, node)
		}
	}
	return
}

type GetChildrenArgs struct {
	Site, Path string
}

func (i *DataService) GetChildren(args GetChildrenArgs,
	reply *[][]byte) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	ret, err := getChildren(site, args.Path)
	*reply = ret
	return err
}

type GetNodeDataArgs struct{ Site, Path, File string }

func (i *DataService) GetNodeData(args *GetNodeDataArgs,
	reply *[]byte) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	path := filepath.Join(site, args.Path[1:], args.File)
	ret, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		*reply = nil
		return nil
	}
	*reply = ret
	return err
}

type WriteNodeDataArgs struct {
	Site, Path, File string
	Content          []byte
}

func (i *DataService) WriteNodeData(args *WriteNodeDataArgs,
	reply *int) error {
	site := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	path := filepath.Join(site, args.Path[1:], args.File)
	err := os.MkdirAll(filepath.Dir(path), 0700)
	if err != nil {
		return fmt.Errorf("Could not create node directory: %v", err)
	}
	err = ioutil.WriteFile(path, []byte(args.Content), 0600)
	if err != nil {
		return fmt.Errorf("Could not write node data: %v", err)
	}
	return nil
}

type RemoveNodeArgs struct {
	Site, Node string
}

func (i *DataService) RemoveNode(args *RemoveNodeArgs, reply *int) error {
	root := i.Settings.Monsti.GetSiteNodesPath(args.Site)
	nodePath := filepath.Join(root, args.Node[1:])
	if err := os.RemoveAll(nodePath); err != nil {
		return fmt.Errorf("Can't remove node: %v", err)
	}
	return nil
}

// getConfig returns the configuration value or section for the given name.
func getConfig(path, name string) ([]byte, error) {
	content, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("Could not read configuration: %v", err)
	}
	var target interface{}
	err = json.Unmarshal(content, &target)
	if err != nil {
		return nil, fmt.Errorf("Could not parse configuration: %v", err)
	}
	subs := strings.Split(name, ".")
	for _, sub := range subs {
		if sub == "" {
			break
		}
		targetT := reflect.TypeOf(target)
		if targetT != reflect.TypeOf(map[string]interface{}{}) {
			target = nil
			break
		}
		var ok bool
		if target, ok = (target.(map[string]interface{}))[sub]; !ok {
			target = nil
			break
		}
	}
	target = map[string]interface{}{"Value": target}
	ret, err := json.Marshal(target)
	if err != nil {
		return nil, fmt.Errorf("Could not encode configuration: %v", err)
	}
	return ret, nil
}

type GetConfigArgs struct{ Site, Module, Name string }

func (i *DataService) GetConfig(args *GetConfigArgs,
	reply *[]byte) error {
	configPath := i.Settings.Monsti.GetSiteConfigPath(args.Site)
	config, err := getConfig(filepath.Join(configPath, args.Module+".json"),
		args.Name)
	if err != nil {
		reply = nil
		return err
	}
	*reply = config
	return nil
}

type settings struct {
	Monsti util.MonstiSettings
}

func main() {
	logger := log.New(os.Stderr, "", 0)

	// Load configuration
	flag.Parse()
	cfgPath := util.GetConfigPath(flag.Arg(0))
	var settings settings
	if err := util.LoadModuleSettings("data", cfgPath, &settings); err != nil {
		logger.Fatal("Could not load settings: ", err)
	}

	// Connect to Info service
	info, err := service.NewInfoConnection(settings.Monsti.GetServicePath(
		service.Info.String()))
	if err != nil {
		logger.Fatalf("Could not connect to Info service: %v", err)
	}

	// Start own Data service
	var waitGroup sync.WaitGroup
	logger.Println("Setting up Data service")
	waitGroup.Add(1)
	dataPath := settings.Monsti.GetServicePath(service.Data.String())
	go func() {
		defer waitGroup.Done()
		var data_ DataService
		data_.Settings = settings
		provider := service.NewProvider("Data", &data_)
		provider.Logger = logger
		if err := provider.Listen(dataPath); err != nil {
			logger.Fatalf("Could not start Data service: %v", err)
		}
		if err := provider.Accept(); err != nil {
			logger.Fatalf("Could not accept at Data service: %v", err)
		}
	}()

	if err := info.PublishService("Data", dataPath); err != nil {
		logger.Fatalf("Could not publish Data service: %v", err)
	}

	waitGroup.Wait()
}
<|endoftext|>


<commit_before>package kafka_avro

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	avro "github.com/elodina/go-avro"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/gliderlabs/logspout/router"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/fsouza/go-dockerclient"
	"github.com/joeshaw/iso8601"
)

var messageSchema = `{
  "type": "record",
  "name": "LogLine",
  "fields": [
    {"name": "timestamp", "type": "string"},
    {"name": "container_name", "type": "string"},
    {"name": "host", "type": "string"},
    {"name": "source", "type": "string"},
    {"name": "line", "type": "string"}
  ]
}`

func init() {
	router.AdapterFactories.Register(NewKafkaAvroAdapter, "kafka_avro")
}

type KafkaAvroAdapter struct {
	route    *router.Route
	brokers  []string
	topic    string
	schema   avro.Schema
	registry *kafkaavro.KafkaAvroEncoder
	producer sarama.AsyncProducer
}

func NewKafkaAvroAdapter(route *router.Route) (router.LogAdapter, error) {
	brokers := readBrokers(route.Address)
	if len(brokers) == 0 {
		return nil, errorf("The Kafka broker host:port is missing. Did you specify it as a route address?")
	}

	topic := readTopic(route.Address, route.Options)
	if topic == "" {
		return nil, errorf("The Kafka topic is missing. Did you specify it as a route option?")
	}

	schemaUrl := readSchemaRegistryUrl(route.Options)
	if schemaUrl == "" {
		return nil, errorf("The schema registry url is missing. Did you specify it as a route option?")
	}

	registry := kafkaavro.NewKafkaAvroEncoder(schemaUrl)

	var schema avro.Schema
	schema, err := avro.ParseSchema(messageSchema)
	if err != nil {
		return nil, errorf("The schema could not be parsed")
	}

	if os.Getenv("DEBUG") != "" {
		log.Printf("Starting Kafka producer for address: %s, topic: %s.\n", brokers, topic)
	}

	var retries int
	retries, err = strconv.Atoi(os.Getenv("KAFKA_CONNECT_RETRIES"))
	if err != nil {
		retries = 3
	}
	var producer sarama.AsyncProducer
	for i := 0; i < retries; i++ {
		producer, err = sarama.NewAsyncProducer(brokers, newConfig())
		if err != nil {
			if os.Getenv("DEBUG") != "" {
				log.Println("Couldn't create Kafka producer. Retrying...", err)
			}
			if i == retries-1 {
				return nil, errorf("Couldn't create Kafka producer. %v", err)
			}
		} else {
			time.Sleep(1 * time.Second)
		}
	}

	return &KafkaAvroAdapter{
		route:    route,
		brokers:  brokers,
		topic:    topic,
		registry: registry,
		schema:   schema,
		producer: producer,
	}, nil
}

func (a *KafkaAvroAdapter) Stream(logstream chan *router.Message) {
	defer a.producer.Close()
	for rm := range logstream {
		message, err := a.formatMessage(rm)
		if err != nil {
			log.Println("kafka:", err)
			a.route.Close()
			break
		}

		a.producer.Input() <- message
	}
}

func newConfig() *sarama.Config {
	config := sarama.NewConfig()
	config.ClientID = "logspout"
	config.Producer.Return.Errors = false
	config.Producer.Return.Successes = false
	config.Producer.Flush.Frequency = 1 * time.Second
	config.Producer.RequiredAcks = sarama.WaitForLocal

	if opt := os.Getenv("KAFKA_COMPRESSION_CODEC"); opt != "" {
		switch opt {
		case "gzip":
			config.Producer.Compression = sarama.CompressionGZIP
		case "snappy":
			config.Producer.Compression = sarama.CompressionSnappy
		}
	}

	return config
}

func (a *KafkaAvroAdapter) formatMessage(message *router.Message) (*sarama.ProducerMessage, error) {
	var encoder sarama.Encoder
	var containerName string
	var env docker.Env = message.Container.Config.Env

	if env != nil {
		containerName = env.Get("MARATHON_APP_ID")
	}
	if containerName == "" {
		containerName = message.Container.Name
	}
	host := env.Get("HOST")
	if host == "" {
		host = getLocalIP()
	}

	record := avro.NewGenericRecord(a.schema)
	record.Set("timestamp", iso8601.Time(message.Time).String())
	record.Set("container_name", containerName)
	record.Set("host", host)
	record.Set("source", message.Source)
	record.Set("line", message.Data)

	b, err := a.registry.Encode(record)
	if err != nil {
		return nil, err
	}
	debug(b)

	encoder = sarama.ByteEncoder(b)

	return &sarama.ProducerMessage{
		Topic: a.topic,
		Value: encoder,
	}, nil
}

func readBrokers(address string) []string {
	if strings.Contains(address, "/") {
		slash := strings.Index(address, "/")
		address = address[:slash]
	}

	return strings.Split(address, ",")
}

func readTopic(address string, options map[string]string) string {
	var topic string
	if !strings.Contains(address, "/") {
		topic = options["topic"]
	} else {
		slash := strings.Index(address, "/")
		topic = address[slash+1:]
	}

	return topic
}

func readSchemaRegistryUrl(options map[string]string) string {
	var url string
	if _, ok := options["schema_registry_url"]; ok {
		url = options["schema_registry_url"]
	}
	return url
}

func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, address := range addrs {
		// check the address type and if it is not a loopback the display it
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	return ""
}

func errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	if os.Getenv("DEBUG") != "" {
		fmt.Println(err.Error())
	}
	return
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}
<commit_msg>Fixing iso8601<commit_after>package kafka_avro

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	avro "github.com/elodina/go-avro"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/gliderlabs/logspout/router"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/fsouza/go-dockerclient"
	"github.com/joeshaw/iso8601"
)

var messageSchema = `{
  "type": "record",
  "name": "LogLine",
  "fields": [
    {"name": "timestamp", "type": "string"},
    {"name": "container_name", "type": "string"},
    {"name": "host", "type": "string"},
    {"name": "source", "type": "string"},
    {"name": "line", "type": "string"}
  ]
}`

func init() {
	router.AdapterFactories.Register(NewKafkaAvroAdapter, "kafka_avro")
}

type KafkaAvroAdapter struct {
	route    *router.Route
	brokers  []string
	topic    string
	schema   avro.Schema
	registry *kafkaavro.KafkaAvroEncoder
	producer sarama.AsyncProducer
}

func NewKafkaAvroAdapter(route *router.Route) (router.LogAdapter, error) {
	brokers := readBrokers(route.Address)
	if len(brokers) == 0 {
		return nil, errorf("The Kafka broker host:port is missing. Did you specify it as a route address?")
	}

	topic := readTopic(route.Address, route.Options)
	if topic == "" {
		return nil, errorf("The Kafka topic is missing. Did you specify it as a route option?")
	}

	schemaUrl := readSchemaRegistryUrl(route.Options)
	if schemaUrl == "" {
		return nil, errorf("The schema registry url is missing. Did you specify it as a route option?")
	}

	registry := kafkaavro.NewKafkaAvroEncoder(schemaUrl)

	var schema avro.Schema
	schema, err := avro.ParseSchema(messageSchema)
	if err != nil {
		return nil, errorf("The schema could not be parsed")
	}

	if os.Getenv("DEBUG") != "" {
		log.Printf("Starting Kafka producer for address: %s, topic: %s.\n", brokers, topic)
	}

	var retries int
	retries, err = strconv.Atoi(os.Getenv("KAFKA_CONNECT_RETRIES"))
	if err != nil {
		retries = 3
	}
	var producer sarama.AsyncProducer
	for i := 0; i < retries; i++ {
		producer, err = sarama.NewAsyncProducer(brokers, newConfig())
		if err != nil {
			if os.Getenv("DEBUG") != "" {
				log.Println("Couldn't create Kafka producer. Retrying...", err)
			}
			if i == retries-1 {
				return nil, errorf("Couldn't create Kafka producer. %v", err)
			}
		} else {
			time.Sleep(1 * time.Second)
		}
	}

	return &KafkaAvroAdapter{
		route:    route,
		brokers:  brokers,
		topic:    topic,
		registry: registry,
		schema:   schema,
		producer: producer,
	}, nil
}

func (a *KafkaAvroAdapter) Stream(logstream chan *router.Message) {
	defer a.producer.Close()
	for rm := range logstream {
		message, err := a.formatMessage(rm)
		if err != nil {
			log.Println("kafka:", err)
			a.route.Close()
			break
		}

		a.producer.Input() <- message
	}
}

func newConfig() *sarama.Config {
	config := sarama.NewConfig()
	config.ClientID = "logspout"
	config.Producer.Return.Errors = false
	config.Producer.Return.Successes = false
	config.Producer.Flush.Frequency = 1 * time.Second
	config.Producer.RequiredAcks = sarama.WaitForLocal

	if opt := os.Getenv("KAFKA_COMPRESSION_CODEC"); opt != "" {
		switch opt {
		case "gzip":
			config.Producer.Compression = sarama.CompressionGZIP
		case "snappy":
			config.Producer.Compression = sarama.CompressionSnappy
		}
	}

	return config
}

func (a *KafkaAvroAdapter) formatMessage(message *router.Message) (*sarama.ProducerMessage, error) {
	var encoder sarama.Encoder
	var containerName string
	var env docker.Env = message.Container.Config.Env

	if env != nil {
		containerName = env.Get("MARATHON_APP_ID")
	}
	if containerName == "" {
		containerName = message.Container.Name
	}
	host := env.Get("HOST")
	if host == "" {
		host = getLocalIP()
	}

	record := avro.NewGenericRecord(a.schema)
	time, _ := json.Marshal(iso8601.Time(message.Time))
	record.Set("timestamp", string(time))
	record.Set("container_name", containerName)
	record.Set("host", host)
	record.Set("source", message.Source)
	record.Set("line", message.Data)

	b, err := a.registry.Encode(record)
	if err != nil {
		return nil, err
	}
	debug(b)

	encoder = sarama.ByteEncoder(b)

	return &sarama.ProducerMessage{
		Topic: a.topic,
		Value: encoder,
	}, nil
}

func readBrokers(address string) []string {
	if strings.Contains(address, "/") {
		slash := strings.Index(address, "/")
		address = address[:slash]
	}

	return strings.Split(address, ",")
}

func readTopic(address string, options map[string]string) string {
	var topic string
	if !strings.Contains(address, "/") {
		topic = options["topic"]
	} else {
		slash := strings.Index(address, "/")
		topic = address[slash+1:]
	}

	return topic
}

func readSchemaRegistryUrl(options map[string]string) string {
	var url string
	if _, ok := options["schema_registry_url"]; ok {
		url = options["schema_registry_url"]
	}
	return url
}

func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, address := range addrs {
		// check the address type and if it is not a loopback the display it
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	return ""
}

func errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	if os.Getenv("DEBUG") != "" {
		fmt.Println(err.Error())
	}
	return
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}
<|endoftext|>


<commit_before>package utility

import (
	"fmt"
	"github.com/FidelityInternational/virgil/Godeps/_workspace/src/github.com/cloudfoundry-community/go-cfclient"
	"reflect"
	"sort"
	"strconv"
	"strings"
)

// FirewallRules - A collection of Firewall Rules with version
type FirewallRules struct {
	SchemaVersion string         `yaml:"schema_version"`
	FirewallRules []FirewallRule `yaml:"firewall_rules"`
}

// FirewallRule struct
type FirewallRule struct {
	Port        string
	Destination []string
	Protocol    string
	Source      []string
}

// ByPort - impliments sort.Interface for []FirewallRule bases on the Port field
type ByPort []FirewallRule

func (p ByPort) Len() int      { return len(p) }
func (p ByPort) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p ByPort) Less(i, j int) bool {
	if strings.EqualFold(p[i].Protocol, "TCP") && strings.EqualFold(p[j].Protocol, "UDP") {
		return true
	} else if strings.EqualFold(p[i].Protocol, "UDP") && strings.EqualFold(p[j].Protocol, "TCP") {
		return false
	}
	portIInt, _ := strconv.Atoi(p[i].Port)
	portJInt, _ := strconv.Atoi(p[j].Port)
	return portIInt < portJInt
}

// PortExpand - serperates port string into array, for example 2,5-7 becomes {2 5 6 7}
func PortExpand(portString string) ([]string, error) {
	var ports []string
	portsBefore := strings.Split(portString, ",")
	for _, port := range portsBefore {
		port = strings.TrimSpace(port)
		if strings.Contains(port, "-") {
			startFinish := strings.Split(port, "-")
			startString := strings.TrimSpace(startFinish[0])
			start, err := strconv.Atoi(startString)
			if err != nil || start <= 0 || start >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid as part of range %s", startString, port)
			}
			endString := strings.TrimSpace(startFinish[1])
			end, err := strconv.Atoi(endString)
			if err != nil || end <= 0 || end >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid as part of range %s", endString, port)
			}
			if len(startFinish) != 2 || start >= end {
				return []string{}, fmt.Errorf("Port range %s was invalid", port)
			}
			for i := start; i <= end; i++ {
				ports = append(ports, strconv.Itoa(i))
			}
		} else {
			portInt, err := strconv.Atoi(port)
			if err != nil || portInt <= 0 || portInt >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid", port)
			}
			ports = append(ports, port)
		}
	}
	return ports, nil
}

// ProcessRule - returns a concise list of firewall rules for one security group rule
func ProcessRule(secGroupRule cfclient.SecGroupRule, firewallRules []FirewallRule, source []string) ([]FirewallRule, error) {
	if strings.EqualFold(secGroupRule.Protocol, "all") {
		newRules := FirewallRule{
			Protocol:    secGroupRule.Protocol,
			Destination: []string{secGroupRule.Destination},
			Source:      source,
		}
		firewallRules = append(firewallRules, newRules)
		return firewallRules, nil
	}
	ports, err := PortExpand(secGroupRule.Ports)
	if err != nil {
		return []FirewallRule{}, err
	}
	for _, port := range ports {
		var newRule = true
		for i, rule := range firewallRules {
			if rule.Port == port && rule.Protocol == secGroupRule.Protocol {
				rule.Destination = append(rule.Destination, secGroupRule.Destination)
				RemoveDuplicates(&rule.Destination)
				firewallRules[i] = rule
				newRule = false
			}
		}
		if newRule {
			newRules := FirewallRule{
				Port:        port,
				Protocol:    secGroupRule.Protocol,
				Destination: []string{secGroupRule.Destination},
				Source:      source,
			}
			firewallRules = append(firewallRules, newRules)
		}
	}
	return firewallRules, nil
}

// RemoveDuplicates - removes duplicated from array of strings
func RemoveDuplicates(xs *[]string) {
	found := make(map[string]bool)
	j := 0
	for i, x := range *xs {
		if !found[x] {
			found[x] = true
			(*xs)[j] = (*xs)[i]
			j++
		}
	}
	*xs = (*xs)[:j]
}

// GetUsedSecGroups - Trims out any security-groups that cannot be used. I.E not running, staging or bound
func GetUsedSecGroups(allSecGroups []cfclient.SecGroup) []cfclient.SecGroup {
	var secGroups []cfclient.SecGroup
	for _, secGroup := range allSecGroups {
		if secGroup.Running || secGroup.Staging || len(secGroup.SpacesData) != 0 {
			secGroups = append(secGroups, secGroup)
		}
	}
	return secGroups
}

// GetFirewallRules - Returns a concise list of firewall rules for all security groups
func GetFirewallRules(source []string, secGroups []cfclient.SecGroup) FirewallRules {
	var (
		firewallRules, fwRules FirewallRules
		err                    error
	)
	firewallRules.SchemaVersion = "1"
	for _, secGroup := range secGroups {
		for _, secGroupRule := range secGroup.Rules {
			fwRules.FirewallRules, err = ProcessRule(secGroupRule, firewallRules.FirewallRules, source)
			if err != nil {
				continue
			}
			firewallRules.FirewallRules = fwRules.FirewallRules
		}
	}
	return compressDuplicateDestinations(firewallRules)
}

func compressDuplicateDestinations(firewallRules FirewallRules) FirewallRules {
	var firewallRulesResult []FirewallRule
	schema := firewallRules.SchemaVersion
	fwRules := firewallRules.FirewallRules
	sort.Sort(ByPort(fwRules))
	for i, fwRule := range fwRules {
		if i == 0 {
			firewallRulesResult = append(firewallRulesResult, fwRule)
			continue
		}
		prevRule := fwRules[i-1]
		if strings.EqualFold(fwRule.Protocol, "ALL") || fwRule.Protocol != prevRule.Protocol {
			firewallRulesResult = append(firewallRulesResult, fwRule)
			continue
		}
		rulePortInt, _ := strconv.Atoi(fwRule.Port)
		prevRulePorts := strings.Split(prevRule.Port, "-")
		prevRuleStartPort := prevRulePorts[0]
		prevRuleEndPort := ""
		if len(prevRulePorts) > 1 {
			prevRuleEndPort = prevRulePorts[1]
		}
		prevRulePortInt, _ := strconv.Atoi(prevRuleStartPort)
		if prevRuleEndPort != "" {
			prevRulePortInt, _ = strconv.Atoi(prevRuleEndPort)
		}
		if rulePortInt == prevRulePortInt+1 {
			if reflect.DeepEqual(fwRule.Destination, prevRule.Destination) {
				previousPort := strings.Split(firewallRulesResult[len(firewallRulesResult)-1].Port, "-")[0]
				firewallRulesResult[len(firewallRulesResult)-1].Port = fmt.Sprintf("%s-%s", previousPort, fwRule.Port)
			} else {
				firewallRulesResult = append(firewallRulesResult, fwRule)
			}
		} else {
			firewallRulesResult = append(firewallRulesResult, fwRule)
		}
	}
	return FirewallRules{SchemaVersion: schema, FirewallRules: firewallRulesResult}
}
<commit_msg>correct typo<commit_after>package utility

import (
	"fmt"
	"github.com/FidelityInternational/virgil/Godeps/_workspace/src/github.com/cloudfoundry-community/go-cfclient"
	"reflect"
	"sort"
	"strconv"
	"strings"
)

// FirewallRules - A collection of Firewall Rules with version
type FirewallRules struct {
	SchemaVersion string         `yaml:"schema_version"`
	FirewallRules []FirewallRule `yaml:"firewall_rules"`
}

// FirewallRule struct
type FirewallRule struct {
	Port        string
	Destination []string
	Protocol    string
	Source      []string
}

// ByPort - implements sort.Interface for []FirewallRule bases on the Port field
type ByPort []FirewallRule

func (p ByPort) Len() int      { return len(p) }
func (p ByPort) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p ByPort) Less(i, j int) bool {
	if strings.EqualFold(p[i].Protocol, "TCP") && strings.EqualFold(p[j].Protocol, "UDP") {
		return true
	} else if strings.EqualFold(p[i].Protocol, "UDP") && strings.EqualFold(p[j].Protocol, "TCP") {
		return false
	}
	portIInt, _ := strconv.Atoi(p[i].Port)
	portJInt, _ := strconv.Atoi(p[j].Port)
	return portIInt < portJInt
}

// PortExpand - serperates port string into array, for example 2,5-7 becomes {2 5 6 7}
func PortExpand(portString string) ([]string, error) {
	var ports []string
	portsBefore := strings.Split(portString, ",")
	for _, port := range portsBefore {
		port = strings.TrimSpace(port)
		if strings.Contains(port, "-") {
			startFinish := strings.Split(port, "-")
			startString := strings.TrimSpace(startFinish[0])
			start, err := strconv.Atoi(startString)
			if err != nil || start <= 0 || start >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid as part of range %s", startString, port)
			}
			endString := strings.TrimSpace(startFinish[1])
			end, err := strconv.Atoi(endString)
			if err != nil || end <= 0 || end >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid as part of range %s", endString, port)
			}
			if len(startFinish) != 2 || start >= end {
				return []string{}, fmt.Errorf("Port range %s was invalid", port)
			}
			for i := start; i <= end; i++ {
				ports = append(ports, strconv.Itoa(i))
			}
		} else {
			portInt, err := strconv.Atoi(port)
			if err != nil || portInt <= 0 || portInt >= 65536 {
				return []string{}, fmt.Errorf("Port %s was invalid", port)
			}
			ports = append(ports, port)
		}
	}
	return ports, nil
}

// ProcessRule - returns a concise list of firewall rules for one security group rule
func ProcessRule(secGroupRule cfclient.SecGroupRule, firewallRules []FirewallRule, source []string) ([]FirewallRule, error) {
	if strings.EqualFold(secGroupRule.Protocol, "all") {
		newRules := FirewallRule{
			Protocol:    secGroupRule.Protocol,
			Destination: []string{secGroupRule.Destination},
			Source:      source,
		}
		firewallRules = append(firewallRules, newRules)
		return firewallRules, nil
	}
	ports, err := PortExpand(secGroupRule.Ports)
	if err != nil {
		return []FirewallRule{}, err
	}
	for _, port := range ports {
		var newRule = true
		for i, rule := range firewallRules {
			if rule.Port == port && rule.Protocol == secGroupRule.Protocol {
				rule.Destination = append(rule.Destination, secGroupRule.Destination)
				RemoveDuplicates(&rule.Destination)
				firewallRules[i] = rule
				newRule = false
			}
		}
		if newRule {
			newRules := FirewallRule{
				Port:        port,
				Protocol:    secGroupRule.Protocol,
				Destination: []string{secGroupRule.Destination},
				Source:      source,
			}
			firewallRules = append(firewallRules, newRules)
		}
	}
	return firewallRules, nil
}

// RemoveDuplicates - removes duplicated from array of strings
func RemoveDuplicates(xs *[]string) {
	found := make(map[string]bool)
	j := 0
	for i, x := range *xs {
		if !found[x] {
			found[x] = true
			(*xs)[j] = (*xs)[i]
			j++
		}
	}
	*xs = (*xs)[:j]
}

// GetUsedSecGroups - Trims out any security-groups that cannot be used. I.E not running, staging or bound
func GetUsedSecGroups(allSecGroups []cfclient.SecGroup) []cfclient.SecGroup {
	var secGroups []cfclient.SecGroup
	for _, secGroup := range allSecGroups {
		if secGroup.Running || secGroup.Staging || len(secGroup.SpacesData) != 0 {
			secGroups = append(secGroups, secGroup)
		}
	}
	return secGroups
}

// GetFirewallRules - Returns a concise list of firewall rules for all security groups
func GetFirewallRules(source []string, secGroups []cfclient.SecGroup) FirewallRules {
	var (
		firewallRules, fwRules FirewallRules
		err                    error
	)
	firewallRules.SchemaVersion = "1"
	for _, secGroup := range secGroups {
		for _, secGroupRule := range secGroup.Rules {
			fwRules.FirewallRules, err = ProcessRule(secGroupRule, firewallRules.FirewallRules, source)
			if err != nil {
				continue
			}
			firewallRules.FirewallRules = fwRules.FirewallRules
		}
	}
	return compressDuplicateDestinations(firewallRules)
}

func compressDuplicateDestinations(firewallRules FirewallRules) FirewallRules {
	var firewallRulesResult []FirewallRule
	schema := firewallRules.SchemaVersion
	fwRules := firewallRules.FirewallRules
	sort.Sort(ByPort(fwRules))
	for i, fwRule := range fwRules {
		if i == 0 {
			firewallRulesResult = append(firewallRulesResult, fwRule)
			continue
		}
		prevRule := fwRules[i-1]
		if strings.EqualFold(fwRule.Protocol, "ALL") || fwRule.Protocol != prevRule.Protocol {
			firewallRulesResult = append(firewallRulesResult, fwRule)
			continue
		}
		rulePortInt, _ := strconv.Atoi(fwRule.Port)
		prevRulePorts := strings.Split(prevRule.Port, "-")
		prevRuleStartPort := prevRulePorts[0]
		prevRuleEndPort := ""
		if len(prevRulePorts) > 1 {
			prevRuleEndPort = prevRulePorts[1]
		}
		prevRulePortInt, _ := strconv.Atoi(prevRuleStartPort)
		if prevRuleEndPort != "" {
			prevRulePortInt, _ = strconv.Atoi(prevRuleEndPort)
		}
		if rulePortInt == prevRulePortInt+1 {
			if reflect.DeepEqual(fwRule.Destination, prevRule.Destination) {
				previousPort := strings.Split(firewallRulesResult[len(firewallRulesResult)-1].Port, "-")[0]
				firewallRulesResult[len(firewallRulesResult)-1].Port = fmt.Sprintf("%s-%s", previousPort, fwRule.Port)
			} else {
				firewallRulesResult = append(firewallRulesResult, fwRule)
			}
		} else {
			firewallRulesResult = append(firewallRulesResult, fwRule)
		}
	}
	return FirewallRules{SchemaVersion: schema, FirewallRules: firewallRulesResult}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2019 shawn1m. All rights reserved.
 * Use of this source code is governed by The MIT License (MIT) that can be
 * found in the LICENSE file..
 */

// Package outbound implements multiple dns client and dispatcher for outbound connection.
package clients

import (
	"bytes"
	"crypto/tls"
	"io/ioutil"
	"net"
	"net/http"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/proxy"

	"github.com/shawn1m/overture/core/cache"
	"github.com/shawn1m/overture/core/common"
)

type RemoteClient struct {
	responseMessage *dns.Msg
	questionMessage *dns.Msg

	dnsUpstream        *common.DNSUpstream
	ednsClientSubnetIP string
	inboundIP          string

	cache *cache.Cache
}

func NewClient(q *dns.Msg, u *common.DNSUpstream, ip string, cache *cache.Cache) *RemoteClient {
	c := &RemoteClient{questionMessage: q.Copy(), dnsUpstream: u, inboundIP: ip, cache: cache}
	c.getEDNSClientSubnetIP()

	return c
}

func (c *RemoteClient) getEDNSClientSubnetIP() {
	switch c.dnsUpstream.EDNSClientSubnet.Policy {
	case "auto":
		if !common.IsIPMatchList(net.ParseIP(c.inboundIP), common.ReservedIPNetworkList, false, "") {
			c.ednsClientSubnetIP = c.inboundIP
		} else {
			c.ednsClientSubnetIP = c.dnsUpstream.EDNSClientSubnet.ExternalIP
		}
	case "manual":
		if c.dnsUpstream.EDNSClientSubnet.ExternalIP != "" &&
			!common.IsIPMatchList(net.ParseIP(c.dnsUpstream.EDNSClientSubnet.ExternalIP), common.ReservedIPNetworkList, false, "") {
			c.ednsClientSubnetIP = c.dnsUpstream.EDNSClientSubnet.ExternalIP
			return
		}
	case "disable":
	}
}

func (c *RemoteClient) ExchangeFromCache() *dns.Msg {
	cacheClient := NewCacheClient(c.questionMessage, c.ednsClientSubnetIP, c.cache)
	c.responseMessage = cacheClient.Exchange()
	if c.responseMessage != nil {
		return c.responseMessage
	}
	return nil
}

func (c *RemoteClient) Exchange(isLog bool) *dns.Msg {
	common.SetEDNSClientSubnet(c.questionMessage, c.ednsClientSubnetIP,
		c.dnsUpstream.EDNSClientSubnet.NoCookie)
	c.ednsClientSubnetIP = common.GetEDNSClientSubnetIP(c.questionMessage)

	if c.responseMessage != nil {
		return c.responseMessage
	}

	var conn net.Conn = nil
	var err error
	if c.dnsUpstream.SOCKS5Address != "" {
		if conn, err = c.createSocks5Conn(); err != nil {
			return nil
		}
	}

	var temp *dns.Msg
	switch c.dnsUpstream.Protocol {
	case "udp":
		temp, err = c.ExchangeByUDP(conn)
	case "tcp":
		temp, err = c.ExchangeByTCP(conn)
	case "tcp-tls":
		temp, err = c.ExchangeByTLS(conn)
	case "https":
		temp, err = c.ExchangeByHTTPS(conn)
	}

	if err != nil {
		log.Debugf("%s Fail: %s", c.dnsUpstream.Name, err)
		return nil
	}
	if temp == nil {
		log.Debugf("%s Fail: Response message returned nil, maybe timeout? Please check your query or DNS configuration")
		return nil
	}

	c.responseMessage = temp

	if isLog {
		c.logAnswer("")
	}

	return c.responseMessage
}

func (c *RemoteClient) logAnswer(indicator string) {

	for _, a := range c.responseMessage.Answer {
		var name string
		if indicator != "" {
			name = indicator
		} else {
			name = c.dnsUpstream.Name
		}
		log.Debugf("Answer from %s: %s", name, a.String())
	}
}

func (c *RemoteClient) createSocks5Conn() (conn net.Conn, err error) {
	socksAddress, err := ExtractSocksAddress(c.dnsUpstream.SOCKS5Address)
	if err != nil {
		return nil, err
	}
	network := ToNetwork(c.dnsUpstream.Protocol)
	s, err := proxy.SOCKS5(network, socksAddress, nil, proxy.Direct)
	if err != nil {
		log.Warnf("Failed to connect to SOCKS5 proxy: %s", err)
		return nil, err
	}
	host, port, err := ExtractDNSAddress(c.dnsUpstream.Address, c.dnsUpstream.Protocol)
	if err != nil {
		return nil, err
	}
	address := net.JoinHostPort(host, port)
	conn, err = s.Dial(network, address)
	if err != nil {
		log.Warnf("Failed to connect to upstream via SOCKS5 proxy: %s", err)
		return nil, err
	}
	return conn, err
}

func (c *RemoteClient) exchangeByDNSClient(conn net.Conn) (msg *dns.Msg, err error) {
	if conn == nil {
		network := ToNetwork(c.dnsUpstream.Protocol)
		host, port, err := ExtractDNSAddress(c.dnsUpstream.Address, c.dnsUpstream.Protocol)
		if err != nil {
			return nil, err
		}
		address := net.JoinHostPort(host, port)
		if conn, err = net.Dial(network, address); err != nil {
			log.Warnf("Failed to connect to DNS upstream: %s", err)
			return nil, err
		}
	}
	c.setTimeout(conn)
	dc := &dns.Conn{Conn: conn, UDPSize: 65535}
	defer dc.Close()
	err = dc.WriteMsg(c.questionMessage)
	if err != nil {
		log.Warnf("%s Fail: Send question message failed", c.dnsUpstream.Name)
		return nil, err
	}
	return dc.ReadMsg()
}

// ExchangeByUDP send dns record by udp protocol
func (c *RemoteClient) ExchangeByUDP(conn net.Conn) (*dns.Msg, error) {
	return c.exchangeByDNSClient(conn)
}

// ExchangeByTCP send dns record by tcp protocol
func (c *RemoteClient) ExchangeByTCP(conn net.Conn) (*dns.Msg, error) {
	return c.exchangeByDNSClient(conn)
}

// ExchangeByTLS send dns record by tcp-tls protocol
func (c *RemoteClient) ExchangeByTLS(conn net.Conn) (msg *dns.Msg, err error) {
	host, port, ip := ExtractTLSDNSAddress(c.dnsUpstream.Address)
	var address string
	if len(ip) > 0 {
		address = net.JoinHostPort(ip, port)
	} else {
		address = net.JoinHostPort(host, port)
	}

	conf := &tls.Config{
		InsecureSkipVerify: false,
		ServerName:         host,
	}
	if conn != nil {
		// crate tls client use the existing connection
		conn = tls.Client(conn, conf)
	} else {
		if conn, err = tls.Dial("tcp", address, conf); err != nil {
			log.Warnf("Failed to connect to DNS-over-TLS upstream: %s", err)
			return nil, err
		}
	}
	c.setTimeout(conn)
	return c.exchangeByDNSClient(conn)
}

// ExchangeByHTTPS send dns record by https protocol
func (c *RemoteClient) ExchangeByHTTPS(conn net.Conn) (*dns.Msg, error) {
	if conn == nil {
		host, port, err := ExtractHTTPSAddress(c.dnsUpstream.Address)
		if err != nil {
			return nil, err
		}
		address := net.JoinHostPort(host, port)
		conn, err = net.Dial("tcp", address)
		if err != nil {
			log.Warnf("Fail connect to dns server %s", address)
			return nil, err
		}
	}
	c.setTimeout(conn)
	client := http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (net.Conn, error) {
				return conn, nil
			},
		},
	}
	defer client.CloseIdleConnections()
	request, err := c.questionMessage.Pack()
	resp, err := client.Post(c.dnsUpstream.Address, "application/dns-message",
		bytes.NewBuffer(request))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	msg := new(dns.Msg)
	err = msg.Unpack(data)
	if err != nil {
		return nil, err
	}
	return msg, nil
}

func (c *RemoteClient) setTimeout(conn net.Conn) {
	dnsTimeout := time.Duration(c.dnsUpstream.Timeout) * time.Second / 3
	conn.SetDeadline(time.Now().Add(dnsTimeout))
	conn.SetReadDeadline(time.Now().Add(dnsTimeout))
	conn.SetWriteDeadline(time.Now().Add(dnsTimeout))
}
<commit_msg>Fix Debugf usage error<commit_after>/*
 * Copyright (c) 2019 shawn1m. All rights reserved.
 * Use of this source code is governed by The MIT License (MIT) that can be
 * found in the LICENSE file..
 */

// Package outbound implements multiple dns client and dispatcher for outbound connection.
package clients

import (
	"bytes"
	"crypto/tls"
	"io/ioutil"
	"net"
	"net/http"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/proxy"

	"github.com/shawn1m/overture/core/cache"
	"github.com/shawn1m/overture/core/common"
)

type RemoteClient struct {
	responseMessage *dns.Msg
	questionMessage *dns.Msg

	dnsUpstream        *common.DNSUpstream
	ednsClientSubnetIP string
	inboundIP          string

	cache *cache.Cache
}

func NewClient(q *dns.Msg, u *common.DNSUpstream, ip string, cache *cache.Cache) *RemoteClient {
	c := &RemoteClient{questionMessage: q.Copy(), dnsUpstream: u, inboundIP: ip, cache: cache}
	c.getEDNSClientSubnetIP()

	return c
}

func (c *RemoteClient) getEDNSClientSubnetIP() {
	switch c.dnsUpstream.EDNSClientSubnet.Policy {
	case "auto":
		if !common.IsIPMatchList(net.ParseIP(c.inboundIP), common.ReservedIPNetworkList, false, "") {
			c.ednsClientSubnetIP = c.inboundIP
		} else {
			c.ednsClientSubnetIP = c.dnsUpstream.EDNSClientSubnet.ExternalIP
		}
	case "manual":
		if c.dnsUpstream.EDNSClientSubnet.ExternalIP != "" &&
			!common.IsIPMatchList(net.ParseIP(c.dnsUpstream.EDNSClientSubnet.ExternalIP), common.ReservedIPNetworkList, false, "") {
			c.ednsClientSubnetIP = c.dnsUpstream.EDNSClientSubnet.ExternalIP
			return
		}
	case "disable":
	}
}

func (c *RemoteClient) ExchangeFromCache() *dns.Msg {
	cacheClient := NewCacheClient(c.questionMessage, c.ednsClientSubnetIP, c.cache)
	c.responseMessage = cacheClient.Exchange()
	if c.responseMessage != nil {
		return c.responseMessage
	}
	return nil
}

func (c *RemoteClient) Exchange(isLog bool) *dns.Msg {
	common.SetEDNSClientSubnet(c.questionMessage, c.ednsClientSubnetIP,
		c.dnsUpstream.EDNSClientSubnet.NoCookie)
	c.ednsClientSubnetIP = common.GetEDNSClientSubnetIP(c.questionMessage)

	if c.responseMessage != nil {
		return c.responseMessage
	}

	var conn net.Conn = nil
	var err error
	if c.dnsUpstream.SOCKS5Address != "" {
		if conn, err = c.createSocks5Conn(); err != nil {
			return nil
		}
	}

	var temp *dns.Msg
	switch c.dnsUpstream.Protocol {
	case "udp":
		temp, err = c.ExchangeByUDP(conn)
	case "tcp":
		temp, err = c.ExchangeByTCP(conn)
	case "tcp-tls":
		temp, err = c.ExchangeByTLS(conn)
	case "https":
		temp, err = c.ExchangeByHTTPS(conn)
	}

	if err != nil {
		log.Debugf("%s Fail: %s", c.dnsUpstream.Name, err)
		return nil
	}
	if temp == nil {
		log.Debugf("%s Fail: Response message returned nil, maybe timeout? Please check your query or DNS configuration", c.dnsUpstream.Name)
		return nil
	}

	c.responseMessage = temp

	if isLog {
		c.logAnswer("")
	}

	return c.responseMessage
}

func (c *RemoteClient) logAnswer(indicator string) {

	for _, a := range c.responseMessage.Answer {
		var name string
		if indicator != "" {
			name = indicator
		} else {
			name = c.dnsUpstream.Name
		}
		log.Debugf("Answer from %s: %s", name, a.String())
	}
}

func (c *RemoteClient) createSocks5Conn() (conn net.Conn, err error) {
	socksAddress, err := ExtractSocksAddress(c.dnsUpstream.SOCKS5Address)
	if err != nil {
		return nil, err
	}
	network := ToNetwork(c.dnsUpstream.Protocol)
	s, err := proxy.SOCKS5(network, socksAddress, nil, proxy.Direct)
	if err != nil {
		log.Warnf("Failed to connect to SOCKS5 proxy: %s", err)
		return nil, err
	}
	host, port, err := ExtractDNSAddress(c.dnsUpstream.Address, c.dnsUpstream.Protocol)
	if err != nil {
		return nil, err
	}
	address := net.JoinHostPort(host, port)
	conn, err = s.Dial(network, address)
	if err != nil {
		log.Warnf("Failed to connect to upstream via SOCKS5 proxy: %s", err)
		return nil, err
	}
	return conn, err
}

func (c *RemoteClient) exchangeByDNSClient(conn net.Conn) (msg *dns.Msg, err error) {
	if conn == nil {
		network := ToNetwork(c.dnsUpstream.Protocol)
		host, port, err := ExtractDNSAddress(c.dnsUpstream.Address, c.dnsUpstream.Protocol)
		if err != nil {
			return nil, err
		}
		address := net.JoinHostPort(host, port)
		if conn, err = net.Dial(network, address); err != nil {
			log.Warnf("Failed to connect to DNS upstream: %s", err)
			return nil, err
		}
	}
	c.setTimeout(conn)
	dc := &dns.Conn{Conn: conn, UDPSize: 65535}
	defer dc.Close()
	err = dc.WriteMsg(c.questionMessage)
	if err != nil {
		log.Warnf("%s Fail: Send question message failed", c.dnsUpstream.Name)
		return nil, err
	}
	return dc.ReadMsg()
}

// ExchangeByUDP send dns record by udp protocol
func (c *RemoteClient) ExchangeByUDP(conn net.Conn) (*dns.Msg, error) {
	return c.exchangeByDNSClient(conn)
}

// ExchangeByTCP send dns record by tcp protocol
func (c *RemoteClient) ExchangeByTCP(conn net.Conn) (*dns.Msg, error) {
	return c.exchangeByDNSClient(conn)
}

// ExchangeByTLS send dns record by tcp-tls protocol
func (c *RemoteClient) ExchangeByTLS(conn net.Conn) (msg *dns.Msg, err error) {
	host, port, ip := ExtractTLSDNSAddress(c.dnsUpstream.Address)
	var address string
	if len(ip) > 0 {
		address = net.JoinHostPort(ip, port)
	} else {
		address = net.JoinHostPort(host, port)
	}

	conf := &tls.Config{
		InsecureSkipVerify: false,
		ServerName:         host,
	}
	if conn != nil {
		// crate tls client use the existing connection
		conn = tls.Client(conn, conf)
	} else {
		if conn, err = tls.Dial("tcp", address, conf); err != nil {
			log.Warnf("Failed to connect to DNS-over-TLS upstream: %s", err)
			return nil, err
		}
	}
	c.setTimeout(conn)
	return c.exchangeByDNSClient(conn)
}

// ExchangeByHTTPS send dns record by https protocol
func (c *RemoteClient) ExchangeByHTTPS(conn net.Conn) (*dns.Msg, error) {
	if conn == nil {
		host, port, err := ExtractHTTPSAddress(c.dnsUpstream.Address)
		if err != nil {
			return nil, err
		}
		address := net.JoinHostPort(host, port)
		conn, err = net.Dial("tcp", address)
		if err != nil {
			log.Warnf("Fail connect to dns server %s", address)
			return nil, err
		}
	}
	c.setTimeout(conn)
	client := http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (net.Conn, error) {
				return conn, nil
			},
		},
	}
	defer client.CloseIdleConnections()
	request, err := c.questionMessage.Pack()
	resp, err := client.Post(c.dnsUpstream.Address, "application/dns-message",
		bytes.NewBuffer(request))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	msg := new(dns.Msg)
	err = msg.Unpack(data)
	if err != nil {
		return nil, err
	}
	return msg, nil
}

func (c *RemoteClient) setTimeout(conn net.Conn) {
	dnsTimeout := time.Duration(c.dnsUpstream.Timeout) * time.Second / 3
	conn.SetDeadline(time.Now().Add(dnsTimeout))
	conn.SetReadDeadline(time.Now().Add(dnsTimeout))
	conn.SetWriteDeadline(time.Now().Add(dnsTimeout))
}
<|endoftext|>