<commit_before>package hclutil

import (
	"github.com/hashicorp/hcl/v2"
	"github.com/hashicorp/hcl/v2/hclsyntax"
	hcls "github.com/hashicorp/hcl/v2/hclsyntax"
)

// BlocksAsAttrs rewrites the hcl.Body so that hcl blocks are treated as
// attributes when schema is unknown.
//
// This conversion is necessary for parsing task driver configs, as they can be
// arbitrary nested without pre-defined schema.
//
// More concretely, it changes the following:
//
// ```
// config {
//   meta { ... }
// }
// ```

// to
//
// ```
// config {
//   meta = { ... } # <- attribute now
// }
// ```
func BlocksAsAttrs(body hcl.Body) hcl.Body {
	if hclb, ok := body.(*hcls.Body); ok {
		return &blockAttrs{body: hclb}
	}
	return body
}

type blockAttrs struct {
	body hcl.Body

	hiddenAttrs  map[string]struct{}
	hiddenBlocks map[string]struct{}
}

func (b *blockAttrs) Content(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Diagnostics) {
	bc, diags := b.body.Content(schema)
	bc.Blocks = expandBlocks(bc.Blocks)
	return bc, diags
}
func (b *blockAttrs) PartialContent(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Body, hcl.Diagnostics) {
	bc, remainBody, diags := b.body.PartialContent(schema)
	bc.Blocks = expandBlocks(bc.Blocks)

	remain := &blockAttrs{
		body:         remainBody,
		hiddenAttrs:  map[string]struct{}{},
		hiddenBlocks: map[string]struct{}{},
	}
	for name := range b.hiddenAttrs {
		remain.hiddenAttrs[name] = struct{}{}
	}
	for typeName := range b.hiddenBlocks {
		remain.hiddenBlocks[typeName] = struct{}{}
	}
	for _, attrS := range schema.Attributes {
		remain.hiddenAttrs[attrS.Name] = struct{}{}
	}
	for _, blockS := range schema.Blocks {
		remain.hiddenBlocks[blockS.Type] = struct{}{}
	}

	return bc, remain, diags
}

func (b *blockAttrs) JustAttributes() (hcl.Attributes, hcl.Diagnostics) {
	body, ok := b.body.(*hcls.Body)
	if !ok {
		return b.body.JustAttributes()
	}

	attrs := make(hcl.Attributes)
	var diags hcl.Diagnostics

	if body.Attributes == nil && len(body.Blocks) == 0 {
		return attrs, diags
	}

	for name, attr := range body.Attributes {
		if _, hidden := b.hiddenAttrs[name]; hidden {
			continue
		}

		na := attr.AsHCLAttribute()
		na.Expr = attrExpr(attr.Expr)
		attrs[name] = na
	}

	for _, blocks := range blocksByType(body.Blocks) {
		if _, hidden := b.hiddenBlocks[blocks[0].Type]; hidden {
			continue
		}

		b := blocks[0]
		attr := &hcls.Attribute{
			Name:        b.Type,
			NameRange:   b.TypeRange,
			EqualsRange: b.OpenBraceRange,
			SrcRange:    b.Body.SrcRange,
			Expr:        blocksToExpr(blocks),
		}

		attrs[blocks[0].Type] = attr.AsHCLAttribute()
	}

	return attrs, diags
}

func (b *blockAttrs) MissingItemRange() hcl.Range {
	return b.body.MissingItemRange()
}

func expandBlocks(blocks hcl.Blocks) hcl.Blocks {
	if len(blocks) == 0 {
		return blocks
	}

	r := make([]*hcl.Block, len(blocks))
	for i, b := range blocks {
		nb := *b
		nb.Body = BlocksAsAttrs(b.Body)
		r[i] = &nb
	}
	return r
}

func blocksByType(blocks hcls.Blocks) map[string]hcls.Blocks {
	r := map[string]hclsyntax.Blocks{}
	for _, b := range blocks {
		r[b.Type] = append(r[b.Type], b)
	}
	return r
}

func blocksToExpr(blocks hcls.Blocks) hcls.Expression {
	if len(blocks) == 0 {
		panic("unexpected empty blocks")
	}

	exprs := make([]hcls.Expression, len(blocks))
	for i, b := range blocks {
		exprs[i] = blockToExpr(b)
	}

	last := blocks[len(blocks)-1]
	return &hcls.TupleConsExpr{
		Exprs: exprs,

		SrcRange:  hcl.RangeBetween(blocks[0].OpenBraceRange, last.CloseBraceRange),
		OpenRange: blocks[0].OpenBraceRange,
	}
}

func blockToExpr(b *hcls.Block) hcls.Expression {
	items := []hcls.ObjectConsItem{}

	for _, attr := range b.Body.Attributes {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     attr.Name,
					SrcRange: attr.NameRange,
				},
			},
			SrcRange: attr.NameRange,
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}

		items = append(items, hcls.ObjectConsItem{
			KeyExpr:   key,
			ValueExpr: attrExpr(attr.Expr),
		})
	}

	for _, blocks := range blocksByType(b.Body.Blocks) {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     blocks[0].Type,
					SrcRange: blocks[0].TypeRange,
				},
			},
			SrcRange: blocks[0].TypeRange,
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}
		item := hcls.ObjectConsItem{
			KeyExpr:   key,
			ValueExpr: blocksToExpr(blocks),
		}

		items = append(items, item)
	}

	v := &hcls.ObjectConsExpr{
		Items: items,
	}

	// Create nested maps, with the labels as keys.
	// Starts wrapping from most inner label to outer
	for i := len(b.Labels) - 1; i >= 0; i-- {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     b.Labels[i],
					SrcRange: b.LabelRanges[i],
				},
			},
			SrcRange: b.LabelRanges[i],
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}
		item := hcls.ObjectConsItem{
			KeyExpr: key,
			ValueExpr: &hcls.TupleConsExpr{
				Exprs: []hcls.Expression{v},
			},
		}

		v = &hcls.ObjectConsExpr{
			Items: []hcls.ObjectConsItem{item},
		}

	}
	return v
}

func attrExpr(expr hcls.Expression) hcls.Expression {
	if _, ok := expr.(*hcls.ObjectConsExpr); ok {
		return &hcls.TupleConsExpr{
			Exprs:     []hcls.Expression{expr},
			SrcRange:  expr.Range(),
			OpenRange: expr.StartRange(),
		}
	}

	return expr
}
<commit_msg>jobspec2: remove duplicate imports statements.<commit_after>package hclutil

import (
	"github.com/hashicorp/hcl/v2"
	hcls "github.com/hashicorp/hcl/v2/hclsyntax"
)

// BlocksAsAttrs rewrites the hcl.Body so that hcl blocks are treated as
// attributes when schema is unknown.
//
// This conversion is necessary for parsing task driver configs, as they can be
// arbitrary nested without pre-defined schema.
//
// More concretely, it changes the following:
//
// ```
// config {
//   meta { ... }
// }
// ```

// to
//
// ```
// config {
//   meta = { ... } # <- attribute now
// }
// ```
func BlocksAsAttrs(body hcl.Body) hcl.Body {
	if hclb, ok := body.(*hcls.Body); ok {
		return &blockAttrs{body: hclb}
	}
	return body
}

type blockAttrs struct {
	body hcl.Body

	hiddenAttrs  map[string]struct{}
	hiddenBlocks map[string]struct{}
}

func (b *blockAttrs) Content(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Diagnostics) {
	bc, diags := b.body.Content(schema)
	bc.Blocks = expandBlocks(bc.Blocks)
	return bc, diags
}
func (b *blockAttrs) PartialContent(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Body, hcl.Diagnostics) {
	bc, remainBody, diags := b.body.PartialContent(schema)
	bc.Blocks = expandBlocks(bc.Blocks)

	remain := &blockAttrs{
		body:         remainBody,
		hiddenAttrs:  map[string]struct{}{},
		hiddenBlocks: map[string]struct{}{},
	}
	for name := range b.hiddenAttrs {
		remain.hiddenAttrs[name] = struct{}{}
	}
	for typeName := range b.hiddenBlocks {
		remain.hiddenBlocks[typeName] = struct{}{}
	}
	for _, attrS := range schema.Attributes {
		remain.hiddenAttrs[attrS.Name] = struct{}{}
	}
	for _, blockS := range schema.Blocks {
		remain.hiddenBlocks[blockS.Type] = struct{}{}
	}

	return bc, remain, diags
}

func (b *blockAttrs) JustAttributes() (hcl.Attributes, hcl.Diagnostics) {
	body, ok := b.body.(*hcls.Body)
	if !ok {
		return b.body.JustAttributes()
	}

	attrs := make(hcl.Attributes)
	var diags hcl.Diagnostics

	if body.Attributes == nil && len(body.Blocks) == 0 {
		return attrs, diags
	}

	for name, attr := range body.Attributes {
		if _, hidden := b.hiddenAttrs[name]; hidden {
			continue
		}

		na := attr.AsHCLAttribute()
		na.Expr = attrExpr(attr.Expr)
		attrs[name] = na
	}

	for _, blocks := range blocksByType(body.Blocks) {
		if _, hidden := b.hiddenBlocks[blocks[0].Type]; hidden {
			continue
		}

		b := blocks[0]
		attr := &hcls.Attribute{
			Name:        b.Type,
			NameRange:   b.TypeRange,
			EqualsRange: b.OpenBraceRange,
			SrcRange:    b.Body.SrcRange,
			Expr:        blocksToExpr(blocks),
		}

		attrs[blocks[0].Type] = attr.AsHCLAttribute()
	}

	return attrs, diags
}

func (b *blockAttrs) MissingItemRange() hcl.Range {
	return b.body.MissingItemRange()
}

func expandBlocks(blocks hcl.Blocks) hcl.Blocks {
	if len(blocks) == 0 {
		return blocks
	}

	r := make([]*hcl.Block, len(blocks))
	for i, b := range blocks {
		nb := *b
		nb.Body = BlocksAsAttrs(b.Body)
		r[i] = &nb
	}
	return r
}

func blocksByType(blocks hcls.Blocks) map[string]hcls.Blocks {
	r := map[string]hcls.Blocks{}
	for _, b := range blocks {
		r[b.Type] = append(r[b.Type], b)
	}
	return r
}

func blocksToExpr(blocks hcls.Blocks) hcls.Expression {
	if len(blocks) == 0 {
		panic("unexpected empty blocks")
	}

	exprs := make([]hcls.Expression, len(blocks))
	for i, b := range blocks {
		exprs[i] = blockToExpr(b)
	}

	last := blocks[len(blocks)-1]
	return &hcls.TupleConsExpr{
		Exprs: exprs,

		SrcRange:  hcl.RangeBetween(blocks[0].OpenBraceRange, last.CloseBraceRange),
		OpenRange: blocks[0].OpenBraceRange,
	}
}

func blockToExpr(b *hcls.Block) hcls.Expression {
	items := []hcls.ObjectConsItem{}

	for _, attr := range b.Body.Attributes {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     attr.Name,
					SrcRange: attr.NameRange,
				},
			},
			SrcRange: attr.NameRange,
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}

		items = append(items, hcls.ObjectConsItem{
			KeyExpr:   key,
			ValueExpr: attrExpr(attr.Expr),
		})
	}

	for _, blocks := range blocksByType(b.Body.Blocks) {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     blocks[0].Type,
					SrcRange: blocks[0].TypeRange,
				},
			},
			SrcRange: blocks[0].TypeRange,
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}
		item := hcls.ObjectConsItem{
			KeyExpr:   key,
			ValueExpr: blocksToExpr(blocks),
		}

		items = append(items, item)
	}

	v := &hcls.ObjectConsExpr{
		Items: items,
	}

	// Create nested maps, with the labels as keys.
	// Starts wrapping from most inner label to outer
	for i := len(b.Labels) - 1; i >= 0; i-- {
		keyExpr := &hcls.ScopeTraversalExpr{
			Traversal: hcl.Traversal{
				hcl.TraverseRoot{
					Name:     b.Labels[i],
					SrcRange: b.LabelRanges[i],
				},
			},
			SrcRange: b.LabelRanges[i],
		}
		key := &hcls.ObjectConsKeyExpr{
			Wrapped: keyExpr,
		}
		item := hcls.ObjectConsItem{
			KeyExpr: key,
			ValueExpr: &hcls.TupleConsExpr{
				Exprs: []hcls.Expression{v},
			},
		}

		v = &hcls.ObjectConsExpr{
			Items: []hcls.ObjectConsItem{item},
		}

	}
	return v
}

func attrExpr(expr hcls.Expression) hcls.Expression {
	if _, ok := expr.(*hcls.ObjectConsExpr); ok {
		return &hcls.TupleConsExpr{
			Exprs:     []hcls.Expression{expr},
			SrcRange:  expr.Range(),
			OpenRange: expr.StartRange(),
		}
	}

	return expr
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"zenhack.net/go/sandstorm-filesystem/filesystem"

	grain_capnp "zenhack.net/go/sandstorm/capnp/grain"
	bridge_capnp "zenhack.net/go/sandstorm/capnp/sandstormhttpbridge"
	"zenhack.net/go/sandstorm/exp/sandstormhttpbridge"

	"zenhack.net/go/sandstorm/exp/util/bytestream"
)

var (
	rootRwDir filesystem.RwDirectory
)

func initZipUploader(bridge bridge_capnp.SandstormHttpBridge) {
	r := mux.NewRouter()

	badReq := func(w http.ResponseWriter, ctx string, err error) {
		log.Print(ctx, ":", err)
		w.WriteHeader(400)
		w.Write([]byte("Bad Request"))
	}

	r.Methods("POST").Path("/filesystem-cap").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			buf, err := ioutil.ReadAll(req.Body)
			if err != nil {
				badReq(w, "read body", err)
				return
			}

			sessionCtx := sandstormhttpbridge.GetSessionContext(bridge, req)
			res, release := sessionCtx.ClaimRequest(
				req.Context(),
				func(p grain_capnp.SessionContext_claimRequest_Params) error {
					p.SetRequestToken(string(buf))
					return nil
				})
			defer release()
			results, err := res.Struct()
			if err != nil {
				badReq(w, "claim request", err)
				return
			}
			capability, err := results.Cap()
			if err != nil {
				log.Print("Error claiming filesystem cap:", err)
				return
			}
			rootRwDir = filesystem.RwDirectory{
				Client: capability.Interface().Client(),
			}
		})

	r.Methods("GET").Path("/").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			tpls.ExecuteTemplate(w, "zip-uploader-index.html", struct{ HaveFS bool }{
				rootRwDir.Client != nil,
			})
		})

	r.Methods("POST").Path("/zipfile").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			ctx := req.Context()
			badReq := func(s string) {
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte(s))
				return
			}
			if rootRwDir.Client == nil {
				badReq("Don't have filesystem cap")
				return
			}
			formFile, _, err := req.FormFile("zipfile")
			if err != nil {
				badReq(err.Error())
				return
			}
			buf := &bytes.Buffer{}
			_, err = io.Copy(buf, formFile)
			if err != nil {
				badReq(err.Error())
				return
			}
			r, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
			if err != nil {
				badReq(err.Error())
				return
			}
			for _, f := range r.File {
				if f.FileInfo().IsDir() {
					// This hasn't happened in my(zenhack) experimentation, but
					// I don't understand the zip format well enough to know if
					// it could. We skip directories just in case. They seem to
					// normally be encoded as files with parent dirs as part of
					// their name.
					continue
				}
				rwDir := rootRwDir
				parts := strings.Split(f.Name, "/")
				if len(parts) > 1 {
					for _, part := range parts[:len(parts)-1] {
						// Mkdir might fail if the directory already exists, so
						// instead of using its return value, we just make
						// another call to walk afterwards.
						_, release := rwDir.Mkdir(
							ctx,
							func(p filesystem.RwDirectory_mkdir_Params) error {
								p.SetName(part)
								return nil
							})
						release()
						res, release := rwDir.Walk(
							ctx,
							func(p filesystem.Directory_walk_Params) error {
								p.SetName(part)
								return nil
							})
						rwDir.Client = res.Node().Client
						// TODO: better place to put this?
						defer release()
					}
				}
				file, err := f.Open()
				if err != nil {
					badReq(err.Error())
					return
				}
				createRes, releaseCreate := rwDir.Create(
					ctx,
					func(p filesystem.RwDirectory_create_Params) error {
						p.SetName(f.Name)
						p.SetExecutable((f.Mode() & 0111) != 0)
						return nil
					})
				writeRes, releaseWrite := createRes.File().Write(
					ctx,
					func(p filesystem.RwFile_write_Params) error {
						p.SetStartAt(0)
						return nil
					})
				out := writeRes.Sink()
				wc := bytestream.ToWriteCloser(ctx, out)
				_, err = io.Copy(wc, file)
				releaseCreate()
				releaseWrite()
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
					log.Print(err)
					return
				}
			}
			w.Header().Set("Location", "/")
			w.WriteHeader(http.StatusSeeOther)
		})

	r.Methods("GET").Path("/pb-request.js").
		Handler(PbRequest(RwDirectoryReq))

	http.Handle("/", withLock(r))
}
<commit_msg>Get rid of superfluous return statement.<commit_after>package main

import (
	"archive/zip"
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"zenhack.net/go/sandstorm-filesystem/filesystem"

	grain_capnp "zenhack.net/go/sandstorm/capnp/grain"
	bridge_capnp "zenhack.net/go/sandstorm/capnp/sandstormhttpbridge"
	"zenhack.net/go/sandstorm/exp/sandstormhttpbridge"

	"zenhack.net/go/sandstorm/exp/util/bytestream"
)

var (
	rootRwDir filesystem.RwDirectory
)

func initZipUploader(bridge bridge_capnp.SandstormHttpBridge) {
	r := mux.NewRouter()

	badReq := func(w http.ResponseWriter, ctx string, err error) {
		log.Print(ctx, ":", err)
		w.WriteHeader(400)
		w.Write([]byte("Bad Request"))
	}

	r.Methods("POST").Path("/filesystem-cap").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			buf, err := ioutil.ReadAll(req.Body)
			if err != nil {
				badReq(w, "read body", err)
				return
			}

			sessionCtx := sandstormhttpbridge.GetSessionContext(bridge, req)
			res, release := sessionCtx.ClaimRequest(
				req.Context(),
				func(p grain_capnp.SessionContext_claimRequest_Params) error {
					p.SetRequestToken(string(buf))
					return nil
				})
			defer release()
			results, err := res.Struct()
			if err != nil {
				badReq(w, "claim request", err)
				return
			}
			capability, err := results.Cap()
			if err != nil {
				log.Print("Error claiming filesystem cap:", err)
				return
			}
			rootRwDir = filesystem.RwDirectory{
				Client: capability.Interface().Client(),
			}
		})

	r.Methods("GET").Path("/").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			tpls.ExecuteTemplate(w, "zip-uploader-index.html", struct{ HaveFS bool }{
				rootRwDir.Client != nil,
			})
		})

	r.Methods("POST").Path("/zipfile").
		HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			ctx := req.Context()
			badReq := func(s string) {
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte(s))
			}
			if rootRwDir.Client == nil {
				badReq("Don't have filesystem cap")
				return
			}
			formFile, _, err := req.FormFile("zipfile")
			if err != nil {
				badReq(err.Error())
				return
			}
			buf := &bytes.Buffer{}
			_, err = io.Copy(buf, formFile)
			if err != nil {
				badReq(err.Error())
				return
			}
			r, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
			if err != nil {
				badReq(err.Error())
				return
			}
			for _, f := range r.File {
				if f.FileInfo().IsDir() {
					// This hasn't happened in my(zenhack) experimentation, but
					// I don't understand the zip format well enough to know if
					// it could. We skip directories just in case. They seem to
					// normally be encoded as files with parent dirs as part of
					// their name.
					continue
				}
				rwDir := rootRwDir
				parts := strings.Split(f.Name, "/")
				if len(parts) > 1 {
					for _, part := range parts[:len(parts)-1] {
						// Mkdir might fail if the directory already exists, so
						// instead of using its return value, we just make
						// another call to walk afterwards.
						_, release := rwDir.Mkdir(
							ctx,
							func(p filesystem.RwDirectory_mkdir_Params) error {
								p.SetName(part)
								return nil
							})
						release()
						res, release := rwDir.Walk(
							ctx,
							func(p filesystem.Directory_walk_Params) error {
								p.SetName(part)
								return nil
							})
						rwDir.Client = res.Node().Client
						// TODO: better place to put this?
						defer release()
					}
				}
				file, err := f.Open()
				if err != nil {
					badReq(err.Error())
					return
				}
				createRes, releaseCreate := rwDir.Create(
					ctx,
					func(p filesystem.RwDirectory_create_Params) error {
						p.SetName(f.Name)
						p.SetExecutable((f.Mode() & 0111) != 0)
						return nil
					})
				writeRes, releaseWrite := createRes.File().Write(
					ctx,
					func(p filesystem.RwFile_write_Params) error {
						p.SetStartAt(0)
						return nil
					})
				out := writeRes.Sink()
				wc := bytestream.ToWriteCloser(ctx, out)
				_, err = io.Copy(wc, file)
				releaseCreate()
				releaseWrite()
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
					log.Print(err)
					return
				}
			}
			w.Header().Set("Location", "/")
			w.WriteHeader(http.StatusSeeOther)
		})

	r.Methods("GET").Path("/pb-request.js").
		Handler(PbRequest(RwDirectoryReq))

	http.Handle("/", withLock(r))
}
<|endoftext|>


<commit_before>// Copyright 2018 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"encoding/binary"
	"fmt"

	"github.com/google/gopacket"
)

const (
	// LCMShortHeaderMagic is the LCM small message header magic number
	LCMShortHeaderMagic uint32 = 0x4c433032
	// LCMFragmentedHeaderMagic is the LCM fragmented message header magic number
	LCMFragmentedHeaderMagic uint32 = 0x4c433033
)

// LCM (Lightweight Communications and Marshalling) is a set of libraries and
// tools for message passing and data marshalling, targeted at real-time systems
// where high-bandwidth and low latency are critical. It provides a
// publish/subscribe message passing model and automatic
// marshalling/unmarshalling code generation with bindings for applications in a
// variety of programming languages.
//
// References
//   https://lcm-proj.github.io/
//   https://github.com/lcm-proj/lcm
type LCM struct {
	// Common (short & fragmented header) fields
	Magic          uint32
	SequenceNumber uint32
	// Fragmented header only fields
	PayloadSize    uint32
	FragmentOffset uint32
	FragmentNumber uint16
	TotalFragments uint16
	// Common field
	ChannelName string
	// Gopacket helper fields
	Fragmented  bool
	fingerprint uint64
	contents    []byte
	payload     []byte
}

var (
	lcmLayerTypes  = map[uint64]gopacket.LayerType{}
	layerTypeIndex = 1001
)

// RegisterLCMLayerType allows users to register decoders for the underlying
// LCM payload. This is done based on the fingerprint that every LCM message
// contains and which identifies it uniquely. If num is not the zero value it
// will be used when registering with RegisterLayerType towards gopacket,
// otherwise an incremental value starting from 1001 will be used.
func RegisterLCMLayerType(num int, fingerprint uint64,
	decoder gopacket.Decoder) gopacket.LayerType {
	name := fmt.Sprintf("%v", fingerprint)
	metadata := gopacket.LayerTypeMetadata{Name: name, Decoder: decoder}

	if num == 0 {
		num = layerTypeIndex
		layerTypeIndex++
	}

	lcmLayerTypes[fingerprint] = gopacket.RegisterLayerType(num, metadata)

	return lcmLayerTypes[fingerprint]
}

// GetLCMLayerType returns the underlying LCM message's LayerType.
// This LayerType has to be registered by using RegisterLCMLayerType.
func GetLCMLayerType(fingerprint uint64) gopacket.LayerType {
	layerType, ok := lcmLayerTypes[fingerprint]
	if !ok {
		return gopacket.LayerTypePayload
	}

	return layerType
}

func decodeLCM(data []byte, p gopacket.PacketBuilder) error {
	lcm := &LCM{}

	err := lcm.DecodeFromBytes(data, p)
	if err != nil {
		return err
	}

	p.AddLayer(lcm)
	p.SetApplicationLayer(lcm)

	return p.NextDecoder(lcm.NextLayerType())
}

// DecodeFromBytes decodes the given bytes into this layer.
func (lcm *LCM) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	offset := 0

	lcm.Magic = binary.BigEndian.Uint32(data[offset:4])
	offset += 4

	if lcm.Magic != LCMShortHeaderMagic && lcm.Magic != LCMFragmentedHeaderMagic {
		return fmt.Errorf("Received LCM header magic %v does not match know "+
			"LCM magic numbers. Dropping packet.", lcm.Magic)
	}

	lcm.SequenceNumber = binary.BigEndian.Uint32(data[offset:8])
	offset += 4

	if lcm.Magic == LCMFragmentedHeaderMagic {
		lcm.Fragmented = true

		lcm.PayloadSize = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4

		lcm.FragmentOffset = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4

		lcm.FragmentNumber = binary.BigEndian.Uint16(data[offset : offset+2])
		offset += 2

		lcm.TotalFragments = binary.BigEndian.Uint16(data[offset : offset+2])
		offset += 2
	} else {
		lcm.Fragmented = false
	}

	if !lcm.Fragmented || (lcm.Fragmented && lcm.FragmentNumber == 0) {
		buffer := make([]byte, 0)
		for _, b := range data[offset:] {
			offset++

			if b == 0 {
				break
			}

			buffer = append(buffer, b)
		}

		lcm.ChannelName = string(buffer)
	}

	lcm.fingerprint = binary.BigEndian.Uint64(data[offset : offset+8])

	lcm.contents = data[:offset]
	lcm.payload = data[offset:]

	return nil
}

// CanDecode returns a set of layers that LCM objects can decode.
// As LCM objects can only decode the LCM layer, we just return that layer.
func (lcm LCM) CanDecode() gopacket.LayerClass {
	return LayerTypeLCM
}

// NextLayerType specifies the LCM payload layer type following this header.
// As LCM packets are serialized structs with uniq fingerprints for each uniq
// combination of data types, lookup of correct layer type is based on that
// fingerprint.
func (lcm LCM) NextLayerType() gopacket.LayerType {
	if !lcm.Fragmented || (lcm.Fragmented && lcm.FragmentNumber == 0) {
		return GetLCMLayerType(lcm.fingerprint)
	}

	return gopacket.LayerTypeFragment
}

// LayerType returns LayerTypeLCM
func (lcm LCM) LayerType() gopacket.LayerType {
	return LayerTypeLCM
}

// LayerContents returns the contents of the LCM header.
func (lcm LCM) LayerContents() []byte {
	return lcm.contents
}

// LayerPayload returns the payload following this LCM header.
func (lcm LCM) LayerPayload() []byte {
	return lcm.payload
}

// Payload returns the payload following this LCM header.
func (lcm LCM) Payload() []byte {
	return lcm.LayerPayload()
}
<commit_msg>layers: Add name & getters to LCM<commit_after>// Copyright 2018 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"encoding/binary"
	"fmt"

	"github.com/google/gopacket"
)

const (
	// LCMShortHeaderMagic is the LCM small message header magic number
	LCMShortHeaderMagic uint32 = 0x4c433032
	// LCMFragmentedHeaderMagic is the LCM fragmented message header magic number
	LCMFragmentedHeaderMagic uint32 = 0x4c433033
)

// LCM (Lightweight Communications and Marshalling) is a set of libraries and
// tools for message passing and data marshalling, targeted at real-time systems
// where high-bandwidth and low latency are critical. It provides a
// publish/subscribe message passing model and automatic
// marshalling/unmarshalling code generation with bindings for applications in a
// variety of programming languages.
//
// References
//   https://lcm-proj.github.io/
//   https://github.com/lcm-proj/lcm
type LCM struct {
	// Common (short & fragmented header) fields
	Magic          uint32
	SequenceNumber uint32
	// Fragmented header only fields
	PayloadSize    uint32
	FragmentOffset uint32
	FragmentNumber uint16
	TotalFragments uint16
	// Common field
	ChannelName string
	// Gopacket helper fields
	Fragmented  bool
	fingerprint LCMFingerprint
	contents    []byte
	payload     []byte
}

// LCMFingerprint is the type of a LCM fingerprint.
type LCMFingerprint uint64

var (
	// lcmLayerTypes contains a map of all LCM fingerprints that we support and
	// their LayerType
	lcmLayerTypes  = map[LCMFingerprint]gopacket.LayerType{}
	layerTypeIndex = 1001
)

// RegisterLCMLayerType allows users to register decoders for the underlying
// LCM payload. This is done based on the fingerprint that every LCM message
// contains and which identifies it uniquely. If num is not the zero value it
// will be used when registering with RegisterLayerType towards gopacket,
// otherwise an incremental value starting from 1001 will be used.
func RegisterLCMLayerType(num int, name string, fingerprint LCMFingerprint,
	decoder gopacket.Decoder) gopacket.LayerType {
	metadata := gopacket.LayerTypeMetadata{Name: name, Decoder: decoder}

	if num == 0 {
		num = layerTypeIndex
		layerTypeIndex++
	}

	lcmLayerTypes[fingerprint] = gopacket.RegisterLayerType(num, metadata)

	return lcmLayerTypes[fingerprint]
}

// SupportedLCMFingerprints returns a slice of all LCM fingerprints that has
// been registered so far.
func SupportedLCMFingerprints() []LCMFingerprint {
	fingerprints := make([]LCMFingerprint, 0, len(lcmLayerTypes))
	for fp := range lcmLayerTypes {
		fingerprints = append(fingerprints, fp)
	}
	return fingerprints
}

// GetLCMLayerType returns the underlying LCM message's LayerType.
// This LayerType has to be registered by using RegisterLCMLayerType.
func GetLCMLayerType(fingerprint LCMFingerprint) gopacket.LayerType {
	layerType, ok := lcmLayerTypes[fingerprint]
	if !ok {
		return gopacket.LayerTypePayload
	}

	return layerType
}

func decodeLCM(data []byte, p gopacket.PacketBuilder) error {
	lcm := &LCM{}

	err := lcm.DecodeFromBytes(data, p)
	if err != nil {
		return err
	}

	p.AddLayer(lcm)
	p.SetApplicationLayer(lcm)

	return p.NextDecoder(lcm.NextLayerType())
}

// DecodeFromBytes decodes the given bytes into this layer.
func (lcm *LCM) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	offset := 0

	lcm.Magic = binary.BigEndian.Uint32(data[offset:4])
	offset += 4

	if lcm.Magic != LCMShortHeaderMagic && lcm.Magic != LCMFragmentedHeaderMagic {
		return fmt.Errorf("Received LCM header magic %v does not match know "+
			"LCM magic numbers. Dropping packet.", lcm.Magic)
	}

	lcm.SequenceNumber = binary.BigEndian.Uint32(data[offset:8])
	offset += 4

	if lcm.Magic == LCMFragmentedHeaderMagic {
		lcm.Fragmented = true

		lcm.PayloadSize = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4

		lcm.FragmentOffset = binary.BigEndian.Uint32(data[offset : offset+4])
		offset += 4

		lcm.FragmentNumber = binary.BigEndian.Uint16(data[offset : offset+2])
		offset += 2

		lcm.TotalFragments = binary.BigEndian.Uint16(data[offset : offset+2])
		offset += 2
	} else {
		lcm.Fragmented = false
	}

	if !lcm.Fragmented || (lcm.Fragmented && lcm.FragmentNumber == 0) {
		buffer := make([]byte, 0)
		for _, b := range data[offset:] {
			offset++

			if b == 0 {
				break
			}

			buffer = append(buffer, b)
		}

		lcm.ChannelName = string(buffer)
	}

	lcm.fingerprint = LCMFingerprint(
		binary.BigEndian.Uint64(data[offset : offset+8]))

	lcm.contents = data[:offset]
	lcm.payload = data[offset:]

	return nil
}

// CanDecode returns a set of layers that LCM objects can decode.
// As LCM objects can only decode the LCM layer, we just return that layer.
func (lcm LCM) CanDecode() gopacket.LayerClass {
	return LayerTypeLCM
}

// NextLayerType specifies the LCM payload layer type following this header.
// As LCM packets are serialized structs with uniq fingerprints for each uniq
// combination of data types, lookup of correct layer type is based on that
// fingerprint.
func (lcm LCM) NextLayerType() gopacket.LayerType {
	if !lcm.Fragmented || (lcm.Fragmented && lcm.FragmentNumber == 0) {
		return GetLCMLayerType(lcm.fingerprint)
	}

	return gopacket.LayerTypeFragment
}

// LayerType returns LayerTypeLCM
func (lcm LCM) LayerType() gopacket.LayerType {
	return LayerTypeLCM
}

// LayerContents returns the contents of the LCM header.
func (lcm LCM) LayerContents() []byte {
	return lcm.contents
}

// LayerPayload returns the payload following this LCM header.
func (lcm LCM) LayerPayload() []byte {
	return lcm.payload
}

// Payload returns the payload following this LCM header.
func (lcm LCM) Payload() []byte {
	return lcm.LayerPayload()
}

// Fingerprint returns the LCM fingerprint of the underlying message.
func (lcm LCM) Fingerprint() LCMFingerprint {
	return lcm.fingerprint
}
<|endoftext|>


<commit_before>package lclip

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	f.Close()
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Fatal(err)
	}
	expect := filepath.Join(h, ".lclip.db")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

func TestCreateStorageFileIfNotExists(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	if _, err := os.Stat(tempPath); err != nil {
		t.Error("not create storage file; want create")
	}
}

type AccessTest struct {
	Label string
	Data  []byte
}

var indexTestsAccess = []AccessTest{
	{Label: "foo", Data: []byte("bar")},
	{Label: "hoge", Data: []byte("piyo")},
	{Label: "日本語", Data: []byte("日本語")},
}

func TestSetText(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	for _, test := range indexTestsAccess {
		if err = c.Set(test.Label, test.Data); err != nil {
			t.Error("Set(%q) returns %q; want nil",
				test.Label, err)
		}
		expect := test.Data
		actual, err := c.Get(test.Label)
		if err != nil {
			t.Error("Get(%q) returns %q; want nil",
				test.Label, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	for _, labels := range indexTestsLabels {
		os.Remove(tempPath)

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		for _, label := range labels {
			if err = c.Set(label, []byte(``)); err != nil {
				t.Fatal(err)
			}
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual := c.Labels()
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
		}
		if err := c.Close(); err != nil {
			t.Fatal(err)
		}
	}
}

func TestSave(t *testing.T) {
	os.Remove(tempPath)

	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		for _, test := range indexTestsAccess {
			if err := c.Set(test.Label, test.Data); err != nil {
				t.Fatal(err)
			}
		}
		if err := c.Close(); err != nil {
			t.Fatal(err)
		}
	}
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		defer c.Close()
		for _, test := range indexTestsAccess {
			expect := test.Data
			actual, err := c.Get(test.Label)
			if err != nil {
				t.Fatal(err)
			}
			if !reflect.DeepEqual(actual, expect) {
				t.Errorf("Get(%q) = %q; want %q",
					test.Label, actual, expect)
			}
		}
	}
}
<commit_msg>Rewrite test for Labels returns err<commit_after>package lclip

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	f.Close()
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Fatal(err)
	}
	expect := filepath.Join(h, ".lclip.db")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

func TestCreateStorageFileIfNotExists(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	if _, err := os.Stat(tempPath); err != nil {
		t.Error("not create storage file; want create")
	}
}

type AccessTest struct {
	Label string
	Data  []byte
}

var indexTestsAccess = []AccessTest{
	{Label: "foo", Data: []byte("bar")},
	{Label: "hoge", Data: []byte("piyo")},
	{Label: "日本語", Data: []byte("日本語")},
}

func TestSetText(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	for _, test := range indexTestsAccess {
		if err = c.Set(test.Label, test.Data); err != nil {
			t.Error("Set(%q) returns %q; want nil",
				test.Label, err)
		}
		expect := test.Data
		actual, err := c.Get(test.Label)
		if err != nil {
			t.Error("Get(%q) returns %q; want nil",
				test.Label, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	for _, labels := range indexTestsLabels {
		os.Remove(tempPath)

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		for _, label := range labels {
			if err = c.Set(label, []byte(``)); err != nil {
				t.Fatal(err)
			}
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual, err := c.Labels()
		if err != nil {
			t.Fatal(err)
		}
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
		}
		if err := c.Close(); err != nil {
			t.Fatal(err)
		}
	}
}

func TestSave(t *testing.T) {
	os.Remove(tempPath)

	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		for _, test := range indexTestsAccess {
			if err := c.Set(test.Label, test.Data); err != nil {
				t.Fatal(err)
			}
		}
		if err := c.Close(); err != nil {
			t.Fatal(err)
		}
	}
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatal(err)
		}
		defer c.Close()
		for _, test := range indexTestsAccess {
			expect := test.Data
			actual, err := c.Get(test.Label)
			if err != nil {
				t.Fatal(err)
			}
			if !reflect.DeepEqual(actual, expect) {
				t.Errorf("Get(%q) = %q; want %q",
					test.Label, actual, expect)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/oschwald/geoip2-golang"
	"html/template"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
)

type ErrorMessage struct {
	Error string
}

type ReturnData struct {
	Ip          net.IP
	CountryCode string
	CountryName string
	City        string
	Latitude    float64
	Longitude   float64
}

// Load the database file
var db, liberr = geoip2.Open("geodata/GeoLite2-City.mmdb")

// Load templates.
var htmlTemplates = template.Must(template.ParseFiles("bin/templates/index.html", "bin/templates/error.html"))

func getData(r *http.Request) interface{} {
	vars := mux.Vars(r)
	requestIp := ""
	if val, ok := vars["ip"]; ok {
		// IP is provided in URL.
		requestIp = val
	} else {
		// Take IP from Remote Address.
		address := r.Header.Get("X-Real-IP")
		if address == "" {
			address = r.Header.Get("X-Forwarded-For")
		}
		if address == "" {
			address = r.RemoteAddr
			idx := strings.LastIndex(address, ":")
			address = address[:idx]
		}
		requestIp = address
	}

	ip := net.ParseIP(requestIp)

	record, err := db.City(ip)
	if err != nil {
		return ErrorMessage{Error: "Unknown IP"}
	}

	return ReturnData{
		Ip:          ip,
		CountryCode: record.Country.IsoCode,
		CountryName: record.Country.Names["en"],
		City:        record.City.Names["en"],
		Latitude:    record.Location.Latitude,
		Longitude:   record.Location.Longitude}
}

func xmlHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/xml")
	retData := getData(r)
	msg, _ := xml.Marshal(retData)
	w.Write(msg)
}

func jsonHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	retData := getData(r)
	msg, _ := json.Marshal(retData)
	w.Write(msg)
}

func htmlHandler(w http.ResponseWriter, r *http.Request) {
	retData := getData(r)

	switch retData.(type) {
	case ReturnData:
		htmlTemplates.ExecuteTemplate(w, "index.html", retData)
	default:
		htmlTemplates.ExecuteTemplate(w, "error.html", "")
	}
}

func validatePortNumber(port string) bool {
	i, _ := strconv.ParseInt(port, 0, 64)

	return i > 0 && i < 65535
}

func main() {

	if liberr != nil {
		// Country database was not loaded successfully
		fmt.Printf("Error: %s\n", liberr.Error())
		return
	}

	r := mux.NewRouter()
	r.HandleFunc("/{ip}/json", jsonHandler)
	r.HandleFunc("/json", jsonHandler)
	r.HandleFunc("/{ip}/xml", xmlHandler)
	r.HandleFunc("/xml", xmlHandler)
	r.HandleFunc("/{ip}", htmlHandler)
	r.HandleFunc("/", htmlHandler)
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("bin/assets")))

	http.Handle("/", r)

	port := ":8080"

	// Check if port was provided via argument and validate port number.
	if len(os.Args) > 1 && validatePortNumber(os.Args[1]) == true {
		port = fmt.Sprintf(":%s", os.Args[1])
	}

	http.ListenAndServe(port, nil)
}
<commit_msg>Adds: Reloading database via SIGHUP.<commit_after>package main

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/oschwald/geoip2-golang"
	"html/template"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
)

type ErrorMessage struct {
	Error string
}

type ReturnData struct {
	Ip          net.IP
	CountryCode string
	CountryName string
	City        string
	Latitude    float64
	Longitude   float64
}

// Load templates.
var htmlTemplates = template.Must(template.ParseFiles("bin/templates/index.html", "bin/templates/error.html"))

// Load the database file
var db, liberr = geoip2.Open("geodata/GeoLite2-City.mmdb")

// Create signal channel for reloading database.
var sigs = make(chan os.Signal, 1)

func getData(r *http.Request) interface{} {
	vars := mux.Vars(r)
	requestIp := ""
	if val, ok := vars["ip"]; ok {
		// IP is provided in URL.
		requestIp = val
	} else {
		// Take IP from Remote Address.
		address := r.Header.Get("X-Real-IP")
		if address == "" {
			address = r.Header.Get("X-Forwarded-For")
		}
		if address == "" {
			address = r.RemoteAddr
			idx := strings.LastIndex(address, ":")
			address = address[:idx]
		}
		requestIp = address
	}

	ip := net.ParseIP(requestIp)

	record, err := db.City(ip)
	if err != nil {
		return ErrorMessage{Error: "Unknown IP"}
	}

	return ReturnData{
		Ip:          ip,
		CountryCode: record.Country.IsoCode,
		CountryName: record.Country.Names["en"],
		City:        record.City.Names["en"],
		Latitude:    record.Location.Latitude,
		Longitude:   record.Location.Longitude}
}

func xmlHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/xml")
	retData := getData(r)
	msg, _ := xml.Marshal(retData)
	w.Write(msg)
}

func jsonHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	retData := getData(r)
	msg, _ := json.Marshal(retData)
	w.Write(msg)
}

func htmlHandler(w http.ResponseWriter, r *http.Request) {
	retData := getData(r)

	switch retData.(type) {
	case ReturnData:
		htmlTemplates.ExecuteTemplate(w, "index.html", retData)
	default:
		htmlTemplates.ExecuteTemplate(w, "error.html", "")
	}
}

func validatePortNumber(port string) bool {
	i, _ := strconv.ParseInt(port, 0, 64)

	return i > 0 && i < 65535
}

func reloadDatabaseListener() {
	// Wait for appropriate signal.
	<-sigs

	// Load the database file.
	db, liberr = geoip2.Open("geodata/GeoLite2-City.mmdb")

	if liberr != nil {
		// Country database was not loaded successfully.
		fmt.Printf("Error: %s\n", liberr.Error())
		return
	}

	fmt.Println("Database reloaded.")
	go reloadDatabaseListener()
}

func main() {
	// Check if data loaded properly.
	if liberr != nil {
		// Country database was not loaded successfully
		fmt.Printf("Error: %s\n", liberr.Error())
		return
	}

	// Use SIGHUP signal for reloading database.
	signal.Notify(sigs, syscall.SIGHUP)
	go reloadDatabaseListener()

	r := mux.NewRouter()
	r.HandleFunc("/{ip}/json", jsonHandler)
	r.HandleFunc("/json", jsonHandler)
	r.HandleFunc("/{ip}/xml", xmlHandler)
	r.HandleFunc("/xml", xmlHandler)
	r.HandleFunc("/{ip}", htmlHandler)
	r.HandleFunc("/", htmlHandler)
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("bin/assets")))

	http.Handle("/", r)

	port := ":8080"

	// Check if port was provided via argument and validate port number.
	if len(os.Args) > 1 && validatePortNumber(os.Args[1]) == true {
		port = fmt.Sprintf(":%s", os.Args[1])
	}

	http.ListenAndServe(port, nil)
}
<|endoftext|>


<commit_before>// Copyright 2016 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

// Package les implements the Light Ethereum Subprotocol.
package les

import (
	"encoding/binary"
	"math"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/ethdb"
	"github.com/ethereum/go-ethereum/les/flowcontrol"
	"github.com/ethereum/go-ethereum/light"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
)

type LesServer struct {
	protocolManager *ProtocolManager
	fcManager       *flowcontrol.ClientManager // nil if our node is client only
	fcCostStats     *requestCostStats
	defParams       *flowcontrol.ServerParams
}

func NewLesServer(eth *eth.Ethereum, config *eth.Config) (*LesServer, error) {
	pm, err := NewProtocolManager(config.ChainConfig, false, config.NetworkId, eth.EventMux(), eth.Pow(), eth.BlockChain(), eth.TxPool(), eth.ChainDb(), nil, nil)
	if err != nil {
		return nil, err
	}
	pm.blockLoop()

	srv := &LesServer{protocolManager: pm}
	pm.server = srv

	srv.defParams = &flowcontrol.ServerParams{
		BufLimit:    300000000,
		MinRecharge: 50000,
	}
	srv.fcManager = flowcontrol.NewClientManager(uint64(config.LightServ), 10, 1000000000)
	srv.fcCostStats = newCostStats(eth.ChainDb())
	return srv, nil
}

func (s *LesServer) Protocols() []p2p.Protocol {
	return s.protocolManager.SubProtocols
}

func (s *LesServer) Start(srvr *p2p.Server) {
	s.protocolManager.Start(srvr)

}

func (s *LesServer) Stop() {
	s.fcCostStats.store()
	s.fcManager.Stop()
	go func() {
		<-s.protocolManager.noMorePeers
	}()
	s.protocolManager.Stop()
}

type requestCosts struct {
	baseCost, reqCost uint64
}

type requestCostTable map[uint64]*requestCosts

type RequestCostList []struct {
	MsgCode, BaseCost, ReqCost uint64
}

func (list RequestCostList) decode() requestCostTable {
	table := make(requestCostTable)
	for _, e := range list {
		table[e.MsgCode] = &requestCosts{
			baseCost: e.BaseCost,
			reqCost:  e.ReqCost,
		}
	}
	return table
}

func (table requestCostTable) encode() RequestCostList {
	list := make(RequestCostList, len(table))
	for idx, code := range reqList {
		list[idx].MsgCode = code
		list[idx].BaseCost = table[code].baseCost
		list[idx].ReqCost = table[code].reqCost
	}
	return list
}

type linReg struct {
	sumX, sumY, sumXX, sumXY float64
	cnt                      uint64
}

const linRegMaxCnt = 100000

func (l *linReg) add(x, y float64) {
	if l.cnt >= linRegMaxCnt {
		sub := float64(l.cnt+1-linRegMaxCnt) / linRegMaxCnt
		l.sumX -= l.sumX * sub
		l.sumY -= l.sumY * sub
		l.sumXX -= l.sumXX * sub
		l.sumXY -= l.sumXY * sub
		l.cnt = linRegMaxCnt - 1
	}
	l.cnt++
	l.sumX += x
	l.sumY += y
	l.sumXX += x * x
	l.sumXY += x * y
}

func (l *linReg) calc() (b, m float64) {
	if l.cnt == 0 {
		return 0, 0
	}
	cnt := float64(l.cnt)
	d := cnt*l.sumXX - l.sumX*l.sumX
	if d < 0.001 {
		return l.sumY / cnt, 0
	}
	m = (cnt*l.sumXY - l.sumX*l.sumY) / d
	b = (l.sumY / cnt) - (m * l.sumX / cnt)
	return b, m
}

func (l *linReg) toBytes() []byte {
	var arr [40]byte
	binary.BigEndian.PutUint64(arr[0:8], math.Float64bits(l.sumX))
	binary.BigEndian.PutUint64(arr[8:16], math.Float64bits(l.sumY))
	binary.BigEndian.PutUint64(arr[16:24], math.Float64bits(l.sumXX))
	binary.BigEndian.PutUint64(arr[24:32], math.Float64bits(l.sumXY))
	binary.BigEndian.PutUint64(arr[32:40], l.cnt)
	return arr[:]
}

func linRegFromBytes(data []byte) *linReg {
	if len(data) != 40 {
		return nil
	}
	l := &linReg{}
	l.sumX = math.Float64frombits(binary.BigEndian.Uint64(data[0:8]))
	l.sumY = math.Float64frombits(binary.BigEndian.Uint64(data[8:16]))
	l.sumXX = math.Float64frombits(binary.BigEndian.Uint64(data[16:24]))
	l.sumXY = math.Float64frombits(binary.BigEndian.Uint64(data[24:32]))
	l.cnt = binary.BigEndian.Uint64(data[32:40])
	return l
}

type requestCostStats struct {
	lock  sync.RWMutex
	db    ethdb.Database
	stats map[uint64]*linReg
}

type requestCostStatsRlp []struct {
	MsgCode uint64
	Data    []byte
}

var rcStatsKey = []byte("_requestCostStats")

func newCostStats(db ethdb.Database) *requestCostStats {
	stats := make(map[uint64]*linReg)
	for _, code := range reqList {
		stats[code] = &linReg{cnt: 100}
	}

	if db != nil {
		data, err := db.Get(rcStatsKey)
		var statsRlp requestCostStatsRlp
		if err == nil {
			err = rlp.DecodeBytes(data, &statsRlp)
		}
		if err == nil {
			for _, r := range statsRlp {
				if stats[r.MsgCode] != nil {
					if l := linRegFromBytes(r.Data); l != nil {
						stats[r.MsgCode] = l
					}
				}
			}
		}
	}

	return &requestCostStats{
		db:    db,
		stats: stats,
	}
}

func (s *requestCostStats) store() {
	s.lock.Lock()
	defer s.lock.Unlock()

	statsRlp := make(requestCostStatsRlp, len(reqList))
	for i, code := range reqList {
		statsRlp[i].MsgCode = code
		statsRlp[i].Data = s.stats[code].toBytes()
	}

	if data, err := rlp.EncodeToBytes(statsRlp); err == nil {
		s.db.Put(rcStatsKey, data)
	}
}

func (s *requestCostStats) getCurrentList() RequestCostList {
	s.lock.Lock()
	defer s.lock.Unlock()

	list := make(RequestCostList, len(reqList))
	//fmt.Println("RequestCostList")
	for idx, code := range reqList {
		b, m := s.stats[code].calc()
		//fmt.Println(code, s.stats[code].cnt, b/1000000, m/1000000)
		if m < 0 {
			b += m
			m = 0
		}
		if b < 0 {
			b = 0
		}

		list[idx].MsgCode = code
		list[idx].BaseCost = uint64(b * 2)
		list[idx].ReqCost = uint64(m * 2)
	}
	return list
}

func (s *requestCostStats) update(msgCode, reqCnt, cost uint64) {
	s.lock.Lock()
	defer s.lock.Unlock()

	c, ok := s.stats[msgCode]
	if !ok || reqCnt == 0 {
		return
	}
	c.add(float64(reqCnt), float64(cost))
}

func (pm *ProtocolManager) blockLoop() {
	pm.wg.Add(1)
	sub := pm.eventMux.Subscribe(core.ChainHeadEvent{})
	newCht := make(chan struct{}, 10)
	newCht <- struct{}{}
	go func() {
		var mu sync.Mutex
		var lastHead *types.Header
		lastBroadcastTd := common.Big0
		for {
			select {
			case ev := <-sub.Chan():
				peers := pm.peers.AllPeers()
				if len(peers) > 0 {
					header := ev.Data.(core.ChainHeadEvent).Block.Header()
					hash := header.Hash()
					number := header.Number.Uint64()
					td := core.GetTd(pm.chainDb, hash, number)
					if td != nil && td.Cmp(lastBroadcastTd) > 0 {
						var reorg uint64
						if lastHead != nil {
							reorg = lastHead.Number.Uint64() - core.FindCommonAncestor(pm.chainDb, header, lastHead).Number.Uint64()
						}
						lastHead = header
						lastBroadcastTd = td

						glog.V(logger.Debug).Infoln("===> ", number, hash, td, reorg)

						announce := announceData{Hash: hash, Number: number, Td: td, ReorgDepth: reorg}
						for _, p := range peers {
							select {
							case p.announceChn <- announce:
							default:
								pm.removePeer(p.id)
							}
						}
					}
				}
				newCht <- struct{}{}
			case <-newCht:
				go func() {
					mu.Lock()
					more := makeCht(pm.chainDb)
					mu.Unlock()
					if more {
						time.Sleep(time.Millisecond * 10)
						newCht <- struct{}{}
					}
				}()
			case <-pm.quitSync:
				sub.Unsubscribe()
				pm.wg.Done()
				return
			}
		}
	}()
}

var (
	lastChtKey       = []byte("LastChtNumber") // chtNum (uint64 big endian)
	chtPrefix        = []byte("cht")           // chtPrefix + chtNum (uint64 big endian) -> trie root hash
	chtConfirmations = light.ChtFrequency / 2
)

func getChtRoot(db ethdb.Database, num uint64) common.Hash {
	var encNumber [8]byte
	binary.BigEndian.PutUint64(encNumber[:], num)
	data, _ := db.Get(append(chtPrefix, encNumber[:]...))
	return common.BytesToHash(data)
}

func storeChtRoot(db ethdb.Database, num uint64, root common.Hash) {
	var encNumber [8]byte
	binary.BigEndian.PutUint64(encNumber[:], num)
	db.Put(append(chtPrefix, encNumber[:]...), root[:])
}

func makeCht(db ethdb.Database) bool {
	headHash := core.GetHeadBlockHash(db)
	headNum := core.GetBlockNumber(db, headHash)

	var newChtNum uint64
	if headNum > chtConfirmations {
		newChtNum = (headNum - chtConfirmations) / light.ChtFrequency
	}

	var lastChtNum uint64
	data, _ := db.Get(lastChtKey)
	if len(data) == 8 {
		lastChtNum = binary.BigEndian.Uint64(data[:])
	}
	if newChtNum <= lastChtNum {
		return false
	}

	var t *trie.Trie
	if lastChtNum > 0 {
		var err error
		t, err = trie.New(getChtRoot(db, lastChtNum), db)
		if err != nil {
			lastChtNum = 0
		}
	}
	if lastChtNum == 0 {
		t, _ = trie.New(common.Hash{}, db)
	}

	for num := lastChtNum * light.ChtFrequency; num < (lastChtNum+1)*light.ChtFrequency; num++ {
		hash := core.GetCanonicalHash(db, num)
		if hash == (common.Hash{}) {
			panic("Canonical hash not found")
		}
		td := core.GetTd(db, hash, num)
		if td == nil {
			panic("TD not found")
		}
		var encNumber [8]byte
		binary.BigEndian.PutUint64(encNumber[:], num)
		var node light.ChtNode
		node.Hash = hash
		node.Td = td
		data, _ := rlp.EncodeToBytes(node)
		t.Update(encNumber[:], data)
	}

	root, err := t.Commit()
	if err != nil {
		lastChtNum = 0
	} else {
		lastChtNum++

		glog.V(logger.Info).Infoln("CHT %d %064x\n", lastChtNum, root)

		storeChtRoot(db, lastChtNum, root)
		var data [8]byte
		binary.BigEndian.PutUint64(data[:], lastChtNum)
		db.Put(lastChtKey, data[:])
	}

	return newChtNum > lastChtNum
}
<commit_msg>les: fixed logging<commit_after>// Copyright 2016 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

// Package les implements the Light Ethereum Subprotocol.
package les

import (
	"encoding/binary"
	"math"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/ethdb"
	"github.com/ethereum/go-ethereum/les/flowcontrol"
	"github.com/ethereum/go-ethereum/light"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/trie"
)

type LesServer struct {
	protocolManager *ProtocolManager
	fcManager       *flowcontrol.ClientManager // nil if our node is client only
	fcCostStats     *requestCostStats
	defParams       *flowcontrol.ServerParams
}

func NewLesServer(eth *eth.Ethereum, config *eth.Config) (*LesServer, error) {
	pm, err := NewProtocolManager(config.ChainConfig, false, config.NetworkId, eth.EventMux(), eth.Pow(), eth.BlockChain(), eth.TxPool(), eth.ChainDb(), nil, nil)
	if err != nil {
		return nil, err
	}
	pm.blockLoop()

	srv := &LesServer{protocolManager: pm}
	pm.server = srv

	srv.defParams = &flowcontrol.ServerParams{
		BufLimit:    300000000,
		MinRecharge: 50000,
	}
	srv.fcManager = flowcontrol.NewClientManager(uint64(config.LightServ), 10, 1000000000)
	srv.fcCostStats = newCostStats(eth.ChainDb())
	return srv, nil
}

func (s *LesServer) Protocols() []p2p.Protocol {
	return s.protocolManager.SubProtocols
}

func (s *LesServer) Start(srvr *p2p.Server) {
	s.protocolManager.Start(srvr)

}

func (s *LesServer) Stop() {
	s.fcCostStats.store()
	s.fcManager.Stop()
	go func() {
		<-s.protocolManager.noMorePeers
	}()
	s.protocolManager.Stop()
}

type requestCosts struct {
	baseCost, reqCost uint64
}

type requestCostTable map[uint64]*requestCosts

type RequestCostList []struct {
	MsgCode, BaseCost, ReqCost uint64
}

func (list RequestCostList) decode() requestCostTable {
	table := make(requestCostTable)
	for _, e := range list {
		table[e.MsgCode] = &requestCosts{
			baseCost: e.BaseCost,
			reqCost:  e.ReqCost,
		}
	}
	return table
}

func (table requestCostTable) encode() RequestCostList {
	list := make(RequestCostList, len(table))
	for idx, code := range reqList {
		list[idx].MsgCode = code
		list[idx].BaseCost = table[code].baseCost
		list[idx].ReqCost = table[code].reqCost
	}
	return list
}

type linReg struct {
	sumX, sumY, sumXX, sumXY float64
	cnt                      uint64
}

const linRegMaxCnt = 100000

func (l *linReg) add(x, y float64) {
	if l.cnt >= linRegMaxCnt {
		sub := float64(l.cnt+1-linRegMaxCnt) / linRegMaxCnt
		l.sumX -= l.sumX * sub
		l.sumY -= l.sumY * sub
		l.sumXX -= l.sumXX * sub
		l.sumXY -= l.sumXY * sub
		l.cnt = linRegMaxCnt - 1
	}
	l.cnt++
	l.sumX += x
	l.sumY += y
	l.sumXX += x * x
	l.sumXY += x * y
}

func (l *linReg) calc() (b, m float64) {
	if l.cnt == 0 {
		return 0, 0
	}
	cnt := float64(l.cnt)
	d := cnt*l.sumXX - l.sumX*l.sumX
	if d < 0.001 {
		return l.sumY / cnt, 0
	}
	m = (cnt*l.sumXY - l.sumX*l.sumY) / d
	b = (l.sumY / cnt) - (m * l.sumX / cnt)
	return b, m
}

func (l *linReg) toBytes() []byte {
	var arr [40]byte
	binary.BigEndian.PutUint64(arr[0:8], math.Float64bits(l.sumX))
	binary.BigEndian.PutUint64(arr[8:16], math.Float64bits(l.sumY))
	binary.BigEndian.PutUint64(arr[16:24], math.Float64bits(l.sumXX))
	binary.BigEndian.PutUint64(arr[24:32], math.Float64bits(l.sumXY))
	binary.BigEndian.PutUint64(arr[32:40], l.cnt)
	return arr[:]
}

func linRegFromBytes(data []byte) *linReg {
	if len(data) != 40 {
		return nil
	}
	l := &linReg{}
	l.sumX = math.Float64frombits(binary.BigEndian.Uint64(data[0:8]))
	l.sumY = math.Float64frombits(binary.BigEndian.Uint64(data[8:16]))
	l.sumXX = math.Float64frombits(binary.BigEndian.Uint64(data[16:24]))
	l.sumXY = math.Float64frombits(binary.BigEndian.Uint64(data[24:32]))
	l.cnt = binary.BigEndian.Uint64(data[32:40])
	return l
}

type requestCostStats struct {
	lock  sync.RWMutex
	db    ethdb.Database
	stats map[uint64]*linReg
}

type requestCostStatsRlp []struct {
	MsgCode uint64
	Data    []byte
}

var rcStatsKey = []byte("_requestCostStats")

func newCostStats(db ethdb.Database) *requestCostStats {
	stats := make(map[uint64]*linReg)
	for _, code := range reqList {
		stats[code] = &linReg{cnt: 100}
	}

	if db != nil {
		data, err := db.Get(rcStatsKey)
		var statsRlp requestCostStatsRlp
		if err == nil {
			err = rlp.DecodeBytes(data, &statsRlp)
		}
		if err == nil {
			for _, r := range statsRlp {
				if stats[r.MsgCode] != nil {
					if l := linRegFromBytes(r.Data); l != nil {
						stats[r.MsgCode] = l
					}
				}
			}
		}
	}

	return &requestCostStats{
		db:    db,
		stats: stats,
	}
}

func (s *requestCostStats) store() {
	s.lock.Lock()
	defer s.lock.Unlock()

	statsRlp := make(requestCostStatsRlp, len(reqList))
	for i, code := range reqList {
		statsRlp[i].MsgCode = code
		statsRlp[i].Data = s.stats[code].toBytes()
	}

	if data, err := rlp.EncodeToBytes(statsRlp); err == nil {
		s.db.Put(rcStatsKey, data)
	}
}

func (s *requestCostStats) getCurrentList() RequestCostList {
	s.lock.Lock()
	defer s.lock.Unlock()

	list := make(RequestCostList, len(reqList))
	//fmt.Println("RequestCostList")
	for idx, code := range reqList {
		b, m := s.stats[code].calc()
		//fmt.Println(code, s.stats[code].cnt, b/1000000, m/1000000)
		if m < 0 {
			b += m
			m = 0
		}
		if b < 0 {
			b = 0
		}

		list[idx].MsgCode = code
		list[idx].BaseCost = uint64(b * 2)
		list[idx].ReqCost = uint64(m * 2)
	}
	return list
}

func (s *requestCostStats) update(msgCode, reqCnt, cost uint64) {
	s.lock.Lock()
	defer s.lock.Unlock()

	c, ok := s.stats[msgCode]
	if !ok || reqCnt == 0 {
		return
	}
	c.add(float64(reqCnt), float64(cost))
}

func (pm *ProtocolManager) blockLoop() {
	pm.wg.Add(1)
	sub := pm.eventMux.Subscribe(core.ChainHeadEvent{})
	newCht := make(chan struct{}, 10)
	newCht <- struct{}{}
	go func() {
		var mu sync.Mutex
		var lastHead *types.Header
		lastBroadcastTd := common.Big0
		for {
			select {
			case ev := <-sub.Chan():
				peers := pm.peers.AllPeers()
				if len(peers) > 0 {
					header := ev.Data.(core.ChainHeadEvent).Block.Header()
					hash := header.Hash()
					number := header.Number.Uint64()
					td := core.GetTd(pm.chainDb, hash, number)
					if td != nil && td.Cmp(lastBroadcastTd) > 0 {
						var reorg uint64
						if lastHead != nil {
							reorg = lastHead.Number.Uint64() - core.FindCommonAncestor(pm.chainDb, header, lastHead).Number.Uint64()
						}
						lastHead = header
						lastBroadcastTd = td

						glog.V(logger.Debug).Infoln("===> ", number, hash, td, reorg)

						announce := announceData{Hash: hash, Number: number, Td: td, ReorgDepth: reorg}
						for _, p := range peers {
							select {
							case p.announceChn <- announce:
							default:
								pm.removePeer(p.id)
							}
						}
					}
				}
				newCht <- struct{}{}
			case <-newCht:
				go func() {
					mu.Lock()
					more := makeCht(pm.chainDb)
					mu.Unlock()
					if more {
						time.Sleep(time.Millisecond * 10)
						newCht <- struct{}{}
					}
				}()
			case <-pm.quitSync:
				sub.Unsubscribe()
				pm.wg.Done()
				return
			}
		}
	}()
}

var (
	lastChtKey       = []byte("LastChtNumber") // chtNum (uint64 big endian)
	chtPrefix        = []byte("cht")           // chtPrefix + chtNum (uint64 big endian) -> trie root hash
	chtConfirmations = light.ChtFrequency / 2
)

func getChtRoot(db ethdb.Database, num uint64) common.Hash {
	var encNumber [8]byte
	binary.BigEndian.PutUint64(encNumber[:], num)
	data, _ := db.Get(append(chtPrefix, encNumber[:]...))
	return common.BytesToHash(data)
}

func storeChtRoot(db ethdb.Database, num uint64, root common.Hash) {
	var encNumber [8]byte
	binary.BigEndian.PutUint64(encNumber[:], num)
	db.Put(append(chtPrefix, encNumber[:]...), root[:])
}

func makeCht(db ethdb.Database) bool {
	headHash := core.GetHeadBlockHash(db)
	headNum := core.GetBlockNumber(db, headHash)

	var newChtNum uint64
	if headNum > chtConfirmations {
		newChtNum = (headNum - chtConfirmations) / light.ChtFrequency
	}

	var lastChtNum uint64
	data, _ := db.Get(lastChtKey)
	if len(data) == 8 {
		lastChtNum = binary.BigEndian.Uint64(data[:])
	}
	if newChtNum <= lastChtNum {
		return false
	}

	var t *trie.Trie
	if lastChtNum > 0 {
		var err error
		t, err = trie.New(getChtRoot(db, lastChtNum), db)
		if err != nil {
			lastChtNum = 0
		}
	}
	if lastChtNum == 0 {
		t, _ = trie.New(common.Hash{}, db)
	}

	for num := lastChtNum * light.ChtFrequency; num < (lastChtNum+1)*light.ChtFrequency; num++ {
		hash := core.GetCanonicalHash(db, num)
		if hash == (common.Hash{}) {
			panic("Canonical hash not found")
		}
		td := core.GetTd(db, hash, num)
		if td == nil {
			panic("TD not found")
		}
		var encNumber [8]byte
		binary.BigEndian.PutUint64(encNumber[:], num)
		var node light.ChtNode
		node.Hash = hash
		node.Td = td
		data, _ := rlp.EncodeToBytes(node)
		t.Update(encNumber[:], data)
	}

	root, err := t.Commit()
	if err != nil {
		lastChtNum = 0
	} else {
		lastChtNum++

		glog.V(logger.Detail).Infof("cht: %d %064x", lastChtNum, root)

		storeChtRoot(db, lastChtNum, root)
		var data [8]byte
		binary.BigEndian.PutUint64(data[:], lastChtNum)
		db.Put(lastChtKey, data[:])
	}

	return newChtNum > lastChtNum
}
<|endoftext|>


<commit_before>// Copyright 2017 David Ackroyd. All Rights Reserved.
// See LICENSE for licensing terms.

package grpc_recovery_test

import (
	"github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/grpc-ecosystem/go-grpc-middleware/recovery"
	"google.golang.org/grpc"
)

var (
	customFunc grpc_recovery.RecoveryHandlerFunc
)

// Initialization shows an initialization sequence with a custom recovery handler func.
func Example_initialization() {
	// Shared options for the logger, with a custom gRPC code to log level function.
	opts := []grpc_recovery.Option{
		grpc_recovery.WithRecoveryHandler(customFunc),
	}
	// Create a server. Recovery handlers should typically be last in the chain so that other middleware
	// (e.g. logging) can operate on the recovered state instead of being directly affected by any panic
	_ = grpc.NewServer(
		grpc_middleware.WithUnaryServerChain(
			grpc_recovery.UnaryServerInterceptor(opts...),
		),
		grpc_middleware.WithStreamServerChain(
			grpc_recovery.StreamServerInterceptor(opts...),
		),
	)
}
<commit_msg>recovery: add custom func on recovery example<commit_after>// Copyright 2017 David Ackroyd. All Rights Reserved.
// See LICENSE for licensing terms.

package grpc_recovery_test

import (
	"github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/grpc-ecosystem/go-grpc-middleware/recovery"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc"
)

var (
	customFunc grpc_recovery.RecoveryHandlerFunc
)

// Initialization shows an initialization sequence with a custom recovery handler func.
func Example_initialization() {
	// Define customfunc to handle panic
	customFunc = func(p interface{}) (err error) {
		return grpc.Errorf(codes.Unknown, "panic triggered: %v", p)
	}
	// Shared options for the logger, with a custom gRPC code to log level function.
	opts := []grpc_recovery.Option{
		grpc_recovery.WithRecoveryHandler(customFunc),
	}
	// Create a server. Recovery handlers should typically be last in the chain so that other middleware
	// (e.g. logging) can operate on the recovered state instead of being directly affected by any panic
	_ = grpc.NewServer(
		grpc_middleware.WithUnaryServerChain(
			grpc_recovery.UnaryServerInterceptor(opts...),
		),
		grpc_middleware.WithStreamServerChain(
			grpc_recovery.StreamServerInterceptor(opts...),
		),
	)
}
<|endoftext|>


<commit_before>package consul

import (
	"crypto/tls"
	"errors"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/king-jam/libreg"
	"github.com/king-jam/libreg/registry"
)

var (
	// ErrMultipleEndpointsUnsupported is thrown when there are
	// multiple endpoints specified for Consul
	ErrMultipleEndpointsUnsupported = errors.New("consul does not support multiple endpoints")

	// ErrSessionRenew is thrown when the session can't be
	// renewed because the Consul version does not support sessions
	ErrSessionRenew = errors.New("cannot set or renew session for ttl, unable to operate on sessions")
)

// Consul is the receiver type for the
// Store interface
type Consul struct {
	sync.Mutex
	config *api.Config
	client *api.Client
}

// Register registers consul to libkv
func Register() {
	libreg.AddRegistry(registry.CONSUL, New)
}

// New creates a new Consul client given a list
// of endpoints and optional tls config
func New(endpoints []string, options *store.Config) (store.Store, error) {
	if len(endpoints) > 1 {
		return nil, ErrMultipleEndpointsUnsupported
	}

	s := &Consul{}

	// Create Consul client
	config := api.DefaultConfig()
	s.config = config
	config.HttpClient = http.DefaultClient
	config.Address = endpoints[0]
	config.Scheme = "http"

	// Set options
	if options != nil {
		if options.TLS != nil {
			s.setTLS(options.TLS)
		}
		if options.ConnectionTimeout != 0 {
			s.setTimeout(options.ConnectionTimeout)
		}
	}

	// Creates a new client
	client, err := api.NewClient(config)
	if err != nil {
		return nil, err
	}
	s.client = client

	return s, nil
}

// SetTLS sets Consul TLS options
func (s *Consul) setTLS(tls *tls.Config) {
	s.config.HttpClient.Transport = &http.Transport{
		TLSClientConfig: tls,
	}
	s.config.Scheme = "https"
}

// SetTimeout sets the timeout for connecting to Consul
func (s *Consul) setTimeout(time time.Duration) {
	s.config.WaitTime = time
}

func (s *Consul) Register(reg *CatalogRegistration, options *WriteOptions) error {
	catalog := s.api.Catalog()
	_, err := catalog.Register(reg, options)
	return err
}

// Deregister removes a node, service or check
func (s *Consul) Deregister(dereg *CatalogDeregistration, options *WriteOptions) error {
	catalog := s.api.Catalog()
	_, err := catalog.Deregister(dereg, options)
	return err
}

// Datacenters lists known datacenters
func (s *Consul) Datacenters() ([]string, error) {
	catalog := s.api.Catalog()
	dc, err := catalog.Datacenters()
	return dc, err
}

// Nodes lists all nodes in a given DC
func (s *Consul) Nodes(options *QueryOptions) ([]*Node, error) {
	catalog := s.api.Catalog()
	nodes, _, err := catalog.Nodes(options)
	return nodes, err
}

// Services lists all services in a given DC
func (s *Consul) Services(options *QueryOptions) (map[string][]string, error) {
	catalog := s.api.Catalog()
	services, _, err := catalog.Services(options)
	return services, err
}

// Service lists the nodes in a given service
func (s *Consul) Service(service, tag string, options *QueryOptions) ([]*CatalogService, error) {
	catalog := s.api.Catalog()
	services, _, err := catalog.Services(service, tag, options)
	return services, err
}

// Node lists the services provided by a given node
func (s *Consul) Node(node string, options *QueryOptions) (*CatalogNode, error) {
	catalog := s.api.Catalog()
	node, _, err := catalog.Node(node, options)
	return node, err
}
<commit_msg>Update consul.go<commit_after>package consul

import (
	"crypto/tls"
	"errors"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/king-jam/libreg"
	"github.com/king-jam/libreg/registry"
)

var (
	// ErrMultipleEndpointsUnsupported is thrown when there are
	// multiple endpoints specified for Consul
	ErrMultipleEndpointsUnsupported = errors.New("consul does not support multiple endpoints")

	// ErrSessionRenew is thrown when the session can't be
	// renewed because the Consul version does not support sessions
	ErrSessionRenew = errors.New("cannot set or renew session for ttl, unable to operate on sessions")
)

// Consul is the receiver type for the
// Store interface
type Consul struct {
	sync.Mutex
	config *api.Config
	client *api.Client
}

// Register registers consul to libreg
func Register() {
	libreg.AddRegistry(registry.CONSUL, New)
}

// New creates a new Consul client given a list
// of endpoints and optional tls config
func New(endpoints []string, options *store.Config) (store.Store, error) {
	if len(endpoints) > 1 {
		return nil, ErrMultipleEndpointsUnsupported
	}

	s := &Consul{}

	// Create Consul client
	config := api.DefaultConfig()
	s.config = config
	config.HttpClient = http.DefaultClient
	config.Address = endpoints[0]
	config.Scheme = "http"

	// Set options
	if options != nil {
		if options.TLS != nil {
			s.setTLS(options.TLS)
		}
		if options.ConnectionTimeout != 0 {
			s.setTimeout(options.ConnectionTimeout)
		}
	}

	// Creates a new client
	client, err := api.NewClient(config)
	if err != nil {
		return nil, err
	}
	s.client = client

	return s, nil
}

// SetTLS sets Consul TLS options
func (s *Consul) setTLS(tls *tls.Config) {
	s.config.HttpClient.Transport = &http.Transport{
		TLSClientConfig: tls,
	}
	s.config.Scheme = "https"
}

// SetTimeout sets the timeout for connecting to Consul
func (s *Consul) setTimeout(time time.Duration) {
	s.config.WaitTime = time
}

func (s *Consul) Register(reg *CatalogRegistration, options *WriteOptions) error {
	catalog := s.api.Catalog()
	_, err := catalog.Register(reg, options)
	return err
}

// Deregister removes a node, service or check
func (s *Consul) Deregister(dereg *CatalogDeregistration, options *WriteOptions) error {
	catalog := s.api.Catalog()
	_, err := catalog.Deregister(dereg, options)
	return err
}

// Datacenters lists known datacenters
func (s *Consul) Datacenters() ([]string, error) {
	catalog := s.api.Catalog()
	dc, err := catalog.Datacenters()
	return dc, err
}

// Nodes lists all nodes in a given DC
func (s *Consul) Nodes(options *QueryOptions) ([]*Node, error) {
	catalog := s.api.Catalog()
	nodes, _, err := catalog.Nodes(options)
	return nodes, err
}

// Services lists all services in a given DC
func (s *Consul) Services(options *QueryOptions) (map[string][]string, error) {
	catalog := s.api.Catalog()
	services, _, err := catalog.Services(options)
	return services, err
}

// Service lists the nodes in a given service
func (s *Consul) Service(service, tag string, options *QueryOptions) ([]*CatalogService, error) {
	catalog := s.api.Catalog()
	services, _, err := catalog.Services(service, tag, options)
	return services, err
}

// Node lists the services provided by a given node
func (s *Consul) Node(node string, options *QueryOptions) (*CatalogNode, error) {
	catalog := s.api.Catalog()
	node, _, err := catalog.Node(node, options)
	return node, err
}
<|endoftext|>


<commit_before>package remote

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/gruntwork-io/terragrunt/aws_helper"
	"github.com/gruntwork-io/terragrunt/dynamodb"
	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/gruntwork-io/terragrunt/options"
	"github.com/gruntwork-io/terragrunt/shell"
	"github.com/mitchellh/mapstructure"
	"time"
)

// A representation of the configuration options available for S3 remote state
type RemoteStateConfigS3 struct {
	Encrypt       bool   `mapstructure:"encrypt"`
	Bucket        string `mapstructure:"bucket"`
	Key           string `mapstructure:"key"`
	Region        string `mapstructure:"region"`
	Profile       string `mapstructure:"profile"`
	RoleArn       string `mapstructure:"role_arn"`
	LockTable     string `mapstructure:"lock_table"`
	DynamoDBTable string `mapstructure:"dynamodb_table"`
}

// The DynamoDB lock table name used to be called lock_table, but has since been renamed to dynamodb_table, and the old
// name deprecated. To maintain backwards compatibility, we support both names.
func (s3Config *RemoteStateConfigS3) GetLockTableName() string {
	if s3Config.DynamoDBTable != "" {
		return s3Config.DynamoDBTable
	}
	return s3Config.LockTable
}

const MAX_RETRIES_WAITING_FOR_S3_BUCKET = 12
const SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET = 5 * time.Second

type S3Initializer struct{}

// Returns true if the S3 bucket or DynamoDB table does not exist
func (s3Initializer S3Initializer) NeedsInitialization(config map[string]interface{}, terragruntOptions *options.TerragruntOptions) (bool, error) {
	s3Config, err := parseS3Config(config)
	if err != nil {
		return false, err
	}

	s3Client, err := CreateS3Client(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return false, err
	}

	if !DoesS3BucketExist(s3Client, s3Config) {
		return true, nil
	}

	if s3Config.GetLockTableName() != "" {
		dynamodbClient, err := dynamodb.CreateDynamoDbClient(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
		if err != nil {
			return false, err
		}

		tableExists, err := dynamodb.LockTableExistsAndIsActive(s3Config.GetLockTableName(), dynamodbClient)
		if err != nil {
			return false, err
		}
		if !tableExists {
			return true, nil
		}
	}

	return false, nil
}

// Initialize the remote state S3 bucket specified in the given config. This function will validate the config
// parameters, create the S3 bucket if it doesn't already exist, and check that versioning is enabled.
func (s3Initializer S3Initializer) Initialize(config map[string]interface{}, terragruntOptions *options.TerragruntOptions) error {
	s3Config, err := parseS3Config(config)
	if err != nil {
		return err
	}

	if err := validateS3Config(s3Config, terragruntOptions); err != nil {
		return err
	}

	s3Client, err := CreateS3Client(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return err
	}

	if err := createS3BucketIfNecessary(s3Client, s3Config, terragruntOptions); err != nil {
		return err
	}

	if err := checkIfVersioningEnabled(s3Client, s3Config, terragruntOptions); err != nil {
		return err
	}

	if err := createLockTableIfNecessary(s3Config, terragruntOptions); err != nil {
		return err
	}

	return nil
}

// Parse the given map into an S3 config
func parseS3Config(config map[string]interface{}) (*RemoteStateConfigS3, error) {
	var s3Config RemoteStateConfigS3
	if err := mapstructure.Decode(config, &s3Config); err != nil {
		return nil, errors.WithStackTrace(err)
	}

	return &s3Config, nil
}

// Validate all the parameters of the given S3 remote state configuration
func validateS3Config(config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if config.Region == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("region"))
	}

	if config.Bucket == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("bucket"))
	}

	if config.Key == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("key"))
	}

	if !config.Encrypt {
		terragruntOptions.Logger.Printf("WARNING: encryption is not enabled on the S3 remote state bucket %s. Terraform state files may contain secrets, so we STRONGLY recommend enabling encryption!", config.Bucket)
	}

	return nil
}

// If the bucket specified in the given config doesn't already exist, prompt the user to create it, and if the user
// confirms, create the bucket and enable versioning for it.
func createS3BucketIfNecessary(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if !DoesS3BucketExist(s3Client, config) {
		prompt := fmt.Sprintf("Remote state S3 bucket %s does not exist or you don't have permissions to access it. Would you like Terragrunt to create it?", config.Bucket)
		shouldCreateBucket, err := shell.PromptUserForYesNo(prompt, terragruntOptions)
		if err != nil {
			return err
		}

		if shouldCreateBucket {
			return CreateS3BucketWithVersioning(s3Client, config, terragruntOptions)
		}
	}

	return nil
}

// Check if versioning is enabled for the S3 bucket specified in the given config and warn the user if it is not
func checkIfVersioningEnabled(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	out, err := s3Client.GetBucketVersioning(&s3.GetBucketVersioningInput{Bucket: aws.String(config.Bucket)})
	if err != nil {
		return errors.WithStackTrace(err)
	}

	// NOTE: There must be a bug in the AWS SDK since out == nil when versioning is not enabled. In the future,
	// check the AWS SDK for updates to see if we can remove "out == nil ||".
	if out == nil || out.Status == nil || *out.Status != s3.BucketVersioningStatusEnabled {
		terragruntOptions.Logger.Printf("WARNING: Versioning is not enabled for the remote state S3 bucket %s. We recommend enabling versioning so that you can roll back to previous versions of your Terraform state in case of error.", config.Bucket)
	}

	return nil
}

// Create the given S3 bucket and enable versioning for it
func CreateS3BucketWithVersioning(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if err := CreateS3Bucket(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	if err := WaitUntilS3BucketExists(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	if err := EnableVersioningForS3Bucket(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	return nil
}

// AWS is eventually consistent, so after creating an S3 bucket, this method can be used to wait until the information
// about that S3 bucket has propagated everywhere
func WaitUntilS3BucketExists(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	for retries := 0; retries < MAX_RETRIES_WAITING_FOR_S3_BUCKET; retries++ {
		if DoesS3BucketExist(s3Client, config) {
			terragruntOptions.Logger.Printf("S3 bucket %s created.", config.Bucket)
			return nil
		} else if retries < MAX_RETRIES_WAITING_FOR_S3_BUCKET-1 {
			terragruntOptions.Logger.Printf("S3 bucket %s has not been created yet. Sleeping for %s and will check again.", config.Bucket, SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET)
			time.Sleep(SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET)
		}
	}

	return errors.WithStackTrace(MaxRetriesWaitingForS3BucketExceeded(config.Bucket))
}

// Create the S3 bucket specified in the given config
func CreateS3Bucket(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	terragruntOptions.Logger.Printf("Creating S3 bucket %s", config.Bucket)
	_, err := s3Client.CreateBucket(&s3.CreateBucketInput{Bucket: aws.String(config.Bucket)})
	return errors.WithStackTrace(err)
}

// Enable versioning for the S3 bucket specified in the given config
func EnableVersioningForS3Bucket(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	terragruntOptions.Logger.Printf("Enabling versioning on S3 bucket %s", config.Bucket)
	input := s3.PutBucketVersioningInput{
		Bucket:                  aws.String(config.Bucket),
		VersioningConfiguration: &s3.VersioningConfiguration{Status: aws.String(s3.BucketVersioningStatusEnabled)},
	}
	_, err := s3Client.PutBucketVersioning(&input)
	return errors.WithStackTrace(err)
}

// Returns true if the S3 bucket specified in the given config exists and the current user has the ability to access
// it.
func DoesS3BucketExist(s3Client *s3.S3, config *RemoteStateConfigS3) bool {
	_, err := s3Client.HeadBucket(&s3.HeadBucketInput{Bucket: aws.String(config.Bucket)})
	return err == nil
}

// Create a table for locks in DynamoDB if the user has configured a lock table and the table doesn't already exist
func createLockTableIfNecessary(s3Config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if s3Config.GetLockTableName() == "" {
		return nil
	}

	dynamodbClient, err := dynamodb.CreateDynamoDbClient(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return err
	}

	return dynamodb.CreateLockTableIfNecessary(s3Config.GetLockTableName(), dynamodbClient, terragruntOptions)
}

// Create an authenticated client for DynamoDB
func CreateS3Client(awsRegion, awsProfile string, iamRoleArn string, terragruntOptions *options.TerragruntOptions) (*s3.S3, error) {
	session, err := aws_helper.CreateAwsSession(awsRegion, awsProfile, iamRoleArn, terragruntOptions)
	if err != nil {
		return nil, err
	}

	return s3.New(session), nil
}

// Custom error types

type MissingRequiredS3RemoteStateConfig string

func (configName MissingRequiredS3RemoteStateConfig) Error() string {
	return fmt.Sprintf("Missing required S3 remote state configuration %s", string(configName))
}

type MaxRetriesWaitingForS3BucketExceeded string

func (err MaxRetriesWaitingForS3BucketExceeded) Error() string {
	return fmt.Sprintf("Exceeded max retries (%d) waiting for bucket S3 bucket %s", MAX_RETRIES_WAITING_FOR_S3_BUCKET, string(err))
}
<commit_msg>Handle errors from creating the S3 bucket concurrently<commit_after>package remote

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/gruntwork-io/terragrunt/aws_helper"
	"github.com/gruntwork-io/terragrunt/dynamodb"
	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/gruntwork-io/terragrunt/options"
	"github.com/gruntwork-io/terragrunt/shell"
	"github.com/mitchellh/mapstructure"
	"time"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

// A representation of the configuration options available for S3 remote state
type RemoteStateConfigS3 struct {
	Encrypt       bool   `mapstructure:"encrypt"`
	Bucket        string `mapstructure:"bucket"`
	Key           string `mapstructure:"key"`
	Region        string `mapstructure:"region"`
	Profile       string `mapstructure:"profile"`
	RoleArn       string `mapstructure:"role_arn"`
	LockTable     string `mapstructure:"lock_table"`
	DynamoDBTable string `mapstructure:"dynamodb_table"`
}

// The DynamoDB lock table name used to be called lock_table, but has since been renamed to dynamodb_table, and the old
// name deprecated. To maintain backwards compatibility, we support both names.
func (s3Config *RemoteStateConfigS3) GetLockTableName() string {
	if s3Config.DynamoDBTable != "" {
		return s3Config.DynamoDBTable
	}
	return s3Config.LockTable
}

const MAX_RETRIES_WAITING_FOR_S3_BUCKET = 12
const SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET = 5 * time.Second

type S3Initializer struct{}

// Returns true if the S3 bucket or DynamoDB table does not exist
func (s3Initializer S3Initializer) NeedsInitialization(config map[string]interface{}, terragruntOptions *options.TerragruntOptions) (bool, error) {
	s3Config, err := parseS3Config(config)
	if err != nil {
		return false, err
	}

	s3Client, err := CreateS3Client(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return false, err
	}

	if !DoesS3BucketExist(s3Client, s3Config) {
		return true, nil
	}

	if s3Config.GetLockTableName() != "" {
		dynamodbClient, err := dynamodb.CreateDynamoDbClient(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
		if err != nil {
			return false, err
		}

		tableExists, err := dynamodb.LockTableExistsAndIsActive(s3Config.GetLockTableName(), dynamodbClient)
		if err != nil {
			return false, err
		}
		if !tableExists {
			return true, nil
		}
	}

	return false, nil
}

// Initialize the remote state S3 bucket specified in the given config. This function will validate the config
// parameters, create the S3 bucket if it doesn't already exist, and check that versioning is enabled.
func (s3Initializer S3Initializer) Initialize(config map[string]interface{}, terragruntOptions *options.TerragruntOptions) error {
	s3Config, err := parseS3Config(config)
	if err != nil {
		return err
	}

	if err := validateS3Config(s3Config, terragruntOptions); err != nil {
		return err
	}

	s3Client, err := CreateS3Client(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return err
	}

	if err := createS3BucketIfNecessary(s3Client, s3Config, terragruntOptions); err != nil {
		return err
	}

	if err := checkIfVersioningEnabled(s3Client, s3Config, terragruntOptions); err != nil {
		return err
	}

	if err := createLockTableIfNecessary(s3Config, terragruntOptions); err != nil {
		return err
	}

	return nil
}

// Parse the given map into an S3 config
func parseS3Config(config map[string]interface{}) (*RemoteStateConfigS3, error) {
	var s3Config RemoteStateConfigS3
	if err := mapstructure.Decode(config, &s3Config); err != nil {
		return nil, errors.WithStackTrace(err)
	}

	return &s3Config, nil
}

// Validate all the parameters of the given S3 remote state configuration
func validateS3Config(config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if config.Region == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("region"))
	}

	if config.Bucket == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("bucket"))
	}

	if config.Key == "" {
		return errors.WithStackTrace(MissingRequiredS3RemoteStateConfig("key"))
	}

	if !config.Encrypt {
		terragruntOptions.Logger.Printf("WARNING: encryption is not enabled on the S3 remote state bucket %s. Terraform state files may contain secrets, so we STRONGLY recommend enabling encryption!", config.Bucket)
	}

	return nil
}

// If the bucket specified in the given config doesn't already exist, prompt the user to create it, and if the user
// confirms, create the bucket and enable versioning for it.
func createS3BucketIfNecessary(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if !DoesS3BucketExist(s3Client, config) {
		prompt := fmt.Sprintf("Remote state S3 bucket %s does not exist or you don't have permissions to access it. Would you like Terragrunt to create it?", config.Bucket)
		shouldCreateBucket, err := shell.PromptUserForYesNo(prompt, terragruntOptions)
		if err != nil {
			return err
		}

		if shouldCreateBucket {
			return CreateS3BucketWithVersioning(s3Client, config, terragruntOptions)
		}
	}

	return nil
}

// Check if versioning is enabled for the S3 bucket specified in the given config and warn the user if it is not
func checkIfVersioningEnabled(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	out, err := s3Client.GetBucketVersioning(&s3.GetBucketVersioningInput{Bucket: aws.String(config.Bucket)})
	if err != nil {
		return errors.WithStackTrace(err)
	}

	// NOTE: There must be a bug in the AWS SDK since out == nil when versioning is not enabled. In the future,
	// check the AWS SDK for updates to see if we can remove "out == nil ||".
	if out == nil || out.Status == nil || *out.Status != s3.BucketVersioningStatusEnabled {
		terragruntOptions.Logger.Printf("WARNING: Versioning is not enabled for the remote state S3 bucket %s. We recommend enabling versioning so that you can roll back to previous versions of your Terraform state in case of error.", config.Bucket)
	}

	return nil
}

// Create the given S3 bucket and enable versioning for it
func CreateS3BucketWithVersioning(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if err := CreateS3Bucket(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	if err := WaitUntilS3BucketExists(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	if err := EnableVersioningForS3Bucket(s3Client, config, terragruntOptions); err != nil {
		return err
	}

	return nil
}

// AWS is eventually consistent, so after creating an S3 bucket, this method can be used to wait until the information
// about that S3 bucket has propagated everywhere
func WaitUntilS3BucketExists(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	for retries := 0; retries < MAX_RETRIES_WAITING_FOR_S3_BUCKET; retries++ {
		if DoesS3BucketExist(s3Client, config) {
			terragruntOptions.Logger.Printf("S3 bucket %s created.", config.Bucket)
			return nil
		} else if retries < MAX_RETRIES_WAITING_FOR_S3_BUCKET-1 {
			terragruntOptions.Logger.Printf("S3 bucket %s has not been created yet. Sleeping for %s and will check again.", config.Bucket, SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET)
			time.Sleep(SLEEP_BETWEEN_RETRIES_WAITING_FOR_S3_BUCKET)
		}
	}

	return errors.WithStackTrace(MaxRetriesWaitingForS3BucketExceeded(config.Bucket))
}

// Create the S3 bucket specified in the given config
func CreateS3Bucket(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	terragruntOptions.Logger.Printf("Creating S3 bucket %s", config.Bucket)
	_, err := s3Client.CreateBucket(&s3.CreateBucketInput{Bucket: aws.String(config.Bucket)})

	if err != nil {
		if isBucketAlreadyOwnedByYourError(err) {
			terragruntOptions.Logger.Printf("Looks like someone created bucket %s at the same time. Will wait for it to be in active state.", config.Bucket)
			return nil
		} else {
			return errors.WithStackTrace(err)
		}
	}

	return nil
}

// Determine if this is an error that implies you've already made a request to create the S3 bucket and it succeeded
// or is in progress. This usually happens when running many tests in parallel or xxx-all commands.
func isBucketAlreadyOwnedByYourError(err error) bool {
	awsErr, isAwsErr := err.(awserr.Error)
	return isAwsErr && (awsErr.Code() == "BucketAlreadyOwnedByYou" || awsErr.Code() == "OperationAborted")
}

// Enable versioning for the S3 bucket specified in the given config
func EnableVersioningForS3Bucket(s3Client *s3.S3, config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	terragruntOptions.Logger.Printf("Enabling versioning on S3 bucket %s", config.Bucket)
	input := s3.PutBucketVersioningInput{
		Bucket:                  aws.String(config.Bucket),
		VersioningConfiguration: &s3.VersioningConfiguration{Status: aws.String(s3.BucketVersioningStatusEnabled)},
	}
	_, err := s3Client.PutBucketVersioning(&input)
	return errors.WithStackTrace(err)
}

// Returns true if the S3 bucket specified in the given config exists and the current user has the ability to access
// it.
func DoesS3BucketExist(s3Client *s3.S3, config *RemoteStateConfigS3) bool {
	_, err := s3Client.HeadBucket(&s3.HeadBucketInput{Bucket: aws.String(config.Bucket)})
	return err == nil
}

// Create a table for locks in DynamoDB if the user has configured a lock table and the table doesn't already exist
func createLockTableIfNecessary(s3Config *RemoteStateConfigS3, terragruntOptions *options.TerragruntOptions) error {
	if s3Config.GetLockTableName() == "" {
		return nil
	}

	dynamodbClient, err := dynamodb.CreateDynamoDbClient(s3Config.Region, s3Config.Profile, s3Config.RoleArn, terragruntOptions)
	if err != nil {
		return err
	}

	return dynamodb.CreateLockTableIfNecessary(s3Config.GetLockTableName(), dynamodbClient, terragruntOptions)
}

// Create an authenticated client for DynamoDB
func CreateS3Client(awsRegion, awsProfile string, iamRoleArn string, terragruntOptions *options.TerragruntOptions) (*s3.S3, error) {
	session, err := aws_helper.CreateAwsSession(awsRegion, awsProfile, iamRoleArn, terragruntOptions)
	if err != nil {
		return nil, err
	}

	return s3.New(session), nil
}

// Custom error types

type MissingRequiredS3RemoteStateConfig string

func (configName MissingRequiredS3RemoteStateConfig) Error() string {
	return fmt.Sprintf("Missing required S3 remote state configuration %s", string(configName))
}

type MaxRetriesWaitingForS3BucketExceeded string

func (err MaxRetriesWaitingForS3BucketExceeded) Error() string {
	return fmt.Sprintf("Exceeded max retries (%d) waiting for bucket S3 bucket %s", MAX_RETRIES_WAITING_FOR_S3_BUCKET, string(err))
}
<|endoftext|>


<commit_before>package helpers

import (
	"encoding/json"
	"fmt"
	"html/template"
	"strings"

	"github.com/aymerick/raymond"
	"github.com/markbates/inflect"
	"github.com/shurcooL/github_flavored_markdown"
)

// Helpers that are automatically injected into templates.
var Helpers = map[string]interface{}{
	"js_escape":   template.JSEscapeString,
	"html_escape": template.HTMLEscapeString,
	"json":        ToJSON,
	"content_for": ContentFor,
	"content_of":  ContentOf,
	"upcase":      strings.ToUpper,
	"downcase":    strings.ToLower,
	"markdown":    Markdown,
	"debug":       Debug,
}

func init() {
	for k, v := range inflect.Helpers {
		Helpers[k] = v
	}
}

// ToJSON converts an interface into a string.
func ToJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

// ContentFor stores a block of templating code to be re-used later in the template.
/*
	{{content_for "buttons"}}
		<button>hi</button>
	{{/content_for}}
*/
func ContentFor(name string, options *raymond.Options) string {
	ctx := options.Ctx().(map[string]interface{})
	body := options.Fn()
	ctx[name] = raymond.SafeString(body)
	return ""
}

// ContentOf retrieves a stored block for templating and renders it.
/*
	{{content_of "buttons"}}
*/
func ContentOf(name string, options *raymond.Options) raymond.SafeString {
	ctx := options.Ctx().(map[string]interface{})
	if s, ok := ctx[name]; ok {
		return s.(raymond.SafeString)
	}
	return raymond.SafeString("")
}

// Markdown converts the string into HTML using GitHub flavored markdown.
func Markdown(body string) raymond.SafeString {
	b := github_flavored_markdown.Markdown([]byte(body))
	return raymond.SafeString(string(b))
}

// Debug by verbosely printing out using 'pre' tags.
func Debug(v interface{}) raymond.SafeString {
	return raymond.SafeString(fmt.Sprintf("
%+v
", v))
}
<commit_msg>fixed a typo<commit_after>package helpers

import (
	"encoding/json"
	"fmt"
	"html/template"
	"strings"

	"github.com/aymerick/raymond"
	"github.com/markbates/inflect"
	"github.com/shurcooL/github_flavored_markdown"
)

// Helpers that are automatically injected into templates.
var Helpers = map[string]interface{}{
	"js_escape":   template.JSEscapeString,
	"html_escape": template.HTMLEscapeString,
	"json":        ToJSON,
	"content_for": ContentFor,
	"content_of":  ContentOf,
	"upcase":      strings.ToUpper,
	"downcase":    strings.ToLower,
	"markdown":    Markdown,
	"debug":       Debug,
}

func init() {
	for k, v := range inflect.Helpers {
		Helpers[k] = v
	}
}

// ToJSON converts an interface into a string.
func ToJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		return err.Error()
	}
	return string(b)
}

// ContentFor stores a block of templating code to be re-used later in the template.
/*
	{{#content_for "buttons"}}
		<button>hi</button>
	{{/content_for}}
*/
func ContentFor(name string, options *raymond.Options) string {
	ctx := options.Ctx().(map[string]interface{})
	body := options.Fn()
	ctx[name] = raymond.SafeString(body)
	return ""
}

// ContentOf retrieves a stored block for templating and renders it.
/*
	{{content_of "buttons"}}
*/
func ContentOf(name string, options *raymond.Options) raymond.SafeString {
	ctx := options.Ctx().(map[string]interface{})
	if s, ok := ctx[name]; ok {
		return s.(raymond.SafeString)
	}
	return raymond.SafeString("")
}

// Markdown converts the string into HTML using GitHub flavored markdown.
func Markdown(body string) raymond.SafeString {
	b := github_flavored_markdown.Markdown([]byte(body))
	return raymond.SafeString(string(b))
}

// Debug by verbosely printing out using 'pre' tags.
func Debug(v interface{}) raymond.SafeString {
	return raymond.SafeString(fmt.Sprintf("
%+v
", v))
}
<|endoftext|>