<commit_before>package shardkv

import (
  "container/list"
  "sync"
  "time"
  "shardmaster"
  "strconv"
  "labix.org/v2/mgo"
  "labix.org/v2/mgo/bson"
  "fmt"
)

//
// A simple LRU cache implementation which supports
// basic put and get operations.
//
type ShardCache struct {
  mu sync.Mutex

  // recent usage list and lookup table
  list *list.List // tracks recent usage
  table map[int]map[string]*list.Element // table for fast lookup

  // approximation of current size
  size uint64

  // cache size limitation
  capacity uint64
}

type entry struct {
  key string
  value string
  time_accessed time.Time
}

func MakeCache(capacity uint64) *ShardCache {
  cache := &ShardCache{
    list:     list.New(),
    table:    make(map[int](map[string]*list.Element)),
    capacity: capacity,
  }
  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
  return cache
}

func (cache *ShardCache) Size() uint64 {
  return cache.size
}

func (cache *ShardCache) Get(key string) (v string, ok bool) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return "", false
  }
  cache.touch(element)
  return element.Value.(*entry).value, true
}

func (cache *ShardCache) Put(key string, value string) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element != nil {
    cache.update(element, value)
  } else {
    cache.add(key, value)
  }
}

func (cache *ShardCache) Delete(key string) bool {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return false
  }

  cache.list.Remove(element)
  delete(cache.table[shard], key)
  cache.size -= uint64(len(element.Value.(*entry).value) + len(element.Value.(*entry).key))
  return true
}

func (cache *ShardCache) Clear() {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  // replace old state with initial state
  // old references are garbage collected
  cache.list.Init()
  cache.table = make(map[int]map[string]*list.Element)
  cache.size = 0

  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
}

func (cache *ShardCache) ClearShard(shard int) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  cache.table[shard] = make(map[string]*list.Element)
}

func (cache *ShardCache) Keys() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  keys := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    keys = append(keys, e.Value.(*entry).key)
  }
  return keys
}

func (cache *ShardCache) Values() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  values := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    values = append(values, e.Value.(*entry).value)
  }
  return values
}

func (cache *ShardCache) KVPairs() []KVPair {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  kvpairs := make([]KVPair, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    kvpairs = append(kvpairs, KVPair{key2shard(e.Value.(*entry).key), e.Value.(*entry).key, e.Value.(*entry).value})
  }
  return kvpairs
}

func (cache *ShardCache) update(element *list.Element, value string) {
  valueSize := len(value)
  sizeDiff := valueSize - len(element.Value.(*entry).value)
  element.Value.(*entry).value = value
  cache.size += uint64(sizeDiff)
  cache.touch(element)
  cache.maintainSize()
}

func (cache *ShardCache) touch(element *list.Element) {
  cache.list.MoveToFront(element)
  element.Value.(*entry).time_accessed = time.Now()
}

func (cache *ShardCache) add(key string, value string) {
  shard := key2shard(key)
  newEntry := &entry{key, value, time.Now()}
  element := cache.list.PushFront(newEntry)
  cache.table[shard][key] = element
  cache.size += uint64(len(newEntry.value) + len(newEntry.key))
  cache.maintainSize()
}

func (cache *ShardCache) maintainSize() {
  //deleted := list.New()
  for cache.size > cache.capacity {
    delElem := cache.list.Back()
    delEntry := delElem.Value.(*entry)
    shard := key2shard(delEntry.key)
    delete(cache.table[shard], delEntry.key)
    cache.list.Remove(delElem)
    cache.size -= uint64(len(delEntry.value) + len(delEntry.key))
  }
}

type Storage struct {
  mu sync.Mutex

  cache *ShardCache

  // disk storage state
  dbSession *mgo.Session
  db *mgo.Collection
  snapshots *mgo.Collection
  dedupsnaps *mgo.Collection

  // logging for background writes
  writeLog map[int]WriteOp
  applied int
}

type WriteOp struct {
  shard int
  key string
  value string
  dbok bool
}

func (st *Storage) makeCache(capacity uint64) {
  st.cache = MakeCache(capacity)
}

func (st *Storage) connectToDiskDB(url string) {
  var err error
  st.dbSession, err = mgo.Dial(url)
  if err != nil {
    panic(err)
  }
  st.db = st.dbSession.DB("db").C("kvstore")
  st.snapshots = st.dbSession.DB("db").C("snapshots")
}

func (st *Storage) DBClear() {
  st.db.RemoveAll(bson.M{})
  st.snapshots.RemoveAll(bson.M{})
  st.dedupsnaps.RemoveAll(bson.M{})
}

func (st *Storage) CacheClear() {
  st.cache.Clear()
}

func (st *Storage) Clear() {
  st.DBClear()
  st.CacheClear()
}

func MakeStorage(capacity uint64, dbURL string) *Storage {
  st := new(Storage)
  st.makeCache(capacity)
  st.connectToDiskDB(dbURL)
  st.writeLog = make(map[int]WriteOp)

  fmt.Printf("Making storage...\n")

  go st.writeInBackground()
  return st
}

func (st *Storage) Get(key string, shardNum int) string {

  value, ok := st.cache.Get(key)
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      ok = false
    } else {
      value = result.Value
    }
  }
  return value
}

func (st *Storage) Put(key string, value string, doHash bool, shardNum int) string {
  prev, ok := st.cache.Get(key)
  var dbok bool
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      dbok = false
    } else {
      prev = result.Value
      dbok = true
    }
  }

  //var deleted *list.List // list of cache values removed to send to disk
  if doHash {
    toBeHashed := prev + value
    hash := strconv.Itoa(int(hash(toBeHashed)))
    st.cache.Put(key, hash)
    //deleted = st.cache.Put(key, hash)
  } else {
    st.cache.Put(key, value)
    //deleted = st.cache.Put(key, value)
  }

  st.writeLog[st.applied] = WriteOp{shardNum, key, value, dbok}
  st.applied++

  // insert removed cache entries one at a time into DB, possibly faster if done together?
  /*for e := st.cache.list.Front(); e != nil; e = e.Next() {
    entry := e.(*entry)
    shard := key2shard(entry.key)
    err := st.db.Insert(&KVPair{shard, entry.key, entry.value})
    if err != nil {
      panic(err)
    }
  }*/
  return prev
}

func (st *Storage) CreateSnapshot(confignum int, dedup map[string]ClientReply) {
  cachedata := st.cache.KVPairs()
  results := []KVPair{}
  index := 0
  for len(results) >= 100 {
    st.db.Find(bson.M{}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
    for i := 0; i < len(results); i++ {
      st.db.Insert(&SnapshotKV{confignum, results[i].Shard, results[i].Key, results[i].Value, false})
    }
    index++
  }
  for i := 0; i < len(cachedata); i++ {
    st.db.Insert(&SnapshotKV{confignum, cachedata[i].Shard, cachedata[i].Key, cachedata[i].Value, true})
  }
  for key, value := range dedup {
    st.dedupsnaps.Insert(&SnapshotDedup{confignum, key, value.Value, value.Err, value.Counter})
  }
}

func (st *Storage) ReadSnapshotDB(confignum int, shardnum int, index int, cache bool) map[string]string, bool {
  piece := make(map[string]string)
  results := []SnapshotKV{}
  if cache {
    st.snapshots.Find(bson.M{"cache": "true", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  } else {
    st.snapshots.Find(bson.M{"cache": "false", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  }
  for i := 0; i < len(results); i++ {
    piece[results[i].Key] = results[i].Value
  }
  if len(results) < GrabSize {
    return piece, true
  }
  return piece, false
}

func (st *Storage) ReadSnapshotDedup(confignum int) map[string]ClientReply {
  dedup := make(map[string]ClientReply)
  results := []SnapshotDedup{}
  st.dedupsnaps.Find(bson.M{"config": confignum}).All(&results)
  for i := 0; i < len(results); i++ {
    dedup[results[i].Key] = ClientReply{Value: results[i].Value, results[i].Err, results[i].Counter}
  }
  return results
}

func (st *Storage) writeInBackground() {
  current := 0
  for {
    if st.applied > current {
      currentWrite := st.writeLog[current]
      var err error
      if !currentWrite.dbok {
        err = st.db.Insert(&KVPair{currentWrite.shard, currentWrite.key, currentWrite.value})
      } else {
        err = st.db.Update(bson.M{"shard": currentWrite.shard, "key": currentWrite.key}, bson.M{"$set": bson.M{"value": currentWrite.value}})
      }
      if err != nil {
        panic(err)
      }
      delete(st.writeLog, current)
      current++
    }
    time.Sleep(25 * time.Millisecond)
  }
}

func (st *Storage) closeDBConnection() {
  st.dbSession.Close()
}
<commit_msg>Fixed syntax errors and Mongo clearing bug in storage.go<commit_after>package shardkv

import (
  "container/list"
  "sync"
  "time"
  "shardmaster"
  "strconv"
  "labix.org/v2/mgo"
  "labix.org/v2/mgo/bson"
  "fmt"
)

//
// A simple LRU cache implementation which supports
// basic put and get operations.
//
type ShardCache struct {
  mu sync.Mutex

  // recent usage list and lookup table
  list *list.List // tracks recent usage
  table map[int]map[string]*list.Element // table for fast lookup

  // approximation of current size
  size uint64

  // cache size limitation
  capacity uint64
}

type entry struct {
  key string
  value string
  time_accessed time.Time
}

func MakeCache(capacity uint64) *ShardCache {
  cache := &ShardCache{
    list:     list.New(),
    table:    make(map[int](map[string]*list.Element)),
    capacity: capacity,
  }
  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
  return cache
}

func (cache *ShardCache) Size() uint64 {
  return cache.size
}

func (cache *ShardCache) Get(key string) (v string, ok bool) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return "", false
  }
  cache.touch(element)
  return element.Value.(*entry).value, true
}

func (cache *ShardCache) Put(key string, value string) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element != nil {
    cache.update(element, value)
  } else {
    cache.add(key, value)
  }
}

func (cache *ShardCache) Delete(key string) bool {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return false
  }

  cache.list.Remove(element)
  delete(cache.table[shard], key)
  cache.size -= uint64(len(element.Value.(*entry).value) + len(element.Value.(*entry).key))
  return true
}

func (cache *ShardCache) Clear() {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  // replace old state with initial state
  // old references are garbage collected
  cache.list.Init()
  cache.table = make(map[int]map[string]*list.Element)
  cache.size = 0

  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
}

func (cache *ShardCache) ClearShard(shard int) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  cache.table[shard] = make(map[string]*list.Element)
}

func (cache *ShardCache) Keys() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  keys := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    keys = append(keys, e.Value.(*entry).key)
  }
  return keys
}

func (cache *ShardCache) Values() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  values := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    values = append(values, e.Value.(*entry).value)
  }
  return values
}

func (cache *ShardCache) KVPairs() []KVPair {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  kvpairs := make([]KVPair, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    kvpairs = append(kvpairs, KVPair{key2shard(e.Value.(*entry).key), e.Value.(*entry).key, e.Value.(*entry).value})
  }
  return kvpairs
}

func (cache *ShardCache) update(element *list.Element, value string) {
  valueSize := len(value)
  sizeDiff := valueSize - len(element.Value.(*entry).value)
  element.Value.(*entry).value = value
  cache.size += uint64(sizeDiff)
  cache.touch(element)
  cache.maintainSize()
}

func (cache *ShardCache) touch(element *list.Element) {
  cache.list.MoveToFront(element)
  element.Value.(*entry).time_accessed = time.Now()
}

func (cache *ShardCache) add(key string, value string) {
  shard := key2shard(key)
  newEntry := &entry{key, value, time.Now()}
  element := cache.list.PushFront(newEntry)
  cache.table[shard][key] = element
  cache.size += uint64(len(newEntry.value) + len(newEntry.key))
  cache.maintainSize()
}

func (cache *ShardCache) maintainSize() {
  //deleted := list.New()
  for cache.size > cache.capacity {
    delElem := cache.list.Back()
    delEntry := delElem.Value.(*entry)
    shard := key2shard(delEntry.key)
    delete(cache.table[shard], delEntry.key)
    cache.list.Remove(delElem)
    cache.size -= uint64(len(delEntry.value) + len(delEntry.key))
  }
}

type Storage struct {
  mu sync.Mutex

  cache *ShardCache

  // disk storage state
  dbSession *mgo.Session
  db *mgo.Collection
  snapshots *mgo.Collection
  dedupsnaps *mgo.Collection

  // logging for background writes
  writeLog map[int]WriteOp
  applied int
}

type WriteOp struct {
  shard int
  key string
  value string
  dbok bool
}

func (st *Storage) makeCache(capacity uint64) {
  st.cache = MakeCache(capacity)
}

func (st *Storage) connectToDiskDB(url string) {
  var err error
  st.dbSession, err = mgo.Dial(url)
  if err != nil {
    panic(err)
  }
  st.db = st.dbSession.DB("db").C("kvstore")
  st.snapshots = st.dbSession.DB("db").C("snapshots")
}

func (st *Storage) DBClear() {
  st.db.RemoveAll(bson.M{})
  check, _ := st.snapshots.Find(bson.M{}).Count()
  if (check > 0) {
    st.snapshots.RemoveAll(bson.M{})
    st.dedupsnaps.RemoveAll(bson.M{})
  }
}

func (st *Storage) CacheClear() {
  st.cache.Clear()
}

func (st *Storage) Clear() {
  st.DBClear()
  st.CacheClear()
}

func MakeStorage(capacity uint64, dbURL string) *Storage {
  st := new(Storage)
  st.makeCache(capacity)
  st.connectToDiskDB(dbURL)
  st.writeLog = make(map[int]WriteOp)

  fmt.Printf("Making storage...\n")

  go st.writeInBackground()
  return st
}

func (st *Storage) Get(key string, shardNum int) string {

  value, ok := st.cache.Get(key)
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      ok = false
    } else {
      value = result.Value
    }
  }
  return value
}

func (st *Storage) Put(key string, value string, doHash bool, shardNum int) string {
  prev, ok := st.cache.Get(key)
  var dbok bool
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      dbok = false
    } else {
      prev = result.Value
      dbok = true
    }
  }

  //var deleted *list.List // list of cache values removed to send to disk
  if doHash {
    toBeHashed := prev + value
    hash := strconv.Itoa(int(hash(toBeHashed)))
    st.cache.Put(key, hash)
    //deleted = st.cache.Put(key, hash)
  } else {
    st.cache.Put(key, value)
    //deleted = st.cache.Put(key, value)
  }

  st.writeLog[st.applied] = WriteOp{shardNum, key, value, dbok}
  st.applied++

  // insert removed cache entries one at a time into DB, possibly faster if done together?
  /*for e := st.cache.list.Front(); e != nil; e = e.Next() {
    entry := e.(*entry)
    shard := key2shard(entry.key)
    err := st.db.Insert(&KVPair{shard, entry.key, entry.value})
    if err != nil {
      panic(err)
    }
  }*/
  return prev
}

func (st *Storage) CreateSnapshot(confignum int, dedup map[string]ClientReply) {
  cachedata := st.cache.KVPairs()
  results := []KVPair{}
  index := 0
  for len(results) >= 100 {
    st.db.Find(bson.M{}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
    for i := 0; i < len(results); i++ {
      st.db.Insert(&SnapshotKV{confignum, results[i].Shard, results[i].Key, results[i].Value, false})
    }
    index++
  }
  for i := 0; i < len(cachedata); i++ {
    st.db.Insert(&SnapshotKV{confignum, cachedata[i].Shard, cachedata[i].Key, cachedata[i].Value, true})
  }
  for key, value := range dedup {
    st.dedupsnaps.Insert(&SnapshotDedup{confignum, key, value.Value, value.Err, value.Counter})
  }
}

func (st *Storage) ReadSnapshotDB(confignum int, shardnum int, index int, cache bool) (p map[string]string, fin bool) {
  piece := make(map[string]string)
  results := []SnapshotKV{}
  if cache {
    st.snapshots.Find(bson.M{"cache": "true", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  } else {
    st.snapshots.Find(bson.M{"cache": "false", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  }
  for i := 0; i < len(results); i++ {
    piece[results[i].Key] = results[i].Value
  }
  if len(results) < GrabSize {
    return piece, true
  }
  return piece, false
}

func (st *Storage) ReadSnapshotDedup(confignum int) map[string]ClientReply {
  dedup := make(map[string]ClientReply)
  results := []SnapshotDedup{}
  st.dedupsnaps.Find(bson.M{"config": confignum}).All(&results)
  for i := 0; i < len(results); i++ {
    dedup[results[i].Key] = ClientReply{results[i].Value, results[i].Err, results[i].Counter}
  }
  return dedup
}

func (st *Storage) writeInBackground() {
  current := 0
  for {
    if st.applied > current {
      currentWrite := st.writeLog[current]
      var err error
      if !currentWrite.dbok {
        err = st.db.Insert(&KVPair{currentWrite.shard, currentWrite.key, currentWrite.value})
      } else {
        err = st.db.Update(bson.M{"shard": currentWrite.shard, "key": currentWrite.key}, bson.M{"$set": bson.M{"value": currentWrite.value}})
      }
      if err != nil {
        panic(err)
      }
      delete(st.writeLog, current)
      current++
    }
    time.Sleep(25 * time.Millisecond)
  }
}

func (st *Storage) closeDBConnection() {
  st.dbSession.Close()
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/google/go-github/github"
	"github.com/spf13/cobra"
	"golang.org/x/oauth2"
)

type tokenCounterFlags struct {
	influx InfluxConfig
	tokens []string
}

func (flags *tokenCounterFlags) AddFlags(cmd *cobra.Command) {
	cmd.Flags().StringSliceVar(&flags.tokens, "token", []string{}, "List of tokens")
	cmd.Flags().AddGoFlagSet(flag.CommandLine)
}

// TokenHandler is refreshing token usage
type TokenHandler struct {
	gClient  *github.Client
	influxdb *InfluxDB
	login    string
}

// GetGithubClient creates a client for each token
func GetGithubClient(token string) *github.Client {
	return github.NewClient(
		oauth2.NewClient(
			oauth2.NoContext,
			oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token}),
		),
	)
}

// GetUsername finds the login for each token
func GetUsername(client *github.Client) (string, error) {
	user, _, err := client.Users.Get("")
	if err != nil {
		return "", err
	}
	if user.Login == nil {
		return "", errors.New("Users.Get(\"\") returned empty login.")
	}

	return *user.Login, nil
}

// CreateTokenHandler parses the token and create a handler
func CreateTokenHandler(tokenStream io.Reader, influxdb *InfluxDB) (*TokenHandler, error) {
	token, err := ioutil.ReadAll(tokenStream)
	if err != nil {
		return nil, err
	}
	client := GetGithubClient(strings.TrimSpace(string(token)))
	login, err := GetUsername(client) // Get user name for token
	if err != nil {
		return nil, err
	}

	return &TokenHandler{
		gClient:  client,
		login:    login,
		influxdb: influxdb,
	}, nil
}

// CreateTokenHandlers goes through the list of token files, and create handlers
func CreateTokenHandlers(tokenFiles []string, influxdb *InfluxDB) ([]TokenHandler, error) {
	tokens := []TokenHandler{}
	for _, tokenFile := range tokenFiles {
		f, err := os.Open(tokenFile)
		if err != nil {
			return nil, fmt.Errorf("Can't open token-file (%s): %s", tokenFile, err)
		}
		token, err := CreateTokenHandler(f, influxdb)
		if err != nil {
			return nil, fmt.Errorf("Failed to create token (%s): %s", tokenFile, err)
		}
		tokens = append(tokens, *token)
	}
	return tokens, nil
}

func (t TokenHandler) getCoreRate() (*github.Rate, error) {
	limits, _, err := t.gClient.RateLimits()
	if err != nil {
		return nil, err
	}
	return limits.Core, nil
}

// Process does the main job:
// It tries to get the value of "Remaining" rate just before the token
// gets reset. It does that more and more often (as the reset date gets
// closer) to get the most accurate value.
func (t TokenHandler) Process() {
	lastRate, err := t.getCoreRate()
	if err != nil {
		glog.Fatalf("%s: Couldn't get rate limits:", t.login, err)
	}

	for {
		halfPeriod := lastRate.Reset.Time.Sub(time.Now()) / 2
		glog.Infof("%s: Current rate: %s. Sleeping for %s.", t.login, lastRate, halfPeriod)
		time.Sleep(halfPeriod)
		newRate, err := t.getCoreRate()
		if err != nil {
			glog.Error("Failed to get CoreRate: ", err)
			continue
		}
		if !newRate.Reset.Time.Equal(lastRate.Reset.Time) {
			glog.Infof(
				"%s: ### TOKEN USAGE: %d",
				t.login,
				lastRate.Limit-lastRate.Remaining,
			)
			t.influxdb.Push(
				"github_token_count",
				map[string]string{"login": t.login},
				map[string]interface{}{"value": lastRate.Limit - lastRate.Remaining},
				lastRate.Reset.Time,
			)
		}
		lastRate = newRate
	}
}

func runProgram(flags *tokenCounterFlags) error {
	influxdb, err := flags.influx.CreateDatabaseClient()
	if err != nil {
		return err
	}

	tokens, err := CreateTokenHandlers(flags.tokens, influxdb)
	if err != nil {
		return err
	}

	if len(tokens) == 0 {
		glog.Warning("No token given, nothing to do. Leaving...")
		return nil
	}

	for _, token := range tokens {
		glog.Infof("Processing token for '%s'", token.login)
		go token.Process()
	}

	select {}

	return nil
}

func main() {
	flags := &tokenCounterFlags{}
	cmd := &cobra.Command{
		Use:   filepath.Base(os.Args[0]),
		Short: "Count usage of github token",
		RunE: func(_ *cobra.Command, _ []string) error {
			return runProgram(flags)
		},
	}
	flags.AddFlags(cmd)
	flags.influx.AddFlags(cmd)

	if err := cmd.Execute(); err != nil {
		glog.Error(err)
	}
}
<commit_msg>velodrome: Reset token-counter proactively<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/google/go-github/github"
	"github.com/spf13/cobra"
	"golang.org/x/oauth2"
)

type tokenCounterFlags struct {
	influx InfluxConfig
	tokens []string
}

func (flags *tokenCounterFlags) AddFlags(cmd *cobra.Command) {
	cmd.Flags().StringSliceVar(&flags.tokens, "token", []string{}, "List of tokens")
	cmd.Flags().AddGoFlagSet(flag.CommandLine)
}

// TokenHandler is refreshing token usage
type TokenHandler struct {
	gClient  *github.Client
	influxdb *InfluxDB
	login    string
}

// GetGithubClient creates a client for each token
func GetGithubClient(token string) *github.Client {
	return github.NewClient(
		oauth2.NewClient(
			oauth2.NoContext,
			oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token}),
		),
	)
}

// GetUsername finds the login for each token
func GetUsername(client *github.Client) (string, error) {
	user, _, err := client.Users.Get("")
	if err != nil {
		return "", err
	}
	if user.Login == nil {
		return "", errors.New("Users.Get(\"\") returned empty login.")
	}

	return *user.Login, nil
}

// CreateTokenHandler parses the token and create a handler
func CreateTokenHandler(tokenStream io.Reader, influxdb *InfluxDB) (*TokenHandler, error) {
	token, err := ioutil.ReadAll(tokenStream)
	if err != nil {
		return nil, err
	}
	client := GetGithubClient(strings.TrimSpace(string(token)))
	login, err := GetUsername(client) // Get user name for token
	if err != nil {
		return nil, err
	}

	return &TokenHandler{
		gClient:  client,
		login:    login,
		influxdb: influxdb,
	}, nil
}

// CreateTokenHandlers goes through the list of token files, and create handlers
func CreateTokenHandlers(tokenFiles []string, influxdb *InfluxDB) ([]TokenHandler, error) {
	tokens := []TokenHandler{}
	for _, tokenFile := range tokenFiles {
		f, err := os.Open(tokenFile)
		if err != nil {
			return nil, fmt.Errorf("Can't open token-file (%s): %s", tokenFile, err)
		}
		token, err := CreateTokenHandler(f, influxdb)
		if err != nil {
			return nil, fmt.Errorf("Failed to create token (%s): %s", tokenFile, err)
		}
		tokens = append(tokens, *token)
	}
	return tokens, nil
}

func (t TokenHandler) getCoreRate() (*github.Rate, error) {
	limits, _, err := t.gClient.RateLimits()
	if err != nil {
		return nil, err
	}
	return limits.Core, nil
}

// Process does the main job:
// It tries to get the value of "Remaining" rate just before the token
// gets reset. It does that more and more often (as the reset date gets
// closer) to get the most accurate value.
func (t TokenHandler) Process() {
	lastRate, err := t.getCoreRate()
	if err != nil {
		glog.Fatalf("%s: Couldn't get rate limits:", t.login, err)
	}

	for {
		halfPeriod := lastRate.Reset.Time.Sub(time.Now()) / 2
		glog.Infof("%s: Current rate: %s. Sleeping for %s.", t.login, lastRate, halfPeriod)
		time.Sleep(halfPeriod)
		newRate, err := t.getCoreRate()
		if err != nil {
			glog.Error("Failed to get CoreRate: ", err)
			continue
		}
		// There is a bug in Github. They seem to reset the Remaining value before reseting the Reset value.
		if !newRate.Reset.Time.Equal(lastRate.Reset.Time) || newRate.Remaining > lastRate.Remaining {
			glog.Infof(
				"%s: ### TOKEN USAGE: %d",
				t.login,
				lastRate.Limit-lastRate.Remaining,
			)
			t.influxdb.Push(
				"github_token_count",
				map[string]string{"login": t.login},
				map[string]interface{}{"value": lastRate.Limit - lastRate.Remaining},
				lastRate.Reset.Time,
			)
			// Make sure the timer is properly reset, and we have time anyway
			time.Sleep(30 * time.Minute)
			for {
				newRate, err = t.getCoreRate()
				if err == nil {
					break
				}
				glog.Error("Failed to get CoreRate: ", err)
				time.Sleep(time.Minute)
			}

		}
		lastRate = newRate
	}
}

func runProgram(flags *tokenCounterFlags) error {
	influxdb, err := flags.influx.CreateDatabaseClient()
	if err != nil {
		return err
	}

	tokens, err := CreateTokenHandlers(flags.tokens, influxdb)
	if err != nil {
		return err
	}

	if len(tokens) == 0 {
		glog.Warning("No token given, nothing to do. Leaving...")
		return nil
	}

	for _, token := range tokens {
		glog.Infof("Processing token for '%s'", token.login)
		go token.Process()
	}

	select {}

	return nil
}

func main() {
	flags := &tokenCounterFlags{}
	cmd := &cobra.Command{
		Use:   filepath.Base(os.Args[0]),
		Short: "Count usage of github token",
		RunE: func(_ *cobra.Command, _ []string) error {
			return runProgram(flags)
		},
	}
	flags.AddFlags(cmd)
	flags.influx.AddFlags(cmd)

	if err := cmd.Execute(); err != nil {
		glog.Error(err)
	}
}
<|endoftext|>


<commit_before>package irc

import (
	"database/sql"
	//"fmt"
	"bufio"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"path/filepath"
	"strings"
)

type RowId uint64

type Queryable interface {
	Exec(string, ...interface{}) (sql.Result, error)
	Query(string, ...interface{}) (*sql.Rows, error)
	QueryRow(string, ...interface{}) *sql.Row
}

type Savable interface {
	Save(q Queryable) bool
}

//
// general
//

func NewDatabase() (db *sql.DB) {
	db, err := sql.Open("sqlite3", "ergonomadic.db")
	if err != nil {
		log.Fatalln("cannot open database")
	}
	return
}

func readLines(filename string) <-chan string {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalln(err)
	}
	reader := bufio.NewReader(file)
	lines := make(chan string)
	go func(lines chan<- string) {
		defer file.Close()
		defer close(lines)
		for {
			line, err := reader.ReadString(';')
			if err != nil {
				break
			}
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			lines <- line
		}
	}(lines)
	return lines
}

func ExecSqlFile(db *sql.DB, filename string) {
	Transact(db, func(q Queryable) bool {
		for line := range readLines(filepath.Join("sql", filename)) {
			log.Println(line)
			_, err := q.Exec(line)
			if err != nil {
				log.Fatalln(err)
			}
		}
		return true
	})
}

func Transact(db *sql.DB, txf func(Queryable) bool) {
	tx, err := db.Begin()
	if err != nil {
		log.Panicln(err)
	}
	if txf(tx) {
		tx.Commit()
	} else {
		tx.Rollback()
	}
}

func Save(db *sql.DB, s Savable) {
	Transact(db, s.Save)
}

//
// general purpose sql
//

func FindId(q Queryable, sql string, args ...interface{}) (rowId RowId, err error) {
	row := q.QueryRow(sql, args...)
	err = row.Scan(&rowId)
	return
}

func Count(q Queryable, sql string, args ...interface{}) (count uint, err error) {
	row := q.QueryRow(sql, args...)
	err = row.Scan(&count)
	return
}

//
// data
//

type UserRow struct {
	id   RowId
	nick string
	hash []byte
}

type ChannelRow struct {
	id   RowId
	name string
}

// user

func FindAllUsers(q Queryable) (urs []*UserRow, err error) {
	var rows *sql.Rows
	rows, err = q.Query("SELECT id, nick, hash FROM user")
	if err != nil {
		return
	}
	urs = make([]*UserRow, 0)
	for rows.Next() {
		ur := &UserRow{}
		err = rows.Scan(&(ur.id), &(ur.nick), &(ur.hash))
		if err != nil {
			return
		}
		urs = append(urs, ur)
	}
	return
}

func FindUserByNick(q Queryable, nick string) (ur *UserRow, err error) {
	ur = &UserRow{}
	row := q.QueryRow("SELECT id, nick, hash FROM user LIMIT 1 WHERE nick = ?",
		nick)
	err = row.Scan(&(ur.id), &(ur.nick), &(ur.hash))
	return
}

func FindUserIdByNick(q Queryable, nick string) (RowId, error) {
	return FindId(q, "SELECT id FROM user WHERE nick = ?", nick)
}

func FindChannelByName(q Queryable, name string) (cr *ChannelRow) {
	cr = new(ChannelRow)
	row := q.QueryRow("SELECT id, name FROM channel LIMIT 1 WHERE name = ?", name)
	err := row.Scan(&(cr.id), &(cr.name))
	if err != nil {
		cr = nil
	}
	return
}

func InsertUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("INSERT INTO user (nick, hash) VALUES (?, ?)",
		user.nick, user.hash)
	return
}

func UpdateUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("UPDATE user SET nick = ?, hash = ? WHERE id = ?",
		user.nick, user.hash, *(user.id))
	return
}

func DeleteUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("DELETE FROM user WHERE id = ?", *(user.id))
	return
}

// user-channel

func DeleteAllUserChannels(q Queryable, rowId RowId) (err error) {
	_, err = q.Exec("DELETE FROM user_channel WHERE user_id = ?", rowId)
	return
}

func DeleteOtherUserChannels(q Queryable, userId RowId, channelIds []RowId) (err error) {
	_, err = q.Exec(`DELETE FROM user_channel WHERE
user_id = ? AND channel_id NOT IN ?`, userId, channelIds)
	return
}

func InsertUserChannels(q Queryable, userId RowId, channelIds []RowId) (err error) {
	ins := "INSERT OR IGNORE INTO user_channel (user_id, channel_id) VALUES "
	vals := strings.Repeat("(?, ?), ", len(channelIds))
	vals = vals[0 : len(vals)-2]
	args := make([]RowId, 2*len(channelIds))
	var i = 0
	for _, channelId := range channelIds {
		args[i] = userId
		args[i+1] = channelId
		i += 2
	}
	_, err = q.Exec(ins+vals, args)
	return
}

// channel

func FindChannelIdByName(q Queryable, name string) (RowId, error) {
	return FindId(q, "SELECT id FROM channel WHERE name = ?", name)
}

func FindChannelsForUser(q Queryable, userId RowId) (crs []*ChannelRow, err error) {
	query := ` FROM channel WHERE id IN
(SELECT channel_id from user_channel WHERE user_id = ?)`
	count, err := Count(q, "SELECT COUNT(id)"+query, userId)
	if err != nil {
		return
	}
	rows, err := q.Query("SELECT id, name"+query, userId)
	if err != nil {
		return
	}
	crs = make([]*ChannelRow, count)
	var i = 0
	for rows.Next() {
		cr := &ChannelRow{}
		err = rows.Scan(&(cr.id), &(cr.name))
		if err != nil {
			return
		}
		crs[i] = cr
		i++
	}
	return
}

func InsertChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("INSERT INTO channel (name) VALUES (?)", channel.name)
	return
}

func UpdateChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("UPDATE channel SET name = ? WHERE id = ?",
		channel.name, *(channel.id))
	return
}

func DeleteChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("DELETE FROM channel WHERE id = ?", *(channel.id))
	return
}
<commit_msg>Make some helper functions private.<commit_after>package irc

import (
	"database/sql"
	//"fmt"
	"bufio"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"path/filepath"
	"strings"
)

type RowId uint64

type Queryable interface {
	Exec(string, ...interface{}) (sql.Result, error)
	Query(string, ...interface{}) (*sql.Rows, error)
	QueryRow(string, ...interface{}) *sql.Row
}

type Savable interface {
	Save(q Queryable) bool
}

//
// general
//

func NewDatabase() (db *sql.DB) {
	db, err := sql.Open("sqlite3", "ergonomadic.db")
	if err != nil {
		log.Fatalln("cannot open database")
	}
	return
}

func readLines(filename string) <-chan string {
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalln(err)
	}
	reader := bufio.NewReader(file)
	lines := make(chan string)
	go func(lines chan<- string) {
		defer file.Close()
		defer close(lines)
		for {
			line, err := reader.ReadString(';')
			if err != nil {
				break
			}
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			lines <- line
		}
	}(lines)
	return lines
}

func ExecSqlFile(db *sql.DB, filename string) {
	Transact(db, func(q Queryable) bool {
		for line := range readLines(filepath.Join("sql", filename)) {
			log.Println(line)
			_, err := q.Exec(line)
			if err != nil {
				log.Fatalln(err)
			}
		}
		return true
	})
}

func Transact(db *sql.DB, txf func(Queryable) bool) {
	tx, err := db.Begin()
	if err != nil {
		log.Panicln(err)
	}
	if txf(tx) {
		tx.Commit()
	} else {
		tx.Rollback()
	}
}

func Save(db *sql.DB, s Savable) {
	Transact(db, s.Save)
}

//
// general purpose sql
//

func findId(q Queryable, sql string, args ...interface{}) (rowId RowId, err error) {
	row := q.QueryRow(sql, args...)
	err = row.Scan(&rowId)
	return
}

func Count(q Queryable, sql string, args ...interface{}) (count uint, err error) {
	row := q.QueryRow(sql, args...)
	err = row.Scan(&count)
	return
}

//
// data
//

type UserRow struct {
	id   RowId
	nick string
	hash []byte
}

type ChannelRow struct {
	id   RowId
	name string
}

// user

func FindAllUsers(q Queryable) (urs []*UserRow, err error) {
	var rows *sql.Rows
	rows, err = q.Query("SELECT id, nick, hash FROM user")
	if err != nil {
		return
	}
	urs = make([]*UserRow, 0)
	for rows.Next() {
		ur := &UserRow{}
		err = rows.Scan(&(ur.id), &(ur.nick), &(ur.hash))
		if err != nil {
			return
		}
		urs = append(urs, ur)
	}
	return
}

func FindUserByNick(q Queryable, nick string) (ur *UserRow, err error) {
	ur = &UserRow{}
	row := q.QueryRow("SELECT id, nick, hash FROM user LIMIT 1 WHERE nick = ?",
		nick)
	err = row.Scan(&(ur.id), &(ur.nick), &(ur.hash))
	return
}

func FindUserIdByNick(q Queryable, nick string) (RowId, error) {
	return findId(q, "SELECT id FROM user WHERE nick = ?", nick)
}

func FindChannelByName(q Queryable, name string) (cr *ChannelRow) {
	cr = new(ChannelRow)
	row := q.QueryRow("SELECT id, name FROM channel LIMIT 1 WHERE name = ?", name)
	err := row.Scan(&(cr.id), &(cr.name))
	if err != nil {
		cr = nil
	}
	return
}

func InsertUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("INSERT INTO user (nick, hash) VALUES (?, ?)",
		user.nick, user.hash)
	return
}

func UpdateUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("UPDATE user SET nick = ?, hash = ? WHERE id = ?",
		user.nick, user.hash, *(user.id))
	return
}

func DeleteUser(q Queryable, user *User) (err error) {
	_, err = q.Exec("DELETE FROM user WHERE id = ?", *(user.id))
	return
}

// user-channel

func DeleteAllUserChannels(q Queryable, rowId RowId) (err error) {
	_, err = q.Exec("DELETE FROM user_channel WHERE user_id = ?", rowId)
	return
}

func DeleteOtherUserChannels(q Queryable, userId RowId, channelIds []RowId) (err error) {
	_, err = q.Exec(`DELETE FROM user_channel WHERE
user_id = ? AND channel_id NOT IN ?`, userId, channelIds)
	return
}

func InsertUserChannels(q Queryable, userId RowId, channelIds []RowId) (err error) {
	ins := "INSERT OR IGNORE INTO user_channel (user_id, channel_id) VALUES "
	vals := strings.Repeat("(?, ?), ", len(channelIds))
	vals = vals[0 : len(vals)-2]
	args := make([]RowId, 2*len(channelIds))
	for i, channelId := range channelIds {
		args[i] = userId
		args[i+1] = channelId
	}
	_, err = q.Exec(ins+vals, args)
	return
}

// channel

func FindChannelIdByName(q Queryable, name string) (RowId, error) {
	return findId(q, "SELECT id FROM channel WHERE name = ?", name)
}

func FindChannelsForUser(q Queryable, userId RowId) (crs []*ChannelRow, err error) {
	query := ` FROM channel WHERE id IN
(SELECT channel_id from user_channel WHERE user_id = ?)`
	count, err := Count(q, "SELECT COUNT(id)"+query, userId)
	if err != nil {
		return
	}
	rows, err := q.Query("SELECT id, name"+query, userId)
	if err != nil {
		return
	}
	crs = make([]*ChannelRow, count)
	var i = 0
	for rows.Next() {
		cr := &ChannelRow{}
		err = rows.Scan(&(cr.id), &(cr.name))
		if err != nil {
			return
		}
		crs[i] = cr
		i++
	}
	return
}

func InsertChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("INSERT INTO channel (name) VALUES (?)", channel.name)
	return
}

func UpdateChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("UPDATE channel SET name = ? WHERE id = ?",
		channel.name, *(channel.id))
	return
}

func DeleteChannel(q Queryable, channel *Channel) (err error) {
	_, err = q.Exec("DELETE FROM channel WHERE id = ?", *(channel.id))
	return
}
<|endoftext|>


<commit_before>// Package pseudohsm provides a pseudo HSM for development environments.
package pseudohsm

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/pborman/uuid"

	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/errors"
	mnem "github.com/bytom/wallet/mnemonic"
)

// pre-define errors for supporting bytom errorFormatter
var (
	ErrDuplicateKeyAlias = errors.New("duplicate key alias")
	ErrLoadKey           = errors.New("key not found or wrong password ")
	ErrDecrypt           = errors.New("could not decrypt key with given passphrase")
)

// HSM type for storing pubkey and privatekey
type HSM struct {
	cacheMu  sync.Mutex
	keyStore keyStore
	cache    *keyCache
	//kdCache  map[chainkd.XPub]chainkd.XPrv
}

// XPub type for pubkey for anyone can see
type XPub struct {
	Alias string       `json:"alias"`
	XPub  chainkd.XPub `json:"xpub"`
	File  string       `json:"file"`
}

// New method for HSM struct
func New(keypath string) (*HSM, error) {
	keydir, _ := filepath.Abs(keypath)
	return &HSM{
		keyStore: &keyStorePassphrase{keydir, LightScryptN, LightScryptP},
		cache:    newKeyCache(keydir),
		//kdCache:  make(map[chainkd.XPub]chainkd.XPrv),
	}, nil
}

// XCreate produces a new random xprv and stores it in the db.
func (h *HSM) XCreate(alias string, auth string, language string) (*XPub, *string, error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	normalizedAlias := strings.ToLower(strings.TrimSpace(alias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return nil, nil, ErrDuplicateKeyAlias
	}

	xpub, mnemonic, err := h.createChainKDKey(normalizedAlias, auth, language)
	if err != nil {
		return nil, nil, err
	}
	h.cache.add(*xpub)
	return xpub, mnemonic, err
}

// ImportKeyFromMnemonic produces a xprv from mnemonic and stores it in the db.
func (h *HSM) ImportKeyFromMnemonic(alias string, auth string, mnemonic string, language string) (*XPub, error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	normalizedAlias := strings.ToLower(strings.TrimSpace(alias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return nil, ErrDuplicateKeyAlias
	}

	// Pre validate that the mnemonic is well formed and only contains words that
	// are present in the word list
	if !mnem.IsMnemonicValid(mnemonic, language) {
		return nil, mnem.ErrInvalidMnemonic
	}

	xpub, err := h.createKeyFromMnemonic(alias, auth, mnemonic)
	if err != nil {
		return nil, err
	}

	h.cache.add(*xpub)
	return xpub, nil
}

func (h *HSM) createKeyFromMnemonic(alias string, auth string, mnemonic string) (*XPub, error) {
	// Generate a Bip32 HD wallet for the mnemonic and a user supplied password
	seed := mnem.NewSeed(mnemonic, "")
	xprv, xpub, err := chainkd.NewXKeys(bytes.NewBuffer(seed))
	if err != nil {
		return nil, err
	}
	id := uuid.NewRandom()
	key := &XKey{
		ID:      id,
		KeyType: "bytom_kd",
		XPub:    xpub,
		XPrv:    xprv,
		Alias:   alias,
	}
	file := h.keyStore.JoinPath(keyFileName(key.ID.String()))
	if err := h.keyStore.StoreKey(file, key, auth); err != nil {
		return nil, errors.Wrap(err, "storing keys")
	}
	return &XPub{XPub: xpub, Alias: alias, File: file}, nil
}

func (h *HSM) createChainKDKey(alias string, auth string, language string) (*XPub, *string, error) {
	// Generate a mnemonic for memorization or user-friendly seeds
	entropy, err := mnem.NewEntropy(256)
	if err != nil {
		return nil, nil, err
	}
	mnemonic, err := mnem.NewMnemonic(entropy, language)
	if err != nil {
		return nil, nil, err
	}
	xpub, err := h.createKeyFromMnemonic(alias, auth, mnemonic)
	if err != nil {
		return nil, nil, err
	}
	return xpub, &mnemonic, nil
}

// UpdateKeyAlias update key alias
func (h *HSM) UpdateKeyAlias(xpub chainkd.XPub, newAlias string) error {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	h.cache.maybeReload()
	h.cache.mu.Lock()
	xpb, err := h.cache.find(XPub{XPub: xpub})
	h.cache.mu.Unlock()
	if err != nil {
		return err
	}

	keyjson, err := ioutil.ReadFile(xpb.File)
	if err != nil {
		return err
	}

	encrptKeyJSON := new(encryptedKeyJSON)
	if err := json.Unmarshal(keyjson, encrptKeyJSON); err != nil {
		return err
	}

	normalizedAlias := strings.ToLower(strings.TrimSpace(newAlias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return ErrDuplicateKeyAlias
	}

	encrptKeyJSON.Alias = normalizedAlias
	keyJSON, err := json.Marshal(encrptKeyJSON)
	if err != nil {
		return err
	}

	if err := writeKeyFile(xpb.File, keyJSON); err != nil {
		return err
	}

	// update key alias
	h.cache.delete(xpb)
	xpb.Alias = normalizedAlias
	h.cache.add(xpb)

	return nil
}

// ListKeys returns a list of all xpubs from the store
func (h *HSM) ListKeys() []XPub {
	xpubs := h.cache.keys()
	return xpubs
}

// XSign looks up the xprv given the xpub, optionally derives a new
// xprv with the given path (but does not store the new xprv), and
// signs the given msg.
func (h *HSM) XSign(xpub chainkd.XPub, path [][]byte, msg []byte, auth string) ([]byte, error) {
	xprv, err := h.LoadChainKDKey(xpub, auth)
	if err != nil {
		return nil, err
	}
	if len(path) > 0 {
		xprv = xprv.Derive(path)
	}
	return xprv.Sign(msg), nil
}

//LoadChainKDKey get xprv from xpub
func (h *HSM) LoadChainKDKey(xpub chainkd.XPub, auth string) (xprv chainkd.XPrv, err error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	//if xprv, ok := h.kdCache[xpub]; ok {
	//	return xprv, nil
	//}

	_, xkey, err := h.loadDecryptedKey(xpub, auth)
	if err != nil {
		return xprv, ErrLoadKey
	}
	//h.kdCache[xpb.XPub] = xkey.XPrv
	return xkey.XPrv, nil
}

// XDelete deletes the key matched by xpub if the passphrase is correct.
// If a contains no filename, the address must match a unique key.
func (h *HSM) XDelete(xpub chainkd.XPub, auth string) error {
	// Decrypting the key isn't really necessary, but we do
	// it anyway to check the password and zero out the key
	// immediately afterwards.

	xpb, xkey, err := h.loadDecryptedKey(xpub, auth)
	if xkey != nil {
		zeroKey(xkey)
	}
	if err != nil {
		return err
	}

	h.cacheMu.Lock()
	// The order is crucial here. The key is dropped from the
	// cache after the file is gone so that a reload happening in
	// between won't insert it into the cache again.
	err = os.Remove(xpb.File)
	if err == nil {
		h.cache.delete(xpb)
	}
	h.cacheMu.Unlock()
	return err
}

func (h *HSM) loadDecryptedKey(xpub chainkd.XPub, auth string) (XPub, *XKey, error) {
	h.cache.maybeReload()
	h.cache.mu.Lock()
	xpb, err := h.cache.find(XPub{XPub: xpub})

	h.cache.mu.Unlock()
	if err != nil {
		return xpb, nil, err
	}
	xkey, err := h.keyStore.GetKey(xpb.Alias, xpb.File, auth)
	return xpb, xkey, err
}

// ResetPassword reset passphrase for an existing xpub
func (h *HSM) ResetPassword(xpub chainkd.XPub, oldAuth, newAuth string) error {
	xpb, xkey, err := h.loadDecryptedKey(xpub, oldAuth)
	if err != nil {
		return err
	}
	return h.keyStore.StoreKey(xpb.File, xkey, newAuth)
}

// HasAlias check whether the key alias exists
func (h *HSM) HasAlias(alias string) bool {
	return h.cache.hasAlias(alias)
}

// HasKey check whether the private key exists
func (h *HSM) HasKey(xprv chainkd.XPrv) bool {
	return h.cache.hasKey(xprv.XPub())
}
<commit_msg>Change mnemonic length to 12<commit_after>// Package pseudohsm provides a pseudo HSM for development environments.
package pseudohsm

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/pborman/uuid"

	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/errors"
	mnem "github.com/bytom/wallet/mnemonic"
)

// pre-define errors for supporting bytom errorFormatter
var (
	ErrDuplicateKeyAlias = errors.New("duplicate key alias")
	ErrLoadKey           = errors.New("key not found or wrong password ")
	ErrDecrypt           = errors.New("could not decrypt key with given passphrase")
	ErrMnemonicLength    = errors.New("mnemonic length error")
)

// EntropyLength random entropy length to generate mnemonics.
const EntropyLength = 128

// HSM type for storing pubkey and privatekey
type HSM struct {
	cacheMu  sync.Mutex
	keyStore keyStore
	cache    *keyCache
	//kdCache  map[chainkd.XPub]chainkd.XPrv
}

// XPub type for pubkey for anyone can see
type XPub struct {
	Alias string       `json:"alias"`
	XPub  chainkd.XPub `json:"xpub"`
	File  string       `json:"file"`
}

// New method for HSM struct
func New(keypath string) (*HSM, error) {
	keydir, _ := filepath.Abs(keypath)
	return &HSM{
		keyStore: &keyStorePassphrase{keydir, LightScryptN, LightScryptP},
		cache:    newKeyCache(keydir),
		//kdCache:  make(map[chainkd.XPub]chainkd.XPrv),
	}, nil
}

// XCreate produces a new random xprv and stores it in the db.
func (h *HSM) XCreate(alias string, auth string, language string) (*XPub, *string, error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	normalizedAlias := strings.ToLower(strings.TrimSpace(alias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return nil, nil, ErrDuplicateKeyAlias
	}

	xpub, mnemonic, err := h.createChainKDKey(normalizedAlias, auth, language)
	if err != nil {
		return nil, nil, err
	}
	h.cache.add(*xpub)
	return xpub, mnemonic, err
}

// ImportKeyFromMnemonic produces a xprv from mnemonic and stores it in the db.
func (h *HSM) ImportKeyFromMnemonic(alias string, auth string, mnemonic string, language string) (*XPub, error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	// checksum length = entropy length /32
	// mnemonic length = (entropy length + checksum length)/11
	if len(strings.Fields(mnemonic)) != (EntropyLength+EntropyLength/32)/11 {
		return nil, ErrMnemonicLength
	}

	normalizedAlias := strings.ToLower(strings.TrimSpace(alias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return nil, ErrDuplicateKeyAlias
	}

	// Pre validate that the mnemonic is well formed and only contains words that
	// are present in the word list
	if !mnem.IsMnemonicValid(mnemonic, language) {
		return nil, mnem.ErrInvalidMnemonic
	}

	xpub, err := h.createKeyFromMnemonic(alias, auth, mnemonic)
	if err != nil {
		return nil, err
	}

	h.cache.add(*xpub)
	return xpub, nil
}

func (h *HSM) createKeyFromMnemonic(alias string, auth string, mnemonic string) (*XPub, error) {
	// Generate a Bip32 HD wallet for the mnemonic and a user supplied password
	seed := mnem.NewSeed(mnemonic, "")
	xprv, xpub, err := chainkd.NewXKeys(bytes.NewBuffer(seed))
	if err != nil {
		return nil, err
	}
	id := uuid.NewRandom()
	key := &XKey{
		ID:      id,
		KeyType: "bytom_kd",
		XPub:    xpub,
		XPrv:    xprv,
		Alias:   alias,
	}
	file := h.keyStore.JoinPath(keyFileName(key.ID.String()))
	if err := h.keyStore.StoreKey(file, key, auth); err != nil {
		return nil, errors.Wrap(err, "storing keys")
	}
	return &XPub{XPub: xpub, Alias: alias, File: file}, nil
}

func (h *HSM) createChainKDKey(alias string, auth string, language string) (*XPub, *string, error) {
	// Generate a mnemonic for memorization or user-friendly seeds
	entropy, err := mnem.NewEntropy(EntropyLength)
	if err != nil {
		return nil, nil, err
	}
	mnemonic, err := mnem.NewMnemonic(entropy, language)
	if err != nil {
		return nil, nil, err
	}
	xpub, err := h.createKeyFromMnemonic(alias, auth, mnemonic)
	if err != nil {
		return nil, nil, err
	}
	return xpub, &mnemonic, nil
}

// UpdateKeyAlias update key alias
func (h *HSM) UpdateKeyAlias(xpub chainkd.XPub, newAlias string) error {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	h.cache.maybeReload()
	h.cache.mu.Lock()
	xpb, err := h.cache.find(XPub{XPub: xpub})
	h.cache.mu.Unlock()
	if err != nil {
		return err
	}

	keyjson, err := ioutil.ReadFile(xpb.File)
	if err != nil {
		return err
	}

	encrptKeyJSON := new(encryptedKeyJSON)
	if err := json.Unmarshal(keyjson, encrptKeyJSON); err != nil {
		return err
	}

	normalizedAlias := strings.ToLower(strings.TrimSpace(newAlias))
	if ok := h.cache.hasAlias(normalizedAlias); ok {
		return ErrDuplicateKeyAlias
	}

	encrptKeyJSON.Alias = normalizedAlias
	keyJSON, err := json.Marshal(encrptKeyJSON)
	if err != nil {
		return err
	}

	if err := writeKeyFile(xpb.File, keyJSON); err != nil {
		return err
	}

	// update key alias
	h.cache.delete(xpb)
	xpb.Alias = normalizedAlias
	h.cache.add(xpb)

	return nil
}

// ListKeys returns a list of all xpubs from the store
func (h *HSM) ListKeys() []XPub {
	xpubs := h.cache.keys()
	return xpubs
}

// XSign looks up the xprv given the xpub, optionally derives a new
// xprv with the given path (but does not store the new xprv), and
// signs the given msg.
func (h *HSM) XSign(xpub chainkd.XPub, path [][]byte, msg []byte, auth string) ([]byte, error) {
	xprv, err := h.LoadChainKDKey(xpub, auth)
	if err != nil {
		return nil, err
	}
	if len(path) > 0 {
		xprv = xprv.Derive(path)
	}
	return xprv.Sign(msg), nil
}

//LoadChainKDKey get xprv from xpub
func (h *HSM) LoadChainKDKey(xpub chainkd.XPub, auth string) (xprv chainkd.XPrv, err error) {
	h.cacheMu.Lock()
	defer h.cacheMu.Unlock()

	//if xprv, ok := h.kdCache[xpub]; ok {
	//	return xprv, nil
	//}

	_, xkey, err := h.loadDecryptedKey(xpub, auth)
	if err != nil {
		return xprv, ErrLoadKey
	}
	//h.kdCache[xpb.XPub] = xkey.XPrv
	return xkey.XPrv, nil
}

// XDelete deletes the key matched by xpub if the passphrase is correct.
// If a contains no filename, the address must match a unique key.
func (h *HSM) XDelete(xpub chainkd.XPub, auth string) error {
	// Decrypting the key isn't really necessary, but we do
	// it anyway to check the password and zero out the key
	// immediately afterwards.

	xpb, xkey, err := h.loadDecryptedKey(xpub, auth)
	if xkey != nil {
		zeroKey(xkey)
	}
	if err != nil {
		return err
	}

	h.cacheMu.Lock()
	// The order is crucial here. The key is dropped from the
	// cache after the file is gone so that a reload happening in
	// between won't insert it into the cache again.
	err = os.Remove(xpb.File)
	if err == nil {
		h.cache.delete(xpb)
	}
	h.cacheMu.Unlock()
	return err
}

func (h *HSM) loadDecryptedKey(xpub chainkd.XPub, auth string) (XPub, *XKey, error) {
	h.cache.maybeReload()
	h.cache.mu.Lock()
	xpb, err := h.cache.find(XPub{XPub: xpub})

	h.cache.mu.Unlock()
	if err != nil {
		return xpb, nil, err
	}
	xkey, err := h.keyStore.GetKey(xpb.Alias, xpb.File, auth)
	return xpb, xkey, err
}

// ResetPassword reset passphrase for an existing xpub
func (h *HSM) ResetPassword(xpub chainkd.XPub, oldAuth, newAuth string) error {
	xpb, xkey, err := h.loadDecryptedKey(xpub, oldAuth)
	if err != nil {
		return err
	}
	return h.keyStore.StoreKey(xpb.File, xkey, newAuth)
}

// HasAlias check whether the key alias exists
func (h *HSM) HasAlias(alias string) bool {
	return h.cache.hasAlias(alias)
}

// HasKey check whether the private key exists
func (h *HSM) HasKey(xprv chainkd.XPrv) bool {
	return h.cache.hasKey(xprv.XPub())
}
<|endoftext|>


<commit_before>// This must be taken into a more general package. A similar
// file already exists in migrators/graphity/common package

package main

import (
	"encoding/json"
	"github.com/streadway/amqp"
	. "koding/db/models"
	"koding/messaging/rabbitmq"
)

var (
	GRAPHITY_CHANNEL *amqp.Channel
	PUBLISHER        *rabbitmq.Producer
)

func init() {
	exchange := rabbitmq.Exchange{
		Name: "graphFeederExchange",
	}
	queue := rabbitmq.Queue{}
	publishingOptions := rabbitmq.PublishingOptions{
		Tag:        "graphityRelationship",
		RoutingKey: "",
	}
	// var err Error
	var err error
	PUBLISHER, err = rabbitmq.NewProducer(exchange, queue, publishingOptions)
	if err != nil {
		panic(err)
	}

	PUBLISHER.RegisterSignalHandler()
}

func CreateGraphRelationship(relationship *Relationship) error {
	return updateRelationship(relationship, "RelationshipSaved")
}

func RemoveGraphRelationship(relationship *Relationship) error {
	return updateRelationship(relationship, "RelationshipRemoved")
}

func updateRelationship(relationship *Relationship, event string) error {
	data := make([]Relationship, 1)
	data[0] = *relationship
	eventData := map[string]interface{}{"event": event, "payload": data}

	neoMessage, err := json.Marshal(eventData)

	if err != nil {
		log.Error("unmarshall error - %v", err)
		return err
	}

	message := amqp.Publishing{
		Body: neoMessage,
	}

	PUBLISHER.NotifyReturn(func(message amqp.Return) {
		log.Info("%v", message)
	})

	err = PUBLISHER.Publish(message)
	if err != nil {
		log.Error(err.Error())
		return err
	}
	return nil
}
<commit_msg>Moderation: Publisher creations are moved to publisher.go<commit_after>// This must be taken into a more general package. A similar
// file already exists in migrators/graphity/common package

package topicmodifier

import (
	. "koding/db/models"
	"koding/messaging/rabbitmq"
)

var (
	GraphPublisher *rabbitmq.Producer
)

func initGraphPublisher() {
	options := &PublisherConfig{
		ExchangeName: "graphFeederExchange",
		Tag:          "graphityRelationship",
		RoutingKey:   "",
	}
	GraphPublisher = createPublisher(options)
}

func CreateGraphRelationship(relationship *Relationship) error {
	return updateRelationship(relationship, "RelationshipSaved")
}

func RemoveGraphRelationship(relationship *Relationship) error {
	return updateRelationship(relationship, "RelationshipRemoved")
}

func updateRelationship(relationship *Relationship, event string) error {
	data := make([]Relationship, 1)
	data[0] = *relationship
	eventData := map[string]interface{}{"event": event, "payload": data}

	return publish(GraphPublisher, eventData)
}
<|endoftext|>


<commit_before>package resource

import (
	"encoding/json"

	"github.com/EngineerBetter/concourse-up/bosh/internal/resource/internal/file"
)

//go:generate go-bindata -o internal/file/file.go -ignore (\.go$)|(\.git) -nometadata -pkg file -prefix=../../../../concourse-up-ops . ../../../../concourse-up-ops/...

// Resource safely exposes the json parameters of a resource
type Resource struct {
	URL     string `json:"url"`
	Version string `json:"version"`
	SHA1    string `json:"sha1"`
}

var resources map[string]Resource

// ID defines the name of a resource in a safer way
type ID struct {
	name string
}

var (
	// AWSCPI statically defines cpi string
	AWSCPI = ID{"cpi"}
	// AWSStemcell statically defines stemcell string
	AWSStemcell = ID{"stemcell"}
	// BOSHRelease statically defines bosh string
	BOSHRelease = ID{"bosh"}
	// BPMRelease statically defines bpm string
	BPMRelease = ID{"bpm"}
)

var (
	// DirectorManifest statically defines director-manifest.yml contents
	DirectorManifest = mustAssetString("director/director-manifest.yml")
	// AWSCPIOps statically defines aws-cpi.yml contents
	AWSCPIOps = mustAssetString("director/aws-cpi.yml")
	// ExternalIPOps statically defines external-ip.yml contents
	ExternalIPOps = mustAssetString("director/external-ip.yml")
	// DirectorCustomOps statically defines custom-ops.yml contents
	DirectorCustomOps = mustAssetString("director/custom-ops.yml")
)

// NOTE(px) remove this in a later version of github.com/mattn/go-bindata
func mustAssetString(name string) string {
	return string(file.MustAsset(name))
}

// Get returns an Resource in a safe way
func Get(id ID) Resource {
	r, ok := resources[id.name]
	if !ok {
		panic("resource " + id.name + " not found")
	}
	return r
}

func init() {
	p := file.MustAsset("director-versions.json")
	err := json.Unmarshal(p, &resources)
	if err != nil {
		panic(err)
	}
}
<commit_msg>setting the correct asset paths<commit_after>package resource

import (
	"encoding/json"

	"github.com/EngineerBetter/concourse-up/bosh/internal/resource/internal/file"
)

//go:generate go-bindata -o internal/file/file.go -ignore (\.go$)|(\.git) -nometadata -pkg file -prefix=../../../../concourse-up-ops . ../../../../concourse-up-ops/...

// Resource safely exposes the json parameters of a resource
type Resource struct {
	URL     string `json:"url"`
	Version string `json:"version"`
	SHA1    string `json:"sha1"`
}

var resources map[string]Resource

// ID defines the name of a resource in a safer way
type ID struct {
	name string
}

var (
	// AWSCPI statically defines cpi string
	AWSCPI = ID{"cpi"}
	// AWSStemcell statically defines stemcell string
	AWSStemcell = ID{"stemcell"}
	// BOSHRelease statically defines bosh string
	BOSHRelease = ID{"bosh"}
	// BPMRelease statically defines bpm string
	BPMRelease = ID{"bpm"}
)

var (
	// DirectorManifest statically defines director-manifest.yml contents
	DirectorManifest = mustAssetString("director/manifest.yml")
	// AWSCPIOps statically defines aws-cpi.yml contents
	AWSCPIOps = mustAssetString("director/aws/cpi.yml")
	// ExternalIPOps statically defines external-ip.yml contents
	ExternalIPOps = mustAssetString("director/external-ip.yml")
	// DirectorCustomOps statically defines custom-ops.yml contents
	DirectorCustomOps = mustAssetString("director/custom-ops.yml")
)

// NOTE(px) remove this in a later version of github.com/mattn/go-bindata
func mustAssetString(name string) string {
	return string(file.MustAsset(name))
}

// Get returns an Resource in a safe way
func Get(id ID) Resource {
	r, ok := resources[id.name]
	if !ok {
		panic("resource " + id.name + " not found")
	}
	return r
}

func init() {
	p := file.MustAsset("director-versions.json")
	err := json.Unmarshal(p, &resources)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>
package main

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"log"
	"mig"
	"os/exec"
	"os"
	"runtime"
	"time"
)

var AMQPBROKER string = "amqp://guest:guest@172.21.1.1:5672/"
var HEARTBEATFREQ string = "10s"

func getCommands(messages <-chan amqp.Delivery, actions chan []byte, terminate chan bool) error {
	// range waits on the channel and returns all incoming messages
	// range will exit when the channel closes
	for m := range messages {
		log.Printf("getCommands: received '%s'", m.Body)
		// Ack this message only
		err := m.Ack(true)
		if err != nil { panic(err) }
		actions <- m.Body
		log.Printf("getCommands: queued in pos. %d", len(actions))
	}
	terminate <- true
	return nil
}

func parseCommands(commands <-chan []byte, fCommandChan chan mig.Command, terminate chan bool) error {
	var cmd mig.Command
	for a := range commands {
		err := json.Unmarshal(a, &cmd)
		if err != nil {
			log.Fatal("parseCommand - json.Unmarshal:", err)
		}
		log.Printf("ParseCommand: Check '%s' Arguments '%s'",
			   cmd.Action.Check, cmd.Action.Arguments)
		switch cmd.Action.Check{
		case "filechecker":
			fCommandChan <- cmd
			log.Println("parseCommands: queued into filechecker",
				    "in pos.", len(fCommandChan))
		}
	}
	terminate <- true
	return nil
}

func runFilechecker(fCommandChan <-chan mig.Command, alertChan chan mig.Alert, resultChan chan mig.Command, terminate chan bool) error {
	for migCmd := range fCommandChan {
		log.Printf("RunFilechecker: running with args '%s'", migCmd.Action.Arguments)
		var cmdArg string
		for _, arg := range migCmd.Action.Arguments {
			cmdArg += arg
		}
		runCmd := exec.Command("./filechecker", cmdArg)
		cmdout, err := runCmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		st := time.Now()
		err = runCmd.Start()
		if err != nil {
			log.Fatal(err)
		}
		results := make(map[string] mig.FileCheckerResult)
		err = json.NewDecoder(cmdout).Decode(&results)
		if err != nil {
			log.Fatal(err)
		}
		cmdDone := make(chan error)
		go func() {
			cmdDone <-runCmd.Wait()
		}()
		select {
		// kill the process when timeout expires
		case <-time.After(30 * time.Second):
			if err := runCmd.Process.Kill(); err != nil {
				log.Fatal("failed to kill:", err)
			}
			log.Fatal("runFileChecker: command '%s' timed out", migCmd)
		// exit normally
		case err := <-cmdDone:
			if err != nil {
				log.Fatal(err)
			}
		}
		for _, r := range results {
			log.Println("runFileChecker: command", migCmd,"tested",
				    r.TestedFiles, "files in", time.Now().Sub(st))
			if r.ResultCount > 0 {
				for _, f := range r.Files {
					alertChan <- mig.Alert{
						Arguments: migCmd.Action.Arguments,
						Item: f,
					}
				}
			}
			migCmd.FCResults = append(migCmd.FCResults, r)
		}
		resultChan <- migCmd
	}
	terminate <- true
	return nil
}

func raiseAlerts(alertChan chan mig.Alert, terminate chan bool) error {
	for a := range alertChan {
		log.Printf("raiseAlerts: IOC '%s' positive match on '%s'",
			   a.Arguments, a.Item)
	}
	return nil
}

func sendResults(c *amqp.Channel, agtQueueLoc string, resultChan <-chan mig.Command, terminate chan bool) error {
	rKey := fmt.Sprintf("mig.agents.%s", agtQueueLoc)
	for r := range resultChan {
		r.AgentQueueLoc = agtQueueLoc
		body, err := json.Marshal(r)
		if err != nil {
			log.Fatalf("sendResults - json.Marshal: %v", err)
		}
		msgXchange(c, "mig", rKey, body)
	}
	return nil
}

func keepAliveAgent(c *amqp.Channel, regMsg mig.KeepAlive) error {
	sleepTime, err := time.ParseDuration(HEARTBEATFREQ)
	if err != nil {
		log.Fatal("sendHeartbeat - time.ParseDuration():", err)
	}
	for {
		body, err := json.Marshal(regMsg)
		if err != nil {
			log.Fatal("sendHeartbeat - json.Marshal:", err)
		}
		msgXchange(c, "mig", "mig.keepalive", body)
		time.Sleep(sleepTime)
	}
	return nil
}

func msgXchange(c *amqp.Channel, excName, routingKey string, body []byte) error {
	msg := amqp.Publishing{
	    DeliveryMode: amqp.Persistent,
	    Timestamp:    time.Now(),
	    ContentType:  "text/plain",
	    Body:         []byte(body),
	}
	err := c.Publish(excName,
			routingKey,
			true,	// is mandatory
			false,	// is immediate
			msg)	// AMQP message
	if err != nil {
		log.Fatalf("msgXchange - ChannelPublish: %v", err)
	}
	log.Printf("msgXchange: published '%s'\n", msg.Body)
	return nil
}

func main() {
	// termChan is used to exit the program
	termChan	:= make(chan bool)
	actionsChan	:= make(chan []byte, 10)
	fCommandChan	:= make(chan mig.Command, 10)
	alertChan	:= make(chan mig.Alert, 10)
	resultChan	:= make(chan mig.Command, 10)
	hostname, err	:= os.Hostname()
	if err != nil {
		log.Fatalf("os.Hostname(): %v", err)
	}
	regMsg := mig.KeepAlive{
		Name: hostname,
		OS: runtime.GOOS,
		QueueLoc: fmt.Sprintf("%s.%s", runtime.GOOS, hostname),
		LastKeepAlive: time.Now(),
	}
	agentQueue := fmt.Sprintf("mig.agt.%s", regMsg.QueueLoc)
	bindings := []mig.Binding{
		mig.Binding{agentQueue, agentQueue},
		mig.Binding{agentQueue, "mig.all"},
	}

	log.Println("MIG agent starting on", hostname)

	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial(AMQPBROKER)
	if err != nil {
		log.Fatalf("amqp.Dial(): %v", err)
	}
	defer conn.Close()
	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("conn.Channel(): %v", err)
	}
	for _, b := range bindings {
		_, err = c.QueueDeclare(b.Queue,	// Queue name
					true,		// is durable
					false,		// is autoDelete
					false,		// is exclusive
					false,		// is noWait
					nil)		// AMQP args
		if err != nil {
			log.Fatalf("QueueDeclare: %v", err)
		}
		err = c.QueueBind(b.Queue,	// Queue name
				b.Key,		// Routing key name
				"mig",		// Exchange name
				false,		// is noWait
				nil)		// AMQP args
		if err != nil {
			log.Fatalf("QueueBind: %v", err)
		}
	}

	// Limit the number of message the channel will receive
	err = c.Qos(2,		// prefetch count (in # of msg)
		    0,		// prefetch size (in bytes)
		    false)	// is global
	if err != nil {
		log.Fatalf("ChannelQoS: %v", err)
	}
	for _, b := range bindings {
		msgChan, err := c.Consume(b.Queue, // queue name
					"",	// some tag
					false,	// is autoAck
					false,	// is exclusive
					false,	// is noLocal
					false,	// is noWait
					nil)	// AMQP args
		if err != nil {
			log.Fatalf("ChannelConsume: %v", err)
		}
		go getCommands(msgChan, actionsChan, termChan)
	}
	go parseCommands(actionsChan, fCommandChan, termChan)
	go runFilechecker(fCommandChan, alertChan, resultChan, termChan)
	go raiseAlerts(alertChan, termChan)
	go sendResults(c, regMsg.QueueLoc, resultChan, termChan)

	// All set, ready to keepAlive
	go keepAliveAgent(c, regMsg)

	// block until terminate chan is called
	<-termChan
}
<commit_msg>Agent: increase heartbeat to 10 minutes<commit_after>
package main

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"log"
	"mig"
	"os/exec"
	"os"
	"runtime"
	"time"
)

var AMQPBROKER string = "amqp://guest:guest@172.21.1.1:5672/"
var HEARTBEATFREQ string = "600s"

func getCommands(messages <-chan amqp.Delivery, actions chan []byte, terminate chan bool) error {
	// range waits on the channel and returns all incoming messages
	// range will exit when the channel closes
	for m := range messages {
		log.Printf("getCommands: received '%s'", m.Body)
		// Ack this message only
		err := m.Ack(true)
		if err != nil { panic(err) }
		actions <- m.Body
		log.Printf("getCommands: queued in pos. %d", len(actions))
	}
	terminate <- true
	return nil
}

func parseCommands(commands <-chan []byte, fCommandChan chan mig.Command, terminate chan bool) error {
	var cmd mig.Command
	for a := range commands {
		err := json.Unmarshal(a, &cmd)
		if err != nil {
			log.Fatal("parseCommand - json.Unmarshal:", err)
		}
		log.Printf("ParseCommand: Check '%s' Arguments '%s'",
			   cmd.Action.Check, cmd.Action.Arguments)
		switch cmd.Action.Check{
		case "filechecker":
			fCommandChan <- cmd
			log.Println("parseCommands: queued into filechecker",
				    "in pos.", len(fCommandChan))
		}
	}
	terminate <- true
	return nil
}

func runFilechecker(fCommandChan <-chan mig.Command, alertChan chan mig.Alert, resultChan chan mig.Command, terminate chan bool) error {
	for migCmd := range fCommandChan {
		log.Printf("RunFilechecker: running with args '%s'", migCmd.Action.Arguments)
		var cmdArg string
		for _, arg := range migCmd.Action.Arguments {
			cmdArg += arg
		}
		runCmd := exec.Command("./filechecker", cmdArg)
		cmdout, err := runCmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		st := time.Now()
		err = runCmd.Start()
		if err != nil {
			log.Fatal(err)
		}
		results := make(map[string] mig.FileCheckerResult)
		err = json.NewDecoder(cmdout).Decode(&results)
		if err != nil {
			log.Fatal(err)
		}
		cmdDone := make(chan error)
		go func() {
			cmdDone <-runCmd.Wait()
		}()
		select {
		// kill the process when timeout expires
		case <-time.After(30 * time.Second):
			if err := runCmd.Process.Kill(); err != nil {
				log.Fatal("failed to kill:", err)
			}
			log.Fatal("runFileChecker: command '%s' timed out", migCmd)
		// exit normally
		case err := <-cmdDone:
			if err != nil {
				log.Fatal(err)
			}
		}
		for _, r := range results {
			log.Println("runFileChecker: command", migCmd,"tested",
				    r.TestedFiles, "files in", time.Now().Sub(st))
			if r.ResultCount > 0 {
				for _, f := range r.Files {
					alertChan <- mig.Alert{
						Arguments: migCmd.Action.Arguments,
						Item: f,
					}
				}
			}
			migCmd.FCResults = append(migCmd.FCResults, r)
		}
		resultChan <- migCmd
	}
	terminate <- true
	return nil
}

func raiseAlerts(alertChan chan mig.Alert, terminate chan bool) error {
	for a := range alertChan {
		log.Printf("raiseAlerts: IOC '%s' positive match on '%s'",
			   a.Arguments, a.Item)
	}
	return nil
}

func sendResults(c *amqp.Channel, agtQueueLoc string, resultChan <-chan mig.Command, terminate chan bool) error {
	rKey := fmt.Sprintf("mig.agents.%s", agtQueueLoc)
	for r := range resultChan {
		r.AgentQueueLoc = agtQueueLoc
		body, err := json.Marshal(r)
		if err != nil {
			log.Fatalf("sendResults - json.Marshal: %v", err)
		}
		msgXchange(c, "mig", rKey, body)
	}
	return nil
}

func keepAliveAgent(c *amqp.Channel, regMsg mig.KeepAlive) error {
	sleepTime, err := time.ParseDuration(HEARTBEATFREQ)
	if err != nil {
		log.Fatal("sendHeartbeat - time.ParseDuration():", err)
	}
	for {
		body, err := json.Marshal(regMsg)
		if err != nil {
			log.Fatal("sendHeartbeat - json.Marshal:", err)
		}
		msgXchange(c, "mig", "mig.keepalive", body)
		time.Sleep(sleepTime)
	}
	return nil
}

func msgXchange(c *amqp.Channel, excName, routingKey string, body []byte) error {
	msg := amqp.Publishing{
	    DeliveryMode: amqp.Persistent,
	    Timestamp:    time.Now(),
	    ContentType:  "text/plain",
	    Body:         []byte(body),
	}
	err := c.Publish(excName,
			routingKey,
			true,	// is mandatory
			false,	// is immediate
			msg)	// AMQP message
	if err != nil {
		log.Fatalf("msgXchange - ChannelPublish: %v", err)
	}
	log.Printf("msgXchange: published '%s'\n", msg.Body)
	return nil
}

func main() {
	// termChan is used to exit the program
	termChan	:= make(chan bool)
	actionsChan	:= make(chan []byte, 10)
	fCommandChan	:= make(chan mig.Command, 10)
	alertChan	:= make(chan mig.Alert, 10)
	resultChan	:= make(chan mig.Command, 10)
	hostname, err	:= os.Hostname()
	if err != nil {
		log.Fatalf("os.Hostname(): %v", err)
	}
	regMsg := mig.KeepAlive{
		Name: hostname,
		OS: runtime.GOOS,
		QueueLoc: fmt.Sprintf("%s.%s", runtime.GOOS, hostname),
		LastKeepAlive: time.Now(),
	}
	agentQueue := fmt.Sprintf("mig.agt.%s", regMsg.QueueLoc)
	bindings := []mig.Binding{
		mig.Binding{agentQueue, agentQueue},
		mig.Binding{agentQueue, "mig.all"},
	}

	log.Println("MIG agent starting on", hostname)

	// Connects opens an AMQP connection from the credentials in the URL.
	conn, err := amqp.Dial(AMQPBROKER)
	if err != nil {
		log.Fatalf("amqp.Dial(): %v", err)
	}
	defer conn.Close()
	c, err := conn.Channel()
	if err != nil {
		log.Fatalf("conn.Channel(): %v", err)
	}
	for _, b := range bindings {
		_, err = c.QueueDeclare(b.Queue,	// Queue name
					true,		// is durable
					false,		// is autoDelete
					false,		// is exclusive
					false,		// is noWait
					nil)		// AMQP args
		if err != nil {
			log.Fatalf("QueueDeclare: %v", err)
		}
		err = c.QueueBind(b.Queue,	// Queue name
				b.Key,		// Routing key name
				"mig",		// Exchange name
				false,		// is noWait
				nil)		// AMQP args
		if err != nil {
			log.Fatalf("QueueBind: %v", err)
		}
	}

	// Limit the number of message the channel will receive
	err = c.Qos(2,		// prefetch count (in # of msg)
		    0,		// prefetch size (in bytes)
		    false)	// is global
	if err != nil {
		log.Fatalf("ChannelQoS: %v", err)
	}
	for _, b := range bindings {
		msgChan, err := c.Consume(b.Queue, // queue name
					"",	// some tag
					false,	// is autoAck
					false,	// is exclusive
					false,	// is noLocal
					false,	// is noWait
					nil)	// AMQP args
		if err != nil {
			log.Fatalf("ChannelConsume: %v", err)
		}
		go getCommands(msgChan, actionsChan, termChan)
	}
	go parseCommands(actionsChan, fCommandChan, termChan)
	go runFilechecker(fCommandChan, alertChan, resultChan, termChan)
	go raiseAlerts(alertChan, termChan)
	go sendResults(c, regMsg.QueueLoc, resultChan, termChan)

	// All set, ready to keepAlive
	go keepAliveAgent(c, regMsg)

	// block until terminate chan is called
	<-termChan
}
<|endoftext|>


<commit_before>package build

import (
	"neon/util"
)

type Task func() error

type TaskConstructor func(target *Target, args util.Object) (Task, error)

type TaskDescriptor struct {
	Constructor TaskConstructor
	Help        string
}

var TaskMap map[string]TaskDescriptor = make(map[string]TaskDescriptor)
<commit_msg>Code cleaning<commit_after>package build

import (
	"neon/util"
)

// A task is a function that returns an error
type Task func() error

// A task constructor is a function that returns a task and an error
type TaskConstructor func(target *Target, args util.Object) (Task, error)

// A task descriptor is made of a task constructor and an help string
type TaskDescriptor struct {
	Constructor TaskConstructor
	Help        string
}

// Map that gives constructor for given task name
var TaskMap map[string]TaskDescriptor = make(map[string]TaskDescriptor)
<|endoftext|>


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

package net

import (
	"flag"
	"fmt"
	"testing"
)

var testDNSFlood = flag.Bool("dnsflood", false, "whether to test dns query flooding")

func TestDNSThreadLimit(t *testing.T) {
	if !*testDNSFlood {
		t.Skip("test disabled; use -dnsflood to enable")
	}

	const N = 10000
	c := make(chan int, N)
	for i := 0; i < N; i++ {
		go func(i int) {
			LookupIP(fmt.Sprintf("%d.net-test.golang.org", i))
			c <- 1
		}(i)
	}
	// Don't bother waiting for the stragglers; stop at 0.9 N.
	for i := 0; i < N*9/10; i++ {
		if i%100 == 0 {
			//println("TestDNSThreadLimit:", i)
		}
		<-c
	}

	// If we're still here, it worked.
}
<commit_msg>net: add test for lookupIPDeadline<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"flag"
	"fmt"
	"testing"
	"time"
)

var testDNSFlood = flag.Bool("dnsflood", false, "whether to test dns query flooding")

func TestDNSThreadLimit(t *testing.T) {
	if !*testDNSFlood {
		t.Skip("test disabled; use -dnsflood to enable")
	}

	const N = 10000
	c := make(chan int, N)
	for i := 0; i < N; i++ {
		go func(i int) {
			LookupIP(fmt.Sprintf("%d.net-test.golang.org", i))
			c <- 1
		}(i)
	}
	// Don't bother waiting for the stragglers; stop at 0.9 N.
	for i := 0; i < N*9/10; i++ {
		if i%100 == 0 {
			//println("TestDNSThreadLimit:", i)
		}
		<-c
	}

	// If we're still here, it worked.
}

func TestLookupIPDeadline(t *testing.T) {
	if !*testDNSFlood {
		t.Skip("test disabled; use -dnsflood to enable")
	}

	const N = 5000
	const timeout = 3 * time.Second
	c := make(chan error, 2*N)
	for i := 0; i < N; i++ {
		name := fmt.Sprintf("%d.net-test.golang.org", i)
		go func() {
			_, err := lookupIPDeadline(name, time.Now().Add(timeout/2))
			c <- err
		}()
		go func() {
			_, err := lookupIPDeadline(name, time.Now().Add(timeout))
			c <- err
		}()
	}
	qstats := struct {
		succeeded, failed         int
		timeout, temporary, other int
		unknown                   int
	}{}
	deadline := time.After(timeout + time.Second)
	for i := 0; i < 2*N; i++ {
		select {
		case <-deadline:
			t.Fatal("deadline exceeded")
		case err := <-c:
			switch err := err.(type) {
			case nil:
				qstats.succeeded++
			case Error:
				qstats.failed++
				if err.Timeout() {
					qstats.timeout++
				}
				if err.Temporary() {
					qstats.temporary++
				}
				if !err.Timeout() && !err.Temporary() {
					qstats.other++
				}
			default:
				qstats.failed++
				qstats.unknown++
			}
		}
	}

	// A high volume of DNS queries for sub-domain of golang.org
	// would be coordinated by authoritative or recursive server,
	// or stub resolver which implements query-response rate
	// limitation, so we can expect some query successes and more
	// failures including timeout, temporary and other here.
	// As a rule, unknown must not be shown but it might possibly
	// happen due to issue 4856 for now.
	t.Logf("%v succeeded, %v failed (%v timeout, %v temporary, %v other, %v unknown)", qstats.succeeded, qstats.failed, qstats.timeout, qstats.temporary, qstats.other, qstats.unknown)
}
<|endoftext|>


<commit_before>// Copyright 2015 The rkt Authors
//
// 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 main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/hashicorp/errwrap"

	"github.com/coreos/rkt/common"
	rktlog "github.com/coreos/rkt/pkg/log"
	"github.com/coreos/rkt/pkg/sys"
	stage1common "github.com/coreos/rkt/stage1/common"
	stage1commontypes "github.com/coreos/rkt/stage1/common/types"
)

const (
	flavor = "fly"
)

type flyMount struct {
	HostPath         string
	TargetPrefixPath string
	RelTargetPath    string
	Fs               string
	Flags            uintptr
}

type volumeMountTuple struct {
	V types.Volume
	M schema.Mount
}

var (
	debug bool

	discardNetlist common.NetList
	discardBool    bool
	discardString  string

	log  *rktlog.Logger
	diag *rktlog.Logger
)

func getHostMounts() (map[string]struct{}, error) {
	hostMounts := map[string]struct{}{}

	mi, err := os.Open("/proc/self/mountinfo")
	if err != nil {
		return nil, err
	}
	defer mi.Close()

	sc := bufio.NewScanner(mi)
	for sc.Scan() {
		var (
			discard    string
			mountPoint string
		)

		_, err := fmt.Sscanf(sc.Text(),
			"%s %s %s %s %s",
			&discard, &discard, &discard, &discard, &mountPoint,
		)
		if err != nil {
			return nil, err
		}

		hostMounts[mountPoint] = struct{}{}
	}
	if sc.Err() != nil {
		return nil, errwrap.Wrap(errors.New("problem parsing mountinfo"), sc.Err())
	}
	return hostMounts, nil
}

func init() {
	flag.BoolVar(&debug, "debug", false, "Run in debug mode")

	// The following flags need to be supported by stage1 according to
	// https://github.com/coreos/rkt/blob/master/Documentation/devel/stage1-implementors-guide.md
	// TODO: either implement functionality or give not implemented warnings
	flag.Var(&discardNetlist, "net", "Setup networking")
	flag.BoolVar(&discardBool, "interactive", true, "The pod is interactive")
	flag.StringVar(&discardString, "mds-token", "", "MDS auth token")
	flag.StringVar(&discardString, "local-config", common.DefaultLocalConfigDir, "Local config path")
}

func evaluateMounts(rfs string, app string, p *stage1commontypes.Pod) ([]flyMount, error) {
	imApp := p.Images[app].App
	namedVolumeMounts := map[types.ACName]volumeMountTuple{}

	for _, m := range p.Manifest.Apps[0].Mounts {
		_, exists := namedVolumeMounts[m.Volume]
		if exists {
			return nil, fmt.Errorf("duplicate mount given: %q", m.Volume)
		}
		namedVolumeMounts[m.Volume] = volumeMountTuple{M: m}
		diag.Printf("adding %+v", namedVolumeMounts[m.Volume])
	}

	// Merge command-line Mounts with ImageManifest's MountPoints
	for _, mp := range imApp.MountPoints {
		tuple, exists := namedVolumeMounts[mp.Name]
		switch {
		case exists && tuple.M.Path != mp.Path:
			return nil, fmt.Errorf("conflicting path information from mount and mountpoint %q", mp.Name)
		case !exists:
			namedVolumeMounts[mp.Name] = volumeMountTuple{M: schema.Mount{Volume: mp.Name, Path: mp.Path}}
			diag.Printf("adding %+v", namedVolumeMounts[mp.Name])
		}
	}

	// Insert the command-line Volumes
	for _, v := range p.Manifest.Volumes {
		// Check if we have a mount for this volume
		tuple, exists := namedVolumeMounts[v.Name]
		if !exists {
			return nil, fmt.Errorf("missing mount for volume %q", v.Name)
		} else if tuple.M.Volume != v.Name {
			// assertion regarding the implementation, should never happen
			return nil, fmt.Errorf("mismatched volume:mount pair: %q != %q", v.Name, tuple.M.Volume)
		}
		namedVolumeMounts[v.Name] = volumeMountTuple{V: v, M: tuple.M}
		diag.Printf("adding %+v", namedVolumeMounts[v.Name])
	}

	// Merge command-line Volumes with ImageManifest's MountPoints
	for _, mp := range imApp.MountPoints {
		// Check if we have a volume for this mountpoint
		tuple, exists := namedVolumeMounts[mp.Name]
		if !exists || tuple.V.Name == "" {
			return nil, fmt.Errorf("missing volume for mountpoint %q", mp.Name)
		}

		// If empty, fill in ReadOnly bit
		if tuple.V.ReadOnly == nil {
			v := tuple.V
			v.ReadOnly = &mp.ReadOnly
			namedVolumeMounts[mp.Name] = volumeMountTuple{M: tuple.M, V: v}
			diag.Printf("adding %+v", namedVolumeMounts[mp.Name])
		}
	}

	// Gather host mounts which we make MS_SHARED if passed as a volume source
	hostMounts, err := getHostMounts()
	if err != nil {
		return nil, errwrap.Wrap(errors.New("can't gather host mounts"), err)
	}

	argFlyMounts := []flyMount{}
	var flags uintptr = syscall.MS_BIND // TODO: allow optional | syscall.MS_REC
	for _, tuple := range namedVolumeMounts {
		if _, isHostMount := hostMounts[tuple.V.Source]; isHostMount {
			// Mark the host mount as SHARED so the container's changes to the mount are propagated to the host
			argFlyMounts = append(argFlyMounts,
				flyMount{"", "", tuple.V.Source, "none", syscall.MS_REC | syscall.MS_SHARED},
			)
		}
		argFlyMounts = append(argFlyMounts,
			flyMount{tuple.V.Source, rfs, tuple.M.Path, "none", flags},
		)

		if tuple.V.ReadOnly != nil && *tuple.V.ReadOnly {
			argFlyMounts = append(argFlyMounts,
				flyMount{"", rfs, tuple.M.Path, "none", flags | syscall.MS_REMOUNT | syscall.MS_RDONLY},
			)
		}
	}
	return argFlyMounts, nil
}

func stage1() int {
	uuid, err := types.NewUUID(flag.Arg(0))
	if err != nil {
		log.Print("UUID is missing or malformed\n")
		return 1
	}

	root := "."
	p, err := stage1commontypes.LoadPod(root, uuid)
	if err != nil {
		log.PrintE("can't load pod", err)
		return 1
	}

	// Sanity checks
	if len(p.Manifest.Apps) != 1 {
		log.Printf("flavor %q only supports 1 application per Pod for now", flavor)
		return 1
	}

	ra := p.Manifest.Apps[0]

	imgName := p.AppNameToImageName(ra.Name)
	args := ra.App.Exec
	if len(args) == 0 {
		log.Printf(`image %q has an empty "exec" (try --exec=BINARY)`, imgName)
		return 1
	}

	lfd, err := common.GetRktLockFD()
	if err != nil {
		log.PrintE("can't get rkt lock fd", err)
		return 1
	}

	// set close-on-exec flag on RKT_LOCK_FD so it gets correctly closed after execution is finished
	if err := sys.CloseOnExec(lfd, true); err != nil {
		log.PrintE("can't set FD_CLOEXEC on rkt lock", err)
		return 1
	}

	env := []string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"}
	for _, e := range ra.App.Environment {
		env = append(env, e.Name+"="+e.Value)
	}

	rfs := filepath.Join(common.AppPath(p.Root, ra.Name), "rootfs")

	argFlyMounts, err := evaluateMounts(rfs, string(ra.Name), p)
	if err != nil {
		log.PrintE("can't evaluate mounts", err)
		return 1
	}

	effectiveMounts := append(
		[]flyMount{
			{"", "", "/dev", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/dev", rfs, "/dev", "none", syscall.MS_BIND | syscall.MS_REC},

			{"", "", "/proc", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/proc", rfs, "/proc", "none", syscall.MS_BIND | syscall.MS_REC},

			{"", "", "/sys", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/sys", rfs, "/sys", "none", syscall.MS_BIND | syscall.MS_REC},

			{"tmpfs", rfs, "/tmp", "tmpfs", 0},
		},
		argFlyMounts...,
	)

	for _, mount := range effectiveMounts {
		var (
			err            error
			hostPathInfo   os.FileInfo
			targetPathInfo os.FileInfo
		)

		if strings.HasPrefix(mount.HostPath, "/") {
			if hostPathInfo, err = os.Stat(mount.HostPath); err != nil {
				log.PrintE(fmt.Sprintf("stat of host directory %s", mount.HostPath), err)
				return 1
			}
		} else {
			hostPathInfo = nil
		}

		absTargetPath := filepath.Join(mount.TargetPrefixPath, mount.RelTargetPath)
		if targetPathInfo, err = os.Stat(absTargetPath); err != nil && !os.IsNotExist(err) {
			log.PrintE(fmt.Sprintf("stat of target directory %s", absTargetPath), err)
			return 1
		}

		switch {
		case targetPathInfo == nil:
			absTargetPathParent, _ := filepath.Split(absTargetPath)
			if err := os.MkdirAll(absTargetPathParent, 0700); err != nil {
				log.PrintE(fmt.Sprintf("can't create directory %q", absTargetPath), err)
				return 1
			}
			switch {
			case hostPathInfo == nil || hostPathInfo.IsDir():
				if err := os.Mkdir(absTargetPath, 0700); err != nil {
					log.PrintE(fmt.Sprintf("can't create directory %q", absTargetPath), err)
					return 1
				}
			case !hostPathInfo.IsDir():
				file, err := os.OpenFile(absTargetPath, os.O_CREATE, 0700)
				if err != nil {
					log.PrintE(fmt.Sprintf("can't create file %q", absTargetPath), err)
					return 1
				}
				file.Close()
			}
		case hostPathInfo != nil:
			switch {
			case hostPathInfo.IsDir() && !targetPathInfo.IsDir():
				log.Printf("can't mount because %q is a directory while %q is not", mount.HostPath, absTargetPath)
				return 1
			case !hostPathInfo.IsDir() && targetPathInfo.IsDir():
				log.Printf("can't mount because %q is not a directory while %q is", mount.HostPath, absTargetPath)
				return 1
			}
		}

		if err := syscall.Mount(mount.HostPath, absTargetPath, mount.Fs, mount.Flags, ""); err != nil {
			log.PrintE(fmt.Sprintf("can't mount %q on %q with flags %v", mount.HostPath, absTargetPath, mount.Flags), err)
			return 1
		}
	}

	if err = stage1common.WritePpid(os.Getpid()); err != nil {
		log.Error(err)
		return 4
	}

	diag.Printf("chroot to %q", rfs)
	if err := syscall.Chroot(rfs); err != nil {
		log.PrintE("can't chroot", err)
		return 1
	}

	if err := os.Chdir("/"); err != nil {
		log.PrintE("can't change to root new directory", err)
		return 1
	}

	diag.Printf("execing %q in %q", args, rfs)
	err = stage1common.WithClearedCloExec(lfd, func() error {
		return syscall.Exec(args[0], args, env)
	})
	if err != nil {
		log.PrintE(fmt.Sprintf("can't execute %q", args[0]), err)
		return 7
	}

	return 0
}

func main() {
	flag.Parse()

	log, diag, _ = rktlog.NewLogSet("run", debug)
	if !debug {
		diag.SetOutput(ioutil.Discard)
	}

	// move code into stage1() helper so defered fns get run
	os.Exit(stage1())
}
<commit_msg>stage1/fly: add workingDirectory support<commit_after>// Copyright 2015 The rkt Authors
//
// 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 main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/hashicorp/errwrap"

	"github.com/coreos/rkt/common"
	rktlog "github.com/coreos/rkt/pkg/log"
	"github.com/coreos/rkt/pkg/sys"
	stage1common "github.com/coreos/rkt/stage1/common"
	stage1commontypes "github.com/coreos/rkt/stage1/common/types"
)

const (
	flavor = "fly"
)

type flyMount struct {
	HostPath         string
	TargetPrefixPath string
	RelTargetPath    string
	Fs               string
	Flags            uintptr
}

type volumeMountTuple struct {
	V types.Volume
	M schema.Mount
}

var (
	debug bool

	discardNetlist common.NetList
	discardBool    bool
	discardString  string

	log  *rktlog.Logger
	diag *rktlog.Logger
)

func getHostMounts() (map[string]struct{}, error) {
	hostMounts := map[string]struct{}{}

	mi, err := os.Open("/proc/self/mountinfo")
	if err != nil {
		return nil, err
	}
	defer mi.Close()

	sc := bufio.NewScanner(mi)
	for sc.Scan() {
		var (
			discard    string
			mountPoint string
		)

		_, err := fmt.Sscanf(sc.Text(),
			"%s %s %s %s %s",
			&discard, &discard, &discard, &discard, &mountPoint,
		)
		if err != nil {
			return nil, err
		}

		hostMounts[mountPoint] = struct{}{}
	}
	if sc.Err() != nil {
		return nil, errwrap.Wrap(errors.New("problem parsing mountinfo"), sc.Err())
	}
	return hostMounts, nil
}

func init() {
	flag.BoolVar(&debug, "debug", false, "Run in debug mode")

	// The following flags need to be supported by stage1 according to
	// https://github.com/coreos/rkt/blob/master/Documentation/devel/stage1-implementors-guide.md
	// TODO: either implement functionality or give not implemented warnings
	flag.Var(&discardNetlist, "net", "Setup networking")
	flag.BoolVar(&discardBool, "interactive", true, "The pod is interactive")
	flag.StringVar(&discardString, "mds-token", "", "MDS auth token")
	flag.StringVar(&discardString, "local-config", common.DefaultLocalConfigDir, "Local config path")
}

func evaluateMounts(rfs string, app string, p *stage1commontypes.Pod) ([]flyMount, error) {
	imApp := p.Images[app].App
	namedVolumeMounts := map[types.ACName]volumeMountTuple{}

	for _, m := range p.Manifest.Apps[0].Mounts {
		_, exists := namedVolumeMounts[m.Volume]
		if exists {
			return nil, fmt.Errorf("duplicate mount given: %q", m.Volume)
		}
		namedVolumeMounts[m.Volume] = volumeMountTuple{M: m}
		diag.Printf("adding %+v", namedVolumeMounts[m.Volume])
	}

	// Merge command-line Mounts with ImageManifest's MountPoints
	for _, mp := range imApp.MountPoints {
		tuple, exists := namedVolumeMounts[mp.Name]
		switch {
		case exists && tuple.M.Path != mp.Path:
			return nil, fmt.Errorf("conflicting path information from mount and mountpoint %q", mp.Name)
		case !exists:
			namedVolumeMounts[mp.Name] = volumeMountTuple{M: schema.Mount{Volume: mp.Name, Path: mp.Path}}
			diag.Printf("adding %+v", namedVolumeMounts[mp.Name])
		}
	}

	// Insert the command-line Volumes
	for _, v := range p.Manifest.Volumes {
		// Check if we have a mount for this volume
		tuple, exists := namedVolumeMounts[v.Name]
		if !exists {
			return nil, fmt.Errorf("missing mount for volume %q", v.Name)
		} else if tuple.M.Volume != v.Name {
			// assertion regarding the implementation, should never happen
			return nil, fmt.Errorf("mismatched volume:mount pair: %q != %q", v.Name, tuple.M.Volume)
		}
		namedVolumeMounts[v.Name] = volumeMountTuple{V: v, M: tuple.M}
		diag.Printf("adding %+v", namedVolumeMounts[v.Name])
	}

	// Merge command-line Volumes with ImageManifest's MountPoints
	for _, mp := range imApp.MountPoints {
		// Check if we have a volume for this mountpoint
		tuple, exists := namedVolumeMounts[mp.Name]
		if !exists || tuple.V.Name == "" {
			return nil, fmt.Errorf("missing volume for mountpoint %q", mp.Name)
		}

		// If empty, fill in ReadOnly bit
		if tuple.V.ReadOnly == nil {
			v := tuple.V
			v.ReadOnly = &mp.ReadOnly
			namedVolumeMounts[mp.Name] = volumeMountTuple{M: tuple.M, V: v}
			diag.Printf("adding %+v", namedVolumeMounts[mp.Name])
		}
	}

	// Gather host mounts which we make MS_SHARED if passed as a volume source
	hostMounts, err := getHostMounts()
	if err != nil {
		return nil, errwrap.Wrap(errors.New("can't gather host mounts"), err)
	}

	argFlyMounts := []flyMount{}
	var flags uintptr = syscall.MS_BIND // TODO: allow optional | syscall.MS_REC
	for _, tuple := range namedVolumeMounts {
		if _, isHostMount := hostMounts[tuple.V.Source]; isHostMount {
			// Mark the host mount as SHARED so the container's changes to the mount are propagated to the host
			argFlyMounts = append(argFlyMounts,
				flyMount{"", "", tuple.V.Source, "none", syscall.MS_REC | syscall.MS_SHARED},
			)
		}
		argFlyMounts = append(argFlyMounts,
			flyMount{tuple.V.Source, rfs, tuple.M.Path, "none", flags},
		)

		if tuple.V.ReadOnly != nil && *tuple.V.ReadOnly {
			argFlyMounts = append(argFlyMounts,
				flyMount{"", rfs, tuple.M.Path, "none", flags | syscall.MS_REMOUNT | syscall.MS_RDONLY},
			)
		}
	}
	return argFlyMounts, nil
}

func stage1() int {
	uuid, err := types.NewUUID(flag.Arg(0))
	if err != nil {
		log.Print("UUID is missing or malformed\n")
		return 1
	}

	root := "."
	p, err := stage1commontypes.LoadPod(root, uuid)
	if err != nil {
		log.PrintE("can't load pod", err)
		return 1
	}

	// Sanity checks
	if len(p.Manifest.Apps) != 1 {
		log.Printf("flavor %q only supports 1 application per Pod for now", flavor)
		return 1
	}

	ra := p.Manifest.Apps[0]

	imgName := p.AppNameToImageName(ra.Name)
	args := ra.App.Exec
	if len(args) == 0 {
		log.Printf(`image %q has an empty "exec" (try --exec=BINARY)`, imgName)
		return 1
	}

	lfd, err := common.GetRktLockFD()
	if err != nil {
		log.PrintE("can't get rkt lock fd", err)
		return 1
	}

	// set close-on-exec flag on RKT_LOCK_FD so it gets correctly closed after execution is finished
	if err := sys.CloseOnExec(lfd, true); err != nil {
		log.PrintE("can't set FD_CLOEXEC on rkt lock", err)
		return 1
	}

	workDir := "/"
	if ra.App.WorkingDirectory != "" {
		workDir = ra.App.WorkingDirectory
	}

	env := []string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"}
	for _, e := range ra.App.Environment {
		env = append(env, e.Name+"="+e.Value)
	}

	rfs := filepath.Join(common.AppPath(p.Root, ra.Name), "rootfs")

	argFlyMounts, err := evaluateMounts(rfs, string(ra.Name), p)
	if err != nil {
		log.PrintE("can't evaluate mounts", err)
		return 1
	}

	effectiveMounts := append(
		[]flyMount{
			{"", "", "/dev", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/dev", rfs, "/dev", "none", syscall.MS_BIND | syscall.MS_REC},

			{"", "", "/proc", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/proc", rfs, "/proc", "none", syscall.MS_BIND | syscall.MS_REC},

			{"", "", "/sys", "none", syscall.MS_REC | syscall.MS_SHARED},
			{"/sys", rfs, "/sys", "none", syscall.MS_BIND | syscall.MS_REC},

			{"tmpfs", rfs, "/tmp", "tmpfs", 0},
		},
		argFlyMounts...,
	)

	for _, mount := range effectiveMounts {
		var (
			err            error
			hostPathInfo   os.FileInfo
			targetPathInfo os.FileInfo
		)

		if strings.HasPrefix(mount.HostPath, "/") {
			if hostPathInfo, err = os.Stat(mount.HostPath); err != nil {
				log.PrintE(fmt.Sprintf("stat of host directory %s", mount.HostPath), err)
				return 1
			}
		} else {
			hostPathInfo = nil
		}

		absTargetPath := filepath.Join(mount.TargetPrefixPath, mount.RelTargetPath)
		if targetPathInfo, err = os.Stat(absTargetPath); err != nil && !os.IsNotExist(err) {
			log.PrintE(fmt.Sprintf("stat of target directory %s", absTargetPath), err)
			return 1
		}

		switch {
		case targetPathInfo == nil:
			absTargetPathParent, _ := filepath.Split(absTargetPath)
			if err := os.MkdirAll(absTargetPathParent, 0700); err != nil {
				log.PrintE(fmt.Sprintf("can't create directory %q", absTargetPath), err)
				return 1
			}
			switch {
			case hostPathInfo == nil || hostPathInfo.IsDir():
				if err := os.Mkdir(absTargetPath, 0700); err != nil {
					log.PrintE(fmt.Sprintf("can't create directory %q", absTargetPath), err)
					return 1
				}
			case !hostPathInfo.IsDir():
				file, err := os.OpenFile(absTargetPath, os.O_CREATE, 0700)
				if err != nil {
					log.PrintE(fmt.Sprintf("can't create file %q", absTargetPath), err)
					return 1
				}
				file.Close()
			}
		case hostPathInfo != nil:
			switch {
			case hostPathInfo.IsDir() && !targetPathInfo.IsDir():
				log.Printf("can't mount because %q is a directory while %q is not", mount.HostPath, absTargetPath)
				return 1
			case !hostPathInfo.IsDir() && targetPathInfo.IsDir():
				log.Printf("can't mount because %q is not a directory while %q is", mount.HostPath, absTargetPath)
				return 1
			}
		}

		if err := syscall.Mount(mount.HostPath, absTargetPath, mount.Fs, mount.Flags, ""); err != nil {
			log.PrintE(fmt.Sprintf("can't mount %q on %q with flags %v", mount.HostPath, absTargetPath, mount.Flags), err)
			return 1
		}
	}

	if err = stage1common.WritePpid(os.Getpid()); err != nil {
		log.Error(err)
		return 4
	}

	diag.Printf("chroot to %q", rfs)
	if err := syscall.Chroot(rfs); err != nil {
		log.PrintE("can't chroot", err)
		return 1
	}

	if err := os.Chdir(workDir); err != nil {
		log.PrintE(fmt.Sprintf("can't change to working directory %q", workDir), err)
		return 1
	}

	diag.Printf("execing %q in %q", args, rfs)
	err = stage1common.WithClearedCloExec(lfd, func() error {
		return syscall.Exec(args[0], args, env)
	})
	if err != nil {
		log.PrintE(fmt.Sprintf("can't execute %q", args[0]), err)
		return 7
	}

	return 0
}

func main() {
	flag.Parse()

	log, diag, _ = rktlog.NewLogSet("run", debug)
	if !debug {
		diag.SetOutput(ioutil.Discard)
	}

	// move code into stage1() helper so defered fns get run
	os.Exit(stage1())
}
<|endoftext|>