<commit_before>// Copyright 2014-2016 The project AUTHORS. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package diff implements diff functions to compare objects.
package diff

import (
	"encoding/json"
	"errors"
	"reflect"
	"strconv"
	"strings"
)

type Diff map[string]interface{}

func (d Diff) PrettyJSON() []byte {
	bs, err := json.MarshalIndent(d, "", "   ")
	if err != nil {
		return []byte{}
	}
	return bs
}

func (d Diff) JSON() []byte {
	bs, err := json.MarshalIndent(d, "", "   ")
	if err != nil {
		return []byte{}
	}
	return bs
}

type Change struct {
	OldVal interface{} `json:"old_value"`
	NewVal interface{} `json:"new_value"`
}

type ChangeType string

const (
	AddType ChangeType = "ADD" // addition
	DelType ChangeType = "DEL" // deletion
	ModType ChangeType = "MOD" // modification
)

type CollectionChange struct {
	Val  interface{} `json:"value"`
	Type ChangeType  `json:"type"`
}

func Compute(x, y interface{}, recursive bool) (Diff, error) {
	vx, vy := reflect.ValueOf(x), reflect.ValueOf(y)
	tx, ty := vx.Type(), vy.Type()

	if !tx.AssignableTo(ty) {
		return nil, errors.New("input objects do not share the same type")
	}

	// since x and y share the same type, there is no need to check them both
	if vx.Kind() != reflect.Struct {
		return nil, errors.New("input values are not struct")
	}

	xNumFields := vx.NumField()

	delta := make(Diff)

	for i := 0; i < xNumFields; i++ {
		fx := vx.Field(i)
		typ := tx.Field(i)
		if !isExported(typ.Name) { // skip non-exported fields
			continue
		}
		fy := vy.FieldByName(typ.Name)

		if d := handleValue(fx, fy); d != nil {
			delta[typ.Name] = d
		}
	}

	return delta, nil
}

func handleValue(fx, fy reflect.Value) interface{} {
	switch fx.Kind() {

	case reflect.Struct:
		if isFullyNonExportedStruct(fx) {
			if fx.String() != fy.String() {
				return Change{OldVal: fx.Interface(), NewVal: fy.Interface()}
			}
			return nil
		}

		delta := make(Diff)
		numFields := fx.NumField()
		for i := 0; i < numFields; i++ {
			newFx := fx.Field(i)
			typ := fx.Type().Field(i)
			if !isExported(typ.Name) { // skip non-exported fields
				continue
			}
			newFy := fy.FieldByName(typ.Name)

			if d := handleValue(newFx, newFy); d != nil {
				delta[typ.Name] = d
			}
		}
		if len(delta) > 0 {
			return delta
		}
		return nil

	case reflect.Array, reflect.Slice:
		xLen, yLen := fx.Len(), fy.Len()

		changes := make(map[string]CollectionChange)
		if xLen == 0 {
			if yLen == 0 {
				return nil
			}
			for i := 0; i < yLen; i++ {
				changes[strconv.Itoa(i)] = CollectionChange{Val: fy.Index(i).Interface(), Type: AddType}
			}
		} else if yLen == 0 {
			for i := 0; i < xLen; i++ {
				changes[strconv.Itoa(i)] = CollectionChange{Val: fx.Index(i).Interface(), Type: DelType}
			}
		} else {
			var maxLen int
			if xLen > yLen {
				maxLen = yLen
				for i := yLen; i < xLen; i++ {
					changes[strconv.Itoa(i)] = CollectionChange{Val: fx.Index(i).Interface(), Type: DelType}
				}
			} else if xLen < yLen {
				maxLen = xLen
				for i := xLen; i < yLen; i++ {
					changes[strconv.Itoa(i)] = CollectionChange{Val: fy.Index(i).Interface(), Type: AddType}
				}
			}
			for i := 0; i < maxLen; i++ {
				if d := handleValue(fx.Index(i), fy.Index(i)); d != nil {
					changes[strconv.Itoa(i)] = CollectionChange{Val: d, Type: ModType}
				}
			}
		}
		if len(changes) > 0 {
			return changes
		}
		return nil

	case reflect.Map:
		return nil

	case reflect.Ptr:
		if fx.IsNil() {
			if fy.IsNil() {
				return nil
			}
			return Change{OldVal: nil, NewVal: fy.Elem().Interface()}
		} else if fy.IsNil() {
			return Change{OldVal: fx.Elem().Interface(), NewVal: nil}
		}
		return handleValue(fx.Elem(), fy.Elem())

	case reflect.Interface, reflect.Func, reflect.Chan, reflect.Invalid, reflect.UnsafePointer, reflect.Complex64, reflect.Complex128:
		// TODO(gilliek): support complex numbers
		return nil

	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ix, iy := fx.Int(), fy.Int()
		if ix != iy {
			return Change{OldVal: ix, NewVal: iy}
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		uix, uiy := fx.Uint(), fy.Uint()
		if uix != uiy {
			return Change{OldVal: uix, NewVal: uiy}
		}
	case reflect.Float32, reflect.Float64:
		flx, fly := fx.Float(), fy.Float()
		if flx-fly < 0.000001 {
			return Change{OldVal: flx, NewVal: fly}
		}
	case reflect.String:
		sx, sy := fx.String(), fy.String()
		if sx != sy {
			return Change{OldVal: sx, NewVal: sy}
		}
	}

	return nil
}

func isExported(fieldName string) bool {
	if fieldName == "" {
		return false
	}
	firstLetter := string(fieldName[0])
	return firstLetter != strings.ToLower(firstLetter)
}

func isFullyNonExportedStruct(s reflect.Value) bool {
	if s.Kind() != reflect.Struct {
		return false
	}
	numFields := s.NumField()
	for i := 0; i < numFields; i++ {
		typ := s.Type().Field(i)
		if isExported(typ.Name) {
			return false
		}
	}
	return true
}
<commit_msg>diff: fix fully non-exported struct comparison<commit_after>// Copyright 2014-2016 The project AUTHORS. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package diff implements diff functions to compare objects.
package diff

import (
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
)

type Diff map[string]interface{}

func (d Diff) PrettyJSON() []byte {
	bs, err := json.MarshalIndent(d, "", "   ")
	if err != nil {
		return []byte{}
	}
	return bs
}

func (d Diff) JSON() []byte {
	bs, err := json.MarshalIndent(d, "", "   ")
	if err != nil {
		return []byte{}
	}
	return bs
}

type Change struct {
	OldVal interface{} `json:"old_value"`
	NewVal interface{} `json:"new_value"`
}

type ChangeType string

const (
	AddType ChangeType = "ADD" // addition
	DelType ChangeType = "DEL" // deletion
	ModType ChangeType = "MOD" // modification
)

type CollectionChange struct {
	Val  interface{} `json:"value"`
	Type ChangeType  `json:"type"`
}

func Compute(x, y interface{}, recursive bool) (Diff, error) {
	vx, vy := reflect.ValueOf(x), reflect.ValueOf(y)
	tx, ty := vx.Type(), vy.Type()

	if !tx.AssignableTo(ty) {
		return nil, errors.New("input objects do not share the same type")
	}

	// since x and y share the same type, there is no need to check them both
	if vx.Kind() != reflect.Struct {
		return nil, errors.New("input values are not struct")
	}

	xNumFields := vx.NumField()

	delta := make(Diff)

	for i := 0; i < xNumFields; i++ {
		fx := vx.Field(i)
		typ := tx.Field(i)
		if !isExported(typ.Name) { // skip non-exported fields
			continue
		}
		fy := vy.FieldByName(typ.Name)

		if d := handleValue(fx, fy); d != nil {
			delta[typ.Name] = d
		}
	}

	return delta, nil
}

func handleValue(fx, fy reflect.Value) interface{} {
	switch fx.Kind() {

	case reflect.Struct:
		if isFullyNonExportedStruct(fx) {
			if !isEqual(fx, fy) {
				return Change{OldVal: fx.Interface(), NewVal: fy.Interface()}
			}
			return nil
		}

		delta := make(Diff)
		numFields := fx.NumField()
		for i := 0; i < numFields; i++ {
			newFx := fx.Field(i)
			typ := fx.Type().Field(i)
			if !isExported(typ.Name) { // skip non-exported fields
				continue
			}
			newFy := fy.FieldByName(typ.Name)

			if d := handleValue(newFx, newFy); d != nil {
				delta[typ.Name] = d
			}
		}
		if len(delta) > 0 {
			return delta
		}
		return nil

	case reflect.Array, reflect.Slice:
		xLen, yLen := fx.Len(), fy.Len()

		changes := make(map[string]CollectionChange)
		if xLen == 0 {
			if yLen == 0 {
				return nil
			}
			for i := 0; i < yLen; i++ {
				changes[strconv.Itoa(i)] = CollectionChange{Val: fy.Index(i).Interface(), Type: AddType}
			}
		} else if yLen == 0 {
			for i := 0; i < xLen; i++ {
				changes[strconv.Itoa(i)] = CollectionChange{Val: fx.Index(i).Interface(), Type: DelType}
			}
		} else {
			var maxLen int
			if xLen > yLen {
				maxLen = yLen
				for i := yLen; i < xLen; i++ {
					changes[strconv.Itoa(i)] = CollectionChange{Val: fx.Index(i).Interface(), Type: DelType}
				}
			} else if xLen < yLen {
				maxLen = xLen
				for i := xLen; i < yLen; i++ {
					changes[strconv.Itoa(i)] = CollectionChange{Val: fy.Index(i).Interface(), Type: AddType}
				}
			}
			for i := 0; i < maxLen; i++ {
				if d := handleValue(fx.Index(i), fy.Index(i)); d != nil {
					changes[strconv.Itoa(i)] = CollectionChange{Val: d, Type: ModType}
				}
			}
		}
		if len(changes) > 0 {
			return changes
		}
		return nil

	case reflect.Map:
		return nil

	case reflect.Ptr:
		if fx.IsNil() {
			if fy.IsNil() {
				return nil
			}
			return Change{OldVal: nil, NewVal: fy.Elem().Interface()}
		} else if fy.IsNil() {
			return Change{OldVal: fx.Elem().Interface(), NewVal: nil}
		}
		return handleValue(fx.Elem(), fy.Elem())

	case reflect.Interface, reflect.Func, reflect.Chan, reflect.Invalid, reflect.UnsafePointer, reflect.Complex64, reflect.Complex128:
		// TODO(gilliek): support complex numbers
		return nil

	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		ix, iy := fx.Int(), fy.Int()
		if ix != iy {
			return Change{OldVal: ix, NewVal: iy}
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		uix, uiy := fx.Uint(), fy.Uint()
		if uix != uiy {
			return Change{OldVal: uix, NewVal: uiy}
		}
	case reflect.Float32, reflect.Float64:
		flx, fly := fx.Float(), fy.Float()
		if flx-fly < 0.000001 {
			return Change{OldVal: flx, NewVal: fly}
		}
	case reflect.String:
		sx, sy := fx.String(), fy.String()
		if sx != sy {
			return Change{OldVal: sx, NewVal: sy}
		}
	}

	return nil
}

func isExported(fieldName string) bool {
	if fieldName == "" {
		return false
	}
	firstLetter := string(fieldName[0])
	return firstLetter != strings.ToLower(firstLetter)
}

func isFullyNonExportedStruct(s reflect.Value) bool {
	if s.Kind() != reflect.Struct {
		return false
	}
	numFields := s.NumField()
	for i := 0; i < numFields; i++ {
		typ := s.Type().Field(i)
		if isExported(typ.Name) {
			return false
		}
	}
	return true
}

func isEqual(x, y reflect.Value) bool {
	return fmt.Sprint(x.Interface()) == fmt.Sprint(y.Interface())
}
<|endoftext|>


<commit_before>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor_test

import (
	"bytes"
	"fmt"
	"syscall"
	"testing"

	"github.com/shiena/ansicolor"
	. "github.com/shiena/ansicolor"
)

func TestWritePlanText(t *testing.T) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)
	expected := "plain text"
	fmt.Fprintf(w, expected)
	actual := inner.String()
	if actual != expected {
		t.Errorf("Get %s, want %s", actual, expected)
	}
}

func TestWriteParseText(t *testing.T) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)

	input_tail := "\x1b[0mtail text"
	expected_tail := "tail text"
	fmt.Fprintf(w, input_tail)
	actual_tail := inner.String()
	inner.Reset()
	if actual_tail != expected_tail {
		t.Errorf("Get %s, want %s", actual_tail, expected_tail)
	}

	input_head := "head text\x1b[0m"
	expected_head := "head text"
	fmt.Fprintf(w, input_head)
	actual_head := inner.String()
	inner.Reset()
	if actual_head != expected_head {
		t.Errorf("Get %s, want %s", actual_head, expected_head)
	}

	input_both_ends := "both ends \x1b[0m text"
	expected_both_ends := "both ends  text"
	fmt.Fprintf(w, input_both_ends)
	actual_both_ends := inner.String()
	inner.Reset()
	if actual_both_ends != expected_both_ends {
		t.Errorf("Get %s, want %s", actual_both_ends, expected_both_ends)
	}

	input_many_esc := "\x1b\x1b\x1b\x1b[0m many esc"
	expected_many_esc := "\x1b\x1b\x1b many esc"
	fmt.Fprintf(w, input_many_esc)
	actual_many_esc := inner.String()
	inner.Reset()
	if actual_many_esc != expected_many_esc {
		t.Errorf("Get %s, want %s", actual_many_esc, expected_many_esc)
	}

	expected_split := "split  text"
	for _, ch := range "split \x1b[0m text" {
		fmt.Fprintf(w, string(ch))
	}
	actual_split := inner.String()
	inner.Reset()
	if actual_split != expected_split {
		t.Errorf("Get %s, want %s", actual_split, expected_split)
	}
}

type screenNotFoundError struct {
	error
}

func writeAnsiColor(expectedText, colorCode string) (actualText string, actualAttributes uint16, err error) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)
	fmt.Fprintf(w, "\x1b[%sm%s", colorCode, expectedText)

	actualText = inner.String()
	screenInfo := GetConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		actualAttributes = screenInfo.WAttributes
	} else {
		err = &screenNotFoundError{}
	}
	return
}

type testParam struct {
	text       string
	attributes uint16
	ansiColor  string
}

func TestWriteAnsiColorText(t *testing.T) {
	screenInfo := GetConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		t.Fatal("Could not get ConsoleScreenBufferInfo")
	}
	defer ChangeColor(screenInfo.WAttributes)
	defaultFgColor := screenInfo.WAttributes & uint16(0x0007)
	defaultBgColor := screenInfo.WAttributes & uint16(0x0070)
	defaultFgIntensity := screenInfo.WAttributes & uint16(0x0008)
	defaultBgIntensity := screenInfo.WAttributes & uint16(0x0080)

	fgParam := []testParam{
		{"foreground black  ", uint16(0x0000 | 0x0000), "30"},
		{"foreground red    ", uint16(0x0004 | 0x0000), "31"},
		{"foreground green  ", uint16(0x0002 | 0x0000), "32"},
		{"foreground yellow ", uint16(0x0006 | 0x0000), "33"},
		{"foreground blue   ", uint16(0x0001 | 0x0000), "34"},
		{"foreground magenta", uint16(0x0005 | 0x0000), "35"},
		{"foreground cyan   ", uint16(0x0003 | 0x0000), "36"},
		{"foreground white  ", uint16(0x0007 | 0x0000), "37"},
		{"foreground default", defaultFgColor | 0x0000, "39"},
	}

	bgParam := []testParam{
		{"background black  ", uint16(0x0007 | 0x0000), "40"},
		{"background red    ", uint16(0x0007 | 0x0040), "41"},
		{"background green  ", uint16(0x0007 | 0x0020), "42"},
		{"background yellow ", uint16(0x0007 | 0x0060), "43"},
		{"background blue   ", uint16(0x0007 | 0x0010), "44"},
		{"background magenta", uint16(0x0007 | 0x0050), "45"},
		{"background cyan   ", uint16(0x0007 | 0x0030), "46"},
		{"background white  ", uint16(0x0007 | 0x0070), "47"},
		{"background default", uint16(0x0007) | defaultBgColor, "49"},
	}

	resetParam := []testParam{
		{"all reset", defaultFgColor | defaultBgColor | defaultFgIntensity | defaultBgIntensity, "0"},
		{"all reset", defaultFgColor | defaultBgColor | defaultFgIntensity | defaultBgIntensity, ""},
	}

	boldParam := []testParam{
		{"bold on", uint16(0x0007 | 0x0008), "1"},
		{"bold off", uint16(0x0007), "21"},
	}

	underscoreParam := []testParam{
		{"underscore on", uint16(0x0007 | 0x8000), "4"},
		{"underscore off", uint16(0x0007), "24"},
	}

	blinkParam := []testParam{
		{"blink on", uint16(0x0007 | 0x0080), "5"},
		{"blink off", uint16(0x0007), "25"},
	}

	mixedParam := []testParam{
		{"both black,   bold, underline, blink", uint16(0x0000 | 0x0000 | 0x0008 | 0x8000 | 0x0080), "30;40;1;4;5"},
		{"both red,     bold, underline, blink", uint16(0x0004 | 0x0040 | 0x0008 | 0x8000 | 0x0080), "31;41;1;4;5"},
		{"both green,   bold, underline, blink", uint16(0x0002 | 0x0020 | 0x0008 | 0x8000 | 0x0080), "32;42;1;4;5"},
		{"both yellow,  bold, underline, blink", uint16(0x0006 | 0x0060 | 0x0008 | 0x8000 | 0x0080), "33;43;1;4;5"},
		{"both blue,    bold, underline, blink", uint16(0x0001 | 0x0010 | 0x0008 | 0x8000 | 0x0080), "34;44;1;4;5"},
		{"both magenta, bold, underline, blink", uint16(0x0005 | 0x0050 | 0x0008 | 0x8000 | 0x0080), "35;45;1;4;5"},
		{"both cyan,    bold, underline, blink", uint16(0x0003 | 0x0030 | 0x0008 | 0x8000 | 0x0080), "36;46;1;4;5"},
		{"both white,   bold, underline, blink", uint16(0x0007 | 0x0070 | 0x0008 | 0x8000 | 0x0080), "37;47;1;4;5"},
		{"both default, bold, underline, blink", uint16(defaultFgColor | defaultBgColor | 0x0008 | 0x8000 | 0x0080), "39;49;1;4;5"},
	}

	assertTextAttribute := func(expectedText string, expectedAttributes uint16, ansiColor string) {
		actualText, actualAttributes, err := writeAnsiColor(expectedText, ansiColor)
		if actualText != expectedText {
			t.Errorf("Get %s, want %s", actualText, expectedText)
		}
		if err != nil {
			t.Fatal("Could not get ConsoleScreenBufferInfo")
		}
		if actualAttributes != expectedAttributes {
			t.Errorf("Text: %s, Get 0x%04x, want 0x%04x", expectedText, actualAttributes, expectedAttributes)
		}
	}

	for _, v := range fgParam {
		ResetColor()
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range bgParam {
		ChangeColor(uint16(0x0070 | 0x0007))
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range resetParam {
		ChangeColor(uint16(0x0000 | 0x0070 | 0x0008))
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range boldParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range underscoreParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range blinkParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range mixedParam {
		ResetColor()
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}
}
<commit_msg>Fix golint suggestions<commit_after>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor_test

import (
	"bytes"
	"fmt"
	"syscall"
	"testing"

	"github.com/shiena/ansicolor"
	. "github.com/shiena/ansicolor"
)

func TestWritePlanText(t *testing.T) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)
	expected := "plain text"
	fmt.Fprintf(w, expected)
	actual := inner.String()
	if actual != expected {
		t.Errorf("Get %s, want %s", actual, expected)
	}
}

func TestWriteParseText(t *testing.T) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)

	inputTail := "\x1b[0mtail text"
	expectedTail := "tail text"
	fmt.Fprintf(w, inputTail)
	actualTail := inner.String()
	inner.Reset()
	if actualTail != expectedTail {
		t.Errorf("Get %s, want %s", actualTail, expectedTail)
	}

	inputHead := "head text\x1b[0m"
	expectedHead := "head text"
	fmt.Fprintf(w, inputHead)
	actualHead := inner.String()
	inner.Reset()
	if actualHead != expectedHead {
		t.Errorf("Get %s, want %s", actualHead, expectedHead)
	}

	inputBothEnds := "both ends \x1b[0m text"
	expectedBothEnds := "both ends  text"
	fmt.Fprintf(w, inputBothEnds)
	actualBothEnds := inner.String()
	inner.Reset()
	if actualBothEnds != expectedBothEnds {
		t.Errorf("Get %s, want %s", actualBothEnds, expectedBothEnds)
	}

	inputManyEsc := "\x1b\x1b\x1b\x1b[0m many esc"
	expectedManyEsc := "\x1b\x1b\x1b many esc"
	fmt.Fprintf(w, inputManyEsc)
	actualManyEsc := inner.String()
	inner.Reset()
	if actualManyEsc != expectedManyEsc {
		t.Errorf("Get %s, want %s", actualManyEsc, expectedManyEsc)
	}

	expectedSplit := "split  text"
	for _, ch := range "split \x1b[0m text" {
		fmt.Fprintf(w, string(ch))
	}
	actualSplit := inner.String()
	inner.Reset()
	if actualSplit != expectedSplit {
		t.Errorf("Get %s, want %s", actualSplit, expectedSplit)
	}
}

type screenNotFoundError struct {
	error
}

func writeAnsiColor(expectedText, colorCode string) (actualText string, actualAttributes uint16, err error) {
	inner := bytes.NewBufferString("")
	w := ansicolor.NewAnsiColorWriter(inner)
	fmt.Fprintf(w, "\x1b[%sm%s", colorCode, expectedText)

	actualText = inner.String()
	screenInfo := GetConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		actualAttributes = screenInfo.WAttributes
	} else {
		err = &screenNotFoundError{}
	}
	return
}

type testParam struct {
	text       string
	attributes uint16
	ansiColor  string
}

func TestWriteAnsiColorText(t *testing.T) {
	screenInfo := GetConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		t.Fatal("Could not get ConsoleScreenBufferInfo")
	}
	defer ChangeColor(screenInfo.WAttributes)
	defaultFgColor := screenInfo.WAttributes & uint16(0x0007)
	defaultBgColor := screenInfo.WAttributes & uint16(0x0070)
	defaultFgIntensity := screenInfo.WAttributes & uint16(0x0008)
	defaultBgIntensity := screenInfo.WAttributes & uint16(0x0080)

	fgParam := []testParam{
		{"foreground black  ", uint16(0x0000 | 0x0000), "30"},
		{"foreground red    ", uint16(0x0004 | 0x0000), "31"},
		{"foreground green  ", uint16(0x0002 | 0x0000), "32"},
		{"foreground yellow ", uint16(0x0006 | 0x0000), "33"},
		{"foreground blue   ", uint16(0x0001 | 0x0000), "34"},
		{"foreground magenta", uint16(0x0005 | 0x0000), "35"},
		{"foreground cyan   ", uint16(0x0003 | 0x0000), "36"},
		{"foreground white  ", uint16(0x0007 | 0x0000), "37"},
		{"foreground default", defaultFgColor | 0x0000, "39"},
	}

	bgParam := []testParam{
		{"background black  ", uint16(0x0007 | 0x0000), "40"},
		{"background red    ", uint16(0x0007 | 0x0040), "41"},
		{"background green  ", uint16(0x0007 | 0x0020), "42"},
		{"background yellow ", uint16(0x0007 | 0x0060), "43"},
		{"background blue   ", uint16(0x0007 | 0x0010), "44"},
		{"background magenta", uint16(0x0007 | 0x0050), "45"},
		{"background cyan   ", uint16(0x0007 | 0x0030), "46"},
		{"background white  ", uint16(0x0007 | 0x0070), "47"},
		{"background default", uint16(0x0007) | defaultBgColor, "49"},
	}

	resetParam := []testParam{
		{"all reset", defaultFgColor | defaultBgColor | defaultFgIntensity | defaultBgIntensity, "0"},
		{"all reset", defaultFgColor | defaultBgColor | defaultFgIntensity | defaultBgIntensity, ""},
	}

	boldParam := []testParam{
		{"bold on", uint16(0x0007 | 0x0008), "1"},
		{"bold off", uint16(0x0007), "21"},
	}

	underscoreParam := []testParam{
		{"underscore on", uint16(0x0007 | 0x8000), "4"},
		{"underscore off", uint16(0x0007), "24"},
	}

	blinkParam := []testParam{
		{"blink on", uint16(0x0007 | 0x0080), "5"},
		{"blink off", uint16(0x0007), "25"},
	}

	mixedParam := []testParam{
		{"both black,   bold, underline, blink", uint16(0x0000 | 0x0000 | 0x0008 | 0x8000 | 0x0080), "30;40;1;4;5"},
		{"both red,     bold, underline, blink", uint16(0x0004 | 0x0040 | 0x0008 | 0x8000 | 0x0080), "31;41;1;4;5"},
		{"both green,   bold, underline, blink", uint16(0x0002 | 0x0020 | 0x0008 | 0x8000 | 0x0080), "32;42;1;4;5"},
		{"both yellow,  bold, underline, blink", uint16(0x0006 | 0x0060 | 0x0008 | 0x8000 | 0x0080), "33;43;1;4;5"},
		{"both blue,    bold, underline, blink", uint16(0x0001 | 0x0010 | 0x0008 | 0x8000 | 0x0080), "34;44;1;4;5"},
		{"both magenta, bold, underline, blink", uint16(0x0005 | 0x0050 | 0x0008 | 0x8000 | 0x0080), "35;45;1;4;5"},
		{"both cyan,    bold, underline, blink", uint16(0x0003 | 0x0030 | 0x0008 | 0x8000 | 0x0080), "36;46;1;4;5"},
		{"both white,   bold, underline, blink", uint16(0x0007 | 0x0070 | 0x0008 | 0x8000 | 0x0080), "37;47;1;4;5"},
		{"both default, bold, underline, blink", uint16(defaultFgColor | defaultBgColor | 0x0008 | 0x8000 | 0x0080), "39;49;1;4;5"},
	}

	assertTextAttribute := func(expectedText string, expectedAttributes uint16, ansiColor string) {
		actualText, actualAttributes, err := writeAnsiColor(expectedText, ansiColor)
		if actualText != expectedText {
			t.Errorf("Get %s, want %s", actualText, expectedText)
		}
		if err != nil {
			t.Fatal("Could not get ConsoleScreenBufferInfo")
		}
		if actualAttributes != expectedAttributes {
			t.Errorf("Text: %s, Get 0x%04x, want 0x%04x", expectedText, actualAttributes, expectedAttributes)
		}
	}

	for _, v := range fgParam {
		ResetColor()
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range bgParam {
		ChangeColor(uint16(0x0070 | 0x0007))
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range resetParam {
		ChangeColor(uint16(0x0000 | 0x0070 | 0x0008))
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range boldParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range underscoreParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	ResetColor()
	for _, v := range blinkParam {
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}

	for _, v := range mixedParam {
		ResetColor()
		assertTextAttribute(v.text, v.attributes, v.ansiColor)
	}
}
<|endoftext|>


<commit_before>package php_serialize

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"strings"
)

func UnSerialize(s string) (PhpValue, error) {
	decoder := NewUnSerializer(s)
	decoder.SetSerializedDecodeFunc(SerializedDecodeFunc(UnSerialize))
	return decoder.Decode()
}

type UnSerializer struct {
	source     string
	r          *strings.Reader
	lastErr    error
	decodeFunc SerializedDecodeFunc
}

func NewUnSerializer(data string) *UnSerializer {
	return &UnSerializer{
		source: data,
	}
}

func (self *UnSerializer) SetReader(r *strings.Reader) {
	self.r = r
}

func (self *UnSerializer) SetSerializedDecodeFunc(f SerializedDecodeFunc) {
	self.decodeFunc = f
}

func (self *UnSerializer) Decode() (PhpValue, error) {
	if self.r == nil {
		self.r = strings.NewReader(self.source)
	}

	var value PhpValue

	if token, _, err := self.r.ReadRune(); err == nil {
		switch token {
		default:
			self.saveError(fmt.Errorf("php_serialize: Unknown token %#U", token))
		case TOKEN_NULL:
			value = self.decodeNull()
		case TOKEN_BOOL:
			value = self.decodeBool()
		case TOKEN_INT:
			value = self.decodeNumber(false)
		case TOKEN_FLOAT:
			value = self.decodeNumber(true)
		case TOKEN_STRING:
			value = self.decodeString(DELIMITER_STRING_LEFT, DELIMITER_STRING_RIGHT, true)
		case TOKEN_ARRAY:
			value = self.decodeArray()
		case TOKEN_OBJECT:
			value = self.decodeObject()
		case TOKEN_OBJECT_SERIALIZED:
			value = self.decodeSerialized()
		case TOKEN_REFERENCE, TOKEN_REFERENCE_OBJECT:
			value = self.decodeReference()

		}
	}

	return value, self.lastErr
}

func (self *UnSerializer) decodeNull() PhpValue {
	self.expect(SEPARATOR_VALUES)
	return nil
}

func (self *UnSerializer) decodeBool() PhpValue {
	var (
		raw rune
		err error
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, _, err = self.r.ReadRune(); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading bool value: %v", err))
	}

	self.expect(SEPARATOR_VALUES)
	return raw == '1'
}

func (self *UnSerializer) decodeNumber(isFloat bool) PhpValue {
	var (
		raw string
		err error
		val PhpValue
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, err = self.readUntil(SEPARATOR_VALUES); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading number value: %v", err))
	} else {
		if isFloat {
			if val, err = strconv.ParseFloat(raw, 64); err != nil {
				self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to float: %v", raw, err))
			}
		} else {
			if val, err = strconv.Atoi(raw); err != nil {
				self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to int: %v", raw, err))
			}
		}
	}

	return val
}

func (self *UnSerializer) decodeString(left, right rune, isFinal bool) PhpValue {
	var (
		err     error
		val     PhpValue
		strLen  int
		readLen int
	)

	strLen = self.readLen()
	self.expect(left)

	if strLen > 0 {
		buf := make([]byte, strLen, strLen)
		if readLen, err = self.r.Read(buf); err != nil {
			self.saveError(fmt.Errorf("php_serialize: Error while reading string value: %v", err))
		} else {
			if readLen != strLen {
				self.saveError(fmt.Errorf("php_serialize: Unable to read string. Expected %d but have got %d bytes", strLen, readLen))
			} else {
				val = string(buf)
			}
		}
	}

	self.expect(right)
	if isFinal {
		self.expect(SEPARATOR_VALUES)
	}
	return val
}

func (self *UnSerializer) decodeArray() PhpValue {
	var arrLen int
	val := make(PhpArray)

	arrLen = self.readLen()
	self.expect(DELIMITER_OBJECT_LEFT)

	for i := 0; i < arrLen; i++ {
		k, errKey := self.Decode()
		v, errVal := self.Decode()

		if errKey == nil && errVal == nil {
			val[k] = v
			/*switch t := k.(type) {
			default:
				self.saveError(fmt.Errorf("php_serialize: Unexpected key type %T", t))
			case string:
				stringKey, _ := k.(string)
				val[stringKey] = v
			case int:
				intKey, _ := k.(int)
				val[strconv.Itoa(intKey)] = v
			}*/
		} else {
			self.saveError(fmt.Errorf("php_serialize: Error while reading key or(and) value of array"))
		}
	}

	self.expect(DELIMITER_OBJECT_RIGHT)
	return val
}

func (self *UnSerializer) decodeObject() PhpValue {
	val := &PhpObject{
		className: self.readClassName(),
	}

	rawMembers := self.decodeArray()
	val.members, _ = rawMembers.(PhpArray)

	return val
}

func (self *UnSerializer) decodeSerialized() PhpValue {
	val := &PhpObjectSerialized{
		className: self.readClassName(),
	}

	rawData := self.decodeString(DELIMITER_OBJECT_LEFT, DELIMITER_OBJECT_RIGHT, false)
	val.data, _ = rawData.(string)

	if self.decodeFunc != nil && val.data != "" {
		var err error
		if val.value, err = self.decodeFunc(val.data); err != nil {
			self.saveError(err)
		}
	}

	return val
}

func (self *UnSerializer) decodeReference() PhpValue {
	self.expect(SEPARATOR_VALUE_TYPE)
	if _, err := self.readUntil(SEPARATOR_VALUES); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading reference value: %v", err))
	}
	return nil
}

func (self *UnSerializer) expect(expected rune) {
	if token, _, err := self.r.ReadRune(); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading expected rune %#U: %v", expected, err))
	} else if token != expected {
		if debugMode {
			log.Printf("php_serialize: source\n%s\n", self.source)
			log.Printf("php_serialize: reader info\n%#v\n", self.r)
		}
		self.saveError(fmt.Errorf("php_serialize: Expected %#U but have got %#U", expected, token))
	}
}

func (self *UnSerializer) readUntil(stop rune) (string, error) {
	var (
		token rune
		err   error
	)
	buf := bytes.NewBuffer([]byte{})

	for {
		if token, _, err = self.r.ReadRune(); err != nil || token == stop {
			break
		} else {
			buf.WriteRune(token)
		}
	}

	return buf.String(), err
}

func (self *UnSerializer) readLen() int {
	var (
		raw string
		err error
		val int
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, err = self.readUntil(SEPARATOR_VALUE_TYPE); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading lenght of value: %v", err))
	} else {
		if val, err = strconv.Atoi(raw); err != nil {
			self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to int: %v", raw, err))
		}
	}
	return val
}

func (self *UnSerializer) readClassName() (res string) {
	rawClass := self.decodeString(DELIMITER_STRING_LEFT, DELIMITER_STRING_RIGHT, false)
	res, _ = rawClass.(string)
	return
}

func (self *UnSerializer) saveError(err error) {
	if self.lastErr == nil {
		self.lastErr = err
	}
}
<commit_msg>fix bug discovered in #15<commit_after>package php_serialize

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"strings"
)

const UNSERIAZABLE_OBJECT_MAX_LEN = 10 * 1024 * 1024 * 1024

func UnSerialize(s string) (PhpValue, error) {
	decoder := NewUnSerializer(s)
	decoder.SetSerializedDecodeFunc(SerializedDecodeFunc(UnSerialize))
	return decoder.Decode()
}

type UnSerializer struct {
	source     string
	r          *strings.Reader
	lastErr    error
	decodeFunc SerializedDecodeFunc
}

func NewUnSerializer(data string) *UnSerializer {
	return &UnSerializer{
		source: data,
	}
}

func (self *UnSerializer) SetReader(r *strings.Reader) {
	self.r = r
}

func (self *UnSerializer) SetSerializedDecodeFunc(f SerializedDecodeFunc) {
	self.decodeFunc = f
}

func (self *UnSerializer) Decode() (PhpValue, error) {
	if self.r == nil {
		self.r = strings.NewReader(self.source)
	}

	var value PhpValue

	if token, _, err := self.r.ReadRune(); err == nil {
		switch token {
		default:
			self.saveError(fmt.Errorf("php_serialize: Unknown token %#U", token))
		case TOKEN_NULL:
			value = self.decodeNull()
		case TOKEN_BOOL:
			value = self.decodeBool()
		case TOKEN_INT:
			value = self.decodeNumber(false)
		case TOKEN_FLOAT:
			value = self.decodeNumber(true)
		case TOKEN_STRING:
			value = self.decodeString(DELIMITER_STRING_LEFT, DELIMITER_STRING_RIGHT, true)
		case TOKEN_ARRAY:
			value = self.decodeArray()
		case TOKEN_OBJECT:
			value = self.decodeObject()
		case TOKEN_OBJECT_SERIALIZED:
			value = self.decodeSerialized()
		case TOKEN_REFERENCE, TOKEN_REFERENCE_OBJECT:
			value = self.decodeReference()

		}
	}

	return value, self.lastErr
}

func (self *UnSerializer) decodeNull() PhpValue {
	self.expect(SEPARATOR_VALUES)
	return nil
}

func (self *UnSerializer) decodeBool() PhpValue {
	var (
		raw rune
		err error
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, _, err = self.r.ReadRune(); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading bool value: %v", err))
	}

	self.expect(SEPARATOR_VALUES)
	return raw == '1'
}

func (self *UnSerializer) decodeNumber(isFloat bool) PhpValue {
	var (
		raw string
		err error
		val PhpValue
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, err = self.readUntil(SEPARATOR_VALUES); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading number value: %v", err))
	} else {
		if isFloat {
			if val, err = strconv.ParseFloat(raw, 64); err != nil {
				self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to float: %v", raw, err))
			}
		} else {
			if val, err = strconv.Atoi(raw); err != nil {
				self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to int: %v", raw, err))
			}
		}
	}

	return val
}

func (self *UnSerializer) decodeString(left, right rune, isFinal bool) PhpValue {
	var (
		err     error
		val     PhpValue
		strLen  int
		readLen int
	)

	strLen = self.readLen()
	self.expect(left)

	if strLen > 0 {
		buf := make([]byte, strLen, strLen)
		if readLen, err = self.r.Read(buf); err != nil {
			self.saveError(fmt.Errorf("php_serialize: Error while reading string value: %v", err))
		} else {
			if readLen != strLen {
				self.saveError(fmt.Errorf("php_serialize: Unable to read string. Expected %d but have got %d bytes", strLen, readLen))
			} else {
				val = string(buf)
			}
		}
	}

	self.expect(right)
	if isFinal {
		self.expect(SEPARATOR_VALUES)
	}
	return val
}

func (self *UnSerializer) decodeArray() PhpValue {
	var arrLen int
	val := make(PhpArray)

	arrLen = self.readLen()
	self.expect(DELIMITER_OBJECT_LEFT)

	for i := 0; i < arrLen; i++ {
		k, errKey := self.Decode()
		v, errVal := self.Decode()

		if errKey == nil && errVal == nil {
			val[k] = v
			/*switch t := k.(type) {
			default:
				self.saveError(fmt.Errorf("php_serialize: Unexpected key type %T", t))
			case string:
				stringKey, _ := k.(string)
				val[stringKey] = v
			case int:
				intKey, _ := k.(int)
				val[strconv.Itoa(intKey)] = v
			}*/
		} else {
			self.saveError(fmt.Errorf("php_serialize: Error while reading key or(and) value of array"))
		}
	}

	self.expect(DELIMITER_OBJECT_RIGHT)
	return val
}

func (self *UnSerializer) decodeObject() PhpValue {
	val := &PhpObject{
		className: self.readClassName(),
	}

	rawMembers := self.decodeArray()
	val.members, _ = rawMembers.(PhpArray)

	return val
}

func (self *UnSerializer) decodeSerialized() PhpValue {
	val := &PhpObjectSerialized{
		className: self.readClassName(),
	}

	rawData := self.decodeString(DELIMITER_OBJECT_LEFT, DELIMITER_OBJECT_RIGHT, false)
	val.data, _ = rawData.(string)

	if self.decodeFunc != nil && val.data != "" {
		var err error
		if val.value, err = self.decodeFunc(val.data); err != nil {
			self.saveError(err)
		}
	}

	return val
}

func (self *UnSerializer) decodeReference() PhpValue {
	self.expect(SEPARATOR_VALUE_TYPE)
	if _, err := self.readUntil(SEPARATOR_VALUES); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading reference value: %v", err))
	}
	return nil
}

func (self *UnSerializer) expect(expected rune) {
	if token, _, err := self.r.ReadRune(); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading expected rune %#U: %v", expected, err))
	} else if token != expected {
		if debugMode {
			log.Printf("php_serialize: source\n%s\n", self.source)
			log.Printf("php_serialize: reader info\n%#v\n", self.r)
		}
		self.saveError(fmt.Errorf("php_serialize: Expected %#U but have got %#U", expected, token))
	}
}

func (self *UnSerializer) readUntil(stop rune) (string, error) {
	var (
		token rune
		err   error
	)
	buf := bytes.NewBuffer([]byte{})

	for {
		if token, _, err = self.r.ReadRune(); err != nil || token == stop {
			break
		} else {
			buf.WriteRune(token)
		}
	}

	return buf.String(), err
}

func (self *UnSerializer) readLen() int {
	var (
		raw string
		err error
		val int
	)
	self.expect(SEPARATOR_VALUE_TYPE)

	if raw, err = self.readUntil(SEPARATOR_VALUE_TYPE); err != nil {
		self.saveError(fmt.Errorf("php_serialize: Error while reading lenght of value: %v", err))
	} else {
		if val, err = strconv.Atoi(raw); err != nil {
			self.saveError(fmt.Errorf("php_serialize: Unable to convert %s to int: %v", raw, err))
		} else if val > UNSERIAZABLE_OBJECT_MAX_LEN {
			self.saveError(fmt.Errorf("php_serialize: Unserializable object length looks too big(%d). If you are sure you wanna unserialise it, please increase UNSERIAZABLE_OBJECT_MAX_LEN const", val, err))
			val = 0
		}
	}
	return val
}

func (self *UnSerializer) readClassName() (res string) {
	rawClass := self.decodeString(DELIMITER_STRING_LEFT, DELIMITER_STRING_RIGHT, false)
	res, _ = rawClass.(string)
	return
}

func (self *UnSerializer) saveError(err error) {
	if self.lastErr == nil {
		self.lastErr = err
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/krolaw/dhcp4"
)

// DHCPService is the DHCP server instance
type DHCPService struct {
	ip                net.IP
	authoritativePool *net.IPNet
	guestPool         *net.IPNet    // must be within authoritativePool (at least for now)
	leaseDuration     time.Duration // FIXME: make a separate duration per pool?
	defaultOptions    dhcp4.Options // FIXME: make different options per pool?
	etcdClient        *etcd.Client
}

func dhcpSetup(etc *etcd.Client) chan bool {
	etc.CreateDir("dhcp", 0)
	exit := make(chan bool, 1)
	serverIP := net.ParseIP("172.16.193.1")
	_, authoritativePool, _ := net.ParseCIDR("172.16.193.0/24")
	guestPool := authoritativePool
	go func() {
		fmt.Println(dhcp4.ListenAndServeIf("vmnet2", &DHCPService{
			ip:                serverIP.To4(),
			leaseDuration:     time.Hour * 12,
			etcdClient:        etc,
			authoritativePool: authoritativePool,
			guestPool:         guestPool,
			defaultOptions: dhcp4.Options{
				dhcp4.OptionSubnetMask:       net.ParseIP("255.255.255.0").To4(),
				dhcp4.OptionRouter:           serverIP.To4(),
				dhcp4.OptionDomainNameServer: serverIP.To4(),
			},
		}))
		exit <- true
	}()
	return exit
}

// ServeDHCP is called by dhcp4.ListenAndServe when the service is started
func (d *DHCPService) ServeDHCP(packet dhcp4.Packet, msgType dhcp4.MessageType, reqOptions dhcp4.Options) (response dhcp4.Packet) {
	switch msgType {
	case dhcp4.Discover:
		// FIXME: send to StatHat and/or increment a counter
		mac := packet.CHAddr()
		fmt.Printf("DHCP Discover from %s\n", mac.String())
		ip := d.getIPFromMAC(mac)
		if ip != nil {
			options := d.getOptionsFromMAC(mac)
			fmt.Printf("DHCP Discover from %s (we return %s)\n", mac.String(), ip.String())
			for x, y := range reqOptions {
				fmt.Printf("\tR[%v] %v %s\n", x, y, y)
			}
			for x, y := range options {
				fmt.Printf("\tO[%v] %v %s\n", x, y, y)
			}
			return dhcp4.ReplyPacket(packet, dhcp4.Offer, d.ip.To4(), ip.To4(), d.leaseDuration, options.SelectOrderOrAll(reqOptions[dhcp4.OptionParameterRequestList]))
		}
		return nil
	case dhcp4.Request:
		// FIXME: send to StatHat and/or increment a counter
		mac := packet.CHAddr()
		fmt.Printf("DHCP Request from %s...\n", mac.String())
		if requestedIP := net.IP(reqOptions[dhcp4.OptionRequestedIPAddress]); len(requestedIP) == 4 { // valid and IPv4
			fmt.Printf("DHCP Request from %s wanting %s\n", mac.String(), requestedIP.String())
			ip := d.getIPFromMAC(mac)
			if ip.Equal(requestedIP) {
				options := d.getOptionsFromMAC(mac)
				fmt.Printf("DHCP Request from %s wanting %s (we agree)\n", mac.String(), requestedIP.String())
				for x, y := range reqOptions {
					fmt.Printf("\tR[%v] %v %s\n", x, y, y)
				}
				for x, y := range options {
					fmt.Printf("\tO[%v] %v %s\n", x, y, y)
				}
				return dhcp4.ReplyPacket(packet, dhcp4.ACK, d.ip.To4(), requestedIP.To4(), d.leaseDuration, options.SelectOrderOrAll(reqOptions[dhcp4.OptionParameterRequestList]))
			}
		}
		return dhcp4.ReplyPacket(packet, dhcp4.NAK, d.ip.To4(), nil, 0, nil)
	case dhcp4.Release:
		// FIXME: release from DB?  tick a flag?  increment a counter?  send to StatHat?
		mac := packet.CHAddr()
		fmt.Printf("DHCP Release from %s\n", mac.String())
	case dhcp4.Decline:
		// FIXME: release from DB?  tick a flag?  increment a counter?  send to StatHat?
		mac := packet.CHAddr()
		fmt.Printf("DHCP Decline from %s\n", mac.String())
	}
	return nil
}

func (d *DHCPService) getIPFromMAC(mac net.HardwareAddr) net.IP {
	response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ip", false, false)
	if response != nil && response.Node != nil {
		ip := net.ParseIP(response.Node.Value)
		if ip != nil {
			d.etcdClient.Set("dhcp/"+ip.String(), mac.String(), uint64(d.leaseDuration.Seconds()+0.5))
			return ip
		}
	}

	// TODO: determine whether or not this MAC should be permitted to get an IP at all (blacklist? whitelist?)

	// locate an unused IP address (can this be more efficient?  yes!  FIXME)
	var ip net.IP
	for testIP := dhcp4.IPAdd(d.guestPool.IP, 1); d.guestPool.Contains(testIP); testIP = dhcp4.IPAdd(testIP, 1) {
		fmt.Println(testIP.String())
		response, _ := d.etcdClient.Get("dhcp/"+testIP.String(), false, false)
		if response == nil || response.Node == nil { // this means that the IP is not already occupied
			ip = testIP
			break
		}
	}

	if ip != nil { // if nil then we're out of IP addresses!
		d.etcdClient.CreateDir("dhcp/"+mac.String(), 0)
		d.etcdClient.Set("dhcp/"+ip.String(), mac.String(), uint64(d.leaseDuration.Seconds()+0.5))
		d.etcdClient.Set("dhcp/"+mac.String()+"/ip", ip.String(), uint64(d.leaseDuration.Seconds()+0.5))

		// TODO: write this new lease to the database under both "dhcp/mac/$MAC" and a MAC pointer to "dhcp/ip/$IP"
		// TODO: determine dhcp4.options based on server default config + MAC config
		// TODO: return valid IP as expected

		return ip
	}

	return nil
}

func (d *DHCPService) getOptionsFromMAC(mac net.HardwareAddr) dhcp4.Options {
	options := d.defaultOptions

	{ // Subnet Mask
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/mask", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionSubnetMask)
			} else {
				options[dhcp4.OptionSubnetMask] = []byte(response.Node.Value)
			}
		}
	}

	{ // Gateway/Router
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/gw", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionRouter)
			} else {
				options[dhcp4.OptionRouter] = []byte(response.Node.Value)
			}
		}
	}

	{ // Name Server
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ns", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionDomainNameServer)
			} else {
				options[dhcp4.OptionDomainNameServer] = []byte(response.Node.Value)
			}
		}
	}

	{ // Host Name
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/name", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionHostName)
			} else {
				options[dhcp4.OptionHostName] = []byte(response.Node.Value)
			}
		}
	}

	{ // Domain Name
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/domain", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionDomainName)
			} else {
				options[dhcp4.OptionDomainName] = []byte(response.Node.Value)
			}
		}
	}

	{ // Domain Name
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/broadcast", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionBroadcastAddress)
			} else {
				options[dhcp4.OptionBroadcastAddress] = []byte(response.Node.Value)
			}
		}
	}

	{ // NTP Server
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ntp", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionNetworkTimeProtocolServers)
			} else {
				options[dhcp4.OptionNetworkTimeProtocolServers] = []byte(response.Node.Value)
			}
		}
	}

	return options
}
<commit_msg>Fixed comment in dhcp.go<commit_after>package main

import (
	"fmt"
	"net"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/krolaw/dhcp4"
)

// DHCPService is the DHCP server instance
type DHCPService struct {
	ip                net.IP
	authoritativePool *net.IPNet
	guestPool         *net.IPNet    // must be within authoritativePool (at least for now)
	leaseDuration     time.Duration // FIXME: make a separate duration per pool?
	defaultOptions    dhcp4.Options // FIXME: make different options per pool?
	etcdClient        *etcd.Client
}

func dhcpSetup(etc *etcd.Client) chan bool {
	etc.CreateDir("dhcp", 0)
	exit := make(chan bool, 1)
	serverIP := net.ParseIP("172.16.193.1")
	_, authoritativePool, _ := net.ParseCIDR("172.16.193.0/24")
	guestPool := authoritativePool
	go func() {
		fmt.Println(dhcp4.ListenAndServeIf("vmnet2", &DHCPService{
			ip:                serverIP.To4(),
			leaseDuration:     time.Hour * 12,
			etcdClient:        etc,
			authoritativePool: authoritativePool,
			guestPool:         guestPool,
			defaultOptions: dhcp4.Options{
				dhcp4.OptionSubnetMask:       net.ParseIP("255.255.255.0").To4(),
				dhcp4.OptionRouter:           serverIP.To4(),
				dhcp4.OptionDomainNameServer: serverIP.To4(),
			},
		}))
		exit <- true
	}()
	return exit
}

// ServeDHCP is called by dhcp4.ListenAndServe when the service is started
func (d *DHCPService) ServeDHCP(packet dhcp4.Packet, msgType dhcp4.MessageType, reqOptions dhcp4.Options) (response dhcp4.Packet) {
	switch msgType {
	case dhcp4.Discover:
		// FIXME: send to StatHat and/or increment a counter
		mac := packet.CHAddr()
		fmt.Printf("DHCP Discover from %s\n", mac.String())
		ip := d.getIPFromMAC(mac)
		if ip != nil {
			options := d.getOptionsFromMAC(mac)
			fmt.Printf("DHCP Discover from %s (we return %s)\n", mac.String(), ip.String())
			for x, y := range reqOptions {
				fmt.Printf("\tR[%v] %v %s\n", x, y, y)
			}
			for x, y := range options {
				fmt.Printf("\tO[%v] %v %s\n", x, y, y)
			}
			return dhcp4.ReplyPacket(packet, dhcp4.Offer, d.ip.To4(), ip.To4(), d.leaseDuration, options.SelectOrderOrAll(reqOptions[dhcp4.OptionParameterRequestList]))
		}
		return nil
	case dhcp4.Request:
		// FIXME: send to StatHat and/or increment a counter
		mac := packet.CHAddr()
		fmt.Printf("DHCP Request from %s...\n", mac.String())
		if requestedIP := net.IP(reqOptions[dhcp4.OptionRequestedIPAddress]); len(requestedIP) == 4 { // valid and IPv4
			fmt.Printf("DHCP Request from %s wanting %s\n", mac.String(), requestedIP.String())
			ip := d.getIPFromMAC(mac)
			if ip.Equal(requestedIP) {
				options := d.getOptionsFromMAC(mac)
				fmt.Printf("DHCP Request from %s wanting %s (we agree)\n", mac.String(), requestedIP.String())
				for x, y := range reqOptions {
					fmt.Printf("\tR[%v] %v %s\n", x, y, y)
				}
				for x, y := range options {
					fmt.Printf("\tO[%v] %v %s\n", x, y, y)
				}
				return dhcp4.ReplyPacket(packet, dhcp4.ACK, d.ip.To4(), requestedIP.To4(), d.leaseDuration, options.SelectOrderOrAll(reqOptions[dhcp4.OptionParameterRequestList]))
			}
		}
		return dhcp4.ReplyPacket(packet, dhcp4.NAK, d.ip.To4(), nil, 0, nil)
	case dhcp4.Release:
		// FIXME: release from DB?  tick a flag?  increment a counter?  send to StatHat?
		mac := packet.CHAddr()
		fmt.Printf("DHCP Release from %s\n", mac.String())
	case dhcp4.Decline:
		// FIXME: release from DB?  tick a flag?  increment a counter?  send to StatHat?
		mac := packet.CHAddr()
		fmt.Printf("DHCP Decline from %s\n", mac.String())
	}
	return nil
}

func (d *DHCPService) getIPFromMAC(mac net.HardwareAddr) net.IP {
	response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ip", false, false)
	if response != nil && response.Node != nil {
		ip := net.ParseIP(response.Node.Value)
		if ip != nil {
			d.etcdClient.Set("dhcp/"+ip.String(), mac.String(), uint64(d.leaseDuration.Seconds()+0.5))
			return ip
		}
	}

	// TODO: determine whether or not this MAC should be permitted to get an IP at all (blacklist? whitelist?)

	// locate an unused IP address (can this be more efficient?  yes!  FIXME)
	var ip net.IP
	for testIP := dhcp4.IPAdd(d.guestPool.IP, 1); d.guestPool.Contains(testIP); testIP = dhcp4.IPAdd(testIP, 1) {
		fmt.Println(testIP.String())
		response, _ := d.etcdClient.Get("dhcp/"+testIP.String(), false, false)
		if response == nil || response.Node == nil { // this means that the IP is not already occupied
			ip = testIP
			break
		}
	}

	if ip != nil { // if nil then we're out of IP addresses!
		d.etcdClient.CreateDir("dhcp/"+mac.String(), 0)
		d.etcdClient.Set("dhcp/"+ip.String(), mac.String(), uint64(d.leaseDuration.Seconds()+0.5))
		d.etcdClient.Set("dhcp/"+mac.String()+"/ip", ip.String(), uint64(d.leaseDuration.Seconds()+0.5))

		// TODO: write this new lease to the database under both "dhcp/mac/$MAC" and a MAC pointer to "dhcp/ip/$IP"
		// TODO: determine dhcp4.options based on server default config + MAC config
		// TODO: return valid IP as expected

		return ip
	}

	return nil
}

func (d *DHCPService) getOptionsFromMAC(mac net.HardwareAddr) dhcp4.Options {
	options := d.defaultOptions

	{ // Subnet Mask
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/mask", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionSubnetMask)
			} else {
				options[dhcp4.OptionSubnetMask] = []byte(response.Node.Value)
			}
		}
	}

	{ // Gateway/Router
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/gw", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionRouter)
			} else {
				options[dhcp4.OptionRouter] = []byte(response.Node.Value)
			}
		}
	}

	{ // Name Server
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ns", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionDomainNameServer)
			} else {
				options[dhcp4.OptionDomainNameServer] = []byte(response.Node.Value)
			}
		}
	}

	{ // Host Name
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/name", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionHostName)
			} else {
				options[dhcp4.OptionHostName] = []byte(response.Node.Value)
			}
		}
	}

	{ // Domain Name
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/domain", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionDomainName)
			} else {
				options[dhcp4.OptionDomainName] = []byte(response.Node.Value)
			}
		}
	}

	{ // Broadcast Address
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/broadcast", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionBroadcastAddress)
			} else {
				options[dhcp4.OptionBroadcastAddress] = []byte(response.Node.Value)
			}
		}
	}

	{ // NTP Server
		response, _ := d.etcdClient.Get("dhcp/"+mac.String()+"/ntp", false, false)
		if response != nil && response.Node != nil {
			if response.Node.Value == "" {
				delete(options, dhcp4.OptionNetworkTimeProtocolServers)
			} else {
				options[dhcp4.OptionNetworkTimeProtocolServers] = []byte(response.Node.Value)
			}
		}
	}

	return options
}
<|endoftext|>


<commit_before>package gophers

import(
	"fmt"
	"testing"
	"net/http"
	"strings"
)

//gopherErrorDAO returns errors for any call
type gopherErrorDAO struct {
}
func (g *gopherErrorDAO) Spawn(gopher *Gopher) error {
	return fmt.Errorf("Invalid input")
}
func (g *gopherErrorDAO) Update(gopher *Gopher) error {
	return fmt.Errorf("invalid gopher")
}
func (g *gopherErrorDAO) GetAll() ([]Gopher, error) {
	return nil, fmt.Errorf("The gophers are revolting")
}
func (g *gopherErrorDAO) Get(gopherId string) (*Gopher, error) {
	return nil, fmt.Errorf("No gopher found")
}

func (g *gopherErrorDAO) Die(gopherId string) error {
	return fmt.Errorf("Gopher already dead")
}

func TestApiPostServerError(t *testing.T) {
	gopherApi := &GopherApi{Dao: &gopherErrorDAO{}}
	rdr := strings.NewReader(`{ "name": "diane" }`)
	r, err := http.NewRequest("POST", "/gophers", rdr)
	r.Header.Add("Content-Type", "application/json")
	if err != nil {
		t.Fatalf("request error: %v", err)
	}
	req, resp, w := wrapRecorder(r)
	gopherApi.Post(req, resp)
	if w.Code != http.StatusInternalServerError {
		t.Errorf("Wrong status code: %d", w.Code)
	}
	t.Logf("Status: %d, Response data: %v", w.Code, w.Body)

}
<commit_msg>converting gophertron to use go-restful + swagger instead of Martini<commit_after>package gophers

import(
	"fmt"
	"testing"
	"net/http"
	"strings"
)

//gopherErrorDAO returns errors for any call
type gopherErrorDAO struct {
}
func (g *gopherErrorDAO) Spawn(gopher *Gopher) error {
	return fmt.Errorf("Invalid input")
}
func (g *gopherErrorDAO) Update(gopher *Gopher) error {
	return fmt.Errorf("invalid gopher")
}
func (g *gopherErrorDAO) GetAll() ([]Gopher, error) {
	return nil, fmt.Errorf("The gophers are revolting")
}
func (g *gopherErrorDAO) Get(gopherId string) (*Gopher, error) {
	return nil, fmt.Errorf("No gopher found")
}

func (g *gopherErrorDAO) Die(gopherId string) error {
	return fmt.Errorf("Gopher already dead")
}

func TestApiPostServerError(t *testing.T) {
	gopherApi := &GopherApi{Dao: &gopherErrorDAO{}}
	rdr := strings.NewReader(`{ "name": "diane" }`)
	r, err := http.NewRequest("POST", "/gophers", rdr)
	r.Header.Add("Content-Type", "application/json")
	if err != nil {
		t.Fatalf("request error: %v", err)
	}
	req, resp, w := newRecorder(r)
	gopherApi.Post(req, resp)
	if w.Code != http.StatusInternalServerError {
		t.Errorf("Wrong status code: %d", w.Code)
	}
	t.Logf("Status: %d, Response data: %v", w.Code, w.Body)

}
<|endoftext|>


<commit_before>package hey

import (
	"os"
	"time"

	uuid "github.com/satori/go.uuid"
	tarantool "github.com/tarantool/go-tarantool"
)

var (
	DefaultTimeout   time.Duration = 500 * time.Millisecond
	DefaultReconnect time.Duration = 1 * time.Second

	DefaultTarantoolPrefix = "chronograph_"

	threadsSpace    = DefaultTarantoolPrefix + "threads"
	observerSpace   = DefaultTarantoolPrefix + "subscriptions"
	threadLineSpace = DefaultTarantoolPrefix + "threadline"
	eventsSpace     = DefaultTarantoolPrefix + "events"
	usersSpace      = DefaultTarantoolPrefix + "users"
	sourcesSpace    = DefaultTarantoolPrefix + "sources"
	relatedSpace    = DefaultTarantoolPrefix + "related"
)

type TarantoolOpts struct {
	Server        string
	Timeout       time.Duration
	Reconnect     time.Duration
	MaxReconnects uint
	User          string
	Pass          string
}

func (t TarantoolOpts) toOpts() tarantool.Opts {
	return tarantool.Opts{
		Timeout:       t.Timeout,
		Reconnect:     t.Reconnect,
		MaxReconnects: t.MaxReconnects,
		User:          t.User,
		Pass:          t.Pass,
	}
}

func setupFromENV() (opts TarantoolOpts, err error) {
	// 127.0.0.1:3013
	return TarantoolOpts{
		Server:        os.Getenv("TARANTOOL_SERVER"),
		Timeout:       DefaultTimeout,
		Reconnect:     DefaultReconnect,
		MaxReconnects: 3,
		User:          os.Getenv("TARANTOOL_USER_NAME"),
		Pass:          os.Getenv("TARANTOOL_USER_PASSWORD"),
	}, err
}

var _ Manager = &TarantoolManager{}

// TarantoolManager main struct
type TarantoolManager struct {
	conn *tarantool.Connection
}

// NewTarantoolManager return manager setupped from env or optsё
func NewTarantoolManager(opts ...TarantoolOpts) (*TarantoolManager, error) {

	var (
		opt TarantoolOpts
		err error
	)

	if len(opts) > 0 {
		opt = opts[0]
	} else {
		opt, err = setupFromENV()
		if err != nil {
			return nil, err
		}
	}

	client, err := tarantool.Connect(opt.Server, opt.toOpts())
	if err != nil {
		return nil, err
	}

	_, err = client.Ping()

	return &TarantoolManager{conn: client}, err
}

// util
func (m *TarantoolManager) get(space string, keyName string, key interface{}, target interface{}) error {
	return m.conn.SelectTyped(space, keyName, 0, 1, tarantool.IterEq, key, &target)
}

// GetThread return thread by ther id
func (m *TarantoolManager) GetThread(threadID string) (*Thread, error) {
	var (
		thread = new(Thread)
	)
	err := m.conn.SelectTyped(threadsSpace, "primary", 0, 1, tarantool.IterEq,
		makeKey(threadID), thread)
	return thread, err
}

// NewThread create new thread
// it return error is thread already exists
// if id empty, will be generated uuid
func (m *TarantoolManager) NewThread(thread *Thread) (err error) {
	// err = m.conn.SelectTyped(threadsSpace, "primary", 0, 1,
	// 	tarantool.IterEq, makeIndex(threadID), &thread)
	// if err != nil {
	// 	return
	// }
	// // thread already exists
	// if thread.ThreadID == threadID {
	// 	return ErrAlreadyExists
	// }
	if thread.ThreadID == "" {
		thread.ThreadID = uuid.NewV4().String()
	}

	_, err = m.conn.Insert(threadsSpace, thread)
	if err != nil {
		return
	}

	return nil
}

func (m *TarantoolManager) UpdateThread(thread *Thread) error {
	_, err := m.conn.Replace(threadsSpace, thread)
	return err
}

// DeleteThread удаляет трэд, подписки и связанные события
// 1. Удаляем все записи из events
// 2. Удаляем все подписки на трэд
func (m *TarantoolManager) DeleteThread(threadID string) error {
	// delete
	m.conn.Delete(threadsSpace, "primary", makeKey(threadID))
	return nil
}

// AddSource подписывает трэд на другой трэд
func (m *TarantoolManager) AddSource(dstThread, sourceThread string) (err error) {
	src := Sources{SourceThreadID: sourceThread, TargetThreadID: dstThread}
	_, err = m.conn.Insert(sourcesSpace, &src)
	if err != nil {
		return
	}
	return nil
}

func (m *TarantoolManager) GetSources(referThreadID string, offset, limit uint32) (threads []Thread, err error) {
	var sources []Sources
	err = m.conn.SelectTyped(sourcesSpace, "primary", offset, limit, tarantool.IterReq, makeKey(referThreadID), &sources)
	if err != nil {
		return
	}
	for _, src := range sources {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(src.SourceThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

// Трэды, у которого данный трэд в источниках
func (m *TarantoolManager) GetRefers(threadID string, offset, limit uint32) (threads []Thread, err error) {
	var sources []Sources
	err = m.conn.SelectTyped(sourcesSpace, "sources_idx", offset, limit, tarantool.IterReq, makeKey(threadID), &sources)
	if err != nil {
		return
	}
	for _, src := range sources {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(src.SourceThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

// RemoveSource отписывает трэд
func (m *TarantoolManager) RemoveSource(dstThread, sourceThread string) (err error) {
	_, err = m.conn.Delete(sourcesSpace, "primary", makeKey(dstThread, sourceThread))
	if err != nil {
		return err
	}
	return nil
}

// Observe - подписка конкретного юзера на трэд
func (m *TarantoolManager) Observe(userID, threadID string) error {
	obs := Observer{
		UserID:            userID,
		ThreadID:          threadID,
		LastDeliveredTime: time.Now(),
	}
	_, err := m.conn.Insert(observerSpace, &obs)
	return err
}

// Ignore удаляет подписку на трэд
func (m *TarantoolManager) Ignore(userID, threadID string) (err error) {

	_, err = m.conn.Delete(observerSpace, "primary", makeKey(threadID, userID))
	if err != nil {
		return err
	}

	return err
}

// Observers возвращает подписчиков трэда
func (m *TarantoolManager) Observers(threadID string, offset, limit uint32) (users []User, err error) {
	var obs []Observer
	err = m.conn.SelectTyped(observerSpace, "primary", offset, limit, tarantool.IterEq, makeKey(threadID), &obs)
	if err != nil {
		return
	}

	for _, o := range obs {
		var u []User
		err = m.get(usersSpace, "primary", makeKey(o.UserID), &u)
		if err != nil {
			return
		}
		users = append(users, u[0])
	}
	return
}

// Observes возвращает подписки юзера
func (m *TarantoolManager) Observes(
	userID string,
	offset,
	limit uint32) (threads []Thread, err error) {
	var obs []Observer
	err = m.conn.SelectTyped(observerSpace, "subscriptions_idx", offset, limit,
		tarantool.IterReq, makeKey(userID), &obs)
	if err != nil {
		return
	}
	for _, o := range obs {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(o.ThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

func (m *TarantoolManager) MarkAsDelivered(userID, threadID string, times ...time.Time) (err error) {
	t := time.Now()
	if len(times) > 0 {
		t = times[0]
	}

	_, err = m.conn.Update(observerSpace, "primary", makeKey(threadID, userID), makeUpdate(newUpdateOp("=", 2, t.UnixNano()))) //makeUpdate(newUpdateOp("=", 3, t.Unix())))

	return
}

// RecentActivityByLastTS возвращает события позже lastts
func (m *TarantoolManager) RecentActivityByLastTS(threadID string,
	limit uint32, lastts time.Time) (events []Event, err error) {
	err = m.conn.Call17Typed("by_last_ts", makeKey(threadID, lastts.UnixNano()), &events)
	//err = m.conn.SelectTyped(eventsSpace, "threadline_idx", 0, limit,
	//	tarantool.IterGe, makeKey(threadID, lastts.Unix()), &events)
	if err != nil {
		return
	}
	return
}

func (m *TarantoolManager) getThreadline(userID, threadID string, lastts time.Time) (events []Event, err error) {
	err = m.conn.Call17Typed("get_threadline", makeKey(userID, threadID, lastts.UnixNano()), &events)
	return
}

// RecentActivity возвращает события позже lastts
func (m *TarantoolManager) RecentActivity(userID, threadID string, limit uint32) (events []Event, err error) {
	var obs []Observer
	err = m.get(observerSpace, "primary", makeKey(threadID, userID), &obs)
	if err != nil {
		return
	}

	return m.getThreadline(userID, threadID, obs[0].LastDeliveredTime) // m.RecentActivityByLastTS(threadID, limit, obs[0].LastDeliveredTime)
}

// Activity события двигаться по limit,offset что предлагает tnt
func (m *TarantoolManager) Activity(threadID string, limit,
	offset uint32) (events []Event, err error) {
	err = m.conn.SelectTyped(eventsSpace, "threadline_idx", limit, offset, tarantool.IterReq, makeKey(threadID), &events)
	if err != nil {
		return
	}
	return
}

// NewEvent cerate new event. if id empty, wiil be generated uuid.
// If CreatedAt zero, it will be setted to time.Now()
func (m *TarantoolManager) NewEvent(ev *Event) (err error) {
	if ev.EventID == "" {
		ev.EventID = uuid.NewV4().String()
	}
	if ev.CreatedAt.IsZero() {
		ev.CreatedAt = time.Now()
	}

	_, err = m.conn.Insert(eventsSpace, ev)
	if err != nil {
		return
	}

	// check is trhreadline enabled
	var thread *Thread
	thread, err = m.GetThread(ev.ThreadID)
	if err != nil {
		return
	}
	if !thread.ThreadlineEnabled {
		return
	}

	_, err = m.conn.Call17("new_event_in_threadline", makeKey(thread.ThreadID, ev.CreatedAt.UnixNano(), ev.EventID))

	// recieve all thread observers
	// var (
	// 	obs    []User
	// 	tmpObs []User
	// 	offset uint32 = 1000
	// 	limit  uint32
	// )
	// for {
	// 	tmpObs, err = m.Observers(thread.ThreadID, offset, limit)
	// 	if err != nil {
	// 		return
	// 	}
	// 	obs = append(obs, tmpObs...)
	// 	if uint32(len(tmpObs)) != offset {
	// 		break
	// 	}
	// }
	//
	// for _,o:=range obs {
	//
	// }

	return
}

// GetEvent return event by their ID
func (m *TarantoolManager) GetEvent(eventID string) (ev *Event, err error) {
	ev = new(Event)
	err = m.conn.SelectTyped(eventsSpace, "primary", 0, 1, tarantool.IterEq,
		makeKey(eventID), ev)
	return
}

// GetEvents returns events by their ids
func (m *TarantoolManager) GetEvents(ids ...string) (events []Event, err error) {
	for _, id := range ids {
		var event *Event
		event, err = m.GetEvent(id)
		if err != nil {
			return
		}
		events = append(events, *event)
	}
	return
}

func (m *TarantoolManager) UpdateEvent(ev *Event) (err error) {
	// todo обновить updatedat связаных трэдов
	ev.UpdatedAt = time.Now()
	_, err = m.conn.Replace(eventsSpace, ev)
	if err != nil {
		return
	}
	return nil
}

func (m *TarantoolManager) DeleteEvent(eventID string) (err error) {
	// todo удалить из threadline связанные записи
	_, err = m.conn.Delete(eventsSpace, "primary", makeKey(eventID))
	if err != nil {
		return
	}
	return nil
}

// SetRelatedData set data to tarantool
func (m *TarantoolManager) SetRelatedData(rel *RelatedData) (err error) {
	_, err = m.conn.Replace(relatedSpace, rel)
	return
}

// GetRelatedDatas возвращает события с кастомными данными юзера
func (m *TarantoolManager) GetRelatedDatas(userID string, events ...Event) (obs []EventObserver, err error) {
	for _, ev := range events {
		var rel []RelatedData
		err = m.get(relatedSpace, "primary", makeKey(userID, ev.EventID), &rel)
		if err != nil {
			return
		}
		obs = append(obs, EventObserver{Event: ev, RelatedData: rel[0]})
	}

	return
}

// NewUser create new user. If u.UserID empty, will be generated uuid
func (m *TarantoolManager) NewUser(u *User) (err error) {
	if u.UserID == "" {
		u.UserID = uuid.NewV4().String()
	}

	_, err = m.conn.Insert(usersSpace, u)
	if err != nil {
		return
	}
	return nil
}

// GetUser return user by their ID
func (m *TarantoolManager) GetUser(userID string) (u *User, err error) {
	u = new(User)
	err = m.get(usersSpace, "primary", makeKey(userID), u)
	return
}
<commit_msg>fix bug<commit_after>package hey

import (
	"os"
	"time"

	uuid "github.com/satori/go.uuid"
	tarantool "github.com/tarantool/go-tarantool"
)

var (
	DefaultTimeout   time.Duration = 500 * time.Millisecond
	DefaultReconnect time.Duration = 1 * time.Second

	DefaultTarantoolPrefix = "chronograph_"

	threadsSpace    = DefaultTarantoolPrefix + "threads"
	observerSpace   = DefaultTarantoolPrefix + "subscriptions"
	threadLineSpace = DefaultTarantoolPrefix + "threadline"
	eventsSpace     = DefaultTarantoolPrefix + "events"
	usersSpace      = DefaultTarantoolPrefix + "users"
	sourcesSpace    = DefaultTarantoolPrefix + "sources"
	relatedSpace    = DefaultTarantoolPrefix + "related"
)

type TarantoolOpts struct {
	Server        string
	Timeout       time.Duration
	Reconnect     time.Duration
	MaxReconnects uint
	User          string
	Pass          string
}

func (t TarantoolOpts) toOpts() tarantool.Opts {
	return tarantool.Opts{
		Timeout:       t.Timeout,
		Reconnect:     t.Reconnect,
		MaxReconnects: t.MaxReconnects,
		User:          t.User,
		Pass:          t.Pass,
	}
}

func setupFromENV() (opts TarantoolOpts, err error) {
	// 127.0.0.1:3013
	return TarantoolOpts{
		Server:        os.Getenv("TARANTOOL_SERVER"),
		Timeout:       DefaultTimeout,
		Reconnect:     DefaultReconnect,
		MaxReconnects: 3,
		User:          os.Getenv("TARANTOOL_USER_NAME"),
		Pass:          os.Getenv("TARANTOOL_USER_PASSWORD"),
	}, err
}

var _ Manager = &TarantoolManager{}

// TarantoolManager main struct
type TarantoolManager struct {
	conn *tarantool.Connection
}

// NewTarantoolManager return manager setupped from env or optsё
func NewTarantoolManager(opts ...TarantoolOpts) (*TarantoolManager, error) {

	var (
		opt TarantoolOpts
		err error
	)

	if len(opts) > 0 {
		opt = opts[0]
	} else {
		opt, err = setupFromENV()
		if err != nil {
			return nil, err
		}
	}

	client, err := tarantool.Connect(opt.Server, opt.toOpts())
	if err != nil {
		return nil, err
	}

	_, err = client.Ping()

	return &TarantoolManager{conn: client}, err
}

// util
func (m *TarantoolManager) get(space string, keyName string, key interface{}, target interface{}) error {
	return m.conn.SelectTyped(space, keyName, 0, 1, tarantool.IterEq, key, &target)
}

// GetThread return thread by ther id
func (m *TarantoolManager) GetThread(threadID string) (*Thread, error) {
	var (
		threads []Thread
	)
	err := m.conn.SelectTyped(threadsSpace, "primary", 0, 1, tarantool.IterEq,
		makeKey(threadID), &threads)
	if err != nil {
		return nil, err
	}
	return &threads[0], err
}

// NewThread create new thread
// it return error is thread already exists
// if id empty, will be generated uuid
func (m *TarantoolManager) NewThread(thread *Thread) (err error) {
	// err = m.conn.SelectTyped(threadsSpace, "primary", 0, 1,
	// 	tarantool.IterEq, makeIndex(threadID), &thread)
	// if err != nil {
	// 	return
	// }
	// // thread already exists
	// if thread.ThreadID == threadID {
	// 	return ErrAlreadyExists
	// }
	if thread.ThreadID == "" {
		thread.ThreadID = uuid.NewV4().String()
	}

	_, err = m.conn.Insert(threadsSpace, thread)
	if err != nil {
		return
	}

	return nil
}

func (m *TarantoolManager) UpdateThread(thread *Thread) error {
	_, err := m.conn.Replace(threadsSpace, thread)
	return err
}

// DeleteThread удаляет трэд, подписки и связанные события
// 1. Удаляем все записи из events
// 2. Удаляем все подписки на трэд
func (m *TarantoolManager) DeleteThread(threadID string) error {
	// delete
	m.conn.Delete(threadsSpace, "primary", makeKey(threadID))
	return nil
}

// AddSource подписывает трэд на другой трэд
func (m *TarantoolManager) AddSource(dstThread, sourceThread string) (err error) {
	src := Sources{SourceThreadID: sourceThread, TargetThreadID: dstThread}
	_, err = m.conn.Insert(sourcesSpace, &src)
	if err != nil {
		return
	}
	return nil
}

func (m *TarantoolManager) GetSources(referThreadID string, offset, limit uint32) (threads []Thread, err error) {
	var sources []Sources
	err = m.conn.SelectTyped(sourcesSpace, "primary", offset, limit, tarantool.IterReq, makeKey(referThreadID), &sources)
	if err != nil {
		return
	}
	for _, src := range sources {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(src.SourceThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

// Трэды, у которого данный трэд в источниках
func (m *TarantoolManager) GetRefers(threadID string, offset, limit uint32) (threads []Thread, err error) {
	var sources []Sources
	err = m.conn.SelectTyped(sourcesSpace, "sources_idx", offset, limit, tarantool.IterReq, makeKey(threadID), &sources)
	if err != nil {
		return
	}
	for _, src := range sources {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(src.SourceThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

// RemoveSource отписывает трэд
func (m *TarantoolManager) RemoveSource(dstThread, sourceThread string) (err error) {
	_, err = m.conn.Delete(sourcesSpace, "primary", makeKey(dstThread, sourceThread))
	if err != nil {
		return err
	}
	return nil
}

// Observe - подписка конкретного юзера на трэд
func (m *TarantoolManager) Observe(userID, threadID string) error {
	obs := Observer{
		UserID:            userID,
		ThreadID:          threadID,
		LastDeliveredTime: time.Now(),
	}
	_, err := m.conn.Insert(observerSpace, &obs)
	return err
}

// Ignore удаляет подписку на трэд
func (m *TarantoolManager) Ignore(userID, threadID string) (err error) {

	_, err = m.conn.Delete(observerSpace, "primary", makeKey(threadID, userID))
	if err != nil {
		return err
	}

	return err
}

// Observers возвращает подписчиков трэда
func (m *TarantoolManager) Observers(threadID string, offset, limit uint32) (users []User, err error) {
	var obs []Observer
	err = m.conn.SelectTyped(observerSpace, "primary", offset, limit, tarantool.IterEq, makeKey(threadID), &obs)
	if err != nil {
		return
	}

	for _, o := range obs {
		var u []User
		err = m.get(usersSpace, "primary", makeKey(o.UserID), &u)
		if err != nil {
			return
		}
		users = append(users, u[0])
	}
	return
}

// Observes возвращает подписки юзера
func (m *TarantoolManager) Observes(
	userID string,
	offset,
	limit uint32) (threads []Thread, err error) {
	var obs []Observer
	err = m.conn.SelectTyped(observerSpace, "subscriptions_idx", offset, limit,
		tarantool.IterReq, makeKey(userID), &obs)
	if err != nil {
		return
	}
	for _, o := range obs {
		var thread Thread
		err = m.get(threadsSpace, "primary", makeKey(o.ThreadID), &thread)
		if err != nil {
			return
		}
		threads = append(threads, thread)
	}
	return
}

func (m *TarantoolManager) MarkAsDelivered(userID, threadID string, times ...time.Time) (err error) {
	t := time.Now()
	if len(times) > 0 {
		t = times[0]
	}

	_, err = m.conn.Update(observerSpace, "primary", makeKey(threadID, userID), makeUpdate(newUpdateOp("=", 2, t.UnixNano()))) //makeUpdate(newUpdateOp("=", 3, t.Unix())))

	return
}

// RecentActivityByLastTS возвращает события позже lastts
func (m *TarantoolManager) RecentActivityByLastTS(threadID string,
	limit uint32, lastts time.Time) (events []Event, err error) {
	err = m.conn.Call17Typed("by_last_ts", makeKey(threadID, lastts.UnixNano()), &events)
	//err = m.conn.SelectTyped(eventsSpace, "threadline_idx", 0, limit,
	//	tarantool.IterGe, makeKey(threadID, lastts.Unix()), &events)
	if err != nil {
		return
	}
	return
}

func (m *TarantoolManager) getThreadline(userID, threadID string, lastts time.Time) (events []Event, err error) {
	err = m.conn.Call17Typed("get_threadline", makeKey(userID, threadID, lastts.UnixNano()), &events)
	return
}

// RecentActivity возвращает события позже lastts
func (m *TarantoolManager) RecentActivity(userID, threadID string, limit uint32) (events []Event, err error) {
	var obs []Observer
	err = m.get(observerSpace, "primary", makeKey(threadID, userID), &obs)
	if err != nil {
		return
	}

	return m.getThreadline(userID, threadID, obs[0].LastDeliveredTime) // m.RecentActivityByLastTS(threadID, limit, obs[0].LastDeliveredTime)
}

// Activity события двигаться по limit,offset что предлагает tnt
func (m *TarantoolManager) Activity(threadID string, limit,
	offset uint32) (events []Event, err error) {
	err = m.conn.SelectTyped(eventsSpace, "threadline_idx", limit, offset, tarantool.IterReq, makeKey(threadID), &events)
	if err != nil {
		return
	}
	return
}

// NewEvent cerate new event. if id empty, wiil be generated uuid.
// If CreatedAt zero, it will be setted to time.Now()
func (m *TarantoolManager) NewEvent(ev *Event) (err error) {
	if ev.EventID == "" {
		ev.EventID = uuid.NewV4().String()
	}
	if ev.CreatedAt.IsZero() {
		ev.CreatedAt = time.Now()
	}

	_, err = m.conn.Insert(eventsSpace, ev)
	if err != nil {
		return
	}

	// check is trhreadline enabled
	var thread *Thread
	thread, err = m.GetThread(ev.ThreadID)
	if err != nil {
		return
	}
	if !thread.ThreadlineEnabled {
		return
	}

	_, err = m.conn.Call17("new_event_in_threadline", makeKey(thread.ThreadID, ev.CreatedAt.UnixNano(), ev.EventID))

	// recieve all thread observers
	// var (
	// 	obs    []User
	// 	tmpObs []User
	// 	offset uint32 = 1000
	// 	limit  uint32
	// )
	// for {
	// 	tmpObs, err = m.Observers(thread.ThreadID, offset, limit)
	// 	if err != nil {
	// 		return
	// 	}
	// 	obs = append(obs, tmpObs...)
	// 	if uint32(len(tmpObs)) != offset {
	// 		break
	// 	}
	// }
	//
	// for _,o:=range obs {
	//
	// }

	return
}

// GetEvent return event by their ID
func (m *TarantoolManager) GetEvent(eventID string) (ev *Event, err error) {
	ev = new(Event)
	err = m.conn.SelectTyped(eventsSpace, "primary", 0, 1, tarantool.IterEq,
		makeKey(eventID), ev)
	return
}

// GetEvents returns events by their ids
func (m *TarantoolManager) GetEvents(ids ...string) (events []Event, err error) {
	for _, id := range ids {
		var event *Event
		event, err = m.GetEvent(id)
		if err != nil {
			return
		}
		events = append(events, *event)
	}
	return
}

func (m *TarantoolManager) UpdateEvent(ev *Event) (err error) {
	// todo обновить updatedat связаных трэдов
	ev.UpdatedAt = time.Now()
	_, err = m.conn.Replace(eventsSpace, ev)
	if err != nil {
		return
	}
	return nil
}

func (m *TarantoolManager) DeleteEvent(eventID string) (err error) {
	// todo удалить из threadline связанные записи
	_, err = m.conn.Delete(eventsSpace, "primary", makeKey(eventID))
	if err != nil {
		return
	}
	return nil
}

// SetRelatedData set data to tarantool
func (m *TarantoolManager) SetRelatedData(rel *RelatedData) (err error) {
	_, err = m.conn.Replace(relatedSpace, rel)
	return
}

// GetRelatedDatas возвращает события с кастомными данными юзера
func (m *TarantoolManager) GetRelatedDatas(userID string, events ...Event) (obs []EventObserver, err error) {
	for _, ev := range events {
		var rel []RelatedData
		err = m.get(relatedSpace, "primary", makeKey(userID, ev.EventID), &rel)
		if err != nil {
			return
		}
		obs = append(obs, EventObserver{Event: ev, RelatedData: rel[0]})
	}

	return
}

// NewUser create new user. If u.UserID empty, will be generated uuid
func (m *TarantoolManager) NewUser(u *User) (err error) {
	if u.UserID == "" {
		u.UserID = uuid.NewV4().String()
	}

	_, err = m.conn.Insert(usersSpace, u)
	if err != nil {
		return
	}
	return nil
}

// GetUser return user by their ID
func (m *TarantoolManager) GetUser(userID string) (u *User, err error) {
	u = new(User)
	err = m.get(usersSpace, "primary", makeKey(userID), u)
	return
}
<|endoftext|>


<commit_before>package env

import (
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"runtime"

	l4g "code.google.com/p/log4go"
	_ "github.com/go-sql-driver/mysql"
	"gopkg.in/yaml.v1"
)

const (
	VERSION = "1.0.0"
)

var ConfigFile string
var ConfigEnvironment string

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	l4g.AddFilter("stdout", l4g.INFO, l4g.NewConsoleLogWriter())
	l4g.AddFilter("logfile", l4g.DEBUG, l4g.NewFileLogWriter("spored.log", false))
}

func ConfigFlags(f *flag.FlagSet) {
	f.StringVar(&ConfigFile, "config", "config.yml", "Configuration file to use.")
	f.StringVar(&ConfigEnvironment, "env", "development", "Environment to use.")
}

type Environment struct {
	Dialect      string `yaml:"dialect"`
	DataSource   string `yaml:"datasource"`
	Dir          string `yaml:"dir"`
	TableName    string `yaml:"table"`
	HTTP         string `yaml:"http"`
	SphinxServer string `yaml:"sphinx_server"`
	SphinxPort   int    `yaml:"sphinx_port"`
	RedisServer  string `yaml:"redis_server"`
	RedisPort    int    `yaml:"redis_port"`
}

func ReadConfig() (map[string]*Environment, error) {
	file, err := ioutil.ReadFile(ConfigFile)
	if err != nil {
		return nil, err
	}

	config := make(map[string]*Environment)
	err = yaml.Unmarshal(file, config)
	if err != nil {
		return nil, err
	}

	return config, nil
}

func GetEnvironment() (*Environment, error) {
	config, err := ReadConfig()
	if err != nil {
		return nil, err
	}

	env := config[ConfigEnvironment]
	if env == nil {
		return nil, errors.New("No environment: " + ConfigEnvironment)
	}

	if env.DataSource == "" {
		return nil, errors.New("No data source specified")
	}

	return env, nil
}

func GetConnection(env *Environment) (*sql.DB, error) {
	db, err := sql.Open("mysql", env.DataSource)
	if err != nil {
		return nil, fmt.Errorf("Cannot connect to database: %s", err)
	}

	return db, nil
}
<commit_msg>release 1.1.0<commit_after>package env

import (
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"runtime"

	l4g "code.google.com/p/log4go"
	_ "github.com/go-sql-driver/mysql"
	"gopkg.in/yaml.v1"
)

const (
	VERSION = "1.1.0"
)

var ConfigFile string
var ConfigEnvironment string

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	l4g.AddFilter("stdout", l4g.INFO, l4g.NewConsoleLogWriter())
	l4g.AddFilter("logfile", l4g.DEBUG, l4g.NewFileLogWriter("spored.log", false))
}

func ConfigFlags(f *flag.FlagSet) {
	f.StringVar(&ConfigFile, "config", "config.yml", "Configuration file to use.")
	f.StringVar(&ConfigEnvironment, "env", "development", "Environment to use.")
}

type Environment struct {
	Dialect      string `yaml:"dialect"`
	DataSource   string `yaml:"datasource"`
	Dir          string `yaml:"dir"`
	TableName    string `yaml:"table"`
	HTTP         string `yaml:"http"`
	SphinxServer string `yaml:"sphinx_server"`
	SphinxPort   int    `yaml:"sphinx_port"`
	RedisServer  string `yaml:"redis_server"`
	RedisPort    int    `yaml:"redis_port"`
}

func ReadConfig() (map[string]*Environment, error) {
	file, err := ioutil.ReadFile(ConfigFile)
	if err != nil {
		return nil, err
	}

	config := make(map[string]*Environment)
	err = yaml.Unmarshal(file, config)
	if err != nil {
		return nil, err
	}

	return config, nil
}

func GetEnvironment() (*Environment, error) {
	config, err := ReadConfig()
	if err != nil {
		return nil, err
	}

	env := config[ConfigEnvironment]
	if env == nil {
		return nil, errors.New("No environment: " + ConfigEnvironment)
	}

	if env.DataSource == "" {
		return nil, errors.New("No data source specified")
	}

	return env, nil
}

func GetConnection(env *Environment) (*sql.DB, error) {
	db, err := sql.Open("mysql", env.DataSource)
	if err != nil {
		return nil, fmt.Errorf("Cannot connect to database: %s", err)
	}

	return db, nil
}
<|endoftext|>


<commit_before>// Copyright 2019 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package metadata

import (
	"fmt"
	"io"
	"net/http"
)

func getMetadata(name string) (string, error) {
	resp, err := http.Get("http://169.254.169.254/latest/meta-data/" + name)
	if err != nil {
		return "", fmt.Errorf("unable to retrieve instance-id from metadata server: %s", err)
	}

	defer resp.Body.Close()
	instanceID, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("unable to read response body: %s", err)
	}

	return string(instanceID), nil
}

// GetInstanceMetadata returns the instance ID and type
func GetInstanceMetadata() (instanceID, instanceType, availabilityZone, vpcID string, err error) {
	instanceID, err = getMetadata("instance-id")
	if err != nil {
		return
	}

	instanceType, err = getMetadata("instance-type")
	if err != nil {
		return
	}

	eth0MAC, err := getMetadata("mac")
	if err != nil {
		return
	}

	vpcIDPath := fmt.Sprintf("network/interfaces/macs/%s/vpc-id", eth0MAC)
	vpcID, err = getMetadata(vpcIDPath)
	if err != nil {
		return
	}

	availabilityZone, err = getMetadata("placement/availability-zone")
	return
}
<commit_msg>eni: Support Instance Metadata Service Version 2<commit_after>// Copyright 2019-2021 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package metadata

import (
	"context"
	"fmt"
	"io"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
)

func newClient() (*imds.Client, error) {
	cfg, err := config.LoadDefaultConfig(context.TODO())
	if err != nil {
		return nil, err
	}

	return imds.NewFromConfig(cfg), nil
}

func getMetadata(client *imds.Client, path string) (string, error) {
	res, err := client.GetMetadata(context.TODO(), &imds.GetMetadataInput{
		Path: path,
	})
	if err != nil {
		return "", fmt.Errorf("unable to retrieve AWS metadata %s: %w", path, err)
	}

	defer res.Content.Close()
	value, err := io.ReadAll(res.Content)
	if err != nil {
		return "", fmt.Errorf("unable to read response content for AWS metadata %q: %w", path, err)
	}

	return string(value), err
}

// GetInstanceMetadata returns required AWS metadatas
func GetInstanceMetadata() (instanceID, instanceType, availabilityZone, vpcID string, err error) {
	client, err := newClient()
	if err != nil {
		return
	}

	instanceID, err = getMetadata(client, "instance-id")
	if err != nil {
		return
	}

	instanceType, err = getMetadata(client, "instance-type")
	if err != nil {
		return
	}

	eth0MAC, err := getMetadata(client, "mac")
	if err != nil {
		return
	}
	vpcIDPath := fmt.Sprintf("network/interfaces/macs/%s/vpc-id", eth0MAC)
	vpcID, err = getMetadata(client, vpcIDPath)
	if err != nil {
		return
	}

	availabilityZone, err = getMetadata(client, "placement/availability-zone")
	if err != nil {
		return
	}

	return
}
<|endoftext|>