<commit_before>package it

import (
	"bytes"
	"net/url"
	"os"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	"github.com/fxnn/deadbox/config"
	"github.com/fxnn/deadbox/crypto"
	"github.com/fxnn/deadbox/daemon"
	"github.com/fxnn/deadbox/drop"
	"github.com/fxnn/deadbox/model"
	"github.com/fxnn/deadbox/rest"
	"github.com/fxnn/deadbox/worker"
)

const (
	workerDbFileName     = "worker.boltdb"
	workerName           = "itWorker"
	dropDbFileName       = "drop.boltdb"
	dropName             = "itDrop"
	port                 = "54123"
	interactionSleepTime = 500 * time.Millisecond
)

func assertWorkerTimeoutInFuture(actualWorker model.Worker, t *testing.T) {
	t.Helper()
	if actualWorker.Timeout.Before(time.Now()) {
		t.Fatalf("expected worker timeout to be in the future, but was %s", actualWorker.Timeout)
	}
}
func assertWorkerName(actualWorker model.Worker, workerName string, t *testing.T) {
	t.Helper()
	if string(actualWorker.Name) != workerName {
		t.Fatalf("expected worker to be %s, but was %v", workerName, actualWorker)
	}
}
func assertNumberOfWorkers(actualWorkers []model.Worker, expectedNumber int, t *testing.T) {
	t.Helper()
	if len(actualWorkers) != expectedNumber {
		t.Fatalf("expected %d workers, but got %v", expectedNumber, actualWorkers)
	}
}
func assertNumberOfRequests(actualRequests []model.WorkerRequest, expectedNumber int, t *testing.T) {
	t.Helper()
	if len(actualRequests) != expectedNumber {
		t.Fatalf("expected %d requests, but got %v", expectedNumber, actualRequests)
	}
}
func assertRequestId(actualRequest model.WorkerRequest, expectedId string, t *testing.T) {
	t.Helper()
	if string(actualRequest.Id) != expectedId {
		t.Fatalf("expected request to have id %s, but got %s", expectedId, actualRequest.Id)
	}
}
func assertRequestEncryptionType(actualRequest model.WorkerRequest, expectedType string, t *testing.T) {
	t.Helper()
	if actualRequest.EncryptionType != expectedType {
		t.Fatalf("expected request to have encryptionType %s, but got %s", expectedType, actualRequest.EncryptionType)
	}
}
func assertRequestContentContains(actualRequest model.WorkerRequest, expectedContentSubstring string, t *testing.T) {
	t.Helper()
	if !bytes.Contains(actualRequest.Content, []byte(expectedContentSubstring)) {
		t.Fatalf("expected request to have content containing '%s', but got %s", expectedContentSubstring, string(actualRequest.Content))
	}
}

func assertResponseContentType(actualResponse model.WorkerResponse, expectedContentType string, t *testing.T) {
	t.Helper()
	if string(actualResponse.ContentType) != expectedContentType {
		t.Fatalf("expected response to have content type '%s', but got '%s' and content '%s'", expectedContentType, actualResponse.ContentType,
			actualResponse.Content)
	}
}
func assertResponseContent(actualResponse model.WorkerResponse, expectedContent string, t *testing.T) {
	t.Helper()
	if string(actualResponse.Content) != expectedContent {
		t.Fatalf("expected response to have content '%s', but got '%s'", expectedContent, actualResponse.Content)
	}
}

func runDropDaemon(t *testing.T) (daemon.Daemon, model.Drop) {
	t.Helper()

	cfg := &config.Drop{
		Name:                       dropName,
		ListenAddress:              ":" + port,
		MaxRequestTimeoutInSeconds: config.DefaultMaxRequestTimeoutInSeconds,
		MaxWorkerTimeoutInSeconds:  config.DefaultMaxWorkerTimeoutInSeconds,
	}
	db, err := bolt.Open(dropDbFileName, 0664, bolt.DefaultOptions)
	if err != nil {
		t.Fatalf("could not open Drop's BoltDB: %s", err)
	}

	dropDaemon := drop.New(cfg, db)
	dropDaemon.OnStop(func() error {
		if err := db.Close(); err != nil {
			return err
		}
		if err := os.Remove(dropDbFileName); err != nil {
			return err
		}
		return nil
	})
	dropDaemon.Start()

	dropClient := rest.NewClient(parseUrlOrPanic("http://localhost:" + port))

	return dropDaemon, dropClient
}

func runWorkerDaemon(t *testing.T) (worker.Daemonized, []byte) {
	t.Helper()

	cfg := &config.Worker{
		Name:    workerName,
		DropUrl: parseUrlOrPanic("http://localhost:" + port),
		RegistrationTimeoutInSeconds:        config.DefaultRegistrationTimeoutInSeconds,
		UpdateRegistrationIntervalInSeconds: config.DefaultUpdateRegistrationIntervalInSeconds,
	}
	db, err := bolt.Open(workerDbFileName, 0664, bolt.DefaultOptions)
	if err != nil {
		t.Fatalf("could not open Worker's BoltDB: %s", err)
	}
	privateKeyBytes, err := worker.GeneratePrivateKeyBytes(2048)
	if err != nil {
		t.Fatalf("couldn't generate private key: %s", err)
	}

	privateKey, err := crypto.UnmarshalPrivateKeyFromPEMBytes(privateKeyBytes)
	if err != nil {
		t.Fatalf("couldn't unmarshal private key: %s", err)
	}
	publicKeyBytes, err := crypto.GeneratePublicKeyBytes(privateKey)
	if err != nil {
		t.Fatalf("couldn't generate public key bytes: %s", err)
	}
	fingerprint, err := crypto.FingerprintPublicKey(&privateKey.PublicKey, 10, 4)
	if err != nil {
		t.Fatalf("couldn't fingerprint public key: %s", err)
	}

	workerDaemon := worker.New(cfg, fingerprint, db, privateKeyBytes)
	workerDaemon.OnStop(func() error {
		if err := db.Close(); err != nil {
			return err
		}
		if err := os.Remove(workerDbFileName); err != nil {
			return err
		}
		return nil
	})
	workerDaemon.Start()

	return workerDaemon, publicKeyBytes
}

func stopDaemon(d daemon.Daemon, t *testing.T) {
	t.Helper()
	err := d.Stop()
	if err != nil {
		t.Error(err)
	}
}

func parseUrlOrPanic(s string) *url.URL {
	result, err := url.Parse(s)
	if err != nil {
		panic(err)
	}

	return result
}
<commit_msg>#27: fix unit test<commit_after>package it

import (
	"bytes"
	"net/url"
	"os"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	"github.com/fxnn/deadbox/config"
	"github.com/fxnn/deadbox/crypto"
	"github.com/fxnn/deadbox/daemon"
	"github.com/fxnn/deadbox/drop"
	"github.com/fxnn/deadbox/model"
	"github.com/fxnn/deadbox/rest"
	"github.com/fxnn/deadbox/worker"
)

const (
	workerDbFileName     = "worker.boltdb"
	workerName           = "itWorker"
	dropDbFileName       = "drop.boltdb"
	dropName             = "itDrop"
	port                 = "54123"
	interactionSleepTime = 500 * time.Millisecond
)

func assertWorkerTimeoutInFuture(actualWorker model.Worker, t *testing.T) {
	t.Helper()
	if actualWorker.Timeout.Before(time.Now()) {
		t.Fatalf("expected worker timeout to be in the future, but was %s", actualWorker.Timeout)
	}
}
func assertWorkerName(actualWorker model.Worker, workerName string, t *testing.T) {
	t.Helper()
	if string(actualWorker.Name) != workerName {
		t.Fatalf("expected worker to be %s, but was %v", workerName, actualWorker)
	}
}
func assertNumberOfWorkers(actualWorkers []model.Worker, expectedNumber int, t *testing.T) {
	t.Helper()
	if len(actualWorkers) != expectedNumber {
		t.Fatalf("expected %d workers, but got %v", expectedNumber, actualWorkers)
	}
}
func assertNumberOfRequests(actualRequests []model.WorkerRequest, expectedNumber int, t *testing.T) {
	t.Helper()
	if len(actualRequests) != expectedNumber {
		t.Fatalf("expected %d requests, but got %v", expectedNumber, actualRequests)
	}
}
func assertRequestId(actualRequest model.WorkerRequest, expectedId string, t *testing.T) {
	t.Helper()
	if string(actualRequest.Id) != expectedId {
		t.Fatalf("expected request to have id %s, but got %s", expectedId, actualRequest.Id)
	}
}
func assertRequestEncryptionType(actualRequest model.WorkerRequest, expectedType string, t *testing.T) {
	t.Helper()
	if actualRequest.EncryptionType != expectedType {
		t.Fatalf("expected request to have encryptionType %s, but got %s", expectedType, actualRequest.EncryptionType)
	}
}
func assertRequestContentContains(actualRequest model.WorkerRequest, expectedContentSubstring string, t *testing.T) {
	t.Helper()
	if !bytes.Contains(actualRequest.Content, []byte(expectedContentSubstring)) {
		t.Fatalf("expected request to have content containing '%s', but got %s", expectedContentSubstring, string(actualRequest.Content))
	}
}

func assertResponseContentType(actualResponse model.WorkerResponse, expectedContentType string, t *testing.T) {
	t.Helper()
	if string(actualResponse.ContentType) != expectedContentType {
		t.Fatalf("expected response to have content type '%s', but got '%s' and content '%s'", expectedContentType, actualResponse.ContentType,
			actualResponse.Content)
	}
}
func assertResponseContent(actualResponse model.WorkerResponse, expectedContent string, t *testing.T) {
	t.Helper()
	if string(actualResponse.Content) != expectedContent {
		t.Fatalf("expected response to have content '%s', but got '%s'", expectedContent, actualResponse.Content)
	}
}

func runDropDaemon(t *testing.T) (daemon.Daemon, model.Drop) {
	t.Helper()

	cfg := &config.Drop{
		Name:                       dropName,
		ListenAddress:              ":" + port,
		MaxRequestTimeoutInSeconds: config.DefaultMaxRequestTimeoutInSeconds,
		MaxWorkerTimeoutInSeconds:  config.DefaultMaxWorkerTimeoutInSeconds,
	}
	db, err := bolt.Open(dropDbFileName, 0664, bolt.DefaultOptions)
	if err != nil {
		t.Fatalf("could not open Drop's BoltDB: %s", err)
	}

	dropDaemon := drop.New(cfg, db)
	dropDaemon.OnStop(func() error {
		if err := db.Close(); err != nil {
			return err
		}
		if err := os.Remove(dropDbFileName); err != nil {
			return err
		}
		return nil
	})
	dropDaemon.Start()

	dropClient := rest.NewClient(parseUrlOrPanic("http://localhost:" + port))

	return dropDaemon, dropClient
}

func runWorkerDaemon(t *testing.T) (worker.Daemonized, []byte) {
	t.Helper()

	cfg := &config.Worker{
		Name:    workerName,
		DropUrl: parseUrlOrPanic("http://localhost:" + port),
		RegistrationTimeoutInSeconds:        config.DefaultRegistrationTimeoutInSeconds,
		UpdateRegistrationIntervalInSeconds: config.DefaultUpdateRegistrationIntervalInSeconds,
	}
	db, err := bolt.Open(workerDbFileName, 0664, bolt.DefaultOptions)
	if err != nil {
		t.Fatalf("could not open Worker's BoltDB: %s", err)
	}
	privateKeyBytes, err := worker.GeneratePrivateKeyBytes(2048)
	if err != nil {
		t.Fatalf("couldn't generate private key: %s", err)
	}

	privateKey, err := crypto.UnmarshalPrivateKeyFromPEMBytes(privateKeyBytes)
	if err != nil {
		t.Fatalf("couldn't unmarshal private key: %s", err)
	}
	publicKeyBytes, err := crypto.GeneratePublicKeyBytes(privateKey)
	if err != nil {
		t.Fatalf("couldn't generate public key bytes: %s", err)
	}
	fingerprint, err := crypto.FingerprintPublicKey(&privateKey.PublicKey, 10, 4)
	if err != nil {
		t.Fatalf("couldn't fingerprint public key: %s", err)
	}

	workerDaemon := worker.New(cfg, fingerprint, db, privateKey)
	workerDaemon.OnStop(func() error {
		if err := db.Close(); err != nil {
			return err
		}
		if err := os.Remove(workerDbFileName); err != nil {
			return err
		}
		return nil
	})
	workerDaemon.Start()

	return workerDaemon, publicKeyBytes
}

func stopDaemon(d daemon.Daemon, t *testing.T) {
	t.Helper()
	err := d.Stop()
	if err != nil {
		t.Error(err)
	}
}

func parseUrlOrPanic(s string) *url.URL {
	result, err := url.Parse(s)
	if err != nil {
		panic(err)
	}

	return result
}
<|endoftext|>


<commit_before>// Copyright 2015, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

// Package bits implements bit-twiddling functionality.
package bits

import "unsafe"
import "bitbucket.org/rawr/golib/errs"

const (
	MinUint     = uint(0)
	MaxUint     = ^MinUint
	NumUintBits = int(8 * unsafe.Sizeof(uint(0)))
)

// ErrAlign is returned if the function called cannot perform the requested
// operation because the internal offset is not aligned to a byte-boundary.
var ErrAlign = errs.New("golib/bits: offset is not byte-aligned")

type BitReader interface {
	// Read a single bit.
	// This returns an error if and only if no bit is read.
	ReadBit() (val bool, err error)
}

type BitsReader interface {
	BitReader

	// Read num bits in LSB order. That is, the first bits read are packed into
	// val as the LSB. The behavior is undefined if a read is attempted on more
	// bits than fits in an uint.
	//
	// If an error is encountered while reading a bit, then an error will be
	// returned along with the number of bits read thus far. The error returned
	// will be io.EOF only if the cnt is 0. Otherwise, io.ErrUnexpectedEOF will
	// be used. This is done to match the behavior of io.ReadFull.
	ReadBits(num int) (val uint, cnt int, err error)
}

type BitWriter interface {
	// Write a single bit.
	// This returns an error if and only if no bit is written.
	WriteBit(bool) error
}

type BitsWriter interface {
	BitWriter

	// Write num bits in LSB order. That is, the LSB bits of val will be the
	// first bits to be written. The behavior is undefined if a read is
	// attempted on more bits than fits in an uint.
	//
	// If an error is encountered while writing a bit, then an error will be
	// returned along with the number of bits written thus far.
	WriteBits(val uint, num int) (cnt int, err error)
}
<commit_msg>Remove dependency on golib/errs<commit_after>// Copyright 2015, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

// Package bits implements bit-twiddling functionality.
package bits

import "unsafe"
import "errors"

const (
	MinUint     = uint(0)
	MaxUint     = ^MinUint
	NumUintBits = int(8 * unsafe.Sizeof(uint(0)))
)

// ErrAlign is returned if the function called cannot perform the requested
// operation because the internal offset is not aligned to a byte-boundary.
var ErrAlign = errors.New("golib/bits: offset is not byte-aligned")

type BitReader interface {
	// Read a single bit.
	// This returns an error if and only if no bit is read.
	ReadBit() (val bool, err error)
}

type BitsReader interface {
	BitReader

	// Read num bits in LSB order. That is, the first bits read are packed into
	// val as the LSB. The behavior is undefined if a read is attempted on more
	// bits than fits in an uint.
	//
	// If an error is encountered while reading a bit, then an error will be
	// returned along with the number of bits read thus far. The error returned
	// will be io.EOF only if the cnt is 0. Otherwise, io.ErrUnexpectedEOF will
	// be used. This is done to match the behavior of io.ReadFull.
	ReadBits(num int) (val uint, cnt int, err error)
}

type BitWriter interface {
	// Write a single bit.
	// This returns an error if and only if no bit is written.
	WriteBit(bool) error
}

type BitsWriter interface {
	BitWriter

	// Write num bits in LSB order. That is, the LSB bits of val will be the
	// first bits to be written. The behavior is undefined if a read is
	// attempted on more bits than fits in an uint.
	//
	// If an error is encountered while writing a bit, then an error will be
	// returned along with the number of bits written thus far.
	WriteBits(val uint, num int) (cnt int, err error)
}
<|endoftext|>


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

package logger

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
)

const schema = "LogData7"

type Service struct {
	sync.Mutex

	stats  Stats
	origin string

	esIndex elasticsearch.IIndex
	id      int
}

func (service *Service) Init(sys *piazza.SystemConfig, esIndex elasticsearch.IIndex) error {
	var err error

	service.stats.CreatedOn = time.Now()
	/***
	err = esIndex.Delete()
	if err != nil {
		log.Fatal(err)
	}
	if esIndex.IndexExists() {
		log.Fatal("index still exists")
	}
	err = esIndex.Create()
	if err != nil {
		log.Fatal(err)
	}
	***/

	ok, err := esIndex.IndexExists()
	if err != nil {
		return err
	}
	if !ok {
		log.Printf("Creating index: %s", esIndex.IndexName())
		err = esIndex.Create("")
		if err != nil {
			log.Fatal(err)
		}
	}

	ok, err = esIndex.TypeExists(schema)
	if err != nil {
		return err
	}
	if !ok {
		log.Printf("Creating type: %s", schema)

		mapping :=
			`{
			"LogData7":{
				"dynamic": "strict",
				"properties": {
					"service": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"address": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"createdOn": {
						"type": "date",
						"store": true,
						"index": "not_analyzed"
					},
					"severity": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"message": {
						"type": "string",
						"store": true,
						"index": "analyzed"
					}
				}
			}
		}`

		err = esIndex.SetMapping(schema, piazza.JsonString(mapping))
		if err != nil {
			log.Printf("LoggerService.Init: %s", err.Error())
			return err
		}
	}

	service.esIndex = esIndex

	service.origin = string(sys.Name)

	return nil
}

func (service *Service) newInternalErrorResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusInternalServerError,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) newBadRequestResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusBadRequest,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) GetRoot() *piazza.JsonResponse {
	resp := &piazza.JsonResponse{
		StatusCode: 200,
		Data:       "Hi. I'm pz-logger.",
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) PostMessage(mssg *Message) *piazza.JsonResponse {
	err := mssg.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	service.Lock()
	idStr := strconv.Itoa(service.id)
	service.id++
	service.Unlock()

	_, err = service.esIndex.PostData(schema, idStr, mssg)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	service.stats.NumMessages++

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       mssg,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetStats() *piazza.JsonResponse {
	service.Lock()
	t := service.stats
	service.Unlock()

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       t,
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetMessage(params *piazza.HttpQueryParams) *piazza.JsonResponse {
	var err error

	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	dsl, err := createQueryDslAsString(pagination, params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	var searchResult *elasticsearch.SearchResult

	if dsl == "" {
		searchResult, err = service.esIndex.FilterByMatchAll(schema, pagination)
	} else {
		searchResult, err = service.esIndex.SearchByJSON(schema, dsl)
	}

	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	var lines = make([]Message, 0)

	if searchResult != nil && searchResult.GetHits() != nil {
		for _, hit := range *searchResult.GetHits() {
			if hit.Source == nil {
				log.Printf("null source hit")
				continue
			}

			var msg Message
			err = json.Unmarshal(*hit.Source, &msg)
			if err != nil {
				log.Printf("UNABLE TO PARSE: %s", string(*hit.Source))
				return service.newInternalErrorResponse(err)
			}

			// just in case
			err = msg.Validate()
			if err != nil {
				log.Printf("UNABLE TO VALIDATE: %s", string(*hit.Source))
				continue
			}

			lines = append(lines, msg)
		}
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       lines,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}
	return resp
}

func createQueryDslAsString(
	pagination *piazza.JsonPagination,
	params *piazza.HttpQueryParams) (string, error) {

	must := []map[string]interface{}{}

	service, err := params.GetAsString("service", "")
	if err != nil {
		return "", err
	}

	contains, err := params.GetAsString("contains", "")
	if err != nil {
		return "", err
	}

	before, err := params.GetBefore(time.Time{})
	if err != nil {
		return "", err
	}

	after, err := params.GetAfter(time.Time{})
	if err != nil {
		return "", err
	}

	if service != "" {
		must = append(must, map[string]interface{}{
			"match": map[string]interface{}{
				"service": service,
			},
		})
	}

	if contains != "" {
		must = append(must, map[string]interface{}{
			"multi_match": map[string]interface{}{
				"query":  contains,
				"fields": []string{"address", "message", "service", "severity"},
			},
		})
	}

	if !after.IsZero() || !before.IsZero() {
		rangeParams := map[string]time.Time{}

		if !after.IsZero() {
			rangeParams["gte"] = after
		}

		if !before.IsZero() {
			rangeParams["lte"] = before
		}

		must = append(must, map[string]interface{}{
			"range": map[string]interface{}{
				"createdOn": rangeParams,
			},
		})
	}

	if len(must) == 0 {
		return "", nil
	}

	dsl := map[string]interface{}{
		"query": map[string]interface{}{
			"filtered": map[string]interface{}{
				"query": map[string]interface{}{
					"bool": map[string]interface{}{
						"must": must,
					},
				},
			},
		},
		"size": pagination.PerPage,
		"from": pagination.PerPage * pagination.Page,
	}

	dsl["sort"] = map[string]string{
		pagination.SortBy: string(pagination.Order),
	}

	output, err := json.Marshal(dsl)
	if err != nil {
		return "", err
	}

	return string(output), nil
}

func (service *Service) PostSyslog(mNew *piazza.SyslogMessage) *piazza.JsonResponse {
	err := mNew.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	rfc := mNew.String()
	mssgOld := Message{CreatedOn: time.Now(), Message: rfc}

	_, err = service.esIndex.PostData(schema, mNew.MessageID, mssgOld)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       rfc,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

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

package logger

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
)

const schema = "LogData7"

type Service struct {
	sync.Mutex

	stats  Stats
	origin string

	esIndex elasticsearch.IIndex
	id      int
}

func (service *Service) Init(sys *piazza.SystemConfig, esIndex elasticsearch.IIndex) error {
	var err error

	service.stats.CreatedOn = time.Now()
	/***
	err = esIndex.Delete()
	if err != nil {
		log.Fatal(err)
	}
	if esIndex.IndexExists() {
		log.Fatal("index still exists")
	}
	err = esIndex.Create()
	if err != nil {
		log.Fatal(err)
	}
	***/

	ok, err := esIndex.IndexExists()
	if err != nil {
		return err
	}
	if !ok {
		log.Printf("Creating index: %s", esIndex.IndexName())
		err = esIndex.Create("")
		if err != nil {
			log.Fatal(err)
		}
	}

	ok, err = esIndex.TypeExists(schema)
	if err != nil {
		return err
	}
	if !ok {
		log.Printf("Creating type: %s", schema)

		mapping :=
			`{
			"LogData7":{
				"dynamic": "strict",
				"properties": {
					"service": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"address": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"createdOn": {
						"type": "date",
						"store": true,
						"index": "not_analyzed"
					},
					"severity": {
						"type": "string",
						"store": true,
						"index": "not_analyzed"
					},
					"message": {
						"type": "string",
						"store": true,
						"index": "analyzed"
					}
				}
			}
		}`

		err = esIndex.SetMapping(schema, piazza.JsonString(mapping))
		if err != nil {
			log.Printf("LoggerService.Init: %s", err.Error())
			return err
		}
	}

	service.esIndex = esIndex

	service.origin = string(sys.Name)

	return nil
}

func (service *Service) newInternalErrorResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusInternalServerError,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) newBadRequestResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusBadRequest,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) GetRoot() *piazza.JsonResponse {
	resp := &piazza.JsonResponse{
		StatusCode: 200,
		Data:       "Hi. I'm pz-logger.",
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) PostMessage(mssg *Message) *piazza.JsonResponse {
	err := mssg.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	service.Lock()
	idStr := strconv.Itoa(service.id)
	service.id++
	service.Unlock()

	_, err = service.esIndex.PostData(schema, idStr, mssg)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	service.stats.NumMessages++

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       mssg,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetStats() *piazza.JsonResponse {
	service.Lock()
	t := service.stats
	service.Unlock()

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       t,
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetMessage(params *piazza.HttpQueryParams) *piazza.JsonResponse {
	var err error

	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	dsl, err := createQueryDslAsString(pagination, params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	var searchResult *elasticsearch.SearchResult

	if dsl == "" {
		searchResult, err = service.esIndex.FilterByMatchAll(schema, pagination)
	} else {
		searchResult, err = service.esIndex.SearchByJSON(schema, dsl)
	}

	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	var lines = make([]Message, 0)

	if searchResult != nil && searchResult.GetHits() != nil {
		for _, hit := range *searchResult.GetHits() {
			if hit.Source == nil {
				log.Printf("null source hit")
				continue
			}

			var msg Message
			err = json.Unmarshal(*hit.Source, &msg)
			if err != nil {
				log.Printf("UNABLE TO PARSE: %s", string(*hit.Source))
				return service.newInternalErrorResponse(err)
			}

			// just in case
			err = msg.Validate()
			if err != nil {
				log.Printf("UNABLE TO VALIDATE: %s", string(*hit.Source))
				continue
			}

			lines = append(lines, msg)
		}
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       lines,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}
	return resp
}

func createQueryDslAsString(
	pagination *piazza.JsonPagination,
	params *piazza.HttpQueryParams) (string, error) {

	must := []map[string]interface{}{}

	service, err := params.GetAsString("service", "")
	if err != nil {
		return "", err
	}

	contains, err := params.GetAsString("contains", "")
	if err != nil {
		return "", err
	}

	before, err := params.GetBefore(time.Time{})
	if err != nil {
		return "", err
	}

	after, err := params.GetAfter(time.Time{})
	if err != nil {
		return "", err
	}

	if service != "" {
		must = append(must, map[string]interface{}{
			"match": map[string]interface{}{
				"service": service,
			},
		})
	}

	if contains != "" {
		must = append(must, map[string]interface{}{
			"multi_match": map[string]interface{}{
				"query":  contains,
				"fields": []string{"address", "message", "service", "severity"},
			},
		})
	}

	if !after.IsZero() || !before.IsZero() {
		rangeParams := map[string]time.Time{}

		if !after.IsZero() {
			rangeParams["gte"] = after
		}

		if !before.IsZero() {
			rangeParams["lte"] = before
		}

		must = append(must, map[string]interface{}{
			"range": map[string]interface{}{
				"createdOn": rangeParams,
			},
		})
	}

	if len(must) == 0 {
		return "", nil
	}

	dsl := map[string]interface{}{
		"query": map[string]interface{}{
			"filtered": map[string]interface{}{
				"query": map[string]interface{}{
					"bool": map[string]interface{}{
						"must": must,
					},
				},
			},
		},
		"size": pagination.PerPage,
		"from": pagination.PerPage * pagination.Page,
	}

	dsl["sort"] = map[string]string{
		pagination.SortBy: string(pagination.Order),
	}

	output, err := json.Marshal(dsl)
	if err != nil {
		return "", err
	}

	return string(output), nil
}

func (service *Service) PostSyslog(mNew *piazza.SyslogMessage) *piazza.JsonResponse {
	err := mNew.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	rfc := mNew.String()
	mssgOld := Message{CreatedOn: time.Now(), Message: rfc}

	service.Lock()
	idStr := strconv.Itoa(service.id)
	service.id++
	service.Unlock()

	_, err = service.esIndex.PostData(schema, idStr, mssgOld)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       rfc,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}
<|endoftext|>


<commit_before>package core

import (
	"bytes"
	"fmt"
	"math/big"
	"sync"

	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/ethutil"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/state"
)

var chainlogger = logger.NewLogger("CHAIN")

/*
func AddTestNetFunds(block *types.Block) {
	for _, addr := range []string{
		"51ba59315b3a95761d0863b05ccc7a7f54703d99",
		"e4157b34ea9615cfbde6b4fda419828124b70c78",
		"b9c015918bdaba24b4ff057a92a3873d6eb201be",
		"6c386a4b26f73c802f34673f7248bb118f97424a",
		"cd2a3d9f938e13cd947ec05abc7fe734df8dd826",
		"2ef47100e0787b915105fd5e3f4ff6752079d5cb",
		"e6716f9544a56c530d868e4bfbacb172315bdead",
		"1a26338f0d905e295fccb71fa9ea849ffa12aaf4",
	} {
		codedAddr := ethutil.Hex2Bytes(addr)
		account := block.State().GetAccount(codedAddr)
		account.SetBalance(ethutil.Big("1606938044258990275541962092341162602522202993782792835301376")) //ethutil.BigPow(2, 200)
		block.State().UpdateStateObject(account)
	}
}
*/

func CalcDifficulty(block, parent *types.Block) *big.Int {
	diff := new(big.Int)

	bh, ph := block.Header(), parent.Header()
	adjust := new(big.Int).Rsh(ph.Difficulty, 10)
	if bh.Time >= ph.Time+5 {
		diff.Sub(ph.Difficulty, adjust)
	} else {
		diff.Add(ph.Difficulty, adjust)
	}

	return diff
}

func CalcGasLimit(parent, block *types.Block) *big.Int {
	if block.Number().Cmp(big.NewInt(0)) == 0 {
		return ethutil.BigPow(10, 6)
	}

	// ((1024-1) * parent.gasLimit + (gasUsed * 6 / 5)) / 1024

	previous := new(big.Int).Mul(big.NewInt(1024-1), parent.GasLimit())
	current := new(big.Rat).Mul(new(big.Rat).SetInt(parent.GasUsed()), big.NewRat(6, 5))
	curInt := new(big.Int).Div(current.Num(), current.Denom())

	result := new(big.Int).Add(previous, curInt)
	result.Div(result, big.NewInt(1024))

	min := big.NewInt(125000)

	return ethutil.BigMax(min, result)
}

type ChainManager struct {
	//eth          EthManager
	processor    types.BlockProcessor
	eventMux     *event.TypeMux
	genesisBlock *types.Block
	// Last known total difficulty
	mu              sync.RWMutex
	td              *big.Int
	lastBlockNumber uint64
	currentBlock    *types.Block
	lastBlockHash   []byte

	transState *state.StateDB
}

func (self *ChainManager) Td() *big.Int {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.td
}

func (self *ChainManager) LastBlockNumber() uint64 {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.lastBlockNumber
}

func (self *ChainManager) LastBlockHash() []byte {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.lastBlockHash
}

func (self *ChainManager) CurrentBlock() *types.Block {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.currentBlock
}

func NewChainManager(mux *event.TypeMux) *ChainManager {
	bc := &ChainManager{}
	bc.genesisBlock = GenesisBlock()
	bc.eventMux = mux

	bc.setLastBlock()

	bc.transState = bc.State().Copy()

	return bc
}

func (self *ChainManager) Status() (td *big.Int, currentBlock []byte, genesisBlock []byte) {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.td, self.currentBlock.Hash(), self.Genesis().Hash()
}

func (self *ChainManager) SetProcessor(proc types.BlockProcessor) {
	self.processor = proc
}

func (self *ChainManager) State() *state.StateDB {
	return state.New(self.CurrentBlock().Trie())
}

func (self *ChainManager) TransState() *state.StateDB {
	return self.transState
}

func (bc *ChainManager) setLastBlock() {
	data, _ := ethutil.Config.Db.Get([]byte("LastBlock"))
	if len(data) != 0 {
		var block types.Block
		rlp.Decode(bytes.NewReader(data), &block)
		bc.currentBlock = &block
		bc.lastBlockHash = block.Hash()
		bc.lastBlockNumber = block.Header().Number.Uint64()

		// Set the last know difficulty (might be 0x0 as initial value, Genesis)
		bc.td = ethutil.BigD(ethutil.Config.Db.LastKnownTD())
	} else {
		bc.Reset()
	}

	chainlogger.Infof("Last block (#%d) %x\n", bc.lastBlockNumber, bc.currentBlock.Hash())
}

// Block creation & chain handling
func (bc *ChainManager) NewBlock(coinbase []byte) *types.Block {
	bc.mu.RLock()
	defer bc.mu.RUnlock()

	var root []byte
	parentHash := ZeroHash256

	if bc.CurrentBlock != nil {
		root = bc.currentBlock.Header().Root
		parentHash = bc.lastBlockHash
	}

	block := types.NewBlock(
		parentHash,
		coinbase,
		root,
		ethutil.BigPow(2, 32),
		nil,
		"")

	parent := bc.currentBlock
	if parent != nil {
		header := block.Header()
		header.Difficulty = CalcDifficulty(block, parent)
		header.Number = new(big.Int).Add(parent.Header().Number, ethutil.Big1)
		header.GasLimit = CalcGasLimit(parent, block)

	}

	return block
}

func (bc *ChainManager) Reset() {
	bc.mu.Lock()
	defer bc.mu.Unlock()

	for block := bc.currentBlock; block != nil; block = bc.GetBlock(block.Header().ParentHash) {
		ethutil.Config.Db.Delete(block.Hash())
	}

	// Prepare the genesis block
	bc.write(bc.genesisBlock)
	bc.insert(bc.genesisBlock)
	bc.currentBlock = bc.genesisBlock

	bc.setTotalDifficulty(ethutil.Big("0"))
}

func (self *ChainManager) Export() []byte {
	self.mu.RLock()
	defer self.mu.RUnlock()

	chainlogger.Infof("exporting %v blocks...\n", self.currentBlock.Header().Number)

	blocks := make([]*types.Block, int(self.currentBlock.NumberU64())+1)
	for block := self.currentBlock; block != nil; block = self.GetBlock(block.Header().ParentHash) {
		blocks[block.NumberU64()] = block
	}

	return ethutil.Encode(blocks)
}

func (bc *ChainManager) insert(block *types.Block) {
	encodedBlock := ethutil.Encode(block)
	ethutil.Config.Db.Put([]byte("LastBlock"), encodedBlock)
	bc.currentBlock = block
	bc.lastBlockHash = block.Hash()
}

func (bc *ChainManager) write(block *types.Block) {
	bc.writeBlockInfo(block)

	encodedBlock := ethutil.Encode(block)
	ethutil.Config.Db.Put(block.Hash(), encodedBlock)
}

// Accessors
func (bc *ChainManager) Genesis() *types.Block {
	return bc.genesisBlock
}

// Block fetching methods
func (bc *ChainManager) HasBlock(hash []byte) bool {
	data, _ := ethutil.Config.Db.Get(hash)
	return len(data) != 0
}

func (self *ChainManager) GetBlockHashesFromHash(hash []byte, max uint64) (chain [][]byte) {
	block := self.GetBlock(hash)
	if block == nil {
		return
	}

	// XXX Could be optimised by using a different database which only holds hashes (i.e., linked list)
	for i := uint64(0); i < max; i++ {
		chain = append(chain, block.Hash())

		if block.Header().Number.Cmp(ethutil.Big0) <= 0 {
			break
		}

		block = self.GetBlock(block.Header().ParentHash)
	}

	return
}

func (self *ChainManager) GetBlock(hash []byte) *types.Block {
	data, _ := ethutil.Config.Db.Get(hash)
	if len(data) == 0 {
		return nil
	}
	var block types.Block
	if err := rlp.Decode(bytes.NewReader(data), &block); err != nil {
		fmt.Println(err)
		return nil
	}

	return &block
}

func (self *ChainManager) GetBlockByNumber(num uint64) *types.Block {
	self.mu.RLock()
	defer self.mu.RUnlock()

	block := self.currentBlock
	for ; block != nil; block = self.GetBlock(block.Header().ParentHash) {
		if block.Header().Number.Uint64() == num {
			break
		}
	}

	if block != nil && block.Header().Number.Uint64() == 0 && num != 0 {
		return nil
	}

	return block
}

func (bc *ChainManager) setTotalDifficulty(td *big.Int) {
	ethutil.Config.Db.Put([]byte("LTD"), td.Bytes())
	bc.td = td
}

func (self *ChainManager) CalcTotalDiff(block *types.Block) (*big.Int, error) {
	parent := self.GetBlock(block.Header().ParentHash)
	if parent == nil {
		return nil, fmt.Errorf("Unable to calculate total diff without known parent %x", block.Header().ParentHash)
	}

	parentTd := parent.Td

	uncleDiff := new(big.Int)
	for _, uncle := range block.Uncles() {
		uncleDiff = uncleDiff.Add(uncleDiff, uncle.Difficulty)
	}

	td := new(big.Int)
	td = td.Add(parentTd, uncleDiff)
	td = td.Add(td, block.Header().Difficulty)

	return td, nil
}

// Unexported method for writing extra non-essential block info to the db
func (bc *ChainManager) writeBlockInfo(block *types.Block) {
	bc.lastBlockNumber++
}

func (bc *ChainManager) Stop() {
	if bc.CurrentBlock != nil {
		chainlogger.Infoln("Stopped")
	}
}

func (self *ChainManager) InsertChain(chain types.Blocks) error {
	for _, block := range chain {
		td, messages, err := self.processor.Process(block)
		if err != nil {
			if IsKnownBlockErr(err) {
				continue
			}

			h := block.Header()
			chainlogger.Infof("block #%v process failed (%x)\n", h.Number, h.Hash()[:4])
			chainlogger.Infoln(block)
			chainlogger.Infoln(err)
			return err
		}
		block.Td = td

		self.mu.Lock()
		{
			self.write(block)
			cblock := self.currentBlock
			if td.Cmp(self.td) > 0 {
				if block.Header().Number.Cmp(new(big.Int).Add(cblock.Header().Number, ethutil.Big1)) < 0 {
					chainlogger.Infof("Split detected. New head #%v (%x), was #%v (%x)\n", block.Header().Number, block.Hash()[:4], cblock.Header().Number, cblock.Hash()[:4])
				}

				self.setTotalDifficulty(td)
				self.insert(block)
				self.transState = state.New(cblock.Trie().Copy())
			}

		}
		self.mu.Unlock()

		self.eventMux.Post(NewBlockEvent{block})
		self.eventMux.Post(messages)
	}

	return nil
}
<commit_msg>Added a query interface for world state<commit_after>package core

import (
	"bytes"
	"fmt"
	"math/big"
	"sync"

	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/ethutil"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/state"
)

var chainlogger = logger.NewLogger("CHAIN")

type StateQuery interface {
	GetAccount(addr []byte) *state.StateObject
}

/*
func AddTestNetFunds(block *types.Block) {
	for _, addr := range []string{
		"51ba59315b3a95761d0863b05ccc7a7f54703d99",
		"e4157b34ea9615cfbde6b4fda419828124b70c78",
		"b9c015918bdaba24b4ff057a92a3873d6eb201be",
		"6c386a4b26f73c802f34673f7248bb118f97424a",
		"cd2a3d9f938e13cd947ec05abc7fe734df8dd826",
		"2ef47100e0787b915105fd5e3f4ff6752079d5cb",
		"e6716f9544a56c530d868e4bfbacb172315bdead",
		"1a26338f0d905e295fccb71fa9ea849ffa12aaf4",
	} {
		codedAddr := ethutil.Hex2Bytes(addr)
		account := block.State().GetAccount(codedAddr)
		account.SetBalance(ethutil.Big("1606938044258990275541962092341162602522202993782792835301376")) //ethutil.BigPow(2, 200)
		block.State().UpdateStateObject(account)
	}
}
*/

func CalcDifficulty(block, parent *types.Block) *big.Int {
	diff := new(big.Int)

	bh, ph := block.Header(), parent.Header()
	adjust := new(big.Int).Rsh(ph.Difficulty, 10)
	if bh.Time >= ph.Time+5 {
		diff.Sub(ph.Difficulty, adjust)
	} else {
		diff.Add(ph.Difficulty, adjust)
	}

	return diff
}

func CalcGasLimit(parent, block *types.Block) *big.Int {
	if block.Number().Cmp(big.NewInt(0)) == 0 {
		return ethutil.BigPow(10, 6)
	}

	// ((1024-1) * parent.gasLimit + (gasUsed * 6 / 5)) / 1024

	previous := new(big.Int).Mul(big.NewInt(1024-1), parent.GasLimit())
	current := new(big.Rat).Mul(new(big.Rat).SetInt(parent.GasUsed()), big.NewRat(6, 5))
	curInt := new(big.Int).Div(current.Num(), current.Denom())

	result := new(big.Int).Add(previous, curInt)
	result.Div(result, big.NewInt(1024))

	min := big.NewInt(125000)

	return ethutil.BigMax(min, result)
}

type ChainManager struct {
	//eth          EthManager
	processor    types.BlockProcessor
	eventMux     *event.TypeMux
	genesisBlock *types.Block
	// Last known total difficulty
	mu              sync.RWMutex
	td              *big.Int
	lastBlockNumber uint64
	currentBlock    *types.Block
	lastBlockHash   []byte

	transState *state.StateDB
}

func (self *ChainManager) Td() *big.Int {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.td
}

func (self *ChainManager) LastBlockNumber() uint64 {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.lastBlockNumber
}

func (self *ChainManager) LastBlockHash() []byte {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.lastBlockHash
}

func (self *ChainManager) CurrentBlock() *types.Block {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.currentBlock
}

func NewChainManager(mux *event.TypeMux) *ChainManager {
	bc := &ChainManager{}
	bc.genesisBlock = GenesisBlock()
	bc.eventMux = mux

	bc.setLastBlock()

	bc.transState = bc.State().Copy()

	return bc
}

func (self *ChainManager) Status() (td *big.Int, currentBlock []byte, genesisBlock []byte) {
	self.mu.RLock()
	defer self.mu.RUnlock()

	return self.td, self.currentBlock.Hash(), self.Genesis().Hash()
}

func (self *ChainManager) SetProcessor(proc types.BlockProcessor) {
	self.processor = proc
}

func (self *ChainManager) State() *state.StateDB {
	return state.New(self.CurrentBlock().Trie())
}

func (self *ChainManager) TransState() *state.StateDB {
	return self.transState
}

func (bc *ChainManager) setLastBlock() {
	data, _ := ethutil.Config.Db.Get([]byte("LastBlock"))
	if len(data) != 0 {
		var block types.Block
		rlp.Decode(bytes.NewReader(data), &block)
		bc.currentBlock = &block
		bc.lastBlockHash = block.Hash()
		bc.lastBlockNumber = block.Header().Number.Uint64()

		// Set the last know difficulty (might be 0x0 as initial value, Genesis)
		bc.td = ethutil.BigD(ethutil.Config.Db.LastKnownTD())
	} else {
		bc.Reset()
	}

	chainlogger.Infof("Last block (#%d) %x\n", bc.lastBlockNumber, bc.currentBlock.Hash())
}

// Block creation & chain handling
func (bc *ChainManager) NewBlock(coinbase []byte) *types.Block {
	bc.mu.RLock()
	defer bc.mu.RUnlock()

	var root []byte
	parentHash := ZeroHash256

	if bc.CurrentBlock != nil {
		root = bc.currentBlock.Header().Root
		parentHash = bc.lastBlockHash
	}

	block := types.NewBlock(
		parentHash,
		coinbase,
		root,
		ethutil.BigPow(2, 32),
		nil,
		"")

	parent := bc.currentBlock
	if parent != nil {
		header := block.Header()
		header.Difficulty = CalcDifficulty(block, parent)
		header.Number = new(big.Int).Add(parent.Header().Number, ethutil.Big1)
		header.GasLimit = CalcGasLimit(parent, block)

	}

	return block
}

func (bc *ChainManager) Reset() {
	bc.mu.Lock()
	defer bc.mu.Unlock()

	for block := bc.currentBlock; block != nil; block = bc.GetBlock(block.Header().ParentHash) {
		ethutil.Config.Db.Delete(block.Hash())
	}

	// Prepare the genesis block
	bc.write(bc.genesisBlock)
	bc.insert(bc.genesisBlock)
	bc.currentBlock = bc.genesisBlock

	bc.setTotalDifficulty(ethutil.Big("0"))
}

func (self *ChainManager) Export() []byte {
	self.mu.RLock()
	defer self.mu.RUnlock()

	chainlogger.Infof("exporting %v blocks...\n", self.currentBlock.Header().Number)

	blocks := make([]*types.Block, int(self.currentBlock.NumberU64())+1)
	for block := self.currentBlock; block != nil; block = self.GetBlock(block.Header().ParentHash) {
		blocks[block.NumberU64()] = block
	}

	return ethutil.Encode(blocks)
}

func (bc *ChainManager) insert(block *types.Block) {
	encodedBlock := ethutil.Encode(block)
	ethutil.Config.Db.Put([]byte("LastBlock"), encodedBlock)
	bc.currentBlock = block
	bc.lastBlockHash = block.Hash()
}

func (bc *ChainManager) write(block *types.Block) {
	bc.writeBlockInfo(block)

	encodedBlock := ethutil.Encode(block)
	ethutil.Config.Db.Put(block.Hash(), encodedBlock)
}

// Accessors
func (bc *ChainManager) Genesis() *types.Block {
	return bc.genesisBlock
}

// Block fetching methods
func (bc *ChainManager) HasBlock(hash []byte) bool {
	data, _ := ethutil.Config.Db.Get(hash)
	return len(data) != 0
}

func (self *ChainManager) GetBlockHashesFromHash(hash []byte, max uint64) (chain [][]byte) {
	block := self.GetBlock(hash)
	if block == nil {
		return
	}

	// XXX Could be optimised by using a different database which only holds hashes (i.e., linked list)
	for i := uint64(0); i < max; i++ {
		chain = append(chain, block.Hash())

		if block.Header().Number.Cmp(ethutil.Big0) <= 0 {
			break
		}

		block = self.GetBlock(block.Header().ParentHash)
	}

	return
}

func (self *ChainManager) GetBlock(hash []byte) *types.Block {
	data, _ := ethutil.Config.Db.Get(hash)
	if len(data) == 0 {
		return nil
	}
	var block types.Block
	if err := rlp.Decode(bytes.NewReader(data), &block); err != nil {
		fmt.Println(err)
		return nil
	}

	return &block
}

func (self *ChainManager) GetBlockByNumber(num uint64) *types.Block {
	self.mu.RLock()
	defer self.mu.RUnlock()

	block := self.currentBlock
	for ; block != nil; block = self.GetBlock(block.Header().ParentHash) {
		if block.Header().Number.Uint64() == num {
			break
		}
	}

	if block != nil && block.Header().Number.Uint64() == 0 && num != 0 {
		return nil
	}

	return block
}

func (bc *ChainManager) setTotalDifficulty(td *big.Int) {
	ethutil.Config.Db.Put([]byte("LTD"), td.Bytes())
	bc.td = td
}

func (self *ChainManager) CalcTotalDiff(block *types.Block) (*big.Int, error) {
	parent := self.GetBlock(block.Header().ParentHash)
	if parent == nil {
		return nil, fmt.Errorf("Unable to calculate total diff without known parent %x", block.Header().ParentHash)
	}

	parentTd := parent.Td

	uncleDiff := new(big.Int)
	for _, uncle := range block.Uncles() {
		uncleDiff = uncleDiff.Add(uncleDiff, uncle.Difficulty)
	}

	td := new(big.Int)
	td = td.Add(parentTd, uncleDiff)
	td = td.Add(td, block.Header().Difficulty)

	return td, nil
}

// Unexported method for writing extra non-essential block info to the db
func (bc *ChainManager) writeBlockInfo(block *types.Block) {
	bc.lastBlockNumber++
}

func (bc *ChainManager) Stop() {
	if bc.CurrentBlock != nil {
		chainlogger.Infoln("Stopped")
	}
}

func (self *ChainManager) InsertChain(chain types.Blocks) error {
	for _, block := range chain {
		td, messages, err := self.processor.Process(block)
		if err != nil {
			if IsKnownBlockErr(err) {
				continue
			}

			h := block.Header()
			chainlogger.Infof("block #%v process failed (%x)\n", h.Number, h.Hash()[:4])
			chainlogger.Infoln(block)
			chainlogger.Infoln(err)
			return err
		}
		block.Td = td

		self.mu.Lock()
		{
			self.write(block)
			cblock := self.currentBlock
			if td.Cmp(self.td) > 0 {
				if block.Header().Number.Cmp(new(big.Int).Add(cblock.Header().Number, ethutil.Big1)) < 0 {
					chainlogger.Infof("Split detected. New head #%v (%x), was #%v (%x)\n", block.Header().Number, block.Hash()[:4], cblock.Header().Number, cblock.Hash()[:4])
				}

				self.setTotalDifficulty(td)
				self.insert(block)
				self.transState = state.New(cblock.Trie().Copy())
			}

		}
		self.mu.Unlock()

		self.eventMux.Post(NewBlockEvent{block})
		self.eventMux.Post(messages)
	}

	return nil
}

// Satisfy state query interface
func (self *ChainManager) GetAccount(addr []byte) *state.StateObject {
	return self.State().GetAccount(addr)
}
<|endoftext|>


<commit_before>package coretest

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"testing"
)

const cloudinitBinPath = "/usr/bin/coreos-cloudinit"

func read(filename string) (string, error) {
	bytes, err := ioutil.ReadFile(filename)
	return string(bytes), err
}

func rmdir(path string) error {
	cmd := exec.Command("sudo", "rm", "-rf", path)
	return cmd.Run()
}

func TestCloudinitCloudConfig(t *testing.T) {
	workspace, err := ioutil.TempDir("", "coretest-cloudinit-")
	if err != nil {
		t.Fatalf("Failed creating workspace: %v", err)
	}
	defer rmdir(workspace)

	keyOne := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC5LaGMGRqZEEvOhHlIEiQgdMJIQ9Qe8L/XSz06GqzcESbEnYLIXar2nou4eW4AGMVC1V0BrcWWnSTxM1/dWeCLOUt5NulKAjtdBUZGhCT83nbimSzbmx3/q2y5bCiS4Zr8ZjYFbi1eLvye2jKPE4xo7cvIfDKc0ztQ9kU7JknUdKNZo3RKXr5EPhJ5UZ8Ff15CI9+hDSvdPwer+HNnEt/psRVC+s29EwNGwUXD4IYqrk3X4ew0YAl/oULHM4cctoBW9GM+kAl40rOuIARlKfe4UdCgDMHYA/whi7Us+cPNgPit9IVJVBU4eo/cF5molD2l+PMSntypuv79obu8sA1H cloudinit-test-key-one"
	keyTwo := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCZw5Ljtt9wlEfyDvmUwu/BeMcIhVarbcM4ajZolxRy9G8vvCa7ODcSjzSyhfG1mLSBB2KfaFFI6zGHBjFX0Gzy9i8m3u7PnZBPX30bb1n0hJCrUhpqUGQUe8OFdoBstf1HIwJU/KoTBL0Ap1WEn0quRT4kNgBLbPrMjYCPbS1q4wJKdIE5rRm/EfTUrmIb0i91gujEGw5oUHDXf0X+/cxwwIVZh1z16YhOgvJBzXhsJ9a0w7kcy/6wPRv03yyMg/r2Ada6ci68LulKz5GLn+xInT0bvIcra/PZ7WE+jyZhZKly239VZyT/1dHkBbTw+kgnGobLMbjOOg5bKaT8NZJ3 cloudinit-test-key-two"

	configTmpl := `#cloud-config
coreos:
    etcd:
        discovery_url: https://discovery.etcd.io/827c73219eeb2fa5530027c37bf18877
ssh_authorized_keys:
    - %s
    - %s
`
	configData := fmt.Sprintf(configTmpl, keyOne, keyTwo)
	configFile, err := ioutil.TempFile(os.TempDir(), "coretest-")
	if err != nil {
		t.Fatalf("Failed creating tempfile: %v", err)
	}
	defer syscall.Unlink(configFile.Name())

	if _, err := io.WriteString(configFile, configData); err != nil {
		t.Fatalf("Failed writing %s: %v", configFile.Name(), err)
	}

	if stdout, stderr, err := Run("sudo", cloudinitBinPath, "--workspace", workspace, "--from-file", configFile.Name(), "--ssh-key-name", "coretest"); err != nil {
		t.Fatalf("coreos-cloudinit failed with error: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	contents, err := read("/var/run/etcd/bootstrap.disco")
	if err != nil {
		t.Errorf("Unable to read etcd bootstrap file: %v", err)
	} else if contents != "https://discovery.etcd.io/827c73219eeb2fa5530027c37bf18877" {
		t.Errorf("Incorrect data written to /var/run/etcd/bootstrap.disco: %s", contents)
	}

	// Attempt to clean up after ourselves
	defer Run("update-ssh-keys", "-d", "coretest")

	authorized_keys, err := read("/home/core/.ssh/authorized_keys")
	if err != nil {
		t.Fatalf("Unable to read authorized_keys file: %v", err)
	}

	if !strings.Contains(authorized_keys, keyOne) {
		t.Errorf("Could not find first key in authorized_keys")
	}

	if !strings.Contains(authorized_keys, keyTwo) {
		t.Errorf("Could not find second key in authorized_keys")
	}
}

func TestCloudinitScript(t *testing.T) {
	workspace, err := ioutil.TempDir("", "coretest-cloudinit-")
	if err != nil {
		t.Fatalf("Failed creating workspace: %v", err)
	}
	defer rmdir(workspace)

	configData := `#!/bin/bash
/bin/sleep 10
`
	configFile, err := ioutil.TempFile(os.TempDir(), "coretest-")
	if err != nil {
		t.Fatalf("Failed creating tempfile: %v", err)
	}
	defer syscall.Unlink(configFile.Name())

	if _, err := io.WriteString(configFile, configData); err != nil {
		t.Fatalf("Failed writing %s: %v", configFile.Name(), err)
	}

	if stdout, stderr, err := Run("sudo", cloudinitBinPath, "--workspace", workspace, "--from-file", configFile.Name()); err != nil {
		t.Fatalf("coreos-cloudinit failed with error: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	unitName, err := read(path.Join(workspace, "scripts", "unit-name"))
	if err != nil {
		t.Fatalf("Unable to read unit name from cloudinit workspace: %v", err)
	}
	defer Run("systemctl", "stop", unitName)

	stdout, stderr, err := Run("systemctl", "status", unitName)
	if err != nil {
		t.Fatalf("Unable to determine if user-data was executed: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	if !strings.Contains(stdout, "Active: active") {
		t.Errorf("User-data unit is not active")
	}
}
<commit_msg>fix(cloudinit): look in etcd.service drop-in for etcd config<commit_after>package coretest

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"testing"
)

const cloudinitBinPath = "/usr/bin/coreos-cloudinit"

func read(filename string) (string, error) {
	bytes, err := ioutil.ReadFile(filename)
	return string(bytes), err
}

func rmdir(path string) error {
	cmd := exec.Command("sudo", "rm", "-rf", path)
	return cmd.Run()
}

func TestCloudinitCloudConfig(t *testing.T) {
	workspace, err := ioutil.TempDir("", "coretest-cloudinit-")
	if err != nil {
		t.Fatalf("Failed creating workspace: %v", err)
	}
	defer rmdir(workspace)

	keyOne := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC5LaGMGRqZEEvOhHlIEiQgdMJIQ9Qe8L/XSz06GqzcESbEnYLIXar2nou4eW4AGMVC1V0BrcWWnSTxM1/dWeCLOUt5NulKAjtdBUZGhCT83nbimSzbmx3/q2y5bCiS4Zr8ZjYFbi1eLvye2jKPE4xo7cvIfDKc0ztQ9kU7JknUdKNZo3RKXr5EPhJ5UZ8Ff15CI9+hDSvdPwer+HNnEt/psRVC+s29EwNGwUXD4IYqrk3X4ew0YAl/oULHM4cctoBW9GM+kAl40rOuIARlKfe4UdCgDMHYA/whi7Us+cPNgPit9IVJVBU4eo/cF5molD2l+PMSntypuv79obu8sA1H cloudinit-test-key-one"
	keyTwo := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCZw5Ljtt9wlEfyDvmUwu/BeMcIhVarbcM4ajZolxRy9G8vvCa7ODcSjzSyhfG1mLSBB2KfaFFI6zGHBjFX0Gzy9i8m3u7PnZBPX30bb1n0hJCrUhpqUGQUe8OFdoBstf1HIwJU/KoTBL0Ap1WEn0quRT4kNgBLbPrMjYCPbS1q4wJKdIE5rRm/EfTUrmIb0i91gujEGw5oUHDXf0X+/cxwwIVZh1z16YhOgvJBzXhsJ9a0w7kcy/6wPRv03yyMg/r2Ada6ci68LulKz5GLn+xInT0bvIcra/PZ7WE+jyZhZKly239VZyT/1dHkBbTw+kgnGobLMbjOOg5bKaT8NZJ3 cloudinit-test-key-two"

	configTmpl := `#cloud-config
coreos:
    etcd:
        discovery: https://discovery.etcd.io/827c73219eeb2fa5530027c37bf18877
ssh_authorized_keys:
    - %s
    - %s
`
	configData := fmt.Sprintf(configTmpl, keyOne, keyTwo)
	configFile, err := ioutil.TempFile(os.TempDir(), "coretest-")
	if err != nil {
		t.Fatalf("Failed creating tempfile: %v", err)
	}
	defer syscall.Unlink(configFile.Name())

	if _, err := io.WriteString(configFile, configData); err != nil {
		t.Fatalf("Failed writing %s: %v", configFile.Name(), err)
	}

	if stdout, stderr, err := Run("sudo", cloudinitBinPath, "--workspace", workspace, "--from-file", configFile.Name(), "--ssh-key-name", "coretest"); err != nil {
		t.Fatalf("coreos-cloudinit failed with error: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	contents, err := read("/run/systemd/system/etcd.service.d/20-cloudinit.conf")
	if err != nil {
		t.Errorf("Unable to read etcd bootstrap file: %v", err)
	} else if !strings.Contains(contents, "Environment=\"ETCD_DISCOVERY=https://discovery.etcd.io/827c73219eeb2fa5530027c37bf18877\"") {
		t.Errorf("Incorrect data written to etcd.service.d/20-cloudinit.conf: %s", contents)
	}

	// Attempt to clean up after ourselves
	defer Run("update-ssh-keys", "-d", "coretest")

	authorized_keys, err := read("/home/core/.ssh/authorized_keys")
	if err != nil {
		t.Fatalf("Unable to read authorized_keys file: %v", err)
	}

	if !strings.Contains(authorized_keys, keyOne) {
		t.Errorf("Could not find first key in authorized_keys")
	}

	if !strings.Contains(authorized_keys, keyTwo) {
		t.Errorf("Could not find second key in authorized_keys")
	}
}

func TestCloudinitScript(t *testing.T) {
	workspace, err := ioutil.TempDir("", "coretest-cloudinit-")
	if err != nil {
		t.Fatalf("Failed creating workspace: %v", err)
	}
	defer rmdir(workspace)

	configData := `#!/bin/bash
/bin/sleep 10
`
	configFile, err := ioutil.TempFile(os.TempDir(), "coretest-")
	if err != nil {
		t.Fatalf("Failed creating tempfile: %v", err)
	}
	defer syscall.Unlink(configFile.Name())

	if _, err := io.WriteString(configFile, configData); err != nil {
		t.Fatalf("Failed writing %s: %v", configFile.Name(), err)
	}

	if stdout, stderr, err := Run("sudo", cloudinitBinPath, "--workspace", workspace, "--from-file", configFile.Name()); err != nil {
		t.Fatalf("coreos-cloudinit failed with error: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	unitName, err := read(path.Join(workspace, "scripts", "unit-name"))
	if err != nil {
		t.Fatalf("Unable to read unit name from cloudinit workspace: %v", err)
	}
	defer Run("systemctl", "stop", unitName)

	stdout, stderr, err := Run("systemctl", "status", unitName)
	if err != nil {
		t.Fatalf("Unable to determine if user-data was executed: %v\nstdout: %s\nstderr: %s", err, stdout, stderr)
	}

	if !strings.Contains(stdout, "Active: active") {
		t.Errorf("User-data unit is not active")
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"net"

	"github.com/ellcrys/util"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/connector"
	"github.com/ncodes/cocoon/core/connector/server"
	"github.com/ncodes/cocoon/core/scheduler"
	logging "github.com/op/go-logging"
	"github.com/spf13/cobra"
)

var (
	log = logging.MustGetLogger("connector")

	// connector RPC API
	defaultConnectorRPCAPI = util.Env("DEV_ADDR_CONNECTOR_RPC", ":8002")

	// Signals channel
	sigs = make(chan os.Signal, 1)
)

func init() {
	config.ConfigureLogger()
}

// creates a deployment request with argument
// fetched from the environment.
func getRequest() (*connector.Request, error) {

	// get cocoon code github link and language
	ccID := os.Getenv("COCOON_ID")
	ccURL := os.Getenv("COCOON_CODE_URL")
	ccTag := os.Getenv("COCOON_CODE_TAG")
	ccLang := os.Getenv("COCOON_CODE_LANG")
	diskLimit := util.Env("COCOON_DISK_LIMIT", "300")
	buildParam := os.Getenv("COCOON_BUILD_PARAMS")
	ccLink := os.Getenv("COCOON_LINK")
	memoryMB, _ := strconv.Atoi(util.Env("ALLOC_MEMORY", "4"))
	cpuShare, _ := strconv.Atoi(util.Env("ALLOC_CPU_SHARE", "100"))

	if ccID == "" {
		return nil, fmt.Errorf("Cocoon code id not set @ $COCOON_ID")
	} else if ccURL == "" {
		return nil, fmt.Errorf("Cocoon code url not set @ $COCOON_CODE_URL")
	} else if ccLang == "" {
		return nil, fmt.Errorf("Cocoon code url not set @ $COCOON_CODE_LANG")
	}

	return &connector.Request{
		ID:          ccID,
		URL:         ccURL,
		Tag:         ccTag,
		Lang:        ccLang,
		DiskLimit:   common.MBToByte(util.ToInt64(diskLimit)),
		BuildParams: buildParam,
		Link:        ccLink,
		Memory:      int64(memoryMB),
		CPUShares:   int64(cpuShare),
	}, nil
}

// onTerminate calls a function when a terminate or interrupt signal is received.
func onTerminate(f func(s os.Signal)) {
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		s := <-sigs
		f(s)
	}()
}

// connectorCmd represents the connector command
var connectorCmd = &cobra.Command{
	Use:   "connector",
	Short: "Start the connector",
	Long:  `Starts the connector and launches a cocoon code.`,
	Run: func(cmd *cobra.Command, args []string) {

		waitCh := make(chan bool, 1)
		serverStartedCh := make(chan bool)

		log.Info("Connector has started")

		// get request
		log.Info("Collecting and validating launch request")
		req, err := getRequest()
		if err != nil {
			log.Error(err.Error())
			return
		}
		log.Infof("Ready to launch cocoon code with id = %s", req.ID)

		connectorRPCAddr := scheduler.Getenv("ADDR_RPC", defaultConnectorRPCAPI)
		cocoonCodeRPCAddr := net.JoinHostPort("", scheduler.Getenv("PORT_code_RPC", "8004"))

		cn := connector.NewConnector(req, waitCh)
		cn.SetAddrs(connectorRPCAddr, cocoonCodeRPCAddr)
		cn.AddLanguage(connector.NewGo(req))
		onTerminate(func(s os.Signal) {
			log.Info("Terminate signal received. Stopping connector")
			cn.Stop(false)
		})

		// start grpc API server
		rpcServer := server.NewRPCServer(cn)
		go rpcServer.Start(connectorRPCAddr, serverStartedCh)

		// wait for rpc server to start then launch the cocoon code
		<-serverStartedCh
		go cn.Launch(connectorRPCAddr, cocoonCodeRPCAddr)

		if <-waitCh {
			rpcServer.Stop(1)
		} else {
			rpcServer.Stop(0)
		}

		log.Info("Connector stopped")
	},
}

func init() {
	RootCmd.AddCommand(connectorCmd)
}
<commit_msg>fix cocoon code addr env var<commit_after>package cmd

import (
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/ellcrys/util"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/connector"
	"github.com/ncodes/cocoon/core/connector/server"
	"github.com/ncodes/cocoon/core/scheduler"
	logging "github.com/op/go-logging"
	"github.com/spf13/cobra"
)

var (
	log = logging.MustGetLogger("connector")

	// connector RPC API
	defaultConnectorRPCAPI = util.Env("DEV_ADDR_CONNECTOR_RPC", ":8002")

	// Signals channel
	sigs = make(chan os.Signal, 1)
)

func init() {
	config.ConfigureLogger()
}

// creates a deployment request with argument
// fetched from the environment.
func getRequest() (*connector.Request, error) {

	// get cocoon code github link and language
	ccID := os.Getenv("COCOON_ID")
	ccURL := os.Getenv("COCOON_CODE_URL")
	ccTag := os.Getenv("COCOON_CODE_TAG")
	ccLang := os.Getenv("COCOON_CODE_LANG")
	diskLimit := util.Env("COCOON_DISK_LIMIT", "300")
	buildParam := os.Getenv("COCOON_BUILD_PARAMS")
	ccLink := os.Getenv("COCOON_LINK")
	memoryMB, _ := strconv.Atoi(util.Env("ALLOC_MEMORY", "4"))
	cpuShare, _ := strconv.Atoi(util.Env("ALLOC_CPU_SHARE", "100"))

	if ccID == "" {
		return nil, fmt.Errorf("Cocoon code id not set @ $COCOON_ID")
	} else if ccURL == "" {
		return nil, fmt.Errorf("Cocoon code url not set @ $COCOON_CODE_URL")
	} else if ccLang == "" {
		return nil, fmt.Errorf("Cocoon code url not set @ $COCOON_CODE_LANG")
	}

	return &connector.Request{
		ID:          ccID,
		URL:         ccURL,
		Tag:         ccTag,
		Lang:        ccLang,
		DiskLimit:   common.MBToByte(util.ToInt64(diskLimit)),
		BuildParams: buildParam,
		Link:        ccLink,
		Memory:      int64(memoryMB),
		CPUShares:   int64(cpuShare),
	}, nil
}

// onTerminate calls a function when a terminate or interrupt signal is received.
func onTerminate(f func(s os.Signal)) {
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		s := <-sigs
		f(s)
	}()
}

// connectorCmd represents the connector command
var connectorCmd = &cobra.Command{
	Use:   "connector",
	Short: "Start the connector",
	Long:  `Starts the connector and launches a cocoon code.`,
	Run: func(cmd *cobra.Command, args []string) {

		waitCh := make(chan bool, 1)
		serverStartedCh := make(chan bool)

		log.Info("Connector has started")

		// get request
		log.Info("Collecting and validating launch request")
		req, err := getRequest()
		if err != nil {
			log.Error(err.Error())
			return
		}
		log.Infof("Ready to launch cocoon code with id = %s", req.ID)

		connectorRPCAddr := scheduler.Getenv("ADDR_RPC", defaultConnectorRPCAPI)
		cocoonCodeRPCAddr := scheduler.Getenv("ADDR_code_RPC", ":8004")

		cn := connector.NewConnector(req, waitCh)
		cn.SetAddrs(connectorRPCAddr, cocoonCodeRPCAddr)
		cn.AddLanguage(connector.NewGo(req))
		onTerminate(func(s os.Signal) {
			log.Info("Terminate signal received. Stopping connector")
			cn.Stop(false)
		})

		// start grpc API server
		rpcServer := server.NewRPCServer(cn)
		go rpcServer.Start(connectorRPCAddr, serverStartedCh)

		// wait for rpc server to start then launch the cocoon code
		<-serverStartedCh
		go cn.Launch(connectorRPCAddr, cocoonCodeRPCAddr)

		if <-waitCh {
			rpcServer.Stop(1)
		} else {
			rpcServer.Stop(0)
		}

		log.Info("Connector stopped")
	},
}

func init() {
	RootCmd.AddCommand(connectorCmd)
}
<|endoftext|>


<commit_before>package commands

import (
	"flag"
	"mayday/core"
)

type PullCommand struct {
	token  *string
	uuid   *string
	server *string
}

func (cmd *PullCommand) Name() string {
	return "pull"
}

func (cmd *PullCommand) Description() string {
	return "Get the current reports files for a case."
}

func (cmd *PullCommand) DefineFlags(fs *flag.FlagSet) {
	cmd.uuid = fs.String("uuid", "", "PGP KeyID to sign the new configuration")
	cmd.token = fs.String("token", "", "PGP KeyID to sign the new configuration")
	cmd.server = fs.String("server", core.DefaultAPIBaseURL, "Mayday server address")
}

func (cmd *PullCommand) Run() {
	// real stuff
}
<commit_msg>Added "mayday pull" command<commit_after>package commands

import (
	"flag"
	"fmt"
	"io/ioutil"
	"mayday/core"
	"os"
	"path"
)

type PullCommand struct {
	token  *string
	id     *string
	server *string
	all    *bool
	fileId *string
	to     *string
}

func (cmd *PullCommand) Name() string {
	return "pull"
}

func (cmd *PullCommand) Description() string {
	return "Get the current reports files for a case."
}

func (cmd *PullCommand) DefineFlags(fs *flag.FlagSet) {
	cmd.id = fs.String("id", "", "Case ID to pull from")
	cmd.all = fs.Bool("all", true, "Pull all files from the case")
	cmd.fileId = fs.String("file-id", "", "File Id to retrieve")
	cmd.to = fs.String("to", "", "Path to store the retrieved files")
	cmd.token = fs.String("token", "", "PGP KeyID to sign the new configuration")
	cmd.server = fs.String("server", core.DefaultAPIBaseURL, "Mayday server address")
}

func (cmd *PullCommand) Run() {
	if *cmd.id == "" {
		fmt.Println("Please specify a valid Case Id --id\n")
		os.Exit(1)
	}

	mayday, err := core.NewClient(*cmd.server, *cmd.id, *cmd.token)
	if err != nil {
		fmt.Println(err)
	}

	var files map[string]string

	if *cmd.fileId != "" {
		files, err = mayday.Pull(*cmd.fileId)
		if err != nil {
			fmt.Println(err)
		}

	} else if !*cmd.all {
		fmt.Println("Please specify a --file-id or --all")
		os.Exit(1)
	} else {
		files, err = mayday.PullAll()
		if err != nil {
			fmt.Println(err)
		}
	}

	for name, content := range files {
		filename := ""
		if *cmd.to == "" {
			base, err := core.GetDefaultDirectory()
			if err != nil {
				fmt.Println(err)
			}

			base = path.Join(base, "pull", *cmd.id)
			if _, err := os.Stat(base); os.IsNotExist(err) {
				err := os.MkdirAll(base, 0700)
				if err != nil {
					fmt.Println(err)
				}
			}

			filename = path.Join(base, name)
		} else {
			filename = path.Join(*cmd.to, name)
		}

		err := ioutil.WriteFile(filename, []byte(content), 0700)
		if err != nil {
			fmt.Println("Error writing file")
		} else {
			fmt.Printf("Pulled report on path: %s\n", filename)
		}
	}

	// real stuff
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	importer "github.com/jbenet/go-ipfs/importer"
	"github.com/jbenet/go-ipfs/importer/chunk"
	dag "github.com/jbenet/go-ipfs/merkledag"
	pinning "github.com/jbenet/go-ipfs/pin"
	u "github.com/jbenet/go-ipfs/util"
)

// Error indicating the max depth has been exceded.
var ErrDepthLimitExceeded = fmt.Errorf("depth limit exceeded")

type AddOutput struct {
	Objects []*Object
	Names   []string
}

var addCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Add an object to ipfs.",
		ShortDescription: `
Adds contents of <path> to ipfs. Use -r to add directories.
Note that directories are added recursively, to form the ipfs
MerkleDAG. A smarter partial add with a staging area (like git)
remains to be implemented.
`,
	},

	Options: []cmds.Option{
		cmds.BoolOption("recursive", "r", "Must be specified when adding directories"),
	},
	Arguments: []cmds.Argument{
		cmds.FileArg("path", true, true, "The path to a file to be added to IPFS"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		added := &AddOutput{}
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		_, _, err = req.Option("r").Bool()
		if err != nil {
			return nil, err
		}

		// returns the last one
		addDagnode := func(name string, dn *dag.Node) error {
			o, err := getOutput(dn)
			if err != nil {
				return err
			}

			added.Objects = append(added.Objects, o)
			added.Names = append(added.Names, name)
			return nil
		}

		addFile := func(file cmds.File) (*dag.Node, error) {
			dns, err := add(n, []io.Reader{file})
			if err != nil {
				return nil, err
			}

			log.Infof("adding file: %s", file.FileName())
			if err := addDagnode(file.FileName(), dns[len(dns)-1]); err != nil {
				return nil, err
			}
			return dns[len(dns)-1], nil // last dag node is the file.
		}

		// TODO: handle directories

		file, err := req.Files().NextFile()
		for file != nil {
			_, err := addFile(file)
			if err != nil {
				return nil, err
			}

			file, err = req.Files().NextFile()
		}
		if err != nil && err != io.EOF {
			return nil, err
		}

		return added, nil

		// readers, err := internal.CastToReaders(req.Arguments())
		// if err != nil {
		// 	return nil, err
		// }
		//
		// dagnodes, err := add(n, readers)
		// if err != nil {
		// 	return nil, err
		// }
		//
		// // TODO: include fs paths in output (will need a way to specify paths in underlying filearg system)
		// added := make([]*Object, 0, len(req.Arguments()))
		// for _, dagnode := range dagnodes {
		// 	object, err := getOutput(dagnode)
		// 	if err != nil {
		// 		return nil, err
		// 	}
		//
		// 	added = append(added, object)
		// }
		//
		// return &AddOutput{added}, nil
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) ([]byte, error) {
			val, ok := res.Output().(*AddOutput)
			if !ok {
				return nil, u.ErrCast()
			}

			var buf bytes.Buffer
			for i, obj := range val.Objects {
				buf.Write([]byte(fmt.Sprintf("added %s %s\n", obj.Hash, val.Names[i])))
			}
			return buf.Bytes(), nil
		},
	},
	Type: &AddOutput{},
}

func add(n *core.IpfsNode, readers []io.Reader) ([]*dag.Node, error) {
	mp, ok := n.Pinning.(pinning.ManualPinner)
	if !ok {
		return nil, errors.New("invalid pinner type! expected manual pinner")
	}

	dagnodes := make([]*dag.Node, 0)

	// TODO: allow adding directories (will need support for multiple files in filearg system)

	for _, reader := range readers {
		node, err := importer.BuildDagFromReader(reader, n.DAG, mp, chunk.DefaultSplitter)
		if err != nil {
			return nil, err
		}
		dagnodes = append(dagnodes, node)
	}

	return dagnodes, nil
}

func addNode(n *core.IpfsNode, node *dag.Node) error {
	err := n.DAG.AddRecursive(node) // add the file to the graph + local storage
	if err != nil {
		return err
	}

	err = n.Pinning.Pin(node, true) // ensure we keep it
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>core/commands2: Updated 'add' command for new file API<commit_after>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	importer "github.com/jbenet/go-ipfs/importer"
	"github.com/jbenet/go-ipfs/importer/chunk"
	dag "github.com/jbenet/go-ipfs/merkledag"
	pinning "github.com/jbenet/go-ipfs/pin"
	ft "github.com/jbenet/go-ipfs/unixfs"
	u "github.com/jbenet/go-ipfs/util"
)

// Error indicating the max depth has been exceded.
var ErrDepthLimitExceeded = fmt.Errorf("depth limit exceeded")

type AddOutput struct {
	Objects []*Object
	Names   []string
}

var addCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Add an object to ipfs.",
		ShortDescription: `
Adds contents of <path> to ipfs. Use -r to add directories.
Note that directories are added recursively, to form the ipfs
MerkleDAG. A smarter partial add with a staging area (like git)
remains to be implemented.
`,
	},

	Arguments: []cmds.Argument{
		cmds.FileArg("path", true, true, "The path to a file to be added to IPFS").EnableRecursive(),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		added := &AddOutput{}
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		for {
			file, err := req.Files().NextFile()
			if err != nil && err != io.EOF {
				return nil, err
			}
			if file == nil {
				break
			}

			_, err = addFile(n, file, added)
			if err != nil {
				return nil, err
			}
		}

		return added, nil
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) ([]byte, error) {
			val, ok := res.Output().(*AddOutput)
			if !ok {
				return nil, u.ErrCast()
			}

			var buf bytes.Buffer
			for i, obj := range val.Objects {
				buf.Write([]byte(fmt.Sprintf("added %s %s\n", obj.Hash, val.Names[i])))
			}
			return buf.Bytes(), nil
		},
	},
	Type: &AddOutput{},
}

func add(n *core.IpfsNode, readers []io.Reader) ([]*dag.Node, error) {
	mp, ok := n.Pinning.(pinning.ManualPinner)
	if !ok {
		return nil, errors.New("invalid pinner type! expected manual pinner")
	}

	dagnodes := make([]*dag.Node, 0)

	for _, reader := range readers {
		node, err := importer.BuildDagFromReader(reader, n.DAG, mp, chunk.DefaultSplitter)
		if err != nil {
			return nil, err
		}
		dagnodes = append(dagnodes, node)
	}

	return dagnodes, nil
}

func addNode(n *core.IpfsNode, node *dag.Node) error {
	err := n.DAG.AddRecursive(node) // add the file to the graph + local storage
	if err != nil {
		return err
	}

	err = n.Pinning.Pin(node, true) // ensure we keep it
	if err != nil {
		return err
	}

	return nil
}

func addFile(n *core.IpfsNode, file cmds.File, added *AddOutput) (*dag.Node, error) {
	if file.IsDirectory() {
		return addDir(n, file, added)
	}

	dns, err := add(n, []io.Reader{file})
	if err != nil {
		return nil, err
	}

	log.Infof("adding file: %s", file.FileName())
	if err := addDagnode(added, file.FileName(), dns[len(dns)-1]); err != nil {
		return nil, err
	}
	return dns[len(dns)-1], nil // last dag node is the file.
}

func addDir(n *core.IpfsNode, dir cmds.File, added *AddOutput) (*dag.Node, error) {
	log.Infof("adding directory: %s", dir.FileName())

	tree := &dag.Node{Data: ft.FolderPBData()}

	for {
		file, err := dir.NextFile()
		if err != nil && err != io.EOF {
			return nil, err
		}
		if file == nil {
			break
		}

		node, err := addFile(n, file, added)
		if err != nil {
			return nil, err
		}

		name := file.FileName()
		slashIndex := strings.LastIndex(name, "/")
		name = name[slashIndex+1:]

		err = tree.AddNodeLink(name, node)
		if err != nil {
			return nil, err
		}
	}

	addDagnode(added, dir.FileName(), tree)

	return tree, addNode(n, tree)
}

// addDagnode adds dagnode info to an output object
func addDagnode(output *AddOutput, name string, dn *dag.Node) error {
	o, err := getOutput(dn)
	if err != nil {
		return err
	}

	output.Objects = append(output.Objects, o)
	output.Names = append(output.Names, name)
	return nil
}
<|endoftext|>


<commit_before>package cluster

import (
	"encoding/json"
	"strings"

	"github.com/samalba/dockerclient"
)

// SwarmLabelNamespace defines the key prefix in all custom labels
const SwarmLabelNamespace = "com.docker.swarm"

// ContainerConfig is exported
// TODO store affinities and constraints in their own fields
type ContainerConfig struct {
	dockerclient.ContainerConfig
}

func parseEnv(e string) (bool, string, string) {
	parts := strings.SplitN(e, ":", 2)
	if len(parts) == 2 {
		return true, parts[0], parts[1]
	}
	return false, "", ""
}

// BuildContainerConfig creates a cluster.ContainerConfig from a dockerclient.ContainerConfig
func BuildContainerConfig(c dockerclient.ContainerConfig) *ContainerConfig {
	var (
		affinities  []string
		constraints []string
		env         []string
	)

	// only for tests
	if c.Labels == nil {
		c.Labels = make(map[string]string)
	}

	// parse affinities from labels (ex. docker run --label 'com.docker.swarm.affinities=["container==redis","image==nginx"]')
	if labels, ok := c.Labels[SwarmLabelNamespace+".affinities"]; ok {
		json.Unmarshal([]byte(labels), &affinities)
	}

	// parse contraints from labels (ex. docker run --label 'com.docker.swarm.constraints=["region==us-east","storage==ssd"]')
	if labels, ok := c.Labels[SwarmLabelNamespace+".constraints"]; ok {
		json.Unmarshal([]byte(labels), &constraints)
	}

	// parse affinities/contraints from env (ex. docker run -e affinity:container==redis -e affinity:image==nginx -e constraint:region==us-east -e constraint:storage==ssd)
	for _, e := range c.Env {
		if ok, key, value := parseEnv(e); ok && key == "affinity" {
			affinities = append(affinities, value)
		} else if ok && key == "constraint" {
			constraints = append(constraints, value)
		} else {
			env = append(env, e)
		}
	}

	// remove affinities/contraints from env
	c.Env = env

	// store affinities in labels
	if len(affinities) > 0 {
		if labels, err := json.Marshal(affinities); err == nil {
			c.Labels[SwarmLabelNamespace+".affinities"] = string(labels)
		}
	}

	// store contraints in labels
	if len(constraints) > 0 {
		if labels, err := json.Marshal(constraints); err == nil {
			c.Labels[SwarmLabelNamespace+".constraints"] = string(labels)
		}
	}

	return &ContainerConfig{c}
}

func (c *ContainerConfig) extractExprs(key string) []string {
	var exprs []string

	if labels, ok := c.Labels[SwarmLabelNamespace+"."+key]; ok {
		json.Unmarshal([]byte(labels), &exprs)
	}

	return exprs
}

// SwarmID extracts the Swarm ID from the Config.
// May return an empty string if not set.
func (c *ContainerConfig) SwarmID() string {
	return c.Labels[SwarmLabelNamespace+".id"]
}

// SetSwarmID sets or overrides the Swarm ID in the Config.
func (c *ContainerConfig) SetSwarmID(id string) {
	c.Labels[SwarmLabelNamespace+".id"] = id
}

// Affinities returns all the affinities from the ContainerConfig
func (c *ContainerConfig) Affinities() []string {
	return c.extractExprs("affinities")
}

// Constraints returns all the constraints from the ContainerConfig
func (c *ContainerConfig) Constraints() []string {
	return c.extractExprs("constraints")
}
<commit_msg>handle compat<commit_after>package cluster

import (
	"encoding/json"
	"strings"

	"github.com/samalba/dockerclient"
)

// SwarmLabelNamespace defines the key prefix in all custom labels
const SwarmLabelNamespace = "com.docker.swarm"

// ContainerConfig is exported
// TODO store affinities and constraints in their own fields
type ContainerConfig struct {
	dockerclient.ContainerConfig
}

func parseEnv(e string) (bool, string, string) {
	parts := strings.SplitN(e, ":", 2)
	if len(parts) == 2 {
		return true, parts[0], parts[1]
	}
	return false, "", ""
}

// FIXME: Temporary fix to handle forward/backward compatibility between Docker <1.6 and >=1.7
// ContainerConfig should be handling converting to/from different docker versions
func consolidateResourceFields(c *dockerclient.ContainerConfig) {
	if c.Memory != c.HostConfig.Memory {
		if c.Memory != 0 {
			c.HostConfig.Memory = c.Memory
		} else {
			c.Memory = c.HostConfig.Memory
		}
	}

	if c.MemorySwap != c.HostConfig.MemorySwap {
		if c.Memory != 0 {
			c.HostConfig.MemorySwap = c.MemorySwap
		} else {
			c.MemorySwap = c.HostConfig.MemorySwap
		}
	}

	if c.CpuShares != c.HostConfig.CpuShares {
		if c.CpuShares != 0 {
			c.HostConfig.CpuShares = c.CpuShares
		} else {
			c.CpuShares = c.HostConfig.CpuShares
		}
	}

	if c.Cpuset != c.HostConfig.CpusetCpus {
		if c.Cpuset != "" {
			c.HostConfig.CpusetCpus = c.Cpuset
		} else {
			c.Cpuset = c.HostConfig.CpusetCpus
		}
	}
}

// BuildContainerConfig creates a cluster.ContainerConfig from a dockerclient.ContainerConfig
func BuildContainerConfig(c dockerclient.ContainerConfig) *ContainerConfig {
	var (
		affinities  []string
		constraints []string
		env         []string
	)

	// only for tests
	if c.Labels == nil {
		c.Labels = make(map[string]string)
	}

	// parse affinities from labels (ex. docker run --label 'com.docker.swarm.affinities=["container==redis","image==nginx"]')
	if labels, ok := c.Labels[SwarmLabelNamespace+".affinities"]; ok {
		json.Unmarshal([]byte(labels), &affinities)
	}

	// parse contraints from labels (ex. docker run --label 'com.docker.swarm.constraints=["region==us-east","storage==ssd"]')
	if labels, ok := c.Labels[SwarmLabelNamespace+".constraints"]; ok {
		json.Unmarshal([]byte(labels), &constraints)
	}

	// parse affinities/contraints from env (ex. docker run -e affinity:container==redis -e affinity:image==nginx -e constraint:region==us-east -e constraint:storage==ssd)
	for _, e := range c.Env {
		if ok, key, value := parseEnv(e); ok && key == "affinity" {
			affinities = append(affinities, value)
		} else if ok && key == "constraint" {
			constraints = append(constraints, value)
		} else {
			env = append(env, e)
		}
	}

	// remove affinities/contraints from env
	c.Env = env

	// store affinities in labels
	if len(affinities) > 0 {
		if labels, err := json.Marshal(affinities); err == nil {
			c.Labels[SwarmLabelNamespace+".affinities"] = string(labels)
		}
	}

	// store contraints in labels
	if len(constraints) > 0 {
		if labels, err := json.Marshal(constraints); err == nil {
			c.Labels[SwarmLabelNamespace+".constraints"] = string(labels)
		}
	}

	consolidateResourceFields(&c)

	return &ContainerConfig{c}
}

func (c *ContainerConfig) extractExprs(key string) []string {
	var exprs []string

	if labels, ok := c.Labels[SwarmLabelNamespace+"."+key]; ok {
		json.Unmarshal([]byte(labels), &exprs)
	}

	return exprs
}

// SwarmID extracts the Swarm ID from the Config.
// May return an empty string if not set.
func (c *ContainerConfig) SwarmID() string {
	return c.Labels[SwarmLabelNamespace+".id"]
}

// SetSwarmID sets or overrides the Swarm ID in the Config.
func (c *ContainerConfig) SetSwarmID(id string) {
	c.Labels[SwarmLabelNamespace+".id"] = id
}

// Affinities returns all the affinities from the ContainerConfig
func (c *ContainerConfig) Affinities() []string {
	return c.extractExprs("affinities")
}

// Constraints returns all the constraints from the ContainerConfig
func (c *ContainerConfig) Constraints() []string {
	return c.extractExprs("constraints")
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package packet_injector

import (
	"errors"
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/flow"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
)

var (
	options = gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}
)

// PacketParams describes the packet parameters to be injected
type PacketParams struct {
	UUID      string
	SrcNodeID graph.Identifier `valid:"nonzero"`
	SrcIP     string           `valid:"nonzero"`
	SrcMAC    string           `valid:"nonzero"`
	SrcPort   int64            `valid:"min=0"`
	DstIP     string           `valid:"nonzero"`
	DstMAC    string           `valid:"nonzero"`
	DstPort   int64            `valid:"min=0"`
	Type      string           `valid:"regexp=^(icmp4|icmp6|tcp4|tcp6|udp4|udp6)$"`
	Count     int64            `valid:"min=1"`
	ID        int64            `valid:"min=0"`
	Interval  int64            `valid:"min=0"`
	Payload   string
}

type Channels struct {
	sync.Mutex
	Pipes map[string](chan bool)
}

// InjectPacket inject some packets based on the graph
func InjectPacket(pp *PacketParams, g *graph.Graph, chnl *Channels) (string, error) {
	srcIP := getIP(pp.SrcIP)
	if srcIP == nil {
		return "", errors.New("Source Node doesn't have proper IP")
	}

	dstIP := getIP(pp.DstIP)
	if dstIP == nil {
		return "", errors.New("Destination Node doesn't have proper IP")
	}

	srcMAC, err := net.ParseMAC(pp.SrcMAC)
	if err != nil || srcMAC == nil {
		return "", errors.New("Source Node doesn't have proper MAC")
	}

	dstMAC, err := net.ParseMAC(pp.DstMAC)
	if err != nil || dstMAC == nil {
		return "", errors.New("Destination Node doesn't have proper MAC")
	}

	g.RLock()

	srcNode := g.GetNode(pp.SrcNodeID)
	if srcNode == nil {
		g.RUnlock()
		return "", errors.New("Unable to find source node")
	}

	tid, err := srcNode.GetFieldString("TID")
	if err != nil {
		g.RUnlock()
		return "", errors.New("Source node has no TID")
	}

	ifName, err := srcNode.GetFieldString("Name")
	if err != nil {
		g.RUnlock()
		return "", errors.New("Source node has no name")
	}

	_, nsPath, err := topology.NamespaceFromNode(g, srcNode)

	g.RUnlock()

	if err != nil {
		return "", err
	}

	var rawSocket *common.RawSocket
	if nsPath != "" {
		rawSocket, err = common.NewRawSocketInNs(nsPath, ifName)
	} else {
		rawSocket, err = common.NewRawSocket(ifName)
	}
	if err != nil {
		return "", err
	}

	var l []gopacket.SerializableLayer
	var layerType gopacket.LayerType
	ethLayer := &layers.Ethernet{SrcMAC: srcMAC, DstMAC: dstMAC}
	payload := gopacket.Payload([]byte(pp.Payload))

	switch pp.Type {
	case "icmp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{Version: 4, SrcIP: srcIP, DstIP: dstIP, Protocol: layers.IPProtocolICMPv4}
		icmpLayer := &layers.ICMPv4{
			TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
			Id:       uint16(pp.ID),
		}
		layerType = layers.LayerTypeEthernet
		l = append(l, ethLayer, ipLayer, icmpLayer, payload)
	case "icmp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{Version: 6, SrcIP: srcIP, DstIP: dstIP, NextHeader: layers.IPProtocolICMPv6}
		icmpLayer := &layers.ICMPv6{
			TypeCode:  layers.CreateICMPv6TypeCode(layers.ICMPv6TypeEchoRequest, 0),
			TypeBytes: []byte{byte(pp.ID & int64(0xFF00) >> 8), byte(pp.ID & int64(0xFF)), 0, 0},
		}
		layerType = layers.LayerTypeEthernet
		icmpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, icmpLayer, payload)
	case "tcp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{SrcIP: srcIP, DstIP: dstIP, Version: 4, Protocol: layers.IPProtocolTCP, TTL: 64}
		srcPort := layers.TCPPort(pp.SrcPort)
		dstPort := layers.TCPPort(pp.DstPort)
		tcpLayer := &layers.TCP{SrcPort: srcPort, DstPort: dstPort, Seq: rand.Uint32(), SYN: true}
		layerType = layers.LayerTypeTCP
		tcpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, tcpLayer, payload)
	case "tcp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{Version: 6, SrcIP: srcIP, DstIP: dstIP, NextHeader: layers.IPProtocolTCP}
		srcPort := layers.TCPPort(pp.SrcPort)
		dstPort := layers.TCPPort(pp.DstPort)
		tcpLayer := &layers.TCP{SrcPort: srcPort, DstPort: dstPort, Seq: rand.Uint32(), SYN: true}
		layerType = layers.LayerTypeTCP
		tcpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, tcpLayer, payload)
	case "udp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{SrcIP: srcIP, DstIP: dstIP, Version: 4, Protocol: layers.IPProtocolUDP, TTL: 64}
		srcPort := layers.UDPPort(pp.SrcPort)
		dstPort := layers.UDPPort(pp.DstPort)
		udpLayer := &layers.UDP{SrcPort: srcPort, DstPort: dstPort}
		layerType = layers.LayerTypeUDP
		udpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, udpLayer, payload)
	case "udp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{SrcIP: srcIP, DstIP: dstIP, Version: 6, NextHeader: layers.IPProtocolUDP}
		srcPort := layers.UDPPort(pp.SrcPort)
		dstPort := layers.UDPPort(pp.DstPort)
		udpLayer := &layers.UDP{SrcPort: srcPort, DstPort: dstPort}
		layerType = layers.LayerTypeUDP
		udpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, udpLayer, payload)
	default:
		rawSocket.Close()
		return "", fmt.Errorf("Unsupported traffic type '%s'", pp.Type)
	}

	buffer := gopacket.NewSerializeBuffer()
	if err := gopacket.SerializeLayers(buffer, options, l...); err != nil {
		rawSocket.Close()
		return "", fmt.Errorf("Error while generating %s packet: %s", pp.Type, err.Error())
	}

	packetData := buffer.Bytes()
	packet := gopacket.NewPacket(packetData, layerType, gopacket.Default)
	flowKey := flow.KeyFromGoPacket(&packet, "").String()
	f := flow.NewFlow()
	f.InitFromGoPacket(flowKey, &packet, int64(len(packetData)), tid, flow.FlowUUIDs{}, flow.FlowOpts{})

	p := make(chan bool)
	chnl.Lock()
	chnl.Pipes[pp.UUID] = p
	chnl.Unlock()

	go func(c chan bool) {
		defer rawSocket.Close()

	stopInjection:
		for i := int64(0); i < pp.Count; i++ {
			select {
			case <-c:
				logging.GetLogger().Debugf("Injection stoped on interface %s", ifName)
				break stopInjection
			default:
				logging.GetLogger().Debugf("Injecting packet on interface %s", ifName)

				if _, err := rawSocket.Write(packetData); err != nil {
					if err == syscall.ENXIO {
						logging.GetLogger().Warningf("Write error: %s", err.Error())
					} else {
						logging.GetLogger().Errorf("Write error: %s", err.Error())
					}

					if i != pp.Count-1 {
						time.Sleep(time.Millisecond * time.Duration(pp.Interval))
					}
				}
			}
		}
		chnl.Lock()
		delete(chnl.Pipes, pp.UUID)
		chnl.Unlock()
	}(p)

	return f.TrackingID, nil
}

func getIP(cidr string) net.IP {
	if len(cidr) <= 0 {
		return nil
	}
	ips := strings.Split(cidr, ",")
	//TODO(masco): currently taking first IP, need to implement to select a proper IP
	ip, _, err := net.ParseCIDR(ips[0])
	if err != nil {
		return nil
	}
	return ip
}
<commit_msg>pi: fix interval between packets<commit_after>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package packet_injector

import (
	"errors"
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/flow"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
)

var (
	options = gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}
)

// PacketParams describes the packet parameters to be injected
type PacketParams struct {
	UUID      string
	SrcNodeID graph.Identifier `valid:"nonzero"`
	SrcIP     string           `valid:"nonzero"`
	SrcMAC    string           `valid:"nonzero"`
	SrcPort   int64            `valid:"min=0"`
	DstIP     string           `valid:"nonzero"`
	DstMAC    string           `valid:"nonzero"`
	DstPort   int64            `valid:"min=0"`
	Type      string           `valid:"regexp=^(icmp4|icmp6|tcp4|tcp6|udp4|udp6)$"`
	Count     int64            `valid:"min=1"`
	ID        int64            `valid:"min=0"`
	Interval  int64            `valid:"min=0"`
	Payload   string
}

type Channels struct {
	sync.Mutex
	Pipes map[string](chan bool)
}

// InjectPacket inject some packets based on the graph
func InjectPacket(pp *PacketParams, g *graph.Graph, chnl *Channels) (string, error) {
	srcIP := getIP(pp.SrcIP)
	if srcIP == nil {
		return "", errors.New("Source Node doesn't have proper IP")
	}

	dstIP := getIP(pp.DstIP)
	if dstIP == nil {
		return "", errors.New("Destination Node doesn't have proper IP")
	}

	srcMAC, err := net.ParseMAC(pp.SrcMAC)
	if err != nil || srcMAC == nil {
		return "", errors.New("Source Node doesn't have proper MAC")
	}

	dstMAC, err := net.ParseMAC(pp.DstMAC)
	if err != nil || dstMAC == nil {
		return "", errors.New("Destination Node doesn't have proper MAC")
	}

	g.RLock()

	srcNode := g.GetNode(pp.SrcNodeID)
	if srcNode == nil {
		g.RUnlock()
		return "", errors.New("Unable to find source node")
	}

	tid, err := srcNode.GetFieldString("TID")
	if err != nil {
		g.RUnlock()
		return "", errors.New("Source node has no TID")
	}

	ifName, err := srcNode.GetFieldString("Name")
	if err != nil {
		g.RUnlock()
		return "", errors.New("Source node has no name")
	}

	_, nsPath, err := topology.NamespaceFromNode(g, srcNode)

	g.RUnlock()

	if err != nil {
		return "", err
	}

	var rawSocket *common.RawSocket
	if nsPath != "" {
		rawSocket, err = common.NewRawSocketInNs(nsPath, ifName)
	} else {
		rawSocket, err = common.NewRawSocket(ifName)
	}
	if err != nil {
		return "", err
	}

	var l []gopacket.SerializableLayer
	var layerType gopacket.LayerType
	ethLayer := &layers.Ethernet{SrcMAC: srcMAC, DstMAC: dstMAC}
	payload := gopacket.Payload([]byte(pp.Payload))

	switch pp.Type {
	case "icmp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{Version: 4, SrcIP: srcIP, DstIP: dstIP, Protocol: layers.IPProtocolICMPv4}
		icmpLayer := &layers.ICMPv4{
			TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
			Id:       uint16(pp.ID),
		}
		layerType = layers.LayerTypeEthernet
		l = append(l, ethLayer, ipLayer, icmpLayer, payload)
	case "icmp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{Version: 6, SrcIP: srcIP, DstIP: dstIP, NextHeader: layers.IPProtocolICMPv6}
		icmpLayer := &layers.ICMPv6{
			TypeCode:  layers.CreateICMPv6TypeCode(layers.ICMPv6TypeEchoRequest, 0),
			TypeBytes: []byte{byte(pp.ID & int64(0xFF00) >> 8), byte(pp.ID & int64(0xFF)), 0, 0},
		}
		layerType = layers.LayerTypeEthernet
		icmpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, icmpLayer, payload)
	case "tcp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{SrcIP: srcIP, DstIP: dstIP, Version: 4, Protocol: layers.IPProtocolTCP, TTL: 64}
		srcPort := layers.TCPPort(pp.SrcPort)
		dstPort := layers.TCPPort(pp.DstPort)
		tcpLayer := &layers.TCP{SrcPort: srcPort, DstPort: dstPort, Seq: rand.Uint32(), SYN: true}
		layerType = layers.LayerTypeTCP
		tcpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, tcpLayer, payload)
	case "tcp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{Version: 6, SrcIP: srcIP, DstIP: dstIP, NextHeader: layers.IPProtocolTCP}
		srcPort := layers.TCPPort(pp.SrcPort)
		dstPort := layers.TCPPort(pp.DstPort)
		tcpLayer := &layers.TCP{SrcPort: srcPort, DstPort: dstPort, Seq: rand.Uint32(), SYN: true}
		layerType = layers.LayerTypeTCP
		tcpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, tcpLayer, payload)
	case "udp4":
		ethLayer.EthernetType = layers.EthernetTypeIPv4
		ipLayer := &layers.IPv4{SrcIP: srcIP, DstIP: dstIP, Version: 4, Protocol: layers.IPProtocolUDP, TTL: 64}
		srcPort := layers.UDPPort(pp.SrcPort)
		dstPort := layers.UDPPort(pp.DstPort)
		udpLayer := &layers.UDP{SrcPort: srcPort, DstPort: dstPort}
		layerType = layers.LayerTypeUDP
		udpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, udpLayer, payload)
	case "udp6":
		ethLayer.EthernetType = layers.EthernetTypeIPv6
		ipLayer := &layers.IPv6{SrcIP: srcIP, DstIP: dstIP, Version: 6, NextHeader: layers.IPProtocolUDP}
		srcPort := layers.UDPPort(pp.SrcPort)
		dstPort := layers.UDPPort(pp.DstPort)
		udpLayer := &layers.UDP{SrcPort: srcPort, DstPort: dstPort}
		layerType = layers.LayerTypeUDP
		udpLayer.SetNetworkLayerForChecksum(ipLayer)
		l = append(l, ethLayer, ipLayer, udpLayer, payload)
	default:
		rawSocket.Close()
		return "", fmt.Errorf("Unsupported traffic type '%s'", pp.Type)
	}

	buffer := gopacket.NewSerializeBuffer()
	if err := gopacket.SerializeLayers(buffer, options, l...); err != nil {
		rawSocket.Close()
		return "", fmt.Errorf("Error while generating %s packet: %s", pp.Type, err.Error())
	}

	packetData := buffer.Bytes()
	packet := gopacket.NewPacket(packetData, layerType, gopacket.Default)
	flowKey := flow.KeyFromGoPacket(&packet, "").String()
	f := flow.NewFlow()
	f.InitFromGoPacket(flowKey, &packet, int64(len(packetData)), tid, flow.FlowUUIDs{}, flow.FlowOpts{})

	p := make(chan bool)
	chnl.Lock()
	chnl.Pipes[pp.UUID] = p
	chnl.Unlock()

	go func(c chan bool) {
		defer rawSocket.Close()

	stopInjection:
		for i := int64(0); i < pp.Count; i++ {
			select {
			case <-c:
				logging.GetLogger().Debugf("Injection stoped on interface %s", ifName)
				break stopInjection
			default:
				logging.GetLogger().Debugf("Injecting packet on interface %s", ifName)

				if _, err := rawSocket.Write(packetData); err != nil {
					if err == syscall.ENXIO {
						logging.GetLogger().Warningf("Write error: %s", err.Error())
					} else {
						logging.GetLogger().Errorf("Write error: %s", err.Error())
					}
				}

				if i != pp.Count-1 {
					time.Sleep(time.Millisecond * time.Duration(pp.Interval))
				}
			}
		}
		chnl.Lock()
		delete(chnl.Pipes, pp.UUID)
		chnl.Unlock()
	}(p)

	return f.TrackingID, nil
}

func getIP(cidr string) net.IP {
	if len(cidr) <= 0 {
		return nil
	}
	ips := strings.Split(cidr, ",")
	//TODO(masco): currently taking first IP, need to implement to select a proper IP
	ip, _, err := net.ParseCIDR(ips[0])
	if err != nil {
		return nil
	}
	return ip
}
<|endoftext|>