<commit_before>//
// mod_http.go
//
// http mod panel
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "github.com/gorilla/sessions"
  "errors"
  "encoding/hex"
  "encoding/json"
  "fmt"
  "io"
  "log"
  "net/http"
  "strings"
)

type httpModUI struct {
  regenAll func()
  regen func (ArticleEntry)
  delete func (string)
  modMessageChan chan NNTPMessage
  database Database
  articles ArticleStore
  store *sessions.CookieStore
  prefix string
  mod_prefix string
}

func createHttpModUI(frontend httpFrontend) httpModUI {
  return httpModUI{frontend.regenAll, frontend.Regen, frontend.deleteThreadMarkup, make(chan NNTPMessage), frontend.daemon.database, frontend.daemon.store, frontend.store, frontend.prefix, frontend.prefix + "mod/"}

}
func (self httpModUI) getAdminFunc(funcname string) AdminFunc {
  if funcname == "template.reload" {
    return func(param map[string]interface{}) (string, error) {
      tname, ok := param["template"]
      if ok {
        t := ""
        switch tname.(type) {
        case string:
          t = tname.(string)
        default:
          return "failed to reload templates", errors.New("invalid parameters")
        }
        template.reloadTemplate(t)
        return "reloaded " + t, nil
      }
      template.reloadAllTemplates()
      return "reloaded all templates", nil
    }
  } else if funcname == "frontend.regen" {
    return func(param map[string]interface{}) (string, error) {
      self.regenAll()
      return "started regeneration", nil
    }
  }
  return nil
}

// handle an admin action
func (self httpModUI) HandleAdminCommand(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(url string) {
    action := strings.Split(url, "/admin/")[1]
    f := self.getAdminFunc(action)
    if f == nil {
      wr.WriteHeader(404)
    } else {
      var msg string
      var err error
      req := make(map[string]interface{})
      if r.Method == "POST" {
        dec := json.NewDecoder(r.Body)
        err = dec.Decode(req)
      }
      if err == nil {
        msg, err = f(req)
      }
      resp := make(map[string]interface{})
      if err == nil {
        resp["error"] = nil
      } else {
        resp["error"] = err.Error()
      }
      resp["result"] = msg
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    }

  }, wr, r)
}

// TODO: check for different levels of permissions
func (self httpModUI) CheckKey(privkey string) (bool, error) {
  privkey_bytes, err := hex.DecodeString(privkey)
  if err == nil {
    pubkey_bytes := nacl.GetSignPubkey(privkey_bytes)
    if pubkey_bytes != nil {
      pubkey := hex.EncodeToString(pubkey_bytes)
      if self.database.CheckModPubkeyGlobal(pubkey) {
        // this user is an admin
        return true, nil
      } else {
        return false, nil
      }
    }
  }
  log.Println("invalid key format for key", privkey)
  return false, err
}

func (self httpModUI) MessageChan() chan NNTPMessage {
  return self.modMessageChan
}

func (self httpModUI) getSession(r *http.Request) *sessions.Session {
  s, _ := self.store.Get(r, "nntpchan-mod")
  return s
}

// get the session's private key as bytes or nil if we don't have it
func (self httpModUI) getSessionPrivkeyBytes(r *http.Request) []byte {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    privkey_bytes, err := hex.DecodeString(k.(string))
    if err == nil {
      return privkey_bytes
    }
    log.Println("failed to decode private key bytes from session", err)
  } else {
    log.Println("failed to get private key from session, no private key in session?")
  }
  return nil
}

// returns true if the session is okay
// otherwise redirect to login page
func (self httpModUI) checkSession(r *http.Request) bool {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    ok, err := self.CheckKey(k.(string))
    if err != nil {
      return false
    }
    return ok
  }
  return false
}

func (self httpModUI) writeTemplate(wr http.ResponseWriter, name string) {
  self.writeTemplateParam(wr, name, nil)
}

func (self httpModUI) writeTemplateParam(wr http.ResponseWriter, name string, param map[string]string) {
  if param == nil {
    param = make(map[string]string)
  }
  param["prefix"] = self.prefix
  param["mod_prefix"] = self.mod_prefix
  io.WriteString(wr, template.renderTemplate(name, param))
}

// do a function as authenticated
// pass in the request path to the handler
func (self httpModUI) asAuthed(handler func(string), wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    handler(r.URL.Path)
  } else {
    wr.WriteHeader(403)
  }
}

// do stuff to a certain message if with have it and are authed
func (self httpModUI) asAuthedWithMessage(handler func(ArticleEntry, *http.Request) map[string]interface{}, wr http.ResponseWriter, req *http.Request) {
  self.asAuthed(func(path string) {
    // get the long hash
    if strings.Count(path, "/") > 2 {
      // TOOD: prefix detection
      longhash := strings.Split(path, "/")[3]
      // get the message id
      msg, err := self.database.GetMessageIDByHash(longhash)
      resp := make(map[string]interface{})
      if err == nil {
        resp = handler(msg, req)
      } else {
        resp["error"] = fmt.Sprintf("don't have message with hash %s, %s", longhash, err.Error())
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, req)
}

func (self httpModUI) HandleAddPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleDelPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleUnbanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(path string) {
    // extract the ip address
    // TODO: ip ranges and prefix detection
    if strings.Count(path, "/") > 2 {
      addr := strings.Split(path, "/")[3]
      resp := make(map[string]interface{})
      banned, err := self.database.CheckIPBanned(addr)
      if err != nil {
        resp["error"] = fmt.Sprintf("cannot tell if %s is banned: %s", addr, err.Error())
      } else if banned {
        // TODO: rangebans
        err = self.database.UnbanAddr(addr)
        if err == nil {
          resp["result"] = fmt.Sprintf("%s was unbanned", addr)
        } else {
          resp["error"] = err.Error()
        }
      } else {
        resp["error"] = fmt.Sprintf("%s was not banned", addr)
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, r)
}

// handle ban logic
func (self httpModUI) handleBanAddress(msg ArticleEntry, r *http.Request) map[string]interface{} {
  // get the article headers
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    // we don't got it?!
    resp["error"] = fmt.Sprintf("message %s not on the filesystem wtf?", msgid)
  } else {
    // get the associated encrypted ip
    encip := hdr.Get("X-Encrypted-Ip", hdr.Get("X-Encrypted-IP", ""))
    encip = strings.Trim(encip, "\t ")

    if len(encip) == 0 {
      // no ip header detected
      resp["error"] = fmt.Sprintf("%s has no IP, ban Tor instead", msgid)
    } else {
      // get the ip address if we have it
      ip, err := self.database.GetIPAddress(encip)
      if len(ip) > 0 {
        // we have it
        // ban the address
        err = self.database.BanAddr(ip)
        // then we tell everyone about it
        var key string
        // TODO: we SHOULD have the key, but what if we do not?
        key, err = self.database.GetEncKey(encip)
        // create mod message
        // TODO: hardcoded ban period
        mm := ModMessage{overchanInetBan(encip, key, -1),}
        privkey_bytes := self.getSessionPrivkeyBytes(r)
        if privkey_bytes == nil {
          // this should not happen
          log.Println("failed to get privkey bytes from session")
          resp["error"] = "failed to get private key from session. wtf?"
        } else {
          // wrap and sign
          nntp := wrapModMessage(mm)
          nntp, err = signArticle(nntp, privkey_bytes)
          if err == nil {
            // federate
            self.modMessageChan <- nntp
          }
        }
      } else {
        // we don't have it
        // ban the encrypted version
        err = self.database.BanEncAddr(encip)
      }
      if err == nil {
        result_msg :=  fmt.Sprintf("We banned %s", encip)
        if len(ip) > 0 {
          result_msg += fmt.Sprintf(" (%s)", ip)
        }
        resp["banned"] = result_msg
      } else {
        resp["error"] = err.Error()
      }

    }
  }
  return resp
}

func (self httpModUI) handleDeletePost(msg ArticleEntry, r *http.Request) map[string]interface{} {
  var mm ModMessage
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  delmsgs := []string{}
  // get headers
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    resp["error"] = fmt.Sprintf("message %s is not on the filesystem? wtf!", msgid)
  } else {
    ref := hdr.Get("References", hdr.Get("Reference", ""))
    ref = strings.Trim(ref, "\t ")
    // is it a root post?
    if ref == "" {
      // load replies
      replies := self.database.GetThreadReplies(msgid, 0)
      if replies != nil {
        for _, repl := range(replies) {
          // append mod line to mod message for reply
          mm = append(mm, overchanDelete(repl))
          // add to delete queue
          delmsgs = append(delmsgs, repl)
        }
      }
    }
    delmsgs = append(delmsgs, msgid)
    // append mod line to mod message
    mm = append(mm, overchanDelete(msgid))

    resp["deleted"] = delmsgs
    // only regen threads when we delete a non root port
    if ref != "" {
      group := hdr.Get("Newsgroups", "")
      self.regen(ArticleEntry{
        ref, group,
      })
    }
    privkey_bytes := self.getSessionPrivkeyBytes(r)
    if privkey_bytes == nil {
      // crap this should never happen
      log.Println("failed to get private keys from session, not federating")
    } else {
      // wrap and sign mod message
      nntp, err := signArticle(wrapModMessage(mm), privkey_bytes)
      if err == nil {
        // send it off to federate
        self.modMessageChan <- nntp
      } else {
        resp["error"] = fmt.Sprintf("signing error: %s", err.Error())
      }
    }
  }
  return resp
}

// ban the address of a poster
func (self httpModUI) HandleBanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleBanAddress, wr, r)
}

// delete a post
func (self httpModUI) HandleDeletePost(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleDeletePost, wr, r)
}

func (self httpModUI) HandleLogin(wr http.ResponseWriter, r *http.Request) {
  privkey := r.FormValue("privkey")
  msg := "failed login: "
  if len(privkey) == 0 {
    msg += "no key"
  } else {
    ok, err := self.CheckKey(privkey)
    if err != nil {
      msg += fmt.Sprintf("%s", err)
    } else if ok {
      msg = "login okay"
      sess := self.getSession(r)
      sess.Values["privkey"] = privkey
      sess.Save(r, wr)
    } else {
      msg += "invalid key"
    }
  }
  self.writeTemplateParam(wr, "modlogin_result.mustache", map[string]string { "message" : msg })
}

func (self httpModUI) HandleKeyGen(wr http.ResponseWriter, r *http.Request) {
  pk, sk := newSignKeypair()
  tripcode := makeTripcode(pk)
  self.writeTemplateParam(wr, "keygen.mustache", map[string]string {"public" : pk, "secret" : sk, "tripcode" : tripcode})
}

func (self httpModUI) ServeModPage(wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    // we are logged in
    // serve mod page
    self.writeTemplate(wr, "modpage.mustache")
  } else {
    // we are not logged in
    // serve login page
    self.writeTemplate(wr, "modlogin.mustache")
  }

}
<commit_msg>fork off regen it takes a bit<commit_after>//
// mod_http.go
//
// http mod panel
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "github.com/gorilla/sessions"
  "errors"
  "encoding/hex"
  "encoding/json"
  "fmt"
  "io"
  "log"
  "net/http"
  "strings"
)

type httpModUI struct {
  regenAll func()
  regen func (ArticleEntry)
  delete func (string)
  modMessageChan chan NNTPMessage
  database Database
  articles ArticleStore
  store *sessions.CookieStore
  prefix string
  mod_prefix string
}

func createHttpModUI(frontend httpFrontend) httpModUI {
  return httpModUI{frontend.regenAll, frontend.Regen, frontend.deleteThreadMarkup, make(chan NNTPMessage), frontend.daemon.database, frontend.daemon.store, frontend.store, frontend.prefix, frontend.prefix + "mod/"}

}
func (self httpModUI) getAdminFunc(funcname string) AdminFunc {
  if funcname == "template.reload" {
    return func(param map[string]interface{}) (string, error) {
      tname, ok := param["template"]
      if ok {
        t := ""
        switch tname.(type) {
        case string:
          t = tname.(string)
        default:
          return "failed to reload templates", errors.New("invalid parameters")
        }
        template.reloadTemplate(t)
        return "reloaded " + t, nil
      }
      template.reloadAllTemplates()
      return "reloaded all templates", nil
    }
  } else if funcname == "frontend.regen" {
    return func(param map[string]interface{}) (string, error) {
      go self.regenAll()
      return "started regeneration", nil
    }
  }
  return nil
}

// handle an admin action
func (self httpModUI) HandleAdminCommand(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(url string) {
    action := strings.Split(url, "/admin/")[1]
    f := self.getAdminFunc(action)
    if f == nil {
      wr.WriteHeader(404)
    } else {
      var msg string
      var err error
      req := make(map[string]interface{})
      if r.Method == "POST" {
        dec := json.NewDecoder(r.Body)
        err = dec.Decode(req)
      }
      if err == nil {
        msg, err = f(req)
      }
      resp := make(map[string]interface{})
      if err == nil {
        resp["error"] = nil
      } else {
        resp["error"] = err.Error()
      }
      resp["result"] = msg
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    }

  }, wr, r)
}

// TODO: check for different levels of permissions
func (self httpModUI) CheckKey(privkey string) (bool, error) {
  privkey_bytes, err := hex.DecodeString(privkey)
  if err == nil {
    pubkey_bytes := nacl.GetSignPubkey(privkey_bytes)
    if pubkey_bytes != nil {
      pubkey := hex.EncodeToString(pubkey_bytes)
      if self.database.CheckModPubkeyGlobal(pubkey) {
        // this user is an admin
        return true, nil
      } else {
        return false, nil
      }
    }
  }
  log.Println("invalid key format for key", privkey)
  return false, err
}

func (self httpModUI) MessageChan() chan NNTPMessage {
  return self.modMessageChan
}

func (self httpModUI) getSession(r *http.Request) *sessions.Session {
  s, _ := self.store.Get(r, "nntpchan-mod")
  return s
}

// get the session's private key as bytes or nil if we don't have it
func (self httpModUI) getSessionPrivkeyBytes(r *http.Request) []byte {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    privkey_bytes, err := hex.DecodeString(k.(string))
    if err == nil {
      return privkey_bytes
    }
    log.Println("failed to decode private key bytes from session", err)
  } else {
    log.Println("failed to get private key from session, no private key in session?")
  }
  return nil
}

// returns true if the session is okay
// otherwise redirect to login page
func (self httpModUI) checkSession(r *http.Request) bool {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    ok, err := self.CheckKey(k.(string))
    if err != nil {
      return false
    }
    return ok
  }
  return false
}

func (self httpModUI) writeTemplate(wr http.ResponseWriter, name string) {
  self.writeTemplateParam(wr, name, nil)
}

func (self httpModUI) writeTemplateParam(wr http.ResponseWriter, name string, param map[string]string) {
  if param == nil {
    param = make(map[string]string)
  }
  param["prefix"] = self.prefix
  param["mod_prefix"] = self.mod_prefix
  io.WriteString(wr, template.renderTemplate(name, param))
}

// do a function as authenticated
// pass in the request path to the handler
func (self httpModUI) asAuthed(handler func(string), wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    handler(r.URL.Path)
  } else {
    wr.WriteHeader(403)
  }
}

// do stuff to a certain message if with have it and are authed
func (self httpModUI) asAuthedWithMessage(handler func(ArticleEntry, *http.Request) map[string]interface{}, wr http.ResponseWriter, req *http.Request) {
  self.asAuthed(func(path string) {
    // get the long hash
    if strings.Count(path, "/") > 2 {
      // TOOD: prefix detection
      longhash := strings.Split(path, "/")[3]
      // get the message id
      msg, err := self.database.GetMessageIDByHash(longhash)
      resp := make(map[string]interface{})
      if err == nil {
        resp = handler(msg, req)
      } else {
        resp["error"] = fmt.Sprintf("don't have message with hash %s, %s", longhash, err.Error())
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, req)
}

func (self httpModUI) HandleAddPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleDelPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleUnbanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(path string) {
    // extract the ip address
    // TODO: ip ranges and prefix detection
    if strings.Count(path, "/") > 2 {
      addr := strings.Split(path, "/")[3]
      resp := make(map[string]interface{})
      banned, err := self.database.CheckIPBanned(addr)
      if err != nil {
        resp["error"] = fmt.Sprintf("cannot tell if %s is banned: %s", addr, err.Error())
      } else if banned {
        // TODO: rangebans
        err = self.database.UnbanAddr(addr)
        if err == nil {
          resp["result"] = fmt.Sprintf("%s was unbanned", addr)
        } else {
          resp["error"] = err.Error()
        }
      } else {
        resp["error"] = fmt.Sprintf("%s was not banned", addr)
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, r)
}

// handle ban logic
func (self httpModUI) handleBanAddress(msg ArticleEntry, r *http.Request) map[string]interface{} {
  // get the article headers
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    // we don't got it?!
    resp["error"] = fmt.Sprintf("message %s not on the filesystem wtf?", msgid)
  } else {
    // get the associated encrypted ip
    encip := hdr.Get("X-Encrypted-Ip", hdr.Get("X-Encrypted-IP", ""))
    encip = strings.Trim(encip, "\t ")

    if len(encip) == 0 {
      // no ip header detected
      resp["error"] = fmt.Sprintf("%s has no IP, ban Tor instead", msgid)
    } else {
      // get the ip address if we have it
      ip, err := self.database.GetIPAddress(encip)
      if len(ip) > 0 {
        // we have it
        // ban the address
        err = self.database.BanAddr(ip)
        // then we tell everyone about it
        var key string
        // TODO: we SHOULD have the key, but what if we do not?
        key, err = self.database.GetEncKey(encip)
        // create mod message
        // TODO: hardcoded ban period
        mm := ModMessage{overchanInetBan(encip, key, -1),}
        privkey_bytes := self.getSessionPrivkeyBytes(r)
        if privkey_bytes == nil {
          // this should not happen
          log.Println("failed to get privkey bytes from session")
          resp["error"] = "failed to get private key from session. wtf?"
        } else {
          // wrap and sign
          nntp := wrapModMessage(mm)
          nntp, err = signArticle(nntp, privkey_bytes)
          if err == nil {
            // federate
            self.modMessageChan <- nntp
          }
        }
      } else {
        // we don't have it
        // ban the encrypted version
        err = self.database.BanEncAddr(encip)
      }
      if err == nil {
        result_msg :=  fmt.Sprintf("We banned %s", encip)
        if len(ip) > 0 {
          result_msg += fmt.Sprintf(" (%s)", ip)
        }
        resp["banned"] = result_msg
      } else {
        resp["error"] = err.Error()
      }

    }
  }
  return resp
}

func (self httpModUI) handleDeletePost(msg ArticleEntry, r *http.Request) map[string]interface{} {
  var mm ModMessage
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  delmsgs := []string{}
  // get headers
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    resp["error"] = fmt.Sprintf("message %s is not on the filesystem? wtf!", msgid)
  } else {
    ref := hdr.Get("References", hdr.Get("Reference", ""))
    ref = strings.Trim(ref, "\t ")
    // is it a root post?
    if ref == "" {
      // load replies
      replies := self.database.GetThreadReplies(msgid, 0)
      if replies != nil {
        for _, repl := range(replies) {
          // append mod line to mod message for reply
          mm = append(mm, overchanDelete(repl))
          // add to delete queue
          delmsgs = append(delmsgs, repl)
        }
      }
    }
    delmsgs = append(delmsgs, msgid)
    // append mod line to mod message
    mm = append(mm, overchanDelete(msgid))

    resp["deleted"] = delmsgs
    // only regen threads when we delete a non root port
    if ref != "" {
      group := hdr.Get("Newsgroups", "")
      self.regen(ArticleEntry{
        ref, group,
      })
    }
    privkey_bytes := self.getSessionPrivkeyBytes(r)
    if privkey_bytes == nil {
      // crap this should never happen
      log.Println("failed to get private keys from session, not federating")
    } else {
      // wrap and sign mod message
      nntp, err := signArticle(wrapModMessage(mm), privkey_bytes)
      if err == nil {
        // send it off to federate
        self.modMessageChan <- nntp
      } else {
        resp["error"] = fmt.Sprintf("signing error: %s", err.Error())
      }
    }
  }
  return resp
}

// ban the address of a poster
func (self httpModUI) HandleBanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleBanAddress, wr, r)
}

// delete a post
func (self httpModUI) HandleDeletePost(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleDeletePost, wr, r)
}

func (self httpModUI) HandleLogin(wr http.ResponseWriter, r *http.Request) {
  privkey := r.FormValue("privkey")
  msg := "failed login: "
  if len(privkey) == 0 {
    msg += "no key"
  } else {
    ok, err := self.CheckKey(privkey)
    if err != nil {
      msg += fmt.Sprintf("%s", err)
    } else if ok {
      msg = "login okay"
      sess := self.getSession(r)
      sess.Values["privkey"] = privkey
      sess.Save(r, wr)
    } else {
      msg += "invalid key"
    }
  }
  self.writeTemplateParam(wr, "modlogin_result.mustache", map[string]string { "message" : msg })
}

func (self httpModUI) HandleKeyGen(wr http.ResponseWriter, r *http.Request) {
  pk, sk := newSignKeypair()
  tripcode := makeTripcode(pk)
  self.writeTemplateParam(wr, "keygen.mustache", map[string]string {"public" : pk, "secret" : sk, "tripcode" : tripcode})
}

func (self httpModUI) ServeModPage(wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    // we are logged in
    // serve mod page
    self.writeTemplate(wr, "modpage.mustache")
  } else {
    // we are not logged in
    // serve login page
    self.writeTemplate(wr, "modlogin.mustache")
  }

}
<|endoftext|>


<commit_before>package driver

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/TrilliumIT/iputil"
	"net"
	"sync"
	"time"
)

type candidateNets struct {
	nets map[string]*candidateList // map of network to slice of IPs
	lock sync.Mutex
	quit <-chan struct{}
}

type candidateList struct {
	candidates [candidateSize]*net.IPNet
	quit       <-chan struct{}
	popCh      chan chan *net.IPNet
	addCh      chan *net.IPNet
	delCh      chan *net.IPNet
}

// Does nothing if net already exists
func (cn *candidateNets) addNet(n *net.IPNet, ns *NeighSubscription) *candidateList {
	cn.lock.Lock()
	defer cn.lock.Unlock()
	if cl, ok := cn.nets[n.String()]; ok {
		return cl
	}
	cl := &candidateList{
		quit:  cn.quit,
		popCh: make(chan chan *net.IPNet),
		addCh: make(chan *net.IPNet),
		delCh: make(chan *net.IPNet),
	}
	go cl.fill(n, ns)
	cn.nets[n.String()] = cl
	return cl
}

func (cl *candidateList) pop(ns *NeighSubscription) *net.IPNet {
	pc := make(chan *net.IPNet)
	defer close(pc)
	cl.popCh <- pc
	return <-pc
}

func (cl *candidateList) fill(n *net.IPNet, ns *NeighSubscription) {
	t := time.NewTicker(3 * time.Second)
	defer t.Stop()
	sub := ns.addSub()
	uch := sub.sub
	defer sub.delSub()

mainLoop:
	for {
		select {
		case pc := <-cl.popCh: // pop a suggested address
			for i, p := range cl.candidates {
				if p != nil {
					log.WithField("ip", p).Debug("Popping address from suggestions")
					cl.candidates[i] = nil
					pc <- p
				}
			}
			go func() {
				r, err := getNewRandomUnusedAddr(n, ns)
				if err != nil {
					log.WithError(err).Error("Error getting new random candidate address for pop")
					pc <- nil
				}
				pc <- r
			}()
			continue mainLoop
		case ip := <-cl.addCh:
			for i, p := range cl.candidates {
				if p == nil {
					cl.candidates[i] = ip
				}
			}
			continue mainLoop
		case ip := <-cl.delCh:
			for i, p := range cl.candidates {
				if p.IP.Equal(ip.IP) {
					cl.candidates[i] = nil
					go func() {
						addr, err := getNewRandomUnusedAddr(n, ns)
						if err != nil {
							log.WithError(err).Error("Error getting new random address.")
							return
						}
						cl.addCh <- addr
					}()
					continue mainLoop
				}
			}
			continue mainLoop
		case <-cl.quit:
			return
		case n := <-uch: // We got an update from the arp table
			for _, p := range cl.candidates { // Only do something if it is an ip we care about
				if p != nil && p.IP.Equal(n.IP) {
					continue mainLoop
				}
			}
			break
		case <-t.C: // need to refresh because the timer ticked
			break
		}

		for _, p := range cl.candidates {
			if p == nil {
				continue
			}
			go func() {
				r, err := ns.probeAndWait(p.IP)
				if err != nil {
					log.WithError(err).WithField("ip", p).Error("Error probing candidate IP")
					cl.delCh <- p
					return
				}
				if r {
					log.WithField("ip", p).Debug("Candidate IP in use")
					cl.delCh <- p
				}
			}()
		}
	}
}

func (d *Driver) getRandomUnusedAddr(n *net.IPNet) (*net.IPNet, error) {
	cl := d.candidates.addNet(n, d.ns)
	r := cl.pop(d.ns)
	if r != nil {
		return r, nil
	}
	r, err := getNewRandomUnusedAddr(n, d.ns)
	if err != nil {
		log.WithError(err).Error("Error getting new random address")
		return nil, err
	}
	return r, nil
}

func getNewRandomUnusedAddr(n *net.IPNet, ns *NeighSubscription) (*net.IPNet, error) {
	log.Debugf("Generating Random Address in network %v", n)
	tried := make(map[string]struct{})
	var e struct{}
	ones, maskSize := n.Mask.Size()
	totalAddresses := 1 << uint8(maskSize-ones)
	if totalAddresses > 2 { // This network is not a /30 exclude first and last
		tried[string(iputil.FirstAddr(n))] = e
		tried[string(iputil.LastAddr(n))] = e
	}
	for len(tried) < totalAddresses {
		ip, err := iputil.RandAddr(n)
		if err != nil {
			log.WithError(err).Error("Error generating random address")
			return nil, err
		}
		if _, ok := tried[string(ip)]; ok { // this address was already tried
			continue
		}
		r, err := ns.probeAndWait(ip)
		if err != nil {
			log.WithError(err).Error("Error probing random address")
			continue
		}
		if !r {
			log.WithField("IP", ip).Debug("Returning Random Address")
			return &net.IPNet{IP: ip, Mask: n.Mask}, nil
		}
		log.WithField("ip", ip).Debug("Random address reachable, retrying")
	}
	return nil, fmt.Errorf("All avaliable addresses are in use")
}
<commit_msg>missing continue<commit_after>package driver

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/TrilliumIT/iputil"
	"net"
	"sync"
	"time"
)

type candidateNets struct {
	nets map[string]*candidateList // map of network to slice of IPs
	lock sync.Mutex
	quit <-chan struct{}
}

type candidateList struct {
	candidates [candidateSize]*net.IPNet
	quit       <-chan struct{}
	popCh      chan chan *net.IPNet
	addCh      chan *net.IPNet
	delCh      chan *net.IPNet
}

// Does nothing if net already exists
func (cn *candidateNets) addNet(n *net.IPNet, ns *NeighSubscription) *candidateList {
	cn.lock.Lock()
	defer cn.lock.Unlock()
	if cl, ok := cn.nets[n.String()]; ok {
		return cl
	}
	cl := &candidateList{
		quit:  cn.quit,
		popCh: make(chan chan *net.IPNet),
		addCh: make(chan *net.IPNet),
		delCh: make(chan *net.IPNet),
	}
	go cl.fill(n, ns)
	cn.nets[n.String()] = cl
	return cl
}

func (cl *candidateList) pop(ns *NeighSubscription) *net.IPNet {
	pc := make(chan *net.IPNet)
	defer close(pc)
	cl.popCh <- pc
	return <-pc
}

func (cl *candidateList) fill(n *net.IPNet, ns *NeighSubscription) {
	t := time.NewTicker(3 * time.Second)
	defer t.Stop()
	sub := ns.addSub()
	uch := sub.sub
	defer sub.delSub()

mainLoop:
	for {
		select {
		case pc := <-cl.popCh: // pop a suggested address
			for i, p := range cl.candidates {
				if p != nil {
					log.WithField("ip", p).Debug("Popping address from suggestions")
					cl.candidates[i] = nil
					pc <- p
					continue mainLoop
				}
			}
			go func() {
				r, err := getNewRandomUnusedAddr(n, ns)
				if err != nil {
					log.WithError(err).Error("Error getting new random candidate address for pop")
					pc <- nil
				}
				pc <- r
			}()
			continue mainLoop
		case ip := <-cl.addCh:
			for i, p := range cl.candidates {
				if p == nil {
					cl.candidates[i] = ip
					continue mainLoop
				}
			}
			continue mainLoop
		case ip := <-cl.delCh:
			for i, p := range cl.candidates {
				if p.IP.Equal(ip.IP) {
					cl.candidates[i] = nil
					go func() {
						addr, err := getNewRandomUnusedAddr(n, ns)
						if err != nil {
							log.WithError(err).Error("Error getting new random address.")
							return
						}
						cl.addCh <- addr
					}()
					continue mainLoop
				}
			}
			continue mainLoop
		case <-cl.quit:
			return
		case n := <-uch: // We got an update from the arp table
			for _, p := range cl.candidates { // Only do something if it is an ip we care about
				if p != nil && p.IP.Equal(n.IP) {
					continue mainLoop
				}
			}
			break
		case <-t.C: // need to refresh because the timer ticked
			break
		}

		for _, p := range cl.candidates {
			if p == nil {
				continue
			}
			go func() {
				r, err := ns.probeAndWait(p.IP)
				if err != nil {
					log.WithError(err).WithField("ip", p).Error("Error probing candidate IP")
					cl.delCh <- p
					return
				}
				if r {
					log.WithField("ip", p).Debug("Candidate IP in use")
					cl.delCh <- p
				}
			}()
		}
	}
}

func (d *Driver) getRandomUnusedAddr(n *net.IPNet) (*net.IPNet, error) {
	cl := d.candidates.addNet(n, d.ns)
	r := cl.pop(d.ns)
	if r != nil {
		return r, nil
	}
	r, err := getNewRandomUnusedAddr(n, d.ns)
	if err != nil {
		log.WithError(err).Error("Error getting new random address")
		return nil, err
	}
	return r, nil
}

func getNewRandomUnusedAddr(n *net.IPNet, ns *NeighSubscription) (*net.IPNet, error) {
	log.Debugf("Generating Random Address in network %v", n)
	tried := make(map[string]struct{})
	var e struct{}
	ones, maskSize := n.Mask.Size()
	totalAddresses := 1 << uint8(maskSize-ones)
	if totalAddresses > 2 { // This network is not a /30 exclude first and last
		tried[string(iputil.FirstAddr(n))] = e
		tried[string(iputil.LastAddr(n))] = e
	}
	for len(tried) < totalAddresses {
		ip, err := iputil.RandAddr(n)
		if err != nil {
			log.WithError(err).Error("Error generating random address")
			return nil, err
		}
		if _, ok := tried[string(ip)]; ok { // this address was already tried
			continue
		}
		r, err := ns.probeAndWait(ip)
		if err != nil {
			log.WithError(err).Error("Error probing random address")
			continue
		}
		if !r {
			log.WithField("IP", ip).Debug("Returning Random Address")
			return &net.IPNet{IP: ip, Mask: n.Mask}, nil
		}
		log.WithField("ip", ip).Debug("Random address reachable, retrying")
	}
	return nil, fmt.Errorf("All avaliable addresses are in use")
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Tigera, Inc. All rights reserved.

// 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 backend

import (
	"fmt"

	"regexp"

	"reflect"

	"github.com/golang/glog"
	. "github.com/tigera/libcalico-go/lib/common"
)

var (
	matchWorkloadEndpoint = regexp.MustCompile("^/?calico/v1/host/([^/]+)/workload/([^/]+)/([^/]+)/endpoint/([^/]+)$")
)

type WorkloadEndpointKey struct {
	Hostname       string `json:"-"`
	OrchestratorID string `json:"-"`
	WorkloadID     string `json:"-"`
	EndpointID     string `json:"-"`
}

func (key WorkloadEndpointKey) asEtcdKey() (string, error) {
	if key.Hostname == "" || key.OrchestratorID == "" || key.WorkloadID == "" || key.EndpointID == "" {
		return "", ErrorInsufficientIdentifiers{}
	}
	return fmt.Sprintf("/calico/v1/host/%s/workload/%s/%s/endpoint/%s",
		key.Hostname, key.OrchestratorID, key.WorkloadID, key.EndpointID), nil
}

func (key WorkloadEndpointKey) asEtcdDeleteKey() (string, error) {
	return key.asEtcdKey()
}

func (key WorkloadEndpointKey) valueType() reflect.Type {
	return reflect.TypeOf(WorkloadEndpoint{})
}

type WorkloadEndpointListOptions struct {
	Hostname       string
	OrchestratorID string
	WorkloadID     string
	EndpointID     string
}

func (options WorkloadEndpointListOptions) asEtcdKeyRoot() string {
	k := "/calico/v1/host"
	if options.Hostname == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s/workload", options.Hostname)
	if options.OrchestratorID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s", options.OrchestratorID)
	if options.WorkloadID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s/endpoint", options.WorkloadID)
	if options.EndpointID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s", options.EndpointID)
	return k
}

func (options WorkloadEndpointListOptions) keyFromEtcdResult(ekey string) KeyInterface {
	glog.V(2).Infof("Get WorkloadEndpoint key from %s", ekey)
	r := matchWorkloadEndpoint.FindAllStringSubmatch(ekey, -1)
	if len(r) != 1 {
		glog.V(2).Infof("Didn't match regex")
		return nil
	}
	hostname := r[0][1]
	orch := r[0][2]
	workload := r[0][3]
	endpointID := r[0][4]
	if options.Hostname != "" && hostname != options.Hostname {
		glog.V(2).Infof("Didn't match hostname %s != %s", options.Hostname, hostname)
		return nil
	}
	if options.OrchestratorID != "" && orch != options.OrchestratorID {
		glog.V(2).Infof("Didn't match orchestrator %s != %s", options.OrchestratorID, orch)
		return nil
	}
	if options.WorkloadID != "" && workload != options.WorkloadID {
		glog.V(2).Infof("Didn't match workload %s != %s", options.WorkloadID, workload)
		return nil
	}
	if options.EndpointID != "" && endpointID != options.EndpointID {
		glog.V(2).Infof("Didn't match endpoint ID %s != %s", options.EndpointID, endpointID)
		return nil
	}
	return WorkloadEndpointKey{Hostname: hostname, EndpointID: endpointID}
}

type WorkloadEndpoint struct {
	// TODO: Validation for workload endpoint.
	State      string            `json:"state"`
	Name       string            `json:"name"`
	Mac        MAC               `json:"mac"`
	ProfileIDs []string          `json:"profile_ids"`
	IPv4Nets   []IPNet           `json:"ipv4_nets"`
	IPv6Nets   []IPNet           `json:"ipv6_nets"`
	Labels     map[string]string `json:"labels,omitempty"`
}
<commit_msg>Add omitempty directives to slice fields in backend workload enpoint structure<commit_after>// Copyright (c) 2016 Tigera, Inc. All rights reserved.

// 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 backend

import (
	"fmt"

	"regexp"

	"reflect"

	"github.com/golang/glog"
	. "github.com/tigera/libcalico-go/lib/common"
)

var (
	matchWorkloadEndpoint = regexp.MustCompile("^/?calico/v1/host/([^/]+)/workload/([^/]+)/([^/]+)/endpoint/([^/]+)$")
)

type WorkloadEndpointKey struct {
	Hostname       string `json:"-"`
	OrchestratorID string `json:"-"`
	WorkloadID     string `json:"-"`
	EndpointID     string `json:"-"`
}

func (key WorkloadEndpointKey) asEtcdKey() (string, error) {
	if key.Hostname == "" || key.OrchestratorID == "" || key.WorkloadID == "" || key.EndpointID == "" {
		return "", ErrorInsufficientIdentifiers{}
	}
	return fmt.Sprintf("/calico/v1/host/%s/workload/%s/%s/endpoint/%s",
		key.Hostname, key.OrchestratorID, key.WorkloadID, key.EndpointID), nil
}

func (key WorkloadEndpointKey) asEtcdDeleteKey() (string, error) {
	return key.asEtcdKey()
}

func (key WorkloadEndpointKey) valueType() reflect.Type {
	return reflect.TypeOf(WorkloadEndpoint{})
}

type WorkloadEndpointListOptions struct {
	Hostname       string
	OrchestratorID string
	WorkloadID     string
	EndpointID     string
}

func (options WorkloadEndpointListOptions) asEtcdKeyRoot() string {
	k := "/calico/v1/host"
	if options.Hostname == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s/workload", options.Hostname)
	if options.OrchestratorID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s", options.OrchestratorID)
	if options.WorkloadID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s/endpoint", options.WorkloadID)
	if options.EndpointID == "" {
		return k
	}
	k = k + fmt.Sprintf("/%s", options.EndpointID)
	return k
}

func (options WorkloadEndpointListOptions) keyFromEtcdResult(ekey string) KeyInterface {
	glog.V(2).Infof("Get WorkloadEndpoint key from %s", ekey)
	r := matchWorkloadEndpoint.FindAllStringSubmatch(ekey, -1)
	if len(r) != 1 {
		glog.V(2).Infof("Didn't match regex")
		return nil
	}
	hostname := r[0][1]
	orch := r[0][2]
	workload := r[0][3]
	endpointID := r[0][4]
	if options.Hostname != "" && hostname != options.Hostname {
		glog.V(2).Infof("Didn't match hostname %s != %s", options.Hostname, hostname)
		return nil
	}
	if options.OrchestratorID != "" && orch != options.OrchestratorID {
		glog.V(2).Infof("Didn't match orchestrator %s != %s", options.OrchestratorID, orch)
		return nil
	}
	if options.WorkloadID != "" && workload != options.WorkloadID {
		glog.V(2).Infof("Didn't match workload %s != %s", options.WorkloadID, workload)
		return nil
	}
	if options.EndpointID != "" && endpointID != options.EndpointID {
		glog.V(2).Infof("Didn't match endpoint ID %s != %s", options.EndpointID, endpointID)
		return nil
	}
	return WorkloadEndpointKey{Hostname: hostname, EndpointID: endpointID}
}

type WorkloadEndpoint struct {
	// TODO: Validation for workload endpoint.
	State      string            `json:"state"`
	Name       string            `json:"name"`
	Mac        MAC               `json:"mac"`
	ProfileIDs []string          `json:"profile_ids,omitempty"`
	IPv4Nets   []IPNet           `json:"ipv4_nets,omitempty"`
	IPv6Nets   []IPNet           `json:"ipv6_nets,omitempty"`
	Labels     map[string]string `json:"labels,omitempty"`
}
<|endoftext|>


<commit_before>package models_test

/* NOTE: main_test.go contains the TestMain, setup, and teardown functions for this suite. */

import (
	"fmt"
	"github.com/APTrust/easy-store/db/models"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestProfileGetId(t *testing.T) {
	profile := FakeBagItProfile()
	assert.Equal(t, profile.Id, profile.GetId())
}

func TestProfileTableName(t *testing.T) {
	profile := FakeBagItProfile()
	assert.Equal(t, "bagit_profiles", profile.TableName())
}

func TestProfileValidate(t *testing.T) {
	profile := FakeBagItProfile()
	assert.True(t, profile.Validate())
	assert.NotNil(t, profile.Errors())
	assert.Empty(t, profile.Errors())
}

func TestProfileErrors(t *testing.T) {
	profile := FakeBagItProfile()
	assert.NotNil(t, profile.Errors())
	assert.Empty(t, profile.Errors())
}

func TestProfileSave(t *testing.T) {
	// Save as insert
	profile := FakeBagItProfile()
	profile.Id = 0
	ok := profile.Save(false)
	assert.True(t, ok)
	assert.NotEqual(t, 0, profile.Id)
	assert.Empty(t, profile.Errors())

	// Save as update
	id := profile.Id
	profile.Name = profile.Name + "updated"
	ok = profile.Save(false)
	assert.True(t, ok)
	assert.Equal(t, id, profile.Id)
	assert.Empty(t, profile.Errors())
}

func TestGetBagItProfile(t *testing.T) {
	// Save first
	profile := FakeBagItProfile()
	profile.Id = 0
	ok := profile.Save(false)
	assert.True(t, ok)
	assert.NotEqual(t, 0, profile.Id)

	// Now get
	retrievedProfile, err := models.GetBagItProfile(profile.Id)
	require.Nil(t, err)
	assert.Equal(t, profile.Id, retrievedProfile.Id)
	assert.Equal(t, profile.Name, retrievedProfile.Name)
	assert.Equal(t, profile.Description, retrievedProfile.Description)
	assert.Equal(t, profile.JSON, retrievedProfile.JSON)
	assert.Empty(t, profile.Errors())
}

func TestGetBagItProfiles(t *testing.T) {
	// Delete profiles created by other tests
	_, err := models.ExecCommand("delete from bagit_profiles", nil)
	require.Nil(t, err)

	// Save some first
	for i := 0; i < 10; i++ {
		profile := FakeBagItProfile()
		profile.Id = 0
		if i%2 == 0 {
			profile.Name = fmt.Sprintf("Even %d", i)
		} else {
			profile.Name = fmt.Sprintf("Odd %d", i)
		}
		if i%3 == 0 {
			profile.Description = ""
		}
		ok := profile.Save(false)
		assert.True(t, ok)
		assert.NotEqual(t, 0, profile.Id)
	}

	// Now select
	// Profiles 3 and 9 have Odd name and empty description
	where := "name like ? and description = ?"
	values := []interface{}{
		"Odd%",
		"",
	}
	profiles, err := models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 2, len(profiles))

	// Should get ten profiles
	where = ""
	values = []interface{}{}
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 10, len(profiles))

	// Should also get ten profiles
	where = "name like 'Even %' or name like 'Odd%'"
	values = make([]interface{}, 0)
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 10, len(profiles))

	// Should get profiles
	where = "name like :name"
	values = []interface{}{
		"Even%",
	}
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 5, len(profiles))
}

func TestProfileUnmarshal(t *testing.T) {
	profile := FakeBagItProfile()
	bagItProfile, err := profile.Profile()
	assert.Nil(t, err)
	assert.NotNil(t, bagItProfile)
	require.NotEmpty(t, bagItProfile.ManifestsRequired)
	assert.Equal(t, "md5", bagItProfile.ManifestsRequired[0])
}
<commit_msg>Added Profile validation test<commit_after>package models_test

/* NOTE: main_test.go contains the TestMain, setup, and teardown functions for this suite. */

import (
	"fmt"
	"github.com/APTrust/easy-store/db/models"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestProfileGetId(t *testing.T) {
	profile := FakeBagItProfile()
	assert.Equal(t, profile.Id, profile.GetId())
}

func TestProfileTableName(t *testing.T) {
	profile := FakeBagItProfile()
	assert.Equal(t, "bagit_profiles", profile.TableName())
}

func TestProfileValidate(t *testing.T) {
	profile := FakeBagItProfile()
	assert.True(t, profile.Validate())
	assert.NotNil(t, profile.Errors())
	assert.Empty(t, profile.Errors())

	profile.JSON = "{unparsable}"
	assert.False(t, profile.Validate())
	errors := profile.Errors()
	assert.NotNil(t, errors)
	require.Equal(t, 5, len(errors))
	assert.Equal(t, "Cannot parse JSON: invalid character 'u' looking for beginning of object key string", errors[0])
	assert.Equal(t, "Accept-BagIt-Version must accept at least one BagIt version.", errors[1])
	assert.Equal(t, "Manifests-Required must require at least one algorithm.", errors[2])
	assert.Equal(t, "Tag-Files-Required is missing bagit.txt.", errors[3])
	assert.Equal(t, "Tag-Files-Required is missing bag-info.txt.", errors[4])

	// JSON is parsable, but profile is not valid.
	// One thing we're verifying here is that the
	// Validate() method resets the list of errors
	// each time it's called.
	profile.JSON = "{}"
	assert.False(t, profile.Validate())
	errors = profile.Errors()
	assert.NotNil(t, errors)
	require.Equal(t, 4, len(errors))
	assert.Equal(t, "Accept-BagIt-Version must accept at least one BagIt version.", errors[0])
	assert.Equal(t, "Manifests-Required must require at least one algorithm.", errors[1])
	assert.Equal(t, "Tag-Files-Required is missing bagit.txt.", errors[2])
	assert.Equal(t, "Tag-Files-Required is missing bag-info.txt.", errors[3])

}

func TestProfileErrors(t *testing.T) {
	profile := FakeBagItProfile()
	assert.NotNil(t, profile.Errors())
	assert.Empty(t, profile.Errors())
}

func TestProfileSave(t *testing.T) {
	// Save as insert
	profile := FakeBagItProfile()
	profile.Id = 0
	ok := profile.Save(false)
	assert.True(t, ok)
	assert.NotEqual(t, 0, profile.Id)
	assert.Empty(t, profile.Errors())

	// Save as update
	id := profile.Id
	profile.Name = profile.Name + "updated"
	ok = profile.Save(false)
	assert.True(t, ok)
	assert.Equal(t, id, profile.Id)
	assert.Empty(t, profile.Errors())
}

func TestGetBagItProfile(t *testing.T) {
	// Save first
	profile := FakeBagItProfile()
	profile.Id = 0
	ok := profile.Save(false)
	assert.True(t, ok)
	assert.NotEqual(t, 0, profile.Id)

	// Now get
	retrievedProfile, err := models.GetBagItProfile(profile.Id)
	require.Nil(t, err)
	assert.Equal(t, profile.Id, retrievedProfile.Id)
	assert.Equal(t, profile.Name, retrievedProfile.Name)
	assert.Equal(t, profile.Description, retrievedProfile.Description)
	assert.Equal(t, profile.JSON, retrievedProfile.JSON)
	assert.Empty(t, profile.Errors())
}

func TestGetBagItProfiles(t *testing.T) {
	// Delete profiles created by other tests
	_, err := models.ExecCommand("delete from bagit_profiles", nil)
	require.Nil(t, err)

	// Save some first
	for i := 0; i < 10; i++ {
		profile := FakeBagItProfile()
		profile.Id = 0
		if i%2 == 0 {
			profile.Name = fmt.Sprintf("Even %d", i)
		} else {
			profile.Name = fmt.Sprintf("Odd %d", i)
		}
		if i%3 == 0 {
			profile.Description = ""
		}
		ok := profile.Save(false)
		assert.True(t, ok)
		assert.NotEqual(t, 0, profile.Id)
	}

	// Now select
	// Profiles 3 and 9 have Odd name and empty description
	where := "name like ? and description = ?"
	values := []interface{}{
		"Odd%",
		"",
	}
	profiles, err := models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 2, len(profiles))

	// Should get ten profiles
	where = ""
	values = []interface{}{}
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 10, len(profiles))

	// Should also get ten profiles
	where = "name like 'Even %' or name like 'Odd%'"
	values = make([]interface{}, 0)
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 10, len(profiles))

	// Should get profiles
	where = "name like :name"
	values = []interface{}{
		"Even%",
	}
	profiles, err = models.GetBagItProfiles(where, values)
	require.Nil(t, err)
	require.NotNil(t, profiles)
	assert.Equal(t, 5, len(profiles))
}

func TestProfileUnmarshal(t *testing.T) {
	profile := FakeBagItProfile()
	bagItProfile, err := profile.Profile()
	assert.Nil(t, err)
	assert.NotNil(t, bagItProfile)
	require.NotEmpty(t, bagItProfile.ManifestsRequired)
	assert.Equal(t, "md5", bagItProfile.ManifestsRequired[0])
}
<|endoftext|>


<commit_before>package available_topics

import (
	"gopkg.in/bulletind/khabar.v1/db"
	"gopkg.in/bulletind/khabar.v1/dbapi/locks"
	"gopkg.in/bulletind/khabar.v1/utils"
)

const trueState = "true"
const falseState = "false"
const disabledState = "disabled"

type TopicDetail struct {
	Locked bool   `json:"locked"`
	Value  string `json:"value"`
}

type ChotaTopic map[string]*TopicDetail

func GetAllTopics() []string {
	session := db.Conn.Session.Copy()
	defer session.Close()

	topics := []string{}

	db.Conn.GetCursor(
		session, db.AvailableTopicCollection, utils.M{},
	).Distinct("ident", &topics)

	return topics
}

func GetAppTopics(app_name, org string) *[]string {
	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{"app_name": app_name}
	topics := []string{}

	var topic struct {
		Ident string `bson:"ident"`
	}

	iter := db.Conn.GetCursor(
		session, db.AvailableTopicCollection, query,
	).Select(utils.M{"ident": 1}).Sort("ident").Iter()

	for iter.Next(&topic) {
		topics = append(topics, topic.Ident)
	}

	return &topics
}

func GetOrgTopics(org string, appTopics, channels *[]string) (map[string]ChotaTopic, error) {
	// Add defaults for org level

	topicMap := map[string]ChotaTopic{}

	for _, ident := range *appTopics {
		ct := ChotaTopic{}
		for _, channel := range *channels {
			ct[channel] = &TopicDetail{Locked: false, Value: trueState}
		}

		topicMap[ident] = ct
	}

	disabled := new(db.Topic)

	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{
		"ident": utils.M{"$in": appTopics},
		"user":  db.BLANK,
		"org":   org,
	}

	pass1 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass1.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = falseState
			}
		}
	}

	//Find Globally Disabled Topics
	query["user"] = db.BLANK
	query["org"] = db.BLANK

	pass3 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass3.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = disabledState
			}
		}
	}

	return topicMap, nil
}

func ApplyLockes(org string, appTopics *[]string, topicMap map[string]ChotaTopic) {
	enabled := locks.GetAll(org)
	for _, pref := range enabled {
		if _, ok := topicMap[pref.Topic]; ok {
			for _, blocked := range pref.Channels {

				if topicMap[pref.Topic][blocked].Value == disabledState {
					continue
				}

				topicMap[pref.Topic][blocked].Locked = true

				if pref.Enabled {
					topicMap[pref.Topic][blocked].Value = trueState
				} else {
					topicMap[pref.Topic][blocked].Value = falseState
				}
			}
		}
	}
}

func GetUserTopics(user, org string, appTopics, channels *[]string) (map[string]ChotaTopic, error) {
	// Add defaults for user level

	topicMap := map[string]ChotaTopic{}

	for _, ident := range *appTopics {
		ct := ChotaTopic{}
		for _, channel := range *channels {
			ct[channel] = &TopicDetail{Locked: false, Value: trueState}
		}

		topicMap[ident] = ct
	}

	disabled := new(db.Topic)

	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{
		"ident": utils.M{"$in": appTopics},
		"user":  user,
		"org":   org,
	}

	pass1 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass1.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = falseState
			}
		}
	}

	//Find all Topics that have been blocked by the Organization
	query["user"] = db.BLANK

	pass2 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass2.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = disabledState
			}
		}
	}

	//Find Globally Disabled Topics
	query["user"] = db.BLANK
	query["org"] = db.BLANK

	pass3 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass3.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = disabledState
			}
		}
	}

	ApplyLockes(org, appTopics, topicMap)

	return topicMap, nil
}

func Get(topic string) (found *db.AvailableTopic, err error) {
	found = new(db.AvailableTopic)

	err = db.Conn.GetOne(db.AvailableTopicCollection, utils.M{"ident": topic}, found)

	if err != nil {
		return nil, err
	}

	return found, nil
}

func Insert(newTopic *db.AvailableTopic) string {
	return db.Conn.Insert(db.AvailableTopicCollection, newTopic)
}

func Delete(doc *utils.M) error {
	return db.Conn.Delete(db.AvailableTopicCollection, *doc)
}
<commit_msg>Do not show global disabled<commit_after>package available_topics

import (
	"gopkg.in/bulletind/khabar.v1/db"
	"gopkg.in/bulletind/khabar.v1/dbapi/locks"
	"gopkg.in/bulletind/khabar.v1/utils"
)

const trueState = "true"
const falseState = "false"
const disabledState = "disabled"

type TopicDetail struct {
	Locked bool   `json:"locked"`
	Value  string `json:"value"`
}

type ChotaTopic map[string]*TopicDetail

func GetAllTopics() []string {
	session := db.Conn.Session.Copy()
	defer session.Close()

	topics := []string{}

	db.Conn.GetCursor(
		session, db.AvailableTopicCollection, utils.M{},
	).Distinct("ident", &topics)

	return topics
}

func GetAppTopics(app_name, org string) *[]string {
	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{"app_name": app_name}
	topics := []string{}

	var topic struct {
		Ident string `bson:"ident"`
	}

	iter := db.Conn.GetCursor(
		session, db.AvailableTopicCollection, query,
	).Select(utils.M{"ident": 1}).Sort("ident").Iter()

	for iter.Next(&topic) {
		topics = append(topics, topic.Ident)
	}

	return &topics
}

func GetOrgTopics(org string, appTopics, channels *[]string) (map[string]ChotaTopic, error) {
	// Add defaults for org level

	topicMap := map[string]ChotaTopic{}

	for _, ident := range *appTopics {
		ct := ChotaTopic{}
		for _, channel := range *channels {
			ct[channel] = &TopicDetail{Locked: false, Value: trueState}
		}

		topicMap[ident] = ct
	}

	disabled := new(db.Topic)

	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{
		"ident": utils.M{"$in": appTopics},
		"user":  db.BLANK,
		"org":   org,
	}

	pass1 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass1.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = falseState
			}
		}
	}

	//Find Globally Disabled Topics
	query["user"] = db.BLANK
	query["org"] = db.BLANK

	pass3 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass3.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				delete(topicMap[disabled.Ident], blocked)
			}
		}
	}

	return topicMap, nil
}

func ApplyLockes(org string, appTopics *[]string, topicMap map[string]ChotaTopic) {
	enabled := locks.GetAll(org)
	for _, pref := range enabled {
		if _, ok := topicMap[pref.Topic]; ok {
			for _, blocked := range pref.Channels {

				if topicMap[pref.Topic][blocked].Value == disabledState {
					continue
				}

				topicMap[pref.Topic][blocked].Locked = true

				if pref.Enabled {
					topicMap[pref.Topic][blocked].Value = trueState
				} else {
					topicMap[pref.Topic][blocked].Value = falseState
				}
			}
		}
	}
}

func GetUserTopics(user, org string, appTopics, channels *[]string) (map[string]ChotaTopic, error) {
	// Add defaults for user level

	topicMap := map[string]ChotaTopic{}

	for _, ident := range *appTopics {
		ct := ChotaTopic{}
		for _, channel := range *channels {
			ct[channel] = &TopicDetail{Locked: false, Value: trueState}
		}

		topicMap[ident] = ct
	}

	disabled := new(db.Topic)

	session := db.Conn.Session.Copy()
	defer session.Close()

	query := utils.M{
		"ident": utils.M{"$in": appTopics},
		"user":  user,
		"org":   org,
	}

	pass1 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass1.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = falseState
			}
		}
	}

	//Find all Topics that have been blocked by the Organization
	query["user"] = db.BLANK

	pass2 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass2.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				topicMap[disabled.Ident][blocked].Value = disabledState
			}
		}
	}

	//Find Globally Disabled Topics
	query["user"] = db.BLANK
	query["org"] = db.BLANK

	pass3 := db.Conn.GetCursor(session, db.TopicCollection, query).Iter()
	for pass3.Next(disabled) {
		if _, ok := topicMap[disabled.Ident]; ok {
			for _, blocked := range disabled.Channels {
				delete(topicMap[disabled.Ident], blocked)
			}
		}
	}

	ApplyLockes(org, appTopics, topicMap)

	return topicMap, nil
}

func Get(topic string) (found *db.AvailableTopic, err error) {
	found = new(db.AvailableTopic)

	err = db.Conn.GetOne(db.AvailableTopicCollection, utils.M{"ident": topic}, found)

	if err != nil {
		return nil, err
	}

	return found, nil
}

func Insert(newTopic *db.AvailableTopic) string {
	return db.Conn.Insert(db.AvailableTopicCollection, newTopic)
}

func Delete(doc *utils.M) error {
	return db.Conn.Delete(db.AvailableTopicCollection, *doc)
}
<|endoftext|>


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

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc/grpclog"
	clientV2 "github.com/coreos/etcd/client"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
)

func init() {
	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))
}

type Stresser interface {
	// Stress starts to stress the etcd cluster
	Stress() error
	// Cancel cancels the stress test on the etcd cluster
	Cancel()
	// Report reports the success and failure of the stress test
	Report() (success int, failure int)
}

type stresser struct {
	Endpoint string

	KeySize        int
	KeySuffixRange int

	N int

	mu      sync.Mutex
	failure int
	success int

	cancel func()
}

func (s *stresser) Stress() error {
	conn, err := grpc.Dial(s.Endpoint, grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
	if err != nil {
		return fmt.Errorf("%v (%s)", err, s.Endpoint)
	}
	kvc := pb.NewKVClient(conn)

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

	for i := 0; i < s.N; i++ {
		go func(i int) {
			for {
				putctx, putcancel := context.WithTimeout(ctx, 5*time.Second)
				_, err := kvc.Put(putctx, &pb.PutRequest{
					Key:   []byte(fmt.Sprintf("foo%d", rand.Intn(s.KeySuffixRange))),
					Value: []byte(randStr(s.KeySize)),
				})
				putcancel()
				if grpc.ErrorDesc(err) == context.Canceled.Error() {
					return
				}
				s.mu.Lock()
				if err != nil {
					s.failure++
				} else {
					s.success++
				}
				s.mu.Unlock()
			}
		}(i)
	}

	<-ctx.Done()
	return nil
}

func (s *stresser) Cancel() {
	s.cancel()
}

func (s *stresser) Report() (success int, failure int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.success, s.failure
}

type stresserV2 struct {
	Endpoint string

	KeySize        int
	KeySuffixRange int

	N int
	// TODO: not implemented
	Interval time.Duration

	mu      sync.Mutex
	failure int
	success int

	cancel func()
}

func (s *stresserV2) Stress() error {
	cfg := clientV2.Config{
		Endpoints: []string{s.Endpoint},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: s.N,
		},
	}
	c, err := clientV2.New(cfg)
	if err != nil {
		return err
	}

	kv := clientV2.NewKeysAPI(c)
	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	for i := 0; i < s.N; i++ {
		go func() {
			for {
				setctx, setcancel := context.WithTimeout(ctx, clientV2.DefaultRequestTimeout)
				key := fmt.Sprintf("foo%d", rand.Intn(s.KeySuffixRange))
				_, err := kv.Set(setctx, key, randStr(s.KeySize), nil)
				setcancel()
				if err == context.Canceled {
					return
				}
				s.mu.Lock()
				if err != nil {
					s.failure++
				} else {
					s.success++
				}
				s.mu.Unlock()
			}
		}()
	}

	<-ctx.Done()
	return nil
}

func (s *stresserV2) Cancel() {
	s.cancel()
}

func (s *stresserV2) Report() (success int, failure int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.success, s.failure
}

func randStr(size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = byte(int('a') + rand.Intn(26))
	}
	return string(data)
}
<commit_msg>etcd-tester: close gRPC connection when canceling<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc/grpclog"
	clientV2 "github.com/coreos/etcd/client"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
)

func init() {
	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))
}

type Stresser interface {
	// Stress starts to stress the etcd cluster
	Stress() error
	// Cancel cancels the stress test on the etcd cluster
	Cancel()
	// Report reports the success and failure of the stress test
	Report() (success int, failure int)
}

type stresser struct {
	Endpoint string

	KeySize        int
	KeySuffixRange int

	N int

	mu      sync.Mutex
	failure int
	success int

	cancel func()
	conn   *grpc.ClientConn
}

func (s *stresser) Stress() error {
	conn, err := grpc.Dial(s.Endpoint, grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
	if err != nil {
		return fmt.Errorf("%v (%s)", err, s.Endpoint)
	}
	ctx, cancel := context.WithCancel(context.Background())

	s.mu.Lock()
	s.conn = conn
	s.cancel = cancel
	s.mu.Unlock()

	kvc := pb.NewKVClient(conn)

	for i := 0; i < s.N; i++ {
		go func(i int) {
			for {
				putctx, putcancel := context.WithTimeout(ctx, 5*time.Second)
				_, err := kvc.Put(putctx, &pb.PutRequest{
					Key:   []byte(fmt.Sprintf("foo%d", rand.Intn(s.KeySuffixRange))),
					Value: []byte(randStr(s.KeySize)),
				})
				putcancel()
				if grpc.ErrorDesc(err) == context.Canceled.Error() {
					return
				}
				s.mu.Lock()
				if err != nil {
					s.failure++
				} else {
					s.success++
				}
				s.mu.Unlock()
			}
		}(i)
	}

	<-ctx.Done()
	return nil
}

func (s *stresser) Cancel() {
	s.mu.Lock()
	cancel, conn := s.cancel, s.conn
	s.mu.Unlock()
	cancel()
	// TODO: wait for all routines to exit by adding a waitGroup before calling conn.Close()
	conn.Close()
}

func (s *stresser) Report() (success int, failure int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.success, s.failure
}

type stresserV2 struct {
	Endpoint string

	KeySize        int
	KeySuffixRange int

	N int
	// TODO: not implemented
	Interval time.Duration

	mu      sync.Mutex
	failure int
	success int

	cancel func()
}

func (s *stresserV2) Stress() error {
	cfg := clientV2.Config{
		Endpoints: []string{s.Endpoint},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: s.N,
		},
	}
	c, err := clientV2.New(cfg)
	if err != nil {
		return err
	}

	kv := clientV2.NewKeysAPI(c)
	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	for i := 0; i < s.N; i++ {
		go func() {
			for {
				setctx, setcancel := context.WithTimeout(ctx, clientV2.DefaultRequestTimeout)
				key := fmt.Sprintf("foo%d", rand.Intn(s.KeySuffixRange))
				_, err := kv.Set(setctx, key, randStr(s.KeySize), nil)
				setcancel()
				if err == context.Canceled {
					return
				}
				s.mu.Lock()
				if err != nil {
					s.failure++
				} else {
					s.success++
				}
				s.mu.Unlock()
			}
		}()
	}

	<-ctx.Done()
	return nil
}

func (s *stresserV2) Cancel() {
	s.cancel()
}

func (s *stresserV2) Report() (success int, failure int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.success, s.failure
}

func randStr(size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = byte(int('a') + rand.Intn(26))
	}
	return string(data)
}
<|endoftext|>


<commit_before>// Block processing functions

package frundis

import (
	"github.com/anaseto/gofrundis/ast"
	"github.com/anaseto/gofrundis/parser"
)

// ProcessFrundisSource processes a frundis file with a given exporter. In
// restricted mode, no #run nor shell filter are allowed.
func ProcessFrundisSource(exp Exporter, filename string, unrestricted bool) error {
	exp.Init()
	ctx := exp.Context()
	ctx.Unrestricted = unrestricted
	err := processFile(exp, filename)
	if err != nil {
		return err
	}
	err = exp.Reset()
	if err != nil {
		return err
	}
	err = processFile(exp, filename)
	if err != nil {
		return err
	}
	if ctx.loc == nil {
		ctx.loc = &location{curBlock: -1, curFile: filename}
	}
	ctx.Macro = "End Of File"
	closeUnclosedBlocks(exp, "Bm")
	closeUnclosedBlocks(exp, "Bl")
	closeUnclosedBlocks(exp, "Bd")
	checkForUnclosedBlock(exp, "#if")
	checkForUnclosedFormatBlock(exp)
	checkForUnclosedDe(exp)
	endParagraph(exp, false)
	exp.PostProcessing()
	return nil
}

// processFile does one pass through a file with a given exporter.
func processFile(exp Exporter, filename string) error {
	ctx := exp.Context()
	blocks, ok := ctx.files[filename]
	if !ok {
		var err error
		p := parser.Parser{}
		blocks, err = p.ParseFile(filename)
		if err != nil {
			return err
		}
		ctx.files[filename] = blocks
	}
	loc := ctx.loc
	defer func() { ctx.loc = loc }()
	ctx.loc = &location{curBlocks: blocks, curFile: filename}
	processBlocks(exp)
	return nil
}

func processBlocks(exp Exporter) {
	ctx := exp.Context()
	for i, b := range ctx.loc.curBlocks {
		ctx.loc.curBlock = i
		switch b := b.(type) {
		case *ast.Macro:
			ctx.Args = b.Args
			ctx.Macro = b.Name
			ctx.line = b.Line
		case *ast.TextBlock:
			ctx.text = b.Text
			ctx.line = b.Line
		}
		processBlock(exp)
	}
}

func processBlock(exp Exporter) {
	ctx := exp.Context()
	b := ctx.block()
	if ctx.ifIgnoreDepth > 0 {
		b, ok := b.(*ast.Macro)
		if ok {
			switch b.Name {
			case "#;":
				macroIfEnd(exp)
			case "#if":
				macroIfStart(exp)
			default:
			}
		}
		return
	}
	if ctx.uMacroDef != nil {
		switch b := b.(type) {
		case *ast.Macro:
			switch b.Name {
			case "#.":
				macroDefEnd(exp)
			default:
				if !ctx.uMacroDef.ignore {
					ctx.uMacroDef.blocks = append(ctx.uMacroDef.blocks, b)
				}
			}
		case *ast.TextBlock:
			if !ctx.uMacroDef.ignore {
				ctx.uMacroDef.blocks = append(ctx.uMacroDef.blocks, b)
			}
		}
		return
	}

	switch b := b.(type) {
	case *ast.Macro:
		var ok bool
		_, ok = ctx.uMacros[b.Name]
		if ok {
			processUserMacro(exp)
			return
		}
		handler, ok := ctx.Macros[b.Name]
		if ok {
			if ctx.bfInfo != nil {
				switch b.Name {
				case "Ef", "#de", "#.", "#if", "#;", "#dv":
				default:
					checkForUnclosedFormatBlock(exp)
				}
			}
			handler(exp)
			ctx.PrevMacro = b.Name
		} else if b.Name != "" {
			ctx.Error("unknown macro:", b.Name)
		}
	case *ast.TextBlock:
		processText(exp)
		ctx.PrevMacro = ""
	}
}

// DefaultExporterMacros returns a mapping from macros to handling functions,
// with the standard set of frundis macros.
func DefaultExporterMacros() map[string]func(Exporter) {
	return map[string]func(Exporter){
		"Bd":   macroBd,
		"Bf":   macroBf,
		"Bl":   macroBl,
		"Bm":   macroBm,
		"Ch":   macroHeader,
		"D":    macroD,
		"Ed":   macroEd,
		"Ef":   macroEf,
		"El":   macroEl,
		"Em":   macroEm,
		"Ft":   macroFt,
		"If":   macroIncludeFile,
		"Im":   macroIm,
		"It":   macroIt,
		"Lk":   macroLk,
		"P":    macroP,
		"Pt":   macroHeader,
		"Sh":   macroHeader,
		"Sm":   macroSm,
		"Ss":   macroHeader,
		"Sx":   macroSx,
		"Ta":   macroTa,
		"Tc":   macroTc,
		"X":    macroX,
		"#de":  macroDefStart,
		"#.":   macroDefEnd,
		"#if":  macroIfStart,
		"#;":   macroIfEnd,
		"#dv":  macroDefVar,
		"#run": macroRun}
}

// MinimalExporterMacros returns a mapping from macros to handling functions,
// with only the following macros: Bd, Bf, Bm, Ed, Ef, Em, Ft, If, Sm, X.
func MinimalExporterMacros() map[string]func(Exporter) {
	return map[string]func(Exporter){
		"Bd":   macroBd,
		"Bf":   macroBf,
		"Bm":   macroBm,
		"Ed":   macroEd,
		"Ef":   macroEf,
		"Em":   macroEm,
		"Ft":   macroFt,
		"If":   macroIncludeFile,
		"Sm":   macroSm,
		"X":    macroX,
		"#de":  macroDefStart,
		"#.":   macroDefEnd,
		"#if":  macroIfStart,
		"#;":   macroIfEnd,
		"#dv":  macroDefVar,
		"#run": macroRun}
}
<commit_msg>disallow more macros inside filter block<commit_after>// Block processing functions

package frundis

import (
	"github.com/anaseto/gofrundis/ast"
	"github.com/anaseto/gofrundis/parser"
)

// ProcessFrundisSource processes a frundis file with a given exporter. In
// restricted mode, no #run nor shell filter are allowed.
func ProcessFrundisSource(exp Exporter, filename string, unrestricted bool) error {
	exp.Init()
	ctx := exp.Context()
	ctx.Unrestricted = unrestricted
	err := processFile(exp, filename)
	if err != nil {
		return err
	}
	err = exp.Reset()
	if err != nil {
		return err
	}
	err = processFile(exp, filename)
	if err != nil {
		return err
	}
	if ctx.loc == nil {
		ctx.loc = &location{curBlock: -1, curFile: filename}
	}
	ctx.Macro = "End Of File"
	closeUnclosedBlocks(exp, "Bm")
	closeUnclosedBlocks(exp, "Bl")
	closeUnclosedBlocks(exp, "Bd")
	checkForUnclosedBlock(exp, "#if")
	checkForUnclosedFormatBlock(exp)
	checkForUnclosedDe(exp)
	endParagraph(exp, false)
	exp.PostProcessing()
	return nil
}

// processFile does one pass through a file with a given exporter.
func processFile(exp Exporter, filename string) error {
	ctx := exp.Context()
	blocks, ok := ctx.files[filename]
	if !ok {
		var err error
		p := parser.Parser{}
		blocks, err = p.ParseFile(filename)
		if err != nil {
			return err
		}
		ctx.files[filename] = blocks
	}
	loc := ctx.loc
	defer func() { ctx.loc = loc }()
	ctx.loc = &location{curBlocks: blocks, curFile: filename}
	processBlocks(exp)
	return nil
}

func processBlocks(exp Exporter) {
	ctx := exp.Context()
	for i, b := range ctx.loc.curBlocks {
		ctx.loc.curBlock = i
		switch b := b.(type) {
		case *ast.Macro:
			ctx.Args = b.Args
			ctx.Macro = b.Name
			ctx.line = b.Line
		case *ast.TextBlock:
			ctx.text = b.Text
			ctx.line = b.Line
		}
		processBlock(exp)
	}
}

func processBlock(exp Exporter) {
	ctx := exp.Context()
	b := ctx.block()
	if ctx.ifIgnoreDepth > 0 {
		b, ok := b.(*ast.Macro)
		if ok {
			switch b.Name {
			case "#;":
				macroIfEnd(exp)
			case "#if":
				macroIfStart(exp)
			default:
			}
		}
		return
	}
	if ctx.uMacroDef != nil {
		switch b := b.(type) {
		case *ast.Macro:
			switch b.Name {
			case "#.":
				macroDefEnd(exp)
			default:
				if !ctx.uMacroDef.ignore {
					ctx.uMacroDef.blocks = append(ctx.uMacroDef.blocks, b)
				}
			}
		case *ast.TextBlock:
			if !ctx.uMacroDef.ignore {
				ctx.uMacroDef.blocks = append(ctx.uMacroDef.blocks, b)
			}
		}
		return
	}

	switch b := b.(type) {
	case *ast.Macro:
		var ok bool
		_, ok = ctx.uMacros[b.Name]
		if ok {
			processUserMacro(exp)
			return
		}
		handler, ok := ctx.Macros[b.Name]
		if ok {
			if ctx.bfInfo != nil {
				switch b.Name {
				case "Ef", "#if", "#;":
				default:
					checkForUnclosedFormatBlock(exp)
				}
			}
			handler(exp)
			ctx.PrevMacro = b.Name
		} else if b.Name != "" {
			ctx.Error("unknown macro:", b.Name)
		}
	case *ast.TextBlock:
		processText(exp)
		ctx.PrevMacro = ""
	}
}

// DefaultExporterMacros returns a mapping from macros to handling functions,
// with the standard set of frundis macros.
func DefaultExporterMacros() map[string]func(Exporter) {
	return map[string]func(Exporter){
		"Bd":   macroBd,
		"Bf":   macroBf,
		"Bl":   macroBl,
		"Bm":   macroBm,
		"Ch":   macroHeader,
		"D":    macroD,
		"Ed":   macroEd,
		"Ef":   macroEf,
		"El":   macroEl,
		"Em":   macroEm,
		"Ft":   macroFt,
		"If":   macroIncludeFile,
		"Im":   macroIm,
		"It":   macroIt,
		"Lk":   macroLk,
		"P":    macroP,
		"Pt":   macroHeader,
		"Sh":   macroHeader,
		"Sm":   macroSm,
		"Ss":   macroHeader,
		"Sx":   macroSx,
		"Ta":   macroTa,
		"Tc":   macroTc,
		"X":    macroX,
		"#de":  macroDefStart,
		"#.":   macroDefEnd,
		"#if":  macroIfStart,
		"#;":   macroIfEnd,
		"#dv":  macroDefVar,
		"#run": macroRun}
}

// MinimalExporterMacros returns a mapping from macros to handling functions,
// with only the following macros: Bd, Bf, Bm, Ed, Ef, Em, Ft, If, Sm, X.
func MinimalExporterMacros() map[string]func(Exporter) {
	return map[string]func(Exporter){
		"Bd":   macroBd,
		"Bf":   macroBf,
		"Bm":   macroBm,
		"Ed":   macroEd,
		"Ef":   macroEf,
		"Em":   macroEm,
		"Ft":   macroFt,
		"If":   macroIncludeFile,
		"Sm":   macroSm,
		"X":    macroX,
		"#de":  macroDefStart,
		"#.":   macroDefEnd,
		"#if":  macroIfStart,
		"#;":   macroIfEnd,
		"#dv":  macroDefVar,
		"#run": macroRun}
}
<|endoftext|>


<commit_before>package settings

import (
	"encoding/json"
	"io/ioutil"
	"os"
)

type Options struct {
	Dir      string `json:"dir" short:"d" long:"dir" default:"data/" description:"Directory to watch and sync"`
	ClientId string `json:"clientid" short:"c" long:"cid" description:"Client Id unique to this machine"`
	UserId   string `json:"userid" short:"u" long:"uid" description:"User Id unique to the user running using the client"`
}

func (o *Options) LoadFromJSON() {
	if _, err := os.Stat(ConfigFile); err == nil {
		if f, err := os.Open(ConfigFile); err == nil {
			config, _ := ioutil.ReadAll(f)
			options := &Options{}
			json.Unmarshal(config, &options)

			if o.ClientId == "" {
				o.ClientId = options.ClientId
			}

			if o.UserId == "" {
				o.UserId = options.UserId
			}
		}
	}
}
<commit_msg>Save json config<commit_after>package settings

import (
	"encoding/json"
	"io/ioutil"
	"os"
)

type Options struct {
	Dir      string `json:"dir" short:"d" long:"dir" default:"data/" description:"Directory to watch and sync"`
	ClientId string `json:"clientid" short:"c" long:"cid" description:"Client Id unique to this machine"`
	UserId   string `json:"userid" short:"u" long:"uid" description:"User Id unique to the user running using the client"`
}

func (o *Options) Load() {
	if _, err := os.Stat(ConfigFile); err == nil {
		if f, err := os.Open(ConfigFile); err == nil {
			config, _ := ioutil.ReadAll(f)
			options := &Options{}
			json.Unmarshal(config, &options)

			if o.ClientId == "" {
				o.ClientId = options.ClientId
			}

			if o.UserId == "" {
				o.UserId = options.UserId
			}
		}
	}
}

func (o *Options) Save() {
	config, _ := json.Marshal(o)
	ioutil.WriteFile(ConfigFile, config, 0666)
}
<|endoftext|>


<commit_before>package agent

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/stretchr/testify/assert"
)

func TestHTTP_ResourcesWithIllegalMethod(t *testing.T) {
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		req, err := http.NewRequest("DELETE", "/v1/resources", nil)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		_, err = s.Server.ResourcesRequest(respW, req)
		assert.NotNil(t, err, "HTTP DELETE should not be accepted for this endpoint")
	})
}

func createJobForTest(jobID string, s *TestAgent, t *testing.T) {
	job := mock.Job()
	job.ID = jobID
	job.TaskGroups[0].Count = 1
	args := structs.JobRegisterRequest{
		Job:          job,
		WriteRequest: structs.WriteRequest{Region: "global"},
	}
	var resp structs.JobRegisterResponse
	if err := s.Agent.RPC("Job.Register", &args, &resp); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestHTTP_ResourcesWithSingleJob(t *testing.T) {
	testJob := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89706"
	testJobPrefix := "aaaaaaaa-e8f7-fd38"
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		createJobForTest(testJob, s, t)

		data := structs.ResourcesRequest{Prefix: testJobPrefix, Context: "job"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))

		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["job"]
		if j == nil || len(j) != 1 {
			t.Fatalf("The number of jobs that were returned does not equal the number of jobs we expected (1)", j)
		}

		assert.Equal(t, j[0], testJob)
		assert.Equal(t, res.Truncations["job"], false)
	})
}

func TestHTTP_ResourcesWithMultipleJobs(t *testing.T) {
	testJobA := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89706"
	testJobB := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89707"
	testJobC := "bbbbbbbb-e8f7-fd38-c855-ab94ceb89707"

	testJobPrefix := "aaaaaaaa-e8f7-fd38"

	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		createJobForTest(testJobA, s, t)
		createJobForTest(testJobB, s, t)
		createJobForTest(testJobC, s, t)

		data := structs.ResourcesRequest{Prefix: testJobPrefix, Context: "job"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))

		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["job"]
		if j == nil || len(j) != 2 {
			t.Fatalf("The number of jobs that were returned does not equal the number of jobs we expected (2)", j)
		}

		assert.Contains(t, j, testJobA)
		assert.Contains(t, j, testJobB)
		assert.NotContains(t, j, testJobC)
	})
}

func TestHTTP_ResoucesListForEvaluations(t *testing.T) {
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		state := s.Agent.server.State()
		eval1 := mock.Eval()
		eval2 := mock.Eval()
		err := state.UpsertEvals(1000,
			[]*structs.Evaluation{eval1, eval2})
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		// Make the HTTP request
		evalPrefix := eval1.ID[:len(eval1.ID)-2]
		data := structs.ResourcesRequest{Prefix: evalPrefix, Context: "eval"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["eval"]
		if len(j) != 1 {
			t.Fatalf("The number of evaluations that were returned does not equal the number we expected (1)", j)
		}

		assert.Contains(t, j, eval1.ID)
		assert.NotContains(t, j, eval2.ID)
		assert.Equal(t, res.Truncations["eval"], false)
	})
}

// TODO
//func TestHTTP_ResourcesWithNoJob(t *testing.T) {
//}
<commit_msg>use upsert as a test helper<commit_after>package agent

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/stretchr/testify/assert"
)

func TestHTTP_ResourcesWithIllegalMethod(t *testing.T) {
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		req, err := http.NewRequest("DELETE", "/v1/resources", nil)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		_, err = s.Server.ResourcesRequest(respW, req)
		assert.NotNil(t, err, "HTTP DELETE should not be accepted for this endpoint")
	})
}

func createJobForTest(jobID string, s *TestAgent, t *testing.T) {
	job := mock.Job()
	job.ID = jobID
	job.TaskGroups[0].Count = 1

	state := s.Agent.server.State()
	err := state.UpsertJob(1000, job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestHTTP_ResourcesWithSingleJob(t *testing.T) {
	testJob := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89706"
	testJobPrefix := "aaaaaaaa-e8f7-fd38"
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		createJobForTest(testJob, s, t)

		data := structs.ResourcesRequest{Prefix: testJobPrefix, Context: "job"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))

		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["job"]
		if j == nil || len(j) != 1 {
			t.Fatalf("The number of jobs that were returned does not equal the number of jobs we expected (1)", j)
		}

		assert.Equal(t, j[0], testJob)
		assert.Equal(t, res.Truncations["job"], false)
	})
}

func TestHTTP_ResourcesWithMultipleJobs(t *testing.T) {
	testJobA := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89706"
	testJobB := "aaaaaaaa-e8f7-fd38-c855-ab94ceb89707"
	testJobC := "bbbbbbbb-e8f7-fd38-c855-ab94ceb89707"

	testJobPrefix := "aaaaaaaa-e8f7-fd38"

	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		createJobForTest(testJobA, s, t)
		createJobForTest(testJobB, s, t)
		createJobForTest(testJobC, s, t)

		data := structs.ResourcesRequest{Prefix: testJobPrefix, Context: "job"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))

		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["job"]
		if j == nil || len(j) != 2 {
			t.Fatalf("The number of jobs that were returned does not equal the number of jobs we expected (2)", j)
		}

		assert.Contains(t, j, testJobA)
		assert.Contains(t, j, testJobB)
		assert.NotContains(t, j, testJobC)
	})
}

func TestHTTP_ResoucesListForEvaluations(t *testing.T) {
	t.Parallel()
	httpTest(t, nil, func(s *TestAgent) {
		state := s.Agent.server.State()
		eval1 := mock.Eval()
		eval2 := mock.Eval()
		err := state.UpsertEvals(1000,
			[]*structs.Evaluation{eval1, eval2})
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		// Make the HTTP request
		evalPrefix := eval1.ID[:len(eval1.ID)-2]
		data := structs.ResourcesRequest{Prefix: evalPrefix, Context: "eval"}
		req, err := http.NewRequest("POST", "/v1/resources", encodeReq(data))
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		respW := httptest.NewRecorder()

		resp, err := s.Server.ResourcesRequest(respW, req)
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		res := resp.(structs.ResourcesResponse)
		if len(res.Matches) != 1 {
			t.Fatalf("No expected key values in resources list")
		}

		j := res.Matches["eval"]
		if len(j) != 1 {
			t.Fatalf("The number of evaluations that were returned does not equal the number we expected (1)", j)
		}

		assert.Contains(t, j, eval1.ID)
		assert.NotContains(t, j, eval2.ID)
		assert.Equal(t, res.Truncations["eval"], false)
	})
}

// TODO
//func TestHTTP_ResourcesWithNoJob(t *testing.T) {
//}
<|endoftext|>


<commit_before>package libvirt_test

import (
	"github.com/mistifyio/mistify-agent-libvirt"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/rpc"
	"testing"
	"time"
)

var port uint = 9001

type TestClient struct {
	rpc      *rpc.Client
	guest    *client.Guest
	request  *rpc.GuestRequest
	response *rpc.GuestResponse
}

func setup(t *testing.T) *TestClient {
	lv, err := libvirt.NewLibvirt("qemu:///system", 1)
	if err != nil {
		t.Fatalf("NewLibvirt failed: %s\n", err.Error())
	}

	go lv.RunHTTP(port)
	time.Sleep(1 * time.Second)

	cli := new(TestClient)

	cli.rpc, err = rpc.NewClient(port)
	if err != nil {
		t.Fatalf("Can't create RPC client: %s\n", err.Error())
	}

	cli.guest = new(client.Guest)
	cli.guest.Id = "testlibvirt"
	cli.guest.Memory = 1024
	cli.guest.Cpu = 1

	disk := client.Disk{
		Bus:    "sata",
		Device: "/dev/hda",
		Size:   1024,
	}
	cli.guest.Disks = append(cli.guest.Disks, disk)

	nic := client.Nic{
		Name: "eth0",
		Mac:  "01:23:45:67:89:0a",
	}
	cli.guest.Nics = append(cli.guest.Nics, nic)

	cli.request = new(rpc.GuestRequest)
	cli.request.Guest = cli.guest

	cli.response = new(rpc.GuestResponse)

	return cli
}

func do(action string, t *testing.T, cli *TestClient) {
	err := cli.rpc.Do(action, cli.request, cli.response)
	if err != nil {
		t.Fatalf("Error running %s: %s\n", action, err.Error())
	}
}

func TestLibvirt(t *testing.T) {
	cli := setup(t)

	do("Libvirt.Create", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After create, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Shutdown", t, cli)
	if cli.response.Guest.State != "Shutoff" {
		t.Fatalf("After shutdown, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Run", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After run, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Restart", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After restart, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Delete", t, cli)
}
<commit_msg>MIST-138: fix some guest options<commit_after>package libvirt_test

import (
	"github.com/mistifyio/mistify-agent-libvirt"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/rpc"
	"testing"
	"time"
)

var port uint = 9001

type TestClient struct {
	rpc      *rpc.Client
	guest    *client.Guest
	request  *rpc.GuestRequest
	response *rpc.GuestResponse
}

func setup(t *testing.T) *TestClient {
	lv, err := libvirt.NewLibvirt("qemu:///system", 1)
	if err != nil {
		t.Fatalf("NewLibvirt failed: %s\n", err.Error())
	}

	go lv.RunHTTP(port)
	time.Sleep(1 * time.Second)

	cli := new(TestClient)

	cli.rpc, err = rpc.NewClient(port)
	if err != nil {
		t.Fatalf("Can't create RPC client: %s\n", err.Error())
	}

	cli.guest = new(client.Guest)
	cli.guest.Id = "testlibvirt"
	cli.guest.Memory = 1024
	cli.guest.Cpu = 1

	disk := client.Disk{
		Bus:    "sata",
		Device: "hda",
		Size:   1024,
		Source: "/dev/zvol/guests/images/testlibvirt",
	}
	cli.guest.Disks = append(cli.guest.Disks, disk)

	nic := client.Nic{
		Name:    "eth0",
		Mac:     "00:0c:29:2f:00:00",
		Network: "default",
	}
	cli.guest.Nics = append(cli.guest.Nics, nic)

	cli.request = new(rpc.GuestRequest)
	cli.request.Guest = cli.guest

	cli.response = new(rpc.GuestResponse)

	return cli
}

func do(action string, t *testing.T, cli *TestClient) {
	err := cli.rpc.Do(action, cli.request, cli.response)
	if err != nil {
		t.Fatalf("Error running %s: %s\n", action, err.Error())
	}
}

func TestLibvirt(t *testing.T) {
	cli := setup(t)

	do("Libvirt.Create", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After create, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Shutdown", t, cli)
	if cli.response.Guest.State != "Shutoff" {
		t.Fatalf("After shutdown, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Run", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After run, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Restart", t, cli)
	if cli.response.Guest.State != "Running" {
		t.Fatalf("After restart, guest state is %s\n", cli.response.Guest.State)
	}

	do("Libvirt.Delete", t, cli)
}
<|endoftext|>