<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/dustinkirkland/golang-petname"
	"github.com/gorilla/websocket"
	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

func createFromImage(d *Daemon, req *containerPostReq) Response {
	var hash string
	var err error
	var run func() shared.OperationResult

	if req.Source.Alias != "" {
		if req.Source.Mode == "pull" && req.Source.Server != "" {
			hash, err = remoteGetImageFingerprint(d, req.Source.Server, req.Source.Alias)
			if err != nil {
				return InternalError(err)
			}
		} else {

			hash, err = dbImageAliasGet(d.db, req.Source.Alias)
			if err != nil {
				return InternalError(err)
			}
		}
	} else if req.Source.Fingerprint != "" {
		hash = req.Source.Fingerprint
	} else {
		return BadRequest(fmt.Errorf("must specify one of alias or fingerprint for init from image"))
	}

	if req.Source.Server != "" {
		err := d.ImageDownload(req.Source.Server, hash, req.Source.Secret, true)
		if err != nil {
			return InternalError(err)
		}
	}

	imgInfo, err := dbImageGet(d.db, hash, false, false)
	if err != nil {
		return SmartError(err)
	}
	hash = imgInfo.Fingerprint

	args := containerLXDArgs{
		Ctype:        cTypeRegular,
		Config:       req.Config,
		Profiles:     req.Profiles,
		Ephemeral:    req.Ephemeral,
		BaseImage:    hash,
		Architecture: imgInfo.Architecture,
	}

	run = shared.OperationWrap(func() error {
		_, err := containerLXDCreateFromImage(d, req.Name, args, hash)
		return err
	})

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromNone(d *Daemon, req *containerPostReq) Response {
	args := containerLXDArgs{
		Ctype:     cTypeRegular,
		Config:    req.Config,
		Profiles:  req.Profiles,
		Ephemeral: req.Ephemeral,
	}

	run := shared.OperationWrap(func() error {
		_, err := containerLXDCreateAsEmpty(d, req.Name, args)
		return err
	})

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromMigration(d *Daemon, req *containerPostReq) Response {
	if req.Source.Mode != "pull" {
		return NotImplemented
	}

	run := func() shared.OperationResult {
		createArgs := containerLXDArgs{
			Ctype:     cTypeRegular,
			Config:    req.Config,
			Profiles:  req.Profiles,
			Ephemeral: req.Ephemeral,
			BaseImage: req.Source.BaseImage,
		}

		var c container
		if _, err := dbImageGet(d.db, req.Source.BaseImage, false, true); err == nil {
			c, err = containerLXDCreateFromImage(
				d, req.Name, createArgs, req.Source.BaseImage)

			if err != nil {
				return shared.OperationError(err)
			}
		} else {
			c, err = containerLXDCreateAsEmpty(d, req.Name, createArgs)
			if err != nil {
				return shared.OperationError(err)
			}
		}

		config, err := shared.GetTLSConfig(d.certf, d.keyf)
		if err != nil {
			c.Delete()
			return shared.OperationError(err)
		}

		args := MigrationSinkArgs{
			Url: req.Source.Operation,
			Dialer: websocket.Dialer{
				TLSClientConfig: config,
				NetDial:         shared.RFC3493Dialer},
			Container: c,
			Secrets:   req.Source.Websockets,
		}

		sink, err := NewMigrationSink(&args)
		if err != nil {
			c.Delete()
			return shared.OperationError(err)
		}

		// Start the storage for this container (LVM mount/umount)
		c.StorageStart()
		defer c.StorageStop()

		// And finaly run the migration.
		err = sink()
		if err != nil {
			c.Delete()
			return shared.OperationError(err)
		}

		err = c.TemplateApply("copy")
		if err != nil {
			return shared.OperationError(err)
		}

		return shared.OperationError(nil)
	}

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromCopy(d *Daemon, req *containerPostReq) Response {
	if req.Source.Source == "" {
		return BadRequest(fmt.Errorf("must specify a source container"))
	}

	source, err := containerLXDLoad(d, req.Source.Source)
	if err != nil {
		return SmartError(err)
	}

	sourceConfig := source.ConfigGet()

	if req.Config == nil {
		req.Config = make(map[string]string)
	}

	for key, value := range sourceConfig {
		if len(key) > 8 && key[0:8] == "volatile" && key[9:] != "base_image" {
			shared.Log.Debug("Skipping volatile key from copy source",
				log.Ctx{"key": key})
			continue
		}
		req.Config[key] = value
	}

	if req.Profiles == nil {
		req.Profiles = source.ProfilesGet()
	}

	args := containerLXDArgs{
		Ctype:     cTypeRegular,
		Config:    req.Config,
		Profiles:  req.Profiles,
		Ephemeral: req.Ephemeral,
		BaseImage: req.Source.BaseImage,
	}

	run := func() shared.OperationResult {
		_, err := containerLXDCreateAsCopy(d, req.Name, args, source)
		if err != nil {
			return shared.OperationError(err)
		}

		return shared.OperationSuccess
	}

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name, req.Source.Source}

	return &asyncResponse{run: run, resources: resources}
}

func containersPost(d *Daemon, r *http.Request) Response {
	shared.Debugf("Responding to container create")

	req := containerPostReq{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	if req.Name == "" {
		req.Name = strings.ToLower(petname.Generate(2, "-"))
		shared.Debugf("No name provided, creating %s", req.Name)
	}

	if strings.Contains(req.Name, shared.SnapshotDelimiter) {
		return BadRequest(fmt.Errorf("Invalid container name: '%s' is reserved for snapshots", shared.SnapshotDelimiter))
	}

	switch req.Source.Type {
	case "image":
		return createFromImage(d, &req)
	case "none":
		return createFromNone(d, &req)
	case "migration":
		return createFromMigration(d, &req)
	case "copy":
		return createFromCopy(d, &req)
	default:
		return BadRequest(fmt.Errorf("unknown source type %s", req.Source.Type))
	}

}
<commit_msg>Cleanup storage before stopping container<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/dustinkirkland/golang-petname"
	"github.com/gorilla/websocket"
	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

func createFromImage(d *Daemon, req *containerPostReq) Response {
	var hash string
	var err error
	var run func() shared.OperationResult

	if req.Source.Alias != "" {
		if req.Source.Mode == "pull" && req.Source.Server != "" {
			hash, err = remoteGetImageFingerprint(d, req.Source.Server, req.Source.Alias)
			if err != nil {
				return InternalError(err)
			}
		} else {

			hash, err = dbImageAliasGet(d.db, req.Source.Alias)
			if err != nil {
				return InternalError(err)
			}
		}
	} else if req.Source.Fingerprint != "" {
		hash = req.Source.Fingerprint
	} else {
		return BadRequest(fmt.Errorf("must specify one of alias or fingerprint for init from image"))
	}

	if req.Source.Server != "" {
		err := d.ImageDownload(req.Source.Server, hash, req.Source.Secret, true)
		if err != nil {
			return InternalError(err)
		}
	}

	imgInfo, err := dbImageGet(d.db, hash, false, false)
	if err != nil {
		return SmartError(err)
	}
	hash = imgInfo.Fingerprint

	args := containerLXDArgs{
		Ctype:        cTypeRegular,
		Config:       req.Config,
		Profiles:     req.Profiles,
		Ephemeral:    req.Ephemeral,
		BaseImage:    hash,
		Architecture: imgInfo.Architecture,
	}

	run = shared.OperationWrap(func() error {
		_, err := containerLXDCreateFromImage(d, req.Name, args, hash)
		return err
	})

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromNone(d *Daemon, req *containerPostReq) Response {
	args := containerLXDArgs{
		Ctype:     cTypeRegular,
		Config:    req.Config,
		Profiles:  req.Profiles,
		Ephemeral: req.Ephemeral,
	}

	run := shared.OperationWrap(func() error {
		_, err := containerLXDCreateAsEmpty(d, req.Name, args)
		return err
	})

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromMigration(d *Daemon, req *containerPostReq) Response {
	if req.Source.Mode != "pull" {
		return NotImplemented
	}

	run := func() shared.OperationResult {
		createArgs := containerLXDArgs{
			Ctype:     cTypeRegular,
			Config:    req.Config,
			Profiles:  req.Profiles,
			Ephemeral: req.Ephemeral,
			BaseImage: req.Source.BaseImage,
		}

		var c container
		if _, err := dbImageGet(d.db, req.Source.BaseImage, false, true); err == nil {
			c, err = containerLXDCreateFromImage(
				d, req.Name, createArgs, req.Source.BaseImage)

			if err != nil {
				return shared.OperationError(err)
			}
		} else {
			c, err = containerLXDCreateAsEmpty(d, req.Name, createArgs)
			if err != nil {
				return shared.OperationError(err)
			}
		}

		config, err := shared.GetTLSConfig(d.certf, d.keyf)
		if err != nil {
			c.Delete()
			return shared.OperationError(err)
		}

		args := MigrationSinkArgs{
			Url: req.Source.Operation,
			Dialer: websocket.Dialer{
				TLSClientConfig: config,
				NetDial:         shared.RFC3493Dialer},
			Container: c,
			Secrets:   req.Source.Websockets,
		}

		sink, err := NewMigrationSink(&args)
		if err != nil {
			c.Delete()
			return shared.OperationError(err)
		}

		// Start the storage for this container (LVM mount/umount)
		c.StorageStart()

		// And finaly run the migration.
		err = sink()
		if err != nil {
			c.StorageStop()
			c.Delete()
			return shared.OperationError(err)
		}

		defer c.StorageStop()

		err = c.TemplateApply("copy")
		if err != nil {
			return shared.OperationError(err)
		}

		return shared.OperationError(nil)
	}

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name}

	return &asyncResponse{run: run, resources: resources}
}

func createFromCopy(d *Daemon, req *containerPostReq) Response {
	if req.Source.Source == "" {
		return BadRequest(fmt.Errorf("must specify a source container"))
	}

	source, err := containerLXDLoad(d, req.Source.Source)
	if err != nil {
		return SmartError(err)
	}

	sourceConfig := source.ConfigGet()

	if req.Config == nil {
		req.Config = make(map[string]string)
	}

	for key, value := range sourceConfig {
		if len(key) > 8 && key[0:8] == "volatile" && key[9:] != "base_image" {
			shared.Log.Debug("Skipping volatile key from copy source",
				log.Ctx{"key": key})
			continue
		}
		req.Config[key] = value
	}

	if req.Profiles == nil {
		req.Profiles = source.ProfilesGet()
	}

	args := containerLXDArgs{
		Ctype:     cTypeRegular,
		Config:    req.Config,
		Profiles:  req.Profiles,
		Ephemeral: req.Ephemeral,
		BaseImage: req.Source.BaseImage,
	}

	run := func() shared.OperationResult {
		_, err := containerLXDCreateAsCopy(d, req.Name, args, source)
		if err != nil {
			return shared.OperationError(err)
		}

		return shared.OperationSuccess
	}

	resources := make(map[string][]string)
	resources["containers"] = []string{req.Name, req.Source.Source}

	return &asyncResponse{run: run, resources: resources}
}

func containersPost(d *Daemon, r *http.Request) Response {
	shared.Debugf("Responding to container create")

	req := containerPostReq{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	if req.Name == "" {
		req.Name = strings.ToLower(petname.Generate(2, "-"))
		shared.Debugf("No name provided, creating %s", req.Name)
	}

	if strings.Contains(req.Name, shared.SnapshotDelimiter) {
		return BadRequest(fmt.Errorf("Invalid container name: '%s' is reserved for snapshots", shared.SnapshotDelimiter))
	}

	switch req.Source.Type {
	case "image":
		return createFromImage(d, &req)
	case "none":
		return createFromNone(d, &req)
	case "migration":
		return createFromMigration(d, &req)
	case "copy":
		return createFromCopy(d, &req)
	default:
		return BadRequest(fmt.Errorf("unknown source type %s", req.Source.Type))
	}

}
<|endoftext|>


<commit_before>// Copyright (C) 2016  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package gohbase

import (
	"bytes"
	"io"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/cznic/b"
	"github.com/tsuna/gohbase/hrpc"
)

// clientRegionCache is client -> region cache. Used to quickly
// look up all the regioninfos that map to a specific client
type clientRegionCache struct {
	m sync.Mutex

	regions map[hrpc.RegionClient]map[hrpc.RegionInfo]struct{}
}

// put caches client and associates a region with it. Returns a client that is in cache.
// TODO: obvious place for optimization (use map with address as key to lookup exisiting clients)
func (rcc *clientRegionCache) put(c hrpc.RegionClient, r hrpc.RegionInfo) hrpc.RegionClient {
	rcc.m.Lock()
	for existingClient, regions := range rcc.regions {
		// check if client already exists, checking by host and port
		// because concurrent callers might try to put the same client
		if c.Host() == existingClient.Host() && c.Port() == existingClient.Port() {
			// check client already knows about the region, checking
			// by pointer is enough because we make sure that there are
			// no regions with the same name around
			if _, ok := regions[r]; !ok {
				regions[r] = struct{}{}
			}
			rcc.m.Unlock()

			log.WithFields(log.Fields{
				"existingClient": existingClient,
				"client":         c,
			}).Debug("region client is already in client's cache")
			return existingClient
		}
	}

	// no such client yet
	rcc.regions[c] = map[hrpc.RegionInfo]struct{}{r: struct{}{}}
	rcc.m.Unlock()

	log.WithField("client", c).Info("added new region client")
	return c
}

func (rcc *clientRegionCache) del(r hrpc.RegionInfo) {
	rcc.m.Lock()
	c := r.Client()
	if c != nil {
		r.SetClient(nil)

		regions := rcc.regions[c]
		delete(regions, r)

		if len(regions) == 0 {
			// close region client if noone is using it
			delete(rcc.regions, c)
			c.Close()
		}
	}
	rcc.m.Unlock()
}

func (rcc *clientRegionCache) closeAll() {
	rcc.m.Lock()
	for client, regions := range rcc.regions {
		for region := range regions {
			region.MarkUnavailable()
			region.SetClient(nil)
		}
		client.Close()
	}
	rcc.m.Unlock()
}

func (rcc *clientRegionCache) clientDown(c hrpc.RegionClient) map[hrpc.RegionInfo]struct{} {
	rcc.m.Lock()
	downregions, ok := rcc.regions[c]
	delete(rcc.regions, c)
	rcc.m.Unlock()

	if ok {
		log.WithField("client", c).Info("removed region client")
	}
	return downregions
}

// TODO: obvious place for optimization (use map with address as key to lookup exisiting clients)
func (rcc *clientRegionCache) checkForClient(host string, port uint16) hrpc.RegionClient {
	rcc.m.Lock()
	defer rcc.m.Unlock()

	for client := range rcc.regions {
		if client.Host() == host && client.Port() == port {
			return client
		}
	}
	return nil
}

// key -> region cache.
type keyRegionCache struct {
	m sync.Mutex

	// Maps a []byte of a region start key to a hrpc.RegionInfo
	regions *b.Tree
}

func (krc *keyRegionCache) get(key []byte) ([]byte, hrpc.RegionInfo) {
	// When seeking - "The Enumerator's position is possibly after the last item in the tree"
	// http://godoc.org/github.com/cznic/b#Tree.Set
	krc.m.Lock()

	enum, ok := krc.regions.Seek(key)
	if ok {
		krc.m.Unlock()
		log.Fatalf("WTF: got exact match for region search key %q", key)
		return nil, nil
	}
	k, v, err := enum.Prev()
	enum.Close()

	krc.m.Unlock()

	if err == io.EOF {
		// we are the beginning of the tree
		return nil, nil
	}
	return k.([]byte), v.(hrpc.RegionInfo)
}

func isRegionOverlap(regA, regB hrpc.RegionInfo) bool {
	// if region's stop key is empty, it's assumed to be the greatest key
	return bytes.Equal(regA.Namespace(), regB.Namespace()) &&
		bytes.Equal(regA.Table(), regB.Table()) &&
		(len(regB.StopKey()) == 0 || bytes.Compare(regA.StartKey(), regB.StopKey()) < 0) &&
		(len(regA.StopKey()) == 0 || bytes.Compare(regA.StopKey(), regB.StartKey()) > 0)
}

func (krc *keyRegionCache) getOverlaps(reg hrpc.RegionInfo) []hrpc.RegionInfo {
	var overlaps []hrpc.RegionInfo
	var v interface{}
	var err error

	// deal with empty tree in the beginning so that we don't have to check
	// EOF errors for enum later
	if krc.regions.Len() == 0 {
		return overlaps
	}

	// check if key created from new region falls into any cached regions
	key := createRegionSearchKey(fullyQualifiedTable(reg), reg.StartKey())
	enum, ok := krc.regions.Seek(key)
	if ok {
		log.Fatalf("WTF: found a region with exact name as the search key %q", key)
	}

	// case 1: landed before the first region in cache
	// enum.Prev() returns io.EOF
	// enum.Next() returns io.EOF
	// SeekFirst() + enum.Next() returns the first region, which has larger start key

	// case 2: landed before the second region in cache
	// enum.Prev() returns the first region X and moves pointer to -infinity
	// enum.Next() returns io.EOF
	// SeekFirst() + enum.Next() returns first region X, which has smaller start key

	// case 3: landed anywhere after the second region
	// enum.Prev() returns the region X before it landed, moves pointer to the region X - 1
	// enum.Next() returns X - 1 and move pointer to X, which has smaller start key

	enum.Prev()
	_, _, err = enum.Next()
	if err == io.EOF {
		// we are in the beginning of tree, get new enum starting
		// from first region
		enum.Close()
		enum, err = krc.regions.SeekFirst()
		if err != nil {
			log.Fatalf(
				"error seeking first region when getting  overlaps for region %v: %v", reg, err)
		}
	}

	_, v, err = enum.Next()
	if isRegionOverlap(v.(hrpc.RegionInfo), reg) {
		overlaps = append(overlaps, v.(hrpc.RegionInfo))
	}
	_, v, err = enum.Next()

	// now append all regions that overlap until the end of the tree
	// or until they don't overlap
	for err != io.EOF && isRegionOverlap(v.(hrpc.RegionInfo), reg) {
		overlaps = append(overlaps, v.(hrpc.RegionInfo))
		_, v, err = enum.Next()
	}
	enum.Close()
	return overlaps
}

// put looks up if there's already region with this name in regions cache
// and if there's, returns it in overlaps and doesn't modify the cache.
// Otherwise, it puts the region and removes all overlaps in case all of
// them are older. Returns a slice of overlapping regions and whether
// passed region was put in the cache.
func (krc *keyRegionCache) put(reg hrpc.RegionInfo) (overlaps []hrpc.RegionInfo, replaced bool) {
	krc.m.Lock()
	krc.regions.Put(reg.Name(), func(v interface{}, exists bool) (interface{}, bool) {
		if exists {
			// region is already in cache,
			// note: regions with the same name have the same age
			overlaps = []hrpc.RegionInfo{v.(hrpc.RegionInfo)}
			return nil, false
		}
		// find all entries that are overlapping with the range of the new region.
		overlaps = krc.getOverlaps(reg)
		for _, o := range overlaps {
			if o.ID() > reg.ID() {
				// overlapping region is younger,
				// don't replace any regions
				// TODO: figure out if there can a case where we might
				// have both older and younger overlapping regions, for
				// now we only replace if all overlaps are older
				return nil, false
			}
		}
		// all overlaps are older, put the new region
		replaced = true
		return reg, true
	})
	if !replaced {
		krc.m.Unlock()

		log.WithFields(log.Fields{
			"region":   reg,
			"overlaps": overlaps,
			"replaced": replaced,
		}).Debug("region is already in cache")
		return
	}
	// delete overlapping regions
	// TODO: in case overlaps are always either younger or older,
	// we can just greedily remove them in Put function
	for _, o := range overlaps {
		krc.regions.Delete(o.Name())
		// let region establishers know that they can give up
		o.MarkDead()
	}
	krc.m.Unlock()

	log.WithFields(log.Fields{
		"region":   reg,
		"overlaps": overlaps,
		"replaced": replaced,
	}).Info("added new region")
	return
}

func (krc *keyRegionCache) del(reg hrpc.RegionInfo) bool {
	krc.m.Lock()
	success := krc.regions.Delete(reg.Name())
	krc.m.Unlock()
	// let region establishers know that they can give up
	reg.MarkDead()

	log.WithFields(log.Fields{
		"region": reg,
	}).Debug("removed region")
	return success
}
<commit_msg>cache: use RWMutex for caches<commit_after>// Copyright (C) 2016  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package gohbase

import (
	"bytes"
	"io"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/cznic/b"
	"github.com/tsuna/gohbase/hrpc"
)

// clientRegionCache is client -> region cache. Used to quickly
// look up all the regioninfos that map to a specific client
type clientRegionCache struct {
	m sync.RWMutex

	regions map[hrpc.RegionClient]map[hrpc.RegionInfo]struct{}
}

// put caches client and associates a region with it. Returns a client that is in cache.
// TODO: obvious place for optimization (use map with address as key to lookup exisiting clients)
func (rcc *clientRegionCache) put(c hrpc.RegionClient, r hrpc.RegionInfo) hrpc.RegionClient {
	rcc.m.Lock()
	for existingClient, regions := range rcc.regions {
		// check if client already exists, checking by host and port
		// because concurrent callers might try to put the same client
		if c.Host() == existingClient.Host() && c.Port() == existingClient.Port() {
			// check client already knows about the region, checking
			// by pointer is enough because we make sure that there are
			// no regions with the same name around
			if _, ok := regions[r]; !ok {
				regions[r] = struct{}{}
			}
			rcc.m.Unlock()

			log.WithFields(log.Fields{
				"existingClient": existingClient,
				"client":         c,
			}).Debug("region client is already in client's cache")
			return existingClient
		}
	}

	// no such client yet
	rcc.regions[c] = map[hrpc.RegionInfo]struct{}{r: struct{}{}}
	rcc.m.Unlock()

	log.WithField("client", c).Info("added new region client")
	return c
}

func (rcc *clientRegionCache) del(r hrpc.RegionInfo) {
	rcc.m.Lock()
	c := r.Client()
	if c != nil {
		r.SetClient(nil)

		regions := rcc.regions[c]
		delete(regions, r)

		if len(regions) == 0 {
			// close region client if noone is using it
			delete(rcc.regions, c)
			c.Close()
		}
	}
	rcc.m.Unlock()
}

func (rcc *clientRegionCache) closeAll() {
	rcc.m.Lock()
	for client, regions := range rcc.regions {
		for region := range regions {
			region.MarkUnavailable()
			region.SetClient(nil)
		}
		client.Close()
	}
	rcc.m.Unlock()
}

func (rcc *clientRegionCache) clientDown(c hrpc.RegionClient) map[hrpc.RegionInfo]struct{} {
	rcc.m.Lock()
	downregions, ok := rcc.regions[c]
	delete(rcc.regions, c)
	rcc.m.Unlock()

	if ok {
		log.WithField("client", c).Info("removed region client")
	}
	return downregions
}

// TODO: obvious place for optimization (use map with address as key to lookup exisiting clients)
func (rcc *clientRegionCache) checkForClient(host string, port uint16) hrpc.RegionClient {
	rcc.m.RLock()

	for client := range rcc.regions {
		if client.Host() == host && client.Port() == port {
			rcc.m.RUnlock()
			return client
		}
	}

	rcc.m.RUnlock()
	return nil
}

// key -> region cache.
type keyRegionCache struct {
	m sync.RWMutex

	// Maps a []byte of a region start key to a hrpc.RegionInfo
	regions *b.Tree
}

func (krc *keyRegionCache) get(key []byte) ([]byte, hrpc.RegionInfo) {
	krc.m.RLock()

	enum, ok := krc.regions.Seek(key)
	if ok {
		krc.m.RUnlock()
		log.Fatalf("WTF: got exact match for region search key %q", key)
		return nil, nil
	}
	k, v, err := enum.Prev()
	enum.Close()

	krc.m.RUnlock()

	if err == io.EOF {
		// we are the beginning of the tree
		return nil, nil
	}
	return k.([]byte), v.(hrpc.RegionInfo)
}

func isRegionOverlap(regA, regB hrpc.RegionInfo) bool {
	// if region's stop key is empty, it's assumed to be the greatest key
	return bytes.Equal(regA.Namespace(), regB.Namespace()) &&
		bytes.Equal(regA.Table(), regB.Table()) &&
		(len(regB.StopKey()) == 0 || bytes.Compare(regA.StartKey(), regB.StopKey()) < 0) &&
		(len(regA.StopKey()) == 0 || bytes.Compare(regA.StopKey(), regB.StartKey()) > 0)
}

func (krc *keyRegionCache) getOverlaps(reg hrpc.RegionInfo) []hrpc.RegionInfo {
	var overlaps []hrpc.RegionInfo
	var v interface{}
	var err error

	// deal with empty tree in the beginning so that we don't have to check
	// EOF errors for enum later
	if krc.regions.Len() == 0 {
		return overlaps
	}

	// check if key created from new region falls into any cached regions
	key := createRegionSearchKey(fullyQualifiedTable(reg), reg.StartKey())
	enum, ok := krc.regions.Seek(key)
	if ok {
		log.Fatalf("WTF: found a region with exact name as the search key %q", key)
	}

	// case 1: landed before the first region in cache
	// enum.Prev() returns io.EOF
	// enum.Next() returns io.EOF
	// SeekFirst() + enum.Next() returns the first region, which has larger start key

	// case 2: landed before the second region in cache
	// enum.Prev() returns the first region X and moves pointer to -infinity
	// enum.Next() returns io.EOF
	// SeekFirst() + enum.Next() returns first region X, which has smaller start key

	// case 3: landed anywhere after the second region
	// enum.Prev() returns the region X before it landed, moves pointer to the region X - 1
	// enum.Next() returns X - 1 and move pointer to X, which has smaller start key

	enum.Prev()
	_, _, err = enum.Next()
	if err == io.EOF {
		// we are in the beginning of tree, get new enum starting
		// from first region
		enum.Close()
		enum, err = krc.regions.SeekFirst()
		if err != nil {
			log.Fatalf(
				"error seeking first region when getting  overlaps for region %v: %v", reg, err)
		}
	}

	_, v, err = enum.Next()
	if isRegionOverlap(v.(hrpc.RegionInfo), reg) {
		overlaps = append(overlaps, v.(hrpc.RegionInfo))
	}
	_, v, err = enum.Next()

	// now append all regions that overlap until the end of the tree
	// or until they don't overlap
	for err != io.EOF && isRegionOverlap(v.(hrpc.RegionInfo), reg) {
		overlaps = append(overlaps, v.(hrpc.RegionInfo))
		_, v, err = enum.Next()
	}
	enum.Close()
	return overlaps
}

// put looks up if there's already region with this name in regions cache
// and if there's, returns it in overlaps and doesn't modify the cache.
// Otherwise, it puts the region and removes all overlaps in case all of
// them are older. Returns a slice of overlapping regions and whether
// passed region was put in the cache.
func (krc *keyRegionCache) put(reg hrpc.RegionInfo) (overlaps []hrpc.RegionInfo, replaced bool) {
	krc.m.Lock()
	krc.regions.Put(reg.Name(), func(v interface{}, exists bool) (interface{}, bool) {
		if exists {
			// region is already in cache,
			// note: regions with the same name have the same age
			overlaps = []hrpc.RegionInfo{v.(hrpc.RegionInfo)}
			return nil, false
		}
		// find all entries that are overlapping with the range of the new region.
		overlaps = krc.getOverlaps(reg)
		for _, o := range overlaps {
			if o.ID() > reg.ID() {
				// overlapping region is younger,
				// don't replace any regions
				// TODO: figure out if there can a case where we might
				// have both older and younger overlapping regions, for
				// now we only replace if all overlaps are older
				return nil, false
			}
		}
		// all overlaps are older, put the new region
		replaced = true
		return reg, true
	})
	if !replaced {
		krc.m.Unlock()

		log.WithFields(log.Fields{
			"region":   reg,
			"overlaps": overlaps,
			"replaced": replaced,
		}).Debug("region is already in cache")
		return
	}
	// delete overlapping regions
	// TODO: in case overlaps are always either younger or older,
	// we can just greedily remove them in Put function
	for _, o := range overlaps {
		krc.regions.Delete(o.Name())
		// let region establishers know that they can give up
		o.MarkDead()
	}
	krc.m.Unlock()

	log.WithFields(log.Fields{
		"region":   reg,
		"overlaps": overlaps,
		"replaced": replaced,
	}).Info("added new region")
	return
}

func (krc *keyRegionCache) del(reg hrpc.RegionInfo) bool {
	krc.m.Lock()
	success := krc.regions.Delete(reg.Name())
	krc.m.Unlock()
	// let region establishers know that they can give up
	reg.MarkDead()

	log.WithFields(log.Fields{
		"region": reg,
	}).Debug("removed region")
	return success
}
<|endoftext|>


<commit_before>// Copyright 2016 Liam Stanley <me@liamstanley.io>. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

package girc

import (
	"errors"
	"fmt"
	"math/rand"
	"strings"
	"sync"
)

// RunCallbacks manually runs callbacks for a given event.
func (c *Client) RunCallbacks(event *Event) {
	// Log the event.
	c.log.Print("<-- " + event.Raw())

	// Regular wildcard callbacks.
	c.Callbacks.exec(ALLEVENTS, c, event)

	// Then regular non-threaded callbacks.
	c.Callbacks.exec(event.Command, c, event)
}

// Callback is lower level implementation of a callback. See
// Caller.AddHandler()
type Callback interface {
	Execute(*Client, Event)
}

// CallbackFunc is a type that represents the function necessary to
// implement Callback.
type CallbackFunc func(c *Client, e Event)

// Execute calls the CallbackFunc with the sender and irc message.
func (f CallbackFunc) Execute(c *Client, e Event) {
	f(c, e)
}

// Caller manages internal and external (user facing) callbacks.
//
// external/internal keys are of structure:
//   map[CALLBACK_TYPE][COMMAND/EVENT][CUID]Callback
type Caller struct {
	// mu is the mutex that should be used when accessing callbacks.
	mu sync.RWMutex
	// external is a map of user facing callbacks.
	external map[string]map[string]map[string]Callback
	// internal is a map of internally used callbacks for the client.
	internal map[string]map[string]map[string]Callback
}

// newCaller creates and initializes a new callback handler.
func newCaller() (c *Caller) {
	c = &Caller{}

	c.external = map[string]map[string]map[string]Callback{}
	c.external["routine"] = map[string]map[string]Callback{}
	c.external["std"] = map[string]map[string]Callback{}
	c.internal = map[string]map[string]map[string]Callback{}
	c.internal["routine"] = map[string]map[string]Callback{}
	c.internal["std"] = map[string]map[string]Callback{}

	return c
}

// Len returns the total amount of user-entered callbacks.
func (c *Caller) Len() int {
	var total int

	c.mu.RLock()
	for ctype := range c.external {
		for cmd := range c.external[ctype] {
			total += len(c.external[ctype][cmd])
		}
	}
	c.mu.RUnlock()

	return total
}

func (c *Caller) String() string {
	var total int
	var ctypes []string

	c.mu.RLock()
	for ctype := range c.internal {
		ctypes = append(ctypes, ctype)
		for cmd := range c.internal[ctype] {
			total += len(c.internal[ctype][cmd])
		}
	}
	c.mu.RUnlock()

	return fmt.Sprintf(
		"<Caller() types[%d]:[%s] client:%d internal:%d>",
		len(c.external), strings.Join(ctypes, ","), c.Len(), len(c.internal),
	)
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

// cuid generates a unique UID string for each callback for ease of removal.
func (c *Caller) cuid(ctype, cmd string, n int) (cuid, uid string) {
	b := make([]byte, n)

	for i := range b {
		b[i] = letterBytes[rand.Int63()%int64(len(letterBytes))]
	}

	return ctype + ":" + cmd + ":" + string(b), string(b)
}

// cuidToID allows easy mapping between a generated cuid and the caller``
// external/internal callback maps.
func (c *Caller) cuidToID(input string) (ctype, cmd, uid string) {
	// Ignore the errors because the strings will default to empty anyway.
	_, _ = fmt.Sscanf(input, "%s:%s:%s", &ctype, &cmd, &uid)
	return ctype, cmd, uid
}

// exec executes all callbacks pertaining to specified event. Internal first,
// then external.
//
// Please note that there is no specific order/priority for which the
// callback types themselves or the callbacks are executed.
func (c *Caller) exec(command string, client *Client, event *Event) {
	c.mu.RLock()
	// Execute internal callbacks first.
	for callbackType := range c.internal {
		if _, ok := c.internal[callbackType][command]; !ok {
			continue
		}

		for cuid := range c.internal[callbackType][command] {
			switch callbackType {
			case "routine":
				go c.internal[callbackType][command][cuid].Execute(client, *event)
			default:
				c.internal[callbackType][command][cuid].Execute(client, *event)
			}
		}
	}

	// Aaand then external callbacks.
	for callbackType := range c.external {
		if _, ok := c.external[callbackType][command]; !ok {
			continue
		}

		for cuid := range c.external[callbackType][command] {
			switch callbackType {
			case "routine":
				go c.external[callbackType][command][cuid].Execute(client, *event)
			default:
				c.external[callbackType][command][cuid].Execute(client, *event)
			}
		}
	}
	c.mu.RUnlock()
}

// ClearAll clears all external callbacks currently setup within the client.
//
// This ignores internal callbacks.
func (c *Caller) ClearAll() {
	c.mu.Lock()
	c.external = map[string]map[string]map[string]Callback{}
	c.external["routine"] = map[string]map[string]Callback{}
	c.external["std"] = map[string]map[string]Callback{}
	c.mu.Unlock()
}

// Clear clears all of the callbacks for the given event.
//
// This ignores internal callbacks.
func (c *Caller) Clear(cmd string) {
	c.mu.Lock()
	for ctype := range c.external {
		if _, ok := c.external[ctype][cmd]; ok {
			delete(c.external[ctype], cmd)
		}
	}
	c.mu.Unlock()
}

// Remove removes the callback with cuid from the callback stack. success
// indicates that it existed, and has been removed. If not success, it
// wasn't a registered callback.
func (c *Caller) Remove(cuid string) (success bool) {
	ctype, cmd, uid := c.cuidToID(cuid)
	if len(ctype) == 0 || len(cmd) == 0 || len(uid) == 0 {
		return false
	}

	c.mu.Lock()
	// Check if the callback type exists.
	if _, ok := c.external[ctype]; !ok {
		c.mu.Unlock()
		return false
	}

	// Check if the irc command/event has any callbacks on it.
	if _, ok := c.external[ctype][cmd]; !ok {
		c.mu.Unlock()
		return false
	}

	// Check to see if it's actually a registered callback.
	if _, ok := c.external[ctype][cmd][cuid]; !ok {
		c.mu.Unlock()
		return false
	}

	delete(c.external[ctype][cmd], uid)
	c.mu.Unlock()

	// Assume success.
	return true
}

func (c *Caller) register(internal bool, ctype, cmd string, callback Callback) (cuid string) {
	var uid string

	c.mu.Lock()
	if internal {
		if _, ok := c.internal[ctype]; !ok {
			panic(errors.New("callback type does not exist: " + ctype))
		}

		if _, ok := c.internal[ctype][cmd]; !ok {
			c.internal[ctype][cmd] = map[string]Callback{}
		}

		cuid, uid = c.cuid(ctype, cmd, 20)

		c.internal[ctype][cmd][uid] = callback
	} else {
		if _, ok := c.external[ctype]; !ok {
			panic(errors.New("callback type does not exist: " + ctype))
		}

		if _, ok := c.external[ctype][cmd]; !ok {
			c.external[ctype][cmd] = map[string]Callback{}
		}

		cuid, uid = c.cuid(ctype, cmd, 20)

		c.external[ctype][cmd][uid] = callback
	}
	c.mu.Unlock()

	return cuid
}

// AddHandler registers a callback (matching the Callback interface) for the
// given event. cuid is the callback uid which can be used to remove the
// callback with Caller.Remove().
func (c *Caller) AddHandler(cmd string, callback Callback) (cuid string) {
	return c.register(false, "std", cmd, callback)
}

// AddBgHandler registers a callback (matching the Callback interface) for
// the given event and executes it in a go-routine. cuid is the callback uid
// which can be used to remove the callback with Caller.Remove().
func (c *Caller) AddBgHandler(cmd string, callback Callback) (cuid string) {
	return c.register(false, "routine", cmd, callback)
}

// Add registers the callback function for the given event. cuid is the
// callback uid which can be used to remove the callback with Caller.Remove().
func (c *Caller) Add(cmd string, callback func(c *Client, e Event)) (cuid string) {
	return c.register(false, "std", cmd, CallbackFunc(callback))
}

// AddBg registers the callback function for the given event and executes it
// in a go-routine. cuid is the callback uid which can be used to remove the
// callback with Caller.Remove().
func (c *Caller) AddBg(cmd string, callback func(c *Client, e Event)) (cuid string) {
	return c.register(false, "routine", cmd, CallbackFunc(callback))
}
<commit_msg>minor doc syntax changes<commit_after>// Copyright 2016 Liam Stanley <me@liamstanley.io>. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

package girc

import (
	"errors"
	"fmt"
	"math/rand"
	"strings"
	"sync"
)

// RunCallbacks manually runs callbacks for a given event.
func (c *Client) RunCallbacks(event *Event) {
	// Log the event.
	c.log.Print("<-- " + event.Raw())

	// Regular wildcard callbacks.
	c.Callbacks.exec(ALLEVENTS, c, event)

	// Then regular non-threaded callbacks.
	c.Callbacks.exec(event.Command, c, event)
}

// Callback is lower level implementation of a callback. See
// Caller.AddHandler()
type Callback interface {
	Execute(*Client, Event)
}

// CallbackFunc is a type that represents the function necessary to
// implement Callback.
type CallbackFunc func(c *Client, e Event)

// Execute calls the CallbackFunc with the sender and irc message.
func (f CallbackFunc) Execute(c *Client, e Event) {
	f(c, e)
}

// Caller manages internal and external (user facing) callbacks.
//
// external/internal keys are of structure:
//   map[CALLBACK_TYPE][COMMAND/EVENT][CUID]Callback
type Caller struct {
	// mu is the mutex that should be used when accessing callbacks.
	mu sync.RWMutex
	// external is a map of user facing callbacks.
	external map[string]map[string]map[string]Callback
	// internal is a map of internally used callbacks for the client.
	internal map[string]map[string]map[string]Callback
}

// newCaller creates and initializes a new callback handler.
func newCaller() *Caller {
	c := &Caller{}

	c.external = map[string]map[string]map[string]Callback{}
	c.external["routine"] = map[string]map[string]Callback{}
	c.external["std"] = map[string]map[string]Callback{}
	c.internal = map[string]map[string]map[string]Callback{}
	c.internal["routine"] = map[string]map[string]Callback{}
	c.internal["std"] = map[string]map[string]Callback{}

	return c
}

// Len returns the total amount of user-entered callbacks.
func (c *Caller) Len() int {
	var total int

	c.mu.RLock()
	for ctype := range c.external {
		for cmd := range c.external[ctype] {
			total += len(c.external[ctype][cmd])
		}
	}
	c.mu.RUnlock()

	return total
}

func (c *Caller) String() string {
	var total int
	var ctypes []string

	c.mu.RLock()
	for ctype := range c.internal {
		ctypes = append(ctypes, ctype)
		for cmd := range c.internal[ctype] {
			total += len(c.internal[ctype][cmd])
		}
	}
	c.mu.RUnlock()

	return fmt.Sprintf(
		"<Caller() types[%d]:[%s] client:%d internal:%d>",
		len(c.external), strings.Join(ctypes, ","), c.Len(), len(c.internal),
	)
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

// cuid generates a unique UID string for each callback for ease of removal.
func (c *Caller) cuid(ctype, cmd string, n int) (cuid, uid string) {
	b := make([]byte, n)

	for i := range b {
		b[i] = letterBytes[rand.Int63()%int64(len(letterBytes))]
	}

	return ctype + ":" + cmd + ":" + string(b), string(b)
}

// cuidToID allows easy mapping between a generated cuid and the caller``
// external/internal callback maps.
func (c *Caller) cuidToID(input string) (ctype, cmd, uid string) {
	// Ignore the errors because the strings will default to empty anyway.
	_, _ = fmt.Sscanf(input, "%s:%s:%s", &ctype, &cmd, &uid)
	return ctype, cmd, uid
}

// exec executes all callbacks pertaining to specified event. Internal first,
// then external.
//
// Please note that there is no specific order/priority for which the
// callback types themselves or the callbacks are executed.
func (c *Caller) exec(command string, client *Client, event *Event) {
	c.mu.RLock()
	// Execute internal callbacks first.
	for callbackType := range c.internal {
		if _, ok := c.internal[callbackType][command]; !ok {
			continue
		}

		for cuid := range c.internal[callbackType][command] {
			switch callbackType {
			case "routine":
				go c.internal[callbackType][command][cuid].Execute(client, *event)
			default:
				c.internal[callbackType][command][cuid].Execute(client, *event)
			}
		}
	}

	// Aaand then external callbacks.
	for callbackType := range c.external {
		if _, ok := c.external[callbackType][command]; !ok {
			continue
		}

		for cuid := range c.external[callbackType][command] {
			switch callbackType {
			case "routine":
				go c.external[callbackType][command][cuid].Execute(client, *event)
			default:
				c.external[callbackType][command][cuid].Execute(client, *event)
			}
		}
	}
	c.mu.RUnlock()
}

// ClearAll clears all external callbacks currently setup within the client.
// This ignores internal callbacks.
func (c *Caller) ClearAll() {
	c.mu.Lock()
	c.external = map[string]map[string]map[string]Callback{}
	c.external["routine"] = map[string]map[string]Callback{}
	c.external["std"] = map[string]map[string]Callback{}
	c.mu.Unlock()
}

// Clear clears all of the callbacks for the given event.
// This ignores internal callbacks.
func (c *Caller) Clear(cmd string) {
	c.mu.Lock()
	for ctype := range c.external {
		if _, ok := c.external[ctype][cmd]; ok {
			delete(c.external[ctype], cmd)
		}
	}
	c.mu.Unlock()
}

// Remove removes the callback with cuid from the callback stack. success
// indicates that it existed, and has been removed. If not success, it
// wasn't a registered callback.
func (c *Caller) Remove(cuid string) (success bool) {
	ctype, cmd, uid := c.cuidToID(cuid)
	if len(ctype) == 0 || len(cmd) == 0 || len(uid) == 0 {
		return false
	}

	c.mu.Lock()
	// Check if the callback type exists.
	if _, ok := c.external[ctype]; !ok {
		c.mu.Unlock()
		return false
	}

	// Check if the irc command/event has any callbacks on it.
	if _, ok := c.external[ctype][cmd]; !ok {
		c.mu.Unlock()
		return false
	}

	// Check to see if it's actually a registered callback.
	if _, ok := c.external[ctype][cmd][cuid]; !ok {
		c.mu.Unlock()
		return false
	}

	delete(c.external[ctype][cmd], uid)
	c.mu.Unlock()

	// Assume success.
	return true
}

func (c *Caller) register(internal bool, ctype, cmd string, callback Callback) (cuid string) {
	var uid string

	c.mu.Lock()
	if internal {
		if _, ok := c.internal[ctype]; !ok {
			panic(errors.New("callback type does not exist: " + ctype))
		}

		if _, ok := c.internal[ctype][cmd]; !ok {
			c.internal[ctype][cmd] = map[string]Callback{}
		}

		cuid, uid = c.cuid(ctype, cmd, 20)

		c.internal[ctype][cmd][uid] = callback
	} else {
		if _, ok := c.external[ctype]; !ok {
			panic(errors.New("callback type does not exist: " + ctype))
		}

		if _, ok := c.external[ctype][cmd]; !ok {
			c.external[ctype][cmd] = map[string]Callback{}
		}

		cuid, uid = c.cuid(ctype, cmd, 20)

		c.external[ctype][cmd][uid] = callback
	}
	c.mu.Unlock()

	return cuid
}

// AddHandler registers a callback (matching the Callback interface) for the
// given event. cuid is the callback uid which can be used to remove the
// callback with Caller.Remove().
func (c *Caller) AddHandler(cmd string, callback Callback) (cuid string) {
	return c.register(false, "std", cmd, callback)
}

// AddBgHandler registers a callback (matching the Callback interface) for
// the given event and executes it in a go-routine. cuid is the callback uid
// which can be used to remove the callback with Caller.Remove().
func (c *Caller) AddBgHandler(cmd string, callback Callback) (cuid string) {
	return c.register(false, "routine", cmd, callback)
}

// Add registers the callback function for the given event. cuid is the
// callback uid which can be used to remove the callback with Caller.Remove().
func (c *Caller) Add(cmd string, callback func(c *Client, e Event)) (cuid string) {
	return c.register(false, "std", cmd, CallbackFunc(callback))
}

// AddBg registers the callback function for the given event and executes it
// in a go-routine. cuid is the callback uid which can be used to remove the
// callback with Caller.Remove().
func (c *Caller) AddBg(cmd string, callback func(c *Client, e Event)) (cuid string) {
	return c.register(false, "routine", cmd, CallbackFunc(callback))
}
<|endoftext|>


<commit_before>package atlas

import (
	"encoding/json"
	"fmt"
	"github.com/sendgrid/rest"
	"log"
)

// MeasurementResp contains all the results of the measurements
type MeasurementResp struct {
	Measurements []int
}

// createMeasurement creates a measurement for all types
func createMeasurement(t string, d MeasurementRequest) (m *MeasurementResp, err error) {
    // Check that all Definition.Type are the same and compliant
    if !checkAllTypesAs(d.Definitions, t) {
        err = ErrInvalidMeasurementType
        return
    }

    req := prepareRequest(fmt.Sprintf("measurements/%s", t))

    body, err := json.Marshal(d)
    if err != nil {
        return
    }

    req.Method = rest.Post
    req.Body = body

    log.Printf("body: %s", body)
    resp, err := rest.API(req)

    m = &MeasurementResp{}
    err = json.Unmarshal([]byte(resp.Body), m)
    //r, err := api.Res(base, &resp).Post(d)
    fmt.Printf("m: %v\nresp: %#v\nd: %v\n", m, string(resp.Body), d)
    if err != nil {
        err = fmt.Errorf("err: %v - m:%v\n", err, m)
        return
    }

    return
}

// DNS creates a measurement
func DNS(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("dns", d)
}

// HTTP creates a measurement
func HTTP(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("http", d)
}

// NTP creates a measurement
func NTP(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("ntp", d)
}

// Ping creates a measurement
func Ping(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("ping", d)
}

// SSLCert creates a measurement
func SSLCert(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("SSL", d)
}

// Traceroute creates a measurement
func Traceroute(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("traceroute", d)
}
<commit_msg>Handle API errors here too.<commit_after>package atlas

import (
	"encoding/json"
	"fmt"
	"github.com/sendgrid/rest"
	"log"
)

// MeasurementResp contains all the results of the measurements
type MeasurementResp struct {
	Measurements []int
}

// createMeasurement creates a measurement for all types
func createMeasurement(t string, d MeasurementRequest) (m *MeasurementResp, err error) {
    // Check that all Definition.Type are the same and compliant
    if !checkAllTypesAs(d.Definitions, t) {
        err = ErrInvalidMeasurementType
        return
    }

    req := prepareRequest(fmt.Sprintf("measurements/%s", t))

    body, err := json.Marshal(d)
    if err != nil {
        return
    }

    req.Method = rest.Post
    req.Body = body

    log.Printf("body: %s", body)
    resp, err := rest.API(req)
    err = handleAPIResponse(resp)
    if err != nil {
        return
    }

    m = &MeasurementResp{}
    err = json.Unmarshal([]byte(resp.Body), m)
    //r, err := api.Res(base, &resp).Post(d)
    fmt.Printf("m: %v\nresp: %#v\nd: %v\n", m, string(resp.Body), d)
    if err != nil {
        err = fmt.Errorf("err: %v - m:%v\n", err, m)
        return
    }

    return
}

// DNS creates a measurement
func DNS(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("dns", d)
}

// HTTP creates a measurement
func HTTP(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("http", d)
}

// NTP creates a measurement
func NTP(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("ntp", d)
}

// Ping creates a measurement
func Ping(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("ping", d)
}

// SSLCert creates a measurement
func SSLCert(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("SSL", d)
}

// Traceroute creates a measurement
func Traceroute(d MeasurementRequest) (m *MeasurementResp, err error) {
    return createMeasurement("traceroute", d)
}
<|endoftext|>


<commit_before>package scan

import (
	"crypto/md5"
	"encoding/hex"
	"io/ioutil"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/cadenzr/cadenzr/db"
	"github.com/cadenzr/cadenzr/log"
	"github.com/cadenzr/cadenzr/models"
	"github.com/cadenzr/cadenzr/probers"
)

var ScanDone = make(chan struct{})

func isImage(mime string) bool {
	mime = strings.ToLower(mime)
	return strings.Contains(mime, "image")
}

func ScanHandler(scanCh chan (chan struct{})) {
	requests := []chan struct{}{}
	scanning := false
	for {
		select {
		case done := <-scanCh:
			requests = append(requests, done)
			if scanning == false {
				scanning = true
				go ScanFilesystem("media")
			}
		case <-ScanDone:
			for _, done := range requests {

				done <- struct{}{}
			}
			requests = []chan struct{}{}
			scanning = false
		}
	}
}

func ScanFilesystem(mediaDir string) {
	defer func() {
		ScanDone <- struct{}{}
	}()

	newFiles := 0

	filepath.Walk(mediaDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.WithFields(log.Fields{"reason": err.Error(), "path": path}).Error("Failed to handle file/dir.")
			return nil
		}

		if info.IsDir() {
			return nil
		}

		mimeType := mime.TypeByExtension(filepath.Ext(path))
		if !probers.HasProber(mimeType) {
			log.WithFields(log.Fields{"path": path, "mime": mimeType}).Debug("Skipping file. Unknown mime.")
			return nil
		}

		log.WithFields(log.Fields{"path": path, "mime": mimeType}).Debug("Found file.")

		// First check if file with same path exists.
		var exists uint
		gormDB := db.DB.Table("songs").Where("path = ?", path).Count(&exists)
		if gormDB.Error != nil {
			log.Errorf("scanFilesystem Could not check if path exists. Database failed %v.", gormDB.Error)
			return nil
		}

		if exists != 0 {
			log.Debugf("Skipping file. Path already in database: %s.", path)
			return nil
		}

		meta, err := probers.ProbeAudioFile(path)
		if err != nil {
			log.WithFields(log.Fields{"reason": err.Error(), "file": path}).Error("Probing file failed.")
			return nil
		}

		var cover *models.Image
		if meta.CoverBufer != nil {
			mimeCover := http.DetectContentType(meta.CoverBufer)
			if isImage(mimeCover) {
				md5Sum := md5.Sum(meta.CoverBufer)
				hash := hex.EncodeToString(md5Sum[:])

				extensions, _ := mime.ExtensionsByType(mimeCover)
				if extensions != nil && len(extensions) > 0 {
					destination := "images/" + hash + extensions[0]
					cover = &models.Image{
						Path: destination,
						Link: "/" + destination,
						Mime: mimeCover,
						Hash: hash,
					}

					if err = ioutil.WriteFile(destination, meta.CoverBufer, 0666); err == nil {
						cover = &models.Image{
							Path: destination,
							Link: "/" + destination,
							Mime: mimeCover,
							Hash: hash,
						}

						gormDB := db.DB.Table("images").Where("hash = ?", cover.Hash).First(cover)
						if gormDB.RecordNotFound() {
							gormDB = db.DB.Create(cover)
						}

						if gormDB.Error != nil {
							log.WithFields(log.Fields{"reason": err.Error(), "file": cover.Path}).Error("Could not get or insert image.")
							cover = nil
						}

					} else {
						log.WithFields(log.Fields{"reason": err.Error(), "destination": destination}).Error("Failed to write cover to disk.")
					}
				}
			}
		}

		song := &models.Song{}
		song.Name = meta.Title
		if len(song.Name) == 0 {
			log.WithFields(log.Fields{"file": path}).Error("Could not get name of song.")
			return nil
		}
		song.Mime = mimeType
		song.Path = path
		if cover != nil {
			song.Cover = cover
		}
		if len(meta.Genre) > 0 {
			song.Genre.Set(meta.Genre)
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No genre found.")
		}
		if meta.Year != 0 {
			song.Year.Set(int64(meta.Year))
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No year found.")
		}
		if len(meta.Album) > 0 {
			album := &models.Album{
				Name: meta.Album,
				Year: song.Year,
			}
			if cover != nil {
				album.Cover = cover
			}
			if gormDB := db.DB.FirstOrCreate(album, "name = ?", album.Name); gormDB.Error != nil {
				log.Errorf("Could not create/get album '%s': %v", album.Name, gormDB.Error)
				return nil
			}

			song.Album = album
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No album found.")
		}
		if len(meta.Artist) > 0 {
			artist := &models.Artist{
				Name: meta.Artist,
			}

			if gormDB := db.DB.FirstOrCreate(artist, "name = ?", artist.Name); gormDB.Error != nil {
				log.Errorf("Could not create/get artist '%s': %v", artist.Name, gormDB.Error)
				return nil
			}

			song.Artist = artist
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No Artist found.")
		}

		if meta.Duration != 0 {
			song.Duration.Set(meta.Duration)
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No duration found.")
		}

		if gormDB := db.DB.Create(song); gormDB.Error != nil {
			log.Errorf("Could not create song '%s': %v", song.Name, gormDB.Error)
			return nil
		}

		newFiles++
		return nil
	})

	log.Infof("Added %d new songs.", newFiles)
}
<commit_msg>Show how long scan took.<commit_after>package scan

import (
	"crypto/md5"
	"encoding/hex"
	"io/ioutil"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/cadenzr/cadenzr/db"
	"github.com/cadenzr/cadenzr/log"
	"github.com/cadenzr/cadenzr/models"
	"github.com/cadenzr/cadenzr/probers"
)

var ScanDone = make(chan struct{})

func isImage(mime string) bool {
	mime = strings.ToLower(mime)
	return strings.Contains(mime, "image")
}

func ScanHandler(scanCh chan (chan struct{})) {
	requests := []chan struct{}{}
	scanning := false
	for {
		select {
		case done := <-scanCh:
			requests = append(requests, done)
			if scanning == false {
				scanning = true
				go ScanFilesystem("media")
			}
		case <-ScanDone:
			for _, done := range requests {

				done <- struct{}{}
			}
			requests = []chan struct{}{}
			scanning = false
		}
	}
}

func ScanFilesystem(mediaDir string) {
	defer func() {
		ScanDone <- struct{}{}
	}()

	start := time.Now()
	newFiles := 0

	filepath.Walk(mediaDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.WithFields(log.Fields{"reason": err.Error(), "path": path}).Error("Failed to handle file/dir.")
			return nil
		}

		if info.IsDir() {
			return nil
		}

		mimeType := mime.TypeByExtension(filepath.Ext(path))
		if !probers.HasProber(mimeType) {
			log.WithFields(log.Fields{"path": path, "mime": mimeType}).Debug("Skipping file. Unknown mime.")
			return nil
		}

		log.WithFields(log.Fields{"path": path, "mime": mimeType}).Debug("Found file.")

		// First check if file with same path exists.
		var exists uint
		gormDB := db.DB.Table("songs").Where("path = ?", path).Count(&exists)
		if gormDB.Error != nil {
			log.Errorf("scanFilesystem Could not check if path exists. Database failed %v.", gormDB.Error)
			return nil
		}

		if exists != 0 {
			log.Debugf("Skipping file. Path already in database: %s.", path)
			return nil
		}

		meta, err := probers.ProbeAudioFile(path)
		if err != nil {
			log.WithFields(log.Fields{"reason": err.Error(), "file": path}).Error("Probing file failed.")
			return nil
		}

		var cover *models.Image
		if meta.CoverBufer != nil {
			mimeCover := http.DetectContentType(meta.CoverBufer)
			if isImage(mimeCover) {
				md5Sum := md5.Sum(meta.CoverBufer)
				hash := hex.EncodeToString(md5Sum[:])

				extensions, _ := mime.ExtensionsByType(mimeCover)
				if extensions != nil && len(extensions) > 0 {
					destination := "images/" + hash + extensions[0]
					cover = &models.Image{
						Path: destination,
						Link: "/" + destination,
						Mime: mimeCover,
						Hash: hash,
					}

					if err = ioutil.WriteFile(destination, meta.CoverBufer, 0666); err == nil {
						cover = &models.Image{
							Path: destination,
							Link: "/" + destination,
							Mime: mimeCover,
							Hash: hash,
						}

						gormDB := db.DB.Table("images").Where("hash = ?", cover.Hash).First(cover)
						if gormDB.RecordNotFound() {
							gormDB = db.DB.Create(cover)
						}

						if gormDB.Error != nil {
							log.WithFields(log.Fields{"reason": err.Error(), "file": cover.Path}).Error("Could not get or insert image.")
							cover = nil
						}

					} else {
						log.WithFields(log.Fields{"reason": err.Error(), "destination": destination}).Error("Failed to write cover to disk.")
					}
				}
			}
		}

		song := &models.Song{}
		song.Name = meta.Title
		if len(song.Name) == 0 {
			log.WithFields(log.Fields{"file": path}).Error("Could not get name of song.")
			return nil
		}
		song.Mime = mimeType
		song.Path = path
		if cover != nil {
			song.Cover = cover
		}
		if len(meta.Genre) > 0 {
			song.Genre.Set(meta.Genre)
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No genre found.")
		}
		if meta.Year != 0 {
			song.Year.Set(int64(meta.Year))
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No year found.")
		}
		if len(meta.Album) > 0 {
			album := &models.Album{
				Name: meta.Album,
				Year: song.Year,
			}
			if cover != nil {
				album.Cover = cover
			}
			if gormDB := db.DB.FirstOrCreate(album, "name = ?", album.Name); gormDB.Error != nil {
				log.Errorf("Could not create/get album '%s': %v", album.Name, gormDB.Error)
				return nil
			}

			song.Album = album
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No album found.")
		}
		if len(meta.Artist) > 0 {
			artist := &models.Artist{
				Name: meta.Artist,
			}

			if gormDB := db.DB.FirstOrCreate(artist, "name = ?", artist.Name); gormDB.Error != nil {
				log.Errorf("Could not create/get artist '%s': %v", artist.Name, gormDB.Error)
				return nil
			}

			song.Artist = artist
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No Artist found.")
		}

		if meta.Duration != 0 {
			song.Duration.Set(meta.Duration)
		} else {
			log.WithFields(log.Fields{"file": path}).Debug("No duration found.")
		}

		if gormDB := db.DB.Create(song); gormDB.Error != nil {
			log.Errorf("Could not create song '%s': %v", song.Name, gormDB.Error)
			return nil
		}

		newFiles++
		return nil
	})

	dur := time.Since(start)
	log.Infof("Added %d new songs in %.2f seconds.", newFiles, dur.Seconds())
}
<|endoftext|>


<commit_before>// Package scan TCP ports
package scan

import (
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/olekukonko/tablewriter"
)

// Scan represents the scan parameters
type Scan struct {
	minPort int
	maxPort int
	target  string
}

// NewScan creats scan object
func NewScan(args string, cfg cli.Config) (Scan, error) {
	var (
		scan Scan
		flag map[string]interface{}
		err  error
	)

	args, flag = cli.Flag(args)
	// help
	if _, ok := flag["help"]; ok || args == "" {
		help(cfg)
		return scan, fmt.Errorf("")
	}

	pRange := cli.SetFlag(flag, "p", cfg.Scan.Port).(string)

	re := regexp.MustCompile(`(\d+)(\-{0,1}(\d*))`)
	f := re.FindStringSubmatch(pRange)

	if len(f) != 4 {
		return scan, fmt.Errorf("error! please try scan help")
	}

	scan.target = args
	if len(f) == 4 && f[2] != "" {
		scan.minPort, err = strconv.Atoi(f[1])
		scan.maxPort, err = strconv.Atoi(f[3])
	} else {
		scan.minPort, err = strconv.Atoi(f[1])
		scan.maxPort, err = strconv.Atoi(f[1])
	}

	if err != nil {
		return scan, err
	}

	if !scan.IsCIDR() {
		ipAddr, err := net.ResolveIPAddr("ip", scan.target)
		if err != nil {
			return scan, err
		}
		scan.target = ipAddr.String()
	} else {
		return scan, fmt.Errorf("it doesn't support CIDR")
	}

	return scan, nil
}

// IsCIDR checks the target if it's CIDR
func (s Scan) IsCIDR() bool {
	_, _, err := net.ParseCIDR(s.target)
	if err != nil {
		return false
	}
	return true
}

// Run tries to scan wide range ports (TCP)
func (s Scan) Run() {
	if !s.IsCIDR() {
		host(s.target, s.minPort, s.maxPort)
	}
}

// host tries to scan a single host
func host(ipAddr string, minPort, maxPort int) {
	var (
		wg      sync.WaitGroup
		tStart  = time.Now()
		counter int
	)

	var ports []int
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Protocol", "Port", "Status"})

	for i := minPort; i <= maxPort; i++ {
		wg.Add(1)
		go func(i int) {
			for {
				conn, err := net.DialTimeout("tcp", net.JoinHostPort(ipAddr, fmt.Sprintf("%d", i)), 2*time.Second)
				if err != nil {
					if strings.Contains(err.Error(), "too many open files") {
						// random back-off
						time.Sleep(time.Duration(10+rand.Int31n(30)) * time.Millisecond)
						continue
					}
					wg.Done()
					return
				}
				conn.Close()
				break
			}
			counter++
			ports = append(ports, i)
			wg.Done()
		}(i)
	}

	wg.Wait()

	sort.Ints(ports)
	for i := range ports {
		table.Append([]string{"TCP", fmt.Sprintf("%d", ports[i]), "Open"})
	}
	table.Render()

	if counter == 0 {
		println("there isn't any opened port")
	} else {
		elapsed := fmt.Sprintf("%.3f seconds", time.Since(tStart).Seconds())
		println("Scan done:", counter, "opened port(s) found in", elapsed)
	}
}

// help represents guide to user
func help(cfg cli.Config) {
	fmt.Printf(`
    usage:
          scan ip/host [option]
    options:
          -p port-range or port number      specified range or port number (default is %s)
    example:
          scan 8.8.8.8 -p 53
          scan www.google.com -p 1-500
	`,
		cfg.Scan.Port)
}
<commit_msg>add TCP.SYN scan<commit_after>// Package scan TCP ports
package scan

import (
	"context"
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

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

	"github.com/mehrdadrad/mylg/cli"
	"github.com/olekukonko/tablewriter"
)

// Scan represents the scan parameters
type Scan struct {
	minPort int
	maxPort int
	target  string
	lport   int
	rport   int
	laddr   net.IP
	raddr   net.IP
	ifName  string
	network string
	forceV4 bool
	forceV6 bool
	synScan bool
	ip      gopacket.NetworkLayer
}

// NewScan creats scan object
func NewScan(args string, cfg cli.Config) (Scan, error) {
	var (
		scan Scan
		flag map[string]interface{}
		err  error
	)

	args, flag = cli.Flag(args)
	// help
	if _, ok := flag["help"]; ok || args == "" {
		help(cfg)
		return scan, fmt.Errorf("")
	}

	scan.forceV4 = cli.SetFlag(flag, "4", true).(bool)
	scan.forceV6 = cli.SetFlag(flag, "6", false).(bool)
	scan.synScan = cli.SetFlag(flag, "s", false).(bool)

	pRange := cli.SetFlag(flag, "p", cfg.Scan.Port).(string)

	re := regexp.MustCompile(`(\d+)(\-{0,1}(\d*))`)
	f := re.FindStringSubmatch(pRange)

	if len(f) != 4 {
		return scan, fmt.Errorf("error! please try scan help")
	}

	scan.target = args
	if len(f) == 4 && f[2] != "" {
		scan.minPort, err = strconv.Atoi(f[1])
		scan.maxPort, err = strconv.Atoi(f[3])
	} else {
		scan.minPort, err = strconv.Atoi(f[1])
		scan.maxPort, err = strconv.Atoi(f[1])
	}

	if err != nil {
		return scan, err
	}

	if scan.IsCIDR() {
		return scan, fmt.Errorf("it doesn't support CIDR")
	}

	scan.setIP()

	return scan, nil
}

// IsCIDR checks the target if it's CIDR
func (s *Scan) IsCIDR() bool {
	_, _, err := net.ParseCIDR(s.target)
	if err != nil {
		return false
	}
	return true
}

func isIPv4(ip net.IP) bool {
	return len(ip.To4()) == net.IPv4len
}

func isIPv6(ip net.IP) bool {
	if r := strings.Index(ip.String(), ":"); r != -1 {
		return true
	}
	return false
}

func (s *Scan) setIP() error {
	ips, err := net.LookupIP(s.target)
	if err != nil {
		return err
	}
	for _, ip := range ips {
		if isIPv4(ip) && !s.forceV6 {
			s.raddr = ip
			s.network = "ip4"
			break
		} else if isIPv6(ip) && !s.forceV4 {
			s.raddr = ip
			s.network = "ip6"
			break
		}
	}
	return nil
}

// Run tries to scan wide range ports (TCP)
func (s *Scan) Run() {
	var (
		openPorts []int
		err       error
	)

	fmt.Printf("Scan %s (%s)\n", s.target, s.raddr)

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Protocol", "Port", "Status"})

	tStart := time.Now()
	if !s.synScan {
		openPorts = s.tcpConnScan()
	} else {
		openPorts, err = s.tcpSYNScan()
	}

	if err != nil {
		println(err.Error())
		return
	}

	for _, p := range openPorts {
		table.Append([]string{"TCP", fmt.Sprintf("%d", p), "Open"})
	}
	table.Render()

	if len(openPorts) == 0 {
		println("there isn't any opened port")
	} else {
		elapsed := fmt.Sprintf("%.3f seconds", time.Since(tStart).Seconds())
		println("Scan done:", len(openPorts), "opened port(s) found in", elapsed)
	}

}

func (s *Scan) packetDataTCP(rport int) (error, []byte) {
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(s.lport),
		DstPort: layers.TCPPort(rport),
		Seq:     1,
		SYN:     true,
		Window:  15000,
	}

	tcp.SetNetworkLayerForChecksum(s.ip)
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}

	if err := gopacket.SerializeLayers(buf, opts, tcp); err != nil {
		return err, []byte{}
	}

	return nil, buf.Bytes()
}

func (s *Scan) setProto(proto string) error {
	var netProto = fmt.Sprintf("%s:%s", s.network, proto)

	switch netProto {
	case "ip4:tcp":
		s.ip = &layers.IPv4{
			SrcIP:    s.laddr,
			DstIP:    s.raddr,
			Protocol: layers.IPProtocolTCP,
		}
	case "ip4:udp":
		s.ip = &layers.IPv4{
			SrcIP:    s.laddr,
			DstIP:    s.raddr,
			Protocol: layers.IPProtocolUDP,
		}
	case "ip6:tcp":
		s.ip = &layers.IPv6{
			SrcIP:      s.laddr,
			DstIP:      s.raddr,
			NextHeader: layers.IPProtocolTCP,
		}
	case "ip6:udp":
		s.ip = &layers.IPv6{
			SrcIP:      s.laddr,
			DstIP:      s.raddr,
			NextHeader: layers.IPProtocolUDP,
		}

	}

	return nil
}

func (s *Scan) setLocalNet() error {
	conn, err := net.Dial("udp", net.JoinHostPort(s.raddr.String(), "80"))
	if err != nil {
		return err
	}
	defer conn.Close()

	if lAddr, ok := conn.LocalAddr().(*net.UDPAddr); ok {
		s.laddr = lAddr.IP
		s.lport = lAddr.Port
	} else {
		return fmt.Errorf("can not find local address/port")
	}

	ifs, _ := net.Interfaces()
	for _, i := range ifs {
		addrs, _ := i.Addrs()
		for _, addr := range addrs {
			ip, _, _ := net.ParseCIDR(addr.String())
			if ip.String() == s.laddr.String() {
				s.ifName = i.Name
				break
			}
		}
	}

	return nil
}

func (s *Scan) tcpSYNScan() ([]int, error) {
	var err error

	ctx, cancel := context.WithCancel(context.Background())

	if err = s.setLocalNet(); err != nil {
		println(err.Error())
	}
	if err = s.setProto("tcp"); err != nil {
		println(err.Error())
	}
	go func() {
		if err := s.sendTCPSYN(); err != nil {
			println(err.Error())
		}
		cancel()
	}()

	openPorts, err := s.pCapture(ctx)
	return openPorts, err
}

func (s *Scan) pCapture(ctx context.Context) ([]int, error) {
	var (
		tcp       *layers.TCP
		timeout   = 100 * time.Nanosecond
		openPorts []int
		ok        bool
	)
	handle, err := pcap.OpenLive(s.ifName, 6*1024, false, timeout)
	if err != nil {
		return nil, err
	}
	defer handle.Close()

	filter := "(tcp or udp) and src host " + s.raddr.String()
	if err := handle.SetBPFFilter(filter); err != nil {
		return nil, err
	}

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
LOOP:
	for {
		select {
		case <-ctx.Done():
			break LOOP
		case packet := <-packetSource.Packets():
			tcpLayer := packet.Layer(layers.LayerTypeTCP)
			if tcpLayer != nil {
				if tcp, ok = tcpLayer.(*layers.TCP); !ok {
					continue
				}
				if tcp.SYN && tcp.ACK {
					p, _ := strconv.Atoi(fmt.Sprintf("%d", tcp.SrcPort))
					openPorts = append(openPorts, p)
				}
				continue
			}
		}
	}
	//TODO: unique .....
	sort.Ints(openPorts)
	return openPorts, nil
}

func (s *Scan) sendTCPSYN() error {
	var (
		buf  []byte
		err  error
		conn net.PacketConn
	)
	if s.network != "ip6" {
		conn, err = net.ListenPacket(s.network+":tcp", "0.0.0.0")
	} else {
		conn, err = net.ListenPacket(s.network+":tcp", "[::]")
	}
	if err != nil {
		return err
	}
	for i := s.minPort; i <= s.maxPort; i++ {
		if err, buf = s.packetDataTCP(i); err != nil {
			return err
		}
		if _, err := conn.WriteTo(buf, &net.IPAddr{IP: s.raddr}); err != nil {
			println(err.Error())
		}
		time.Sleep(10 * time.Millisecond)
	}
	time.Sleep(2 * time.Second)
	return nil
}

// tcpConnScan tries to scan a single host
func (s *Scan) tcpConnScan() []int {
	var wg sync.WaitGroup

	var ports []int
	for i := s.minPort; i <= s.maxPort; i++ {
		wg.Add(1)
		go func(i int) {
			for {
				host := net.JoinHostPort(s.raddr.String(), fmt.Sprintf("%d", i))
				conn, err := net.DialTimeout("tcp", host, 2*time.Second)
				if err != nil {
					if strings.Contains(err.Error(), "too many open files") {
						// random back-off
						time.Sleep(time.Duration(10+rand.Int31n(30)) * time.Millisecond)
						continue
					}
					wg.Done()
					return
				}
				conn.Close()
				break
			}
			ports = append(ports, i)
			wg.Done()
		}(i)
	}

	wg.Wait()
	sort.Ints(ports)
	return ports
}

// help represents guide to user
func help(cfg cli.Config) {
	fmt.Printf(`
    usage:
          scan ip/host [option]
    options:
          -p port-range or port number      specified range or port number (default is %s)
    example:
          scan 8.8.8.8 -p 53
          scan www.google.com -p 1-500
	`,
		cfg.Scan.Port)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"strings"

	"github.com/coreos/rocket/Godeps/_workspace/src/github.com/gorilla/mux"
	"github.com/coreos/rocket/app-container/schema"
	"github.com/coreos/rocket/app-container/schema/types"
)

type metadata struct {
	manifest schema.ContainerRuntimeManifest
	apps     map[string]*schema.AppManifest
}

var (
	metadataByIP  = make(map[string]*metadata)
	metadataByUID = make(map[types.UUID]*metadata)
	hmacKey       [sha1.Size]byte
)

const (
	myPort   = "4444"
	metaIP   = "169.254.169.255"
	metaPort = "80"
)

func setupIPTables() error {
	args := []string{"-t", "nat", "-A", "PREROUTING",
		"-p", "tcp", "-d", metaIP, "--dport", metaPort,
		"-j", "REDIRECT", "--to-port", myPort}

	return exec.Command("iptables", args...).Run()
}

func antiSpoof(brPort, ipAddr string) error {
	args := []string{"-t", "filter", "-I", "INPUT", "-i", brPort, "-p", "IPV4", "!", "--ip-source", ipAddr, "-j", "DROP"}
	return exec.Command("ebtables", args...).Run()
}

func queryValue(u *url.URL, key string) string {
	vals, ok := u.Query()[key]
	if !ok || len(vals) != 1 {
		return ""
	}
	return vals[0]
}

func handleRegisterContainer(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	containerIP := queryValue(r.URL, "container_ip")
	if containerIP == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_ip missing")
		return
	}
	containerBrPort := queryValue(r.URL, "container_brport")
	if containerBrPort == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_brport missing")
		return
	}

	m := &metadata{
		apps: make(map[string]*schema.AppManifest),
	}

	if err := json.NewDecoder(r.Body).Decode(&m.manifest); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	if err := antiSpoof(containerBrPort, containerIP); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to set anti-spoofing: %v", err)
		return
	}

	metadataByIP[containerIP] = m
	metadataByUID[m.manifest.UUID] = m

	w.WriteHeader(http.StatusOK)
}

func handleRegisterApp(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	uid, err := types.NewUUID(mux.Vars(r)["uid"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "UUID is missing or mulformed: %v", err)
		return
	}

	m, ok := metadataByUID[*uid]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "Container with given UUID not found")
		return
	}

	an := mux.Vars(r)["app"]

	app := &schema.AppManifest{}
	if err := json.NewDecoder(r.Body).Decode(&app); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	m.apps[an] = app

	w.WriteHeader(http.StatusOK)
}

func containerGet(h func(w http.ResponseWriter, r *http.Request, m *metadata)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		remoteIP := strings.Split(r.RemoteAddr, ":")[0]
		m, ok := metadataByIP[remoteIP]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "metadata by remoteIP (%v) not found", remoteIP)
			return
		}

		h(w, r, m)
	}
}

func appGet(h func(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest)) func(http.ResponseWriter, *http.Request) {
	return containerGet(func(w http.ResponseWriter, r *http.Request, m *metadata) {
		appname := mux.Vars(r)["app"]

		if am, ok := m.apps[appname]; ok {
			h(w, r, m, am)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "App (%v) not found", appname)
		}
	})
}

func handleContainerAnnotations(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range m.manifest.Annotations {
		fmt.Fprintln(w, k)
	}
}

func handleContainerAnnotation(w http.ResponseWriter, r *http.Request, m *metadata) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation is not a valid AC Label")
		return
	}

	v, ok := m.manifest.Annotations[*k]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleContainerManifest(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(m.manifest); err != nil {
		fmt.Println(err)
	}
}

func handleContainerUID(w http.ResponseWriter, r *http.Request, m *metadata) {
	uid := m.manifest.UUID.String()

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(uid))
}

func mergeAppAnnotations(am *schema.AppManifest, cm *schema.ContainerRuntimeManifest) types.Annotations {
	merged := make(types.Annotations)

	for k, v := range am.Annotations {
		merged[k] = v
	}

	if app := cm.Apps.Get(am.Name); app != nil {
		for k, v := range app.Annotations {
			merged[k] = v
		}
	}

	return merged
}

func handleAppAnnotations(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range mergeAppAnnotations(am, &m.manifest) {
		fmt.Fprintln(w, k)
	}
}

func handleAppAnnotation(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation is not a valid AC Label")
		return
	}

	merged := mergeAppAnnotations(am, &m.manifest)

	v, ok := merged[*k]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleAppManifest(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(*am); err != nil {
		fmt.Println(err)
	}
}

func handleAppID(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	a := m.manifest.Apps.Get(am.Name)
	if a == nil {
		panic("could not find app in manifest!")
	}
	w.Write([]byte(a.ImageID.String()))
}

func initCrypto() error {
	if n, err := rand.Reader.Read(hmacKey[:]); err != nil || n != len(hmacKey) {
		return fmt.Errorf("failed to generate HMAC Key")
	}
	return nil
}

func digest(r io.Reader) ([]byte, error) {
	digest := sha1.New()
	if _, err := io.Copy(digest, r); err != nil {
		return nil, err
	}
	return digest.Sum(nil), nil
}

func handleContainerSign(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	m, ok := metadataByIP[remoteIP]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Metadata by remoteIP (%v) not found", remoteIP)
		return
	}

	// compute message digest
	d, err := digest(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "Digest computation failed: %v", err)
		return
	}

	// HMAC(UID:digest)
	h := hmac.New(sha1.New, hmacKey[:])
	h.Write(m.manifest.UUID[:])
	h.Write(d)

	// Send back digest:HMAC as the signature
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	enc := base64.NewEncoder(base64.StdEncoding, w)
	enc.Write(d)
	enc.Write(h.Sum(nil))
	enc.Close()
}

func handleContainerVerify(w http.ResponseWriter, r *http.Request) {
	uid, err := types.NewUUID(r.FormValue("uid"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "uid field missing or malformed: %v", err)
		return
	}

	sig, err := base64.StdEncoding.DecodeString(r.FormValue("signature"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "signature field missing or corrupt: %v", err)
		return
	}

	digest := sig[:sha1.Size]
	sum := sig[sha1.Size:]

	h := hmac.New(sha1.New, hmacKey[:])
	h.Write(uid[:])
	h.Write(digest)

	if hmac.Equal(sum, h.Sum(nil)) {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusForbidden)
	}
}

type httpResp struct {
	writer http.ResponseWriter
	status int
}

func (r *httpResp) Header() http.Header {
	return r.writer.Header()
}

func (r *httpResp) Write(d []byte) (int, error) {
	return r.writer.Write(d)
}

func (r *httpResp) WriteHeader(status int) {
	r.status = status
	r.writer.WriteHeader(status)
}

func logReq(h func(w http.ResponseWriter, r *http.Request)) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		resp := &httpResp{w, 0}
		h(resp, r)
		fmt.Printf("%v %v - %v\n", r.Method, r.RequestURI, resp.status)
	}
}

func main() {
	if err := setupIPTables(); err != nil {
		fmt.Fatal(err)
	}

	if err := initCrypto(); err != nil {
		log.Fatal(err)
	}

	r := mux.NewRouter()
	r.HandleFunc("/containers/", logReq(handleRegisterContainer)).Methods("POST")
	r.HandleFunc("/containers/{uid}/{app:.*}", logReq(handleRegisterApp)).Methods("PUT")

	acRtr := r.Headers("Metadata-Flavor", "AppContainer header").
		PathPrefix("/acMetadata/v1").Subrouter()

	mr := acRtr.Methods("GET").Subrouter()

	mr.HandleFunc("/container/annotations/", logReq(containerGet(handleContainerAnnotations)))
	mr.HandleFunc("/container/annotations/{name}", logReq(containerGet(handleContainerAnnotation)))
	mr.HandleFunc("/container/manifest", logReq(containerGet(handleContainerManifest)))
	mr.HandleFunc("/container/uid", logReq(containerGet(handleContainerUID)))

	mr.HandleFunc("/apps/{app:.*}/annotations/", logReq(appGet(handleAppAnnotations)))
	mr.HandleFunc("/apps/{app:.*}/annotations/{name}", logReq(appGet(handleAppAnnotation)))
	mr.HandleFunc("/apps/{app:.*}/image/manifest", logReq(appGet(handleAppManifest)))
	mr.HandleFunc("/apps/{app:.*}/image/id", logReq(appGet(handleAppID)))

	acRtr.HandleFunc("/container/hmac/sign", logReq(handleContainerSign)).Methods("POST")
	acRtr.HandleFunc("/container/hmac/verify", logReq(handleContainerVerify)).Methods("POST")

	log.Fatal(http.ListenAndServe(":4444", r))
}
<commit_msg>metadatasvc: changed from fmt to log the setupIPTables<commit_after>package main

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"strings"

	"github.com/coreos/rocket/Godeps/_workspace/src/github.com/gorilla/mux"
	"github.com/coreos/rocket/app-container/schema"
	"github.com/coreos/rocket/app-container/schema/types"
)

type metadata struct {
	manifest schema.ContainerRuntimeManifest
	apps     map[string]*schema.AppManifest
}

var (
	metadataByIP  = make(map[string]*metadata)
	metadataByUID = make(map[types.UUID]*metadata)
	hmacKey       [sha1.Size]byte
)

const (
	myPort   = "4444"
	metaIP   = "169.254.169.255"
	metaPort = "80"
)

func setupIPTables() error {
	args := []string{"-t", "nat", "-A", "PREROUTING",
		"-p", "tcp", "-d", metaIP, "--dport", metaPort,
		"-j", "REDIRECT", "--to-port", myPort}

	return exec.Command("iptables", args...).Run()
}

func antiSpoof(brPort, ipAddr string) error {
	args := []string{"-t", "filter", "-I", "INPUT", "-i", brPort, "-p", "IPV4", "!", "--ip-source", ipAddr, "-j", "DROP"}
	return exec.Command("ebtables", args...).Run()
}

func queryValue(u *url.URL, key string) string {
	vals, ok := u.Query()[key]
	if !ok || len(vals) != 1 {
		return ""
	}
	return vals[0]
}

func handleRegisterContainer(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	containerIP := queryValue(r.URL, "container_ip")
	if containerIP == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_ip missing")
		return
	}
	containerBrPort := queryValue(r.URL, "container_brport")
	if containerBrPort == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_brport missing")
		return
	}

	m := &metadata{
		apps: make(map[string]*schema.AppManifest),
	}

	if err := json.NewDecoder(r.Body).Decode(&m.manifest); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	if err := antiSpoof(containerBrPort, containerIP); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to set anti-spoofing: %v", err)
		return
	}

	metadataByIP[containerIP] = m
	metadataByUID[m.manifest.UUID] = m

	w.WriteHeader(http.StatusOK)
}

func handleRegisterApp(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	uid, err := types.NewUUID(mux.Vars(r)["uid"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "UUID is missing or mulformed: %v", err)
		return
	}

	m, ok := metadataByUID[*uid]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "Container with given UUID not found")
		return
	}

	an := mux.Vars(r)["app"]

	app := &schema.AppManifest{}
	if err := json.NewDecoder(r.Body).Decode(&app); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	m.apps[an] = app

	w.WriteHeader(http.StatusOK)
}

func containerGet(h func(w http.ResponseWriter, r *http.Request, m *metadata)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		remoteIP := strings.Split(r.RemoteAddr, ":")[0]
		m, ok := metadataByIP[remoteIP]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "metadata by remoteIP (%v) not found", remoteIP)
			return
		}

		h(w, r, m)
	}
}

func appGet(h func(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest)) func(http.ResponseWriter, *http.Request) {
	return containerGet(func(w http.ResponseWriter, r *http.Request, m *metadata) {
		appname := mux.Vars(r)["app"]

		if am, ok := m.apps[appname]; ok {
			h(w, r, m, am)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "App (%v) not found", appname)
		}
	})
}

func handleContainerAnnotations(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range m.manifest.Annotations {
		fmt.Fprintln(w, k)
	}
}

func handleContainerAnnotation(w http.ResponseWriter, r *http.Request, m *metadata) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation is not a valid AC Label")
		return
	}

	v, ok := m.manifest.Annotations[*k]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleContainerManifest(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(m.manifest); err != nil {
		fmt.Println(err)
	}
}

func handleContainerUID(w http.ResponseWriter, r *http.Request, m *metadata) {
	uid := m.manifest.UUID.String()

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(uid))
}

func mergeAppAnnotations(am *schema.AppManifest, cm *schema.ContainerRuntimeManifest) types.Annotations {
	merged := make(types.Annotations)

	for k, v := range am.Annotations {
		merged[k] = v
	}

	if app := cm.Apps.Get(am.Name); app != nil {
		for k, v := range app.Annotations {
			merged[k] = v
		}
	}

	return merged
}

func handleAppAnnotations(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range mergeAppAnnotations(am, &m.manifest) {
		fmt.Fprintln(w, k)
	}
}

func handleAppAnnotation(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation is not a valid AC Label")
		return
	}

	merged := mergeAppAnnotations(am, &m.manifest)

	v, ok := merged[*k]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleAppManifest(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(*am); err != nil {
		fmt.Println(err)
	}
}

func handleAppID(w http.ResponseWriter, r *http.Request, m *metadata, am *schema.AppManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	a := m.manifest.Apps.Get(am.Name)
	if a == nil {
		panic("could not find app in manifest!")
	}
	w.Write([]byte(a.ImageID.String()))
}

func initCrypto() error {
	if n, err := rand.Reader.Read(hmacKey[:]); err != nil || n != len(hmacKey) {
		return fmt.Errorf("failed to generate HMAC Key")
	}
	return nil
}

func digest(r io.Reader) ([]byte, error) {
	digest := sha1.New()
	if _, err := io.Copy(digest, r); err != nil {
		return nil, err
	}
	return digest.Sum(nil), nil
}

func handleContainerSign(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	m, ok := metadataByIP[remoteIP]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Metadata by remoteIP (%v) not found", remoteIP)
		return
	}

	// compute message digest
	d, err := digest(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "Digest computation failed: %v", err)
		return
	}

	// HMAC(UID:digest)
	h := hmac.New(sha1.New, hmacKey[:])
	h.Write(m.manifest.UUID[:])
	h.Write(d)

	// Send back digest:HMAC as the signature
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	enc := base64.NewEncoder(base64.StdEncoding, w)
	enc.Write(d)
	enc.Write(h.Sum(nil))
	enc.Close()
}

func handleContainerVerify(w http.ResponseWriter, r *http.Request) {
	uid, err := types.NewUUID(r.FormValue("uid"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "uid field missing or malformed: %v", err)
		return
	}

	sig, err := base64.StdEncoding.DecodeString(r.FormValue("signature"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "signature field missing or corrupt: %v", err)
		return
	}

	digest := sig[:sha1.Size]
	sum := sig[sha1.Size:]

	h := hmac.New(sha1.New, hmacKey[:])
	h.Write(uid[:])
	h.Write(digest)

	if hmac.Equal(sum, h.Sum(nil)) {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusForbidden)
	}
}

type httpResp struct {
	writer http.ResponseWriter
	status int
}

func (r *httpResp) Header() http.Header {
	return r.writer.Header()
}

func (r *httpResp) Write(d []byte) (int, error) {
	return r.writer.Write(d)
}

func (r *httpResp) WriteHeader(status int) {
	r.status = status
	r.writer.WriteHeader(status)
}

func logReq(h func(w http.ResponseWriter, r *http.Request)) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		resp := &httpResp{w, 0}
		h(resp, r)
		fmt.Printf("%v %v - %v\n", r.Method, r.RequestURI, resp.status)
	}
}

func main() {
	if err := setupIPTables(); err != nil {
		log.Fatal(err)
	}

	if err := initCrypto(); err != nil {
		log.Fatal(err)
	}

	r := mux.NewRouter()
	r.HandleFunc("/containers/", logReq(handleRegisterContainer)).Methods("POST")
	r.HandleFunc("/containers/{uid}/{app:.*}", logReq(handleRegisterApp)).Methods("PUT")

	acRtr := r.Headers("Metadata-Flavor", "AppContainer header").
		PathPrefix("/acMetadata/v1").Subrouter()

	mr := acRtr.Methods("GET").Subrouter()

	mr.HandleFunc("/container/annotations/", logReq(containerGet(handleContainerAnnotations)))
	mr.HandleFunc("/container/annotations/{name}", logReq(containerGet(handleContainerAnnotation)))
	mr.HandleFunc("/container/manifest", logReq(containerGet(handleContainerManifest)))
	mr.HandleFunc("/container/uid", logReq(containerGet(handleContainerUID)))

	mr.HandleFunc("/apps/{app:.*}/annotations/", logReq(appGet(handleAppAnnotations)))
	mr.HandleFunc("/apps/{app:.*}/annotations/{name}", logReq(appGet(handleAppAnnotation)))
	mr.HandleFunc("/apps/{app:.*}/image/manifest", logReq(appGet(handleAppManifest)))
	mr.HandleFunc("/apps/{app:.*}/image/id", logReq(appGet(handleAppID)))

	acRtr.HandleFunc("/container/hmac/sign", logReq(handleContainerSign)).Methods("POST")
	acRtr.HandleFunc("/container/hmac/verify", logReq(handleContainerVerify)).Methods("POST")

	log.Fatal(http.ListenAndServe(":4444", r))
}
<|endoftext|>