<commit_before>package web

import (
  "reflect"
  "fmt"
)

type HttpMethod string
const (
  HttpMethodGet = HttpMethod("GET")
  HttpMethodPost = HttpMethod("POST")
  HttpMethodPut = HttpMethod("PUT")
  HttpMethodDelete = HttpMethod("DELETE")
  HttpMethodPatch = HttpMethod("PATCH")
)
var HttpMethods = []HttpMethod{HttpMethodGet, HttpMethodPost, HttpMethodPut, HttpMethodDelete, HttpMethodPatch}

type Router struct {

  // Hierarchy:
  parent *Router      // nil if root router.
  children []*Router

  // For each request we'll create one of these objects
  contextType reflect.Type

  // Eg, "/" or "/admin". Any routes added to this router will be prefixed with this.
  pathPrefix string

  // Routeset contents:
  middleware []reflect.Value
  routes []*Route

  // The root pathnode is the same for a tree of Routers
  root map[HttpMethod]*PathNode

  // This can can be set on any router. The target's ErrorHandler will be invoked if it exists
  errorHandler reflect.Value

  // This can only be set on the root handler, since by virtue of not finding a route, we don't have a target.
  // (That being said, in the future we could investigate namespace matches)
  NotFoundHandler func(*ResponseWriter, *Request)
}

type Route struct {
  Router *Router
  Method HttpMethod
  Path string
  Handler reflect.Value // Dynamic method sig.
}

type NextMiddlewareFunc func()

func New(ctx interface{}) *Router {
  validateContext(ctx, nil)

  r := &Router{}
  r.contextType = reflect.TypeOf(ctx)
  r.pathPrefix = "/"
  r.root = make(map[HttpMethod]*PathNode)
  for _, method := range HttpMethods {
    r.root[method] = newPathNode()
  }
  return r
}

func (r *Router) NewSubrouter(ctx interface{}) *Router {

  // First, we need to make sure that ctx includes a pointer to the parent context in the first slot
  validateContext(ctx, r.contextType)

  // Create new router, link up hierarchy
  newRouter := &Router{parent: r}
  r.children = append(r.children, newRouter)

  newRouter.contextType = reflect.TypeOf(ctx)
  newRouter.pathPrefix = r.pathPrefix
  newRouter.root = r.root

  fmt.Println("newRouter: ", newRouter) // Keep this to allow fmt

  return newRouter
}

func (r *Router) AddMiddleware(fn interface{}) *Router {
  fnv := reflect.ValueOf(fn)
  validateMiddleware(fnv, r.contextType)
  r.middleware = append(r.middleware, fnv)
  return r
}

func (r *Router) SetNamespace(ns string) *Router {
  // TODO: do we need to re-eval all the routes ?
  // TODO: validate pathPrefix
  r.pathPrefix = ns
  return r
}

func (r *Router) ErrorHandler(fn interface{}) {
  vfn := reflect.ValueOf(fn)
  validateErrorHandler(vfn, r.contextType)
  r.errorHandler = vfn
}

func (r *Router) Get(path string, fn interface{}) {
  r.addRoute(HttpMethodGet, path, fn)
}

func (r *Router) Post(path string, fn interface{}) {
  r.addRoute(HttpMethodPost, path, fn)
}

func (r *Router) Put(path string, fn interface{}) {
  r.addRoute(HttpMethodPut, path, fn)
}

func (r *Router) Delete(path string, fn interface{}) {
  r.addRoute(HttpMethodDelete, path, fn)
}

func (r *Router) Patch(path string, fn interface{}) {
  r.addRoute(HttpMethodPatch, path, fn)
}

//
//
//
func (r *Router) addRoute(method HttpMethod, path string, fn interface{}) {
  fnv := reflect.ValueOf(fn)
  validateHandler(fnv, r.contextType)

  fullPath := appendPath(r.pathPrefix, path)

  route := &Route{Method: method, Path: fullPath, Handler: fnv, Router: r}
  r.routes = append(r.routes, route)

  r.root[method].add(fullPath, route)
}

//
// Private methods:
//

// Panics unless validation is correct
func validateContext(ctx interface{}, parentCtxType reflect.Type) {
  ctxType := reflect.TypeOf(ctx)

  if ctxType.Kind() != reflect.Struct {
    panic("web: Context needs to be a struct type")
  }

  if parentCtxType != nil {
    if ctxType.NumField() == 0 {
      panic("web: Context needs to have first field be a pointer to parent context")
    }

    fldType := ctxType.Field(0).Type

    // Ensure fld is a pointer to parentCtxType
    if fldType != reflect.PtrTo(parentCtxType) {
      panic("web: Context needs to have first field be a pointer to parent context")
    }
  }
}

// Panics unless fn is a proper handler wrt ctxType
// eg, func(ctx *ctxType, writer, request)
func validateHandler(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req)) {
    panic("web: handler be a function with signature TODO")
  }
}

func validateErrorHandler(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  var wat func() interface{}  // This is weird. I need to get an interface{} reflect.Type; var x interface{}; TypeOf(x) doesn't work, because it returns nil
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req), reflect.TypeOf(wat).Out(0)) {
    panic("web: error handler be a function with signature TODO")
  }
}

// Either of:
//    f(*context, *web.ResponseWriter, *web.Request, NextMiddlewareFunc)
//    f(*web.ResponseWriter, *web.Request, NextMiddlewareFunc)
func validateMiddleware(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  var n NextMiddlewareFunc
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req), reflect.TypeOf(n)) {
    panic("web: middlware must be a function with signature TODO")
  }
}

// Ensures fnv is a function, that optionally takes a ctxType as the first argument, followed by the specified types. Handles have no return value.
// Returns true if valid, false otherwise.
func isValidateHandler(fnv reflect.Value, ctxType reflect.Type, types ...reflect.Type) bool {
  fnType := fnv.Type()

  if fnType.Kind() != reflect.Func {
    fmt.Println("1")
    return false
  }

  typesStartIdx := 0
  typesLen := len(types)
  numIn := fnType.NumIn()
  numOut := fnType.NumOut()

  if numOut != 0 {
    fmt.Println("2")
    return false
  }

  if numIn == typesLen {
    // No context
  } else if numIn == (typesLen + 1) {
    // context, types
    if fnType.In(0) != reflect.PtrTo(ctxType) {
      fmt.Println("3")
      return false
    }
    typesStartIdx = 1
  } else {
    fmt.Println("4")
    return false
  }
  fmt.Println(types)
  for _, typeArg := range types {
    fmt.Println("comparing ", fnType.In(typesStartIdx), " vs ", typeArg)
    if fnType.In(typesStartIdx) != typeArg {
      fmt.Println("5")
      return false
    }
    typesStartIdx += 1
  }

  return true
}

// Both rootPath/childPath are like "/" and "/users"
// Assumption is that both are well-formed paths.
func appendPath(rootPath, childPath string) string {
  if rootPath == "/" {
    return childPath
  }

  return rootPath + childPath
}

<commit_msg>Remove fmt<commit_after>package web

import (
  "reflect"
  "fmt"
)

type HttpMethod string
const (
  HttpMethodGet = HttpMethod("GET")
  HttpMethodPost = HttpMethod("POST")
  HttpMethodPut = HttpMethod("PUT")
  HttpMethodDelete = HttpMethod("DELETE")
  HttpMethodPatch = HttpMethod("PATCH")
)
var HttpMethods = []HttpMethod{HttpMethodGet, HttpMethodPost, HttpMethodPut, HttpMethodDelete, HttpMethodPatch}

type Router struct {

  // Hierarchy:
  parent *Router      // nil if root router.
  children []*Router

  // For each request we'll create one of these objects
  contextType reflect.Type

  // Eg, "/" or "/admin". Any routes added to this router will be prefixed with this.
  pathPrefix string

  // Routeset contents:
  middleware []reflect.Value
  routes []*Route

  // The root pathnode is the same for a tree of Routers
  root map[HttpMethod]*PathNode

  // This can can be set on any router. The target's ErrorHandler will be invoked if it exists
  errorHandler reflect.Value

  // This can only be set on the root handler, since by virtue of not finding a route, we don't have a target.
  // (That being said, in the future we could investigate namespace matches)
  NotFoundHandler func(*ResponseWriter, *Request)
}

type Route struct {
  Router *Router
  Method HttpMethod
  Path string
  Handler reflect.Value // Dynamic method sig.
}

type NextMiddlewareFunc func()

func New(ctx interface{}) *Router {
  validateContext(ctx, nil)

  r := &Router{}
  r.contextType = reflect.TypeOf(ctx)
  r.pathPrefix = "/"
  r.root = make(map[HttpMethod]*PathNode)
  for _, method := range HttpMethods {
    r.root[method] = newPathNode()
  }
  return r
}

func (r *Router) NewSubrouter(ctx interface{}) *Router {

  // First, we need to make sure that ctx includes a pointer to the parent context in the first slot
  validateContext(ctx, r.contextType)

  // Create new router, link up hierarchy
  newRouter := &Router{parent: r}
  r.children = append(r.children, newRouter)

  newRouter.contextType = reflect.TypeOf(ctx)
  newRouter.pathPrefix = r.pathPrefix
  newRouter.root = r.root

  fmt.Println("newRouter: ", newRouter) // Keep this to allow fmt

  return newRouter
}

func (r *Router) AddMiddleware(fn interface{}) *Router {
  fnv := reflect.ValueOf(fn)
  validateMiddleware(fnv, r.contextType)
  r.middleware = append(r.middleware, fnv)
  return r
}

func (r *Router) SetNamespace(ns string) *Router {
  // TODO: do we need to re-eval all the routes ?
  // TODO: validate pathPrefix
  r.pathPrefix = ns
  return r
}

func (r *Router) ErrorHandler(fn interface{}) {
  vfn := reflect.ValueOf(fn)
  validateErrorHandler(vfn, r.contextType)
  r.errorHandler = vfn
}

func (r *Router) Get(path string, fn interface{}) {
  r.addRoute(HttpMethodGet, path, fn)
}

func (r *Router) Post(path string, fn interface{}) {
  r.addRoute(HttpMethodPost, path, fn)
}

func (r *Router) Put(path string, fn interface{}) {
  r.addRoute(HttpMethodPut, path, fn)
}

func (r *Router) Delete(path string, fn interface{}) {
  r.addRoute(HttpMethodDelete, path, fn)
}

func (r *Router) Patch(path string, fn interface{}) {
  r.addRoute(HttpMethodPatch, path, fn)
}

//
//
//
func (r *Router) addRoute(method HttpMethod, path string, fn interface{}) {
  fnv := reflect.ValueOf(fn)
  validateHandler(fnv, r.contextType)

  fullPath := appendPath(r.pathPrefix, path)

  route := &Route{Method: method, Path: fullPath, Handler: fnv, Router: r}
  r.routes = append(r.routes, route)

  r.root[method].add(fullPath, route)
}

//
// Private methods:
//

// Panics unless validation is correct
func validateContext(ctx interface{}, parentCtxType reflect.Type) {
  ctxType := reflect.TypeOf(ctx)

  if ctxType.Kind() != reflect.Struct {
    panic("web: Context needs to be a struct type")
  }

  if parentCtxType != nil {
    if ctxType.NumField() == 0 {
      panic("web: Context needs to have first field be a pointer to parent context")
    }

    fldType := ctxType.Field(0).Type

    // Ensure fld is a pointer to parentCtxType
    if fldType != reflect.PtrTo(parentCtxType) {
      panic("web: Context needs to have first field be a pointer to parent context")
    }
  }
}

// Panics unless fn is a proper handler wrt ctxType
// eg, func(ctx *ctxType, writer, request)
func validateHandler(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req)) {
    panic("web: handler be a function with signature TODO")
  }
}

func validateErrorHandler(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  var wat func() interface{}  // This is weird. I need to get an interface{} reflect.Type; var x interface{}; TypeOf(x) doesn't work, because it returns nil
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req), reflect.TypeOf(wat).Out(0)) {
    panic("web: error handler be a function with signature TODO")
  }
}

// Either of:
//    f(*context, *web.ResponseWriter, *web.Request, NextMiddlewareFunc)
//    f(*web.ResponseWriter, *web.Request, NextMiddlewareFunc)
func validateMiddleware(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp *ResponseWriter
  var n NextMiddlewareFunc
  if !isValidateHandler(fnv, ctxType, reflect.TypeOf(resp), reflect.TypeOf(req), reflect.TypeOf(n)) {
    panic("web: middlware must be a function with signature TODO")
  }
}

// Ensures fnv is a function, that optionally takes a ctxType as the first argument, followed by the specified types. Handles have no return value.
// Returns true if valid, false otherwise.
func isValidateHandler(fnv reflect.Value, ctxType reflect.Type, types ...reflect.Type) bool {
  fnType := fnv.Type()

  if fnType.Kind() != reflect.Func {
    return false
  }

  typesStartIdx := 0
  typesLen := len(types)
  numIn := fnType.NumIn()
  numOut := fnType.NumOut()

  if numOut != 0 {
    return false
  }

  if numIn == typesLen {
    // No context
  } else if numIn == (typesLen + 1) {
    // context, types
    if fnType.In(0) != reflect.PtrTo(ctxType) {
      return false
    }
    typesStartIdx = 1
  } else {
    return false
  }

  for _, typeArg := range types {
    if fnType.In(typesStartIdx) != typeArg {
      return false
    }
    typesStartIdx += 1
  }

  return true
}

// Both rootPath/childPath are like "/" and "/users"
// Assumption is that both are well-formed paths.
func appendPath(rootPath, childPath string) string {
  if rootPath == "/" {
    return childPath
  }

  return rootPath + childPath
}

<|endoftext|>


<commit_before>package web

import (
  "reflect"
)

type HttpMethod string
const (
  HttpMethodGet = HttpMethod("GET")
  HttpMethodPost = HttpMethod("POST")
  HttpMethodPut = HttpMethod("PUT")
  HttpMethodDelete = HttpMethod("DELETE")
  HttpMethodPatch = HttpMethod("PATCH")
)
var HttpMethods = []HttpMethod{HttpMethodGet, HttpMethodPost, HttpMethodPut, HttpMethodDelete, HttpMethodPatch}

type Router struct {

  // Hierarchy:
  parent *Router      // nil if root router.
  children []*Router

  // For each request we'll create one of these objects
  contextType reflect.Type

  // Eg, "/" or "/admin". Any routes added to this router will be prefixed with this.
  pathPrefix string

  // Routeset contents:
  middleware []reflect.Value
  routes []*Route

  // The root pathnode is the same for a tree of Routers
  root map[HttpMethod]*PathNode

  // This can can be set on any router. The target's ErrorHandler will be invoked if it exists
  errorHandler reflect.Value

  // This can only be set on the root handler, since by virtue of not finding a route, we don't have a target.
  // (That being said, in the future we could investigate namespace matches)
  notFoundHandler reflect.Value
}

type Route struct {
  Router *Router
  Method HttpMethod
  Path string
  Handler reflect.Value // Dynamic method sig.
}

type NextMiddlewareFunc func(ResponseWriter, *Request)

func New(ctx interface{}) *Router {
  validateContext(ctx, nil)

  r := &Router{}
  r.contextType = reflect.TypeOf(ctx)
  r.pathPrefix = "/"
  r.root = make(map[HttpMethod]*PathNode)
  for _, method := range HttpMethods {
    r.root[method] = newPathNode()
  }
  return r
}

func (r *Router) Subrouter(ctx interface{}, pathPrefix string) *Router {
  validateContext(ctx, r.contextType)

  // Create new router, link up hierarchy
  newRouter := &Router{parent: r}
  r.children = append(r.children, newRouter)

  newRouter.contextType = reflect.TypeOf(ctx)
  newRouter.pathPrefix = appendPath(r.pathPrefix, pathPrefix)
  newRouter.root = r.root

  return newRouter
}

func (r *Router) Middleware(fn interface{}) *Router {
  vfn := reflect.ValueOf(fn)
  validateMiddleware(vfn, r.contextType)
  r.middleware = append(r.middleware, vfn)
  return r
}

func (r *Router) ErrorHandler(fn interface{}) {
  vfn := reflect.ValueOf(fn)
  validateErrorHandler(vfn, r.contextType)
  r.errorHandler = vfn
}

func (r *Router) NotFoundHandler(fn interface{}) {
  if r.parent != nil {
    panic("You can only set a NotFoundHandler on the root router.")
  }
  vfn := reflect.ValueOf(fn)
  validateNotFoundHandler(vfn, r.contextType)
  r.notFoundHandler = vfn
}

func (r *Router) Get(path string, fn interface{}) {
  r.addRoute(HttpMethodGet, path, fn)
}

func (r *Router) Post(path string, fn interface{}) {
  r.addRoute(HttpMethodPost, path, fn)
}

func (r *Router) Put(path string, fn interface{}) {
  r.addRoute(HttpMethodPut, path, fn)
}

func (r *Router) Delete(path string, fn interface{}) {
  r.addRoute(HttpMethodDelete, path, fn)
}

func (r *Router) Patch(path string, fn interface{}) {
  r.addRoute(HttpMethodPatch, path, fn)
}

//
//
//
func (r *Router) addRoute(method HttpMethod, path string, fn interface{}) {
  fnv := reflect.ValueOf(fn)
  validateHandler(fnv, r.contextType)
  fullPath := appendPath(r.pathPrefix, path)
  route := &Route{Method: method, Path: fullPath, Handler: fnv, Router: r}
  r.routes = append(r.routes, route)
  r.root[method].add(fullPath, route)
}

//
// Private methods:
//

// Panics unless validation is correct
func validateContext(ctx interface{}, parentCtxType reflect.Type) {
  ctxType := reflect.TypeOf(ctx)

  if ctxType.Kind() != reflect.Struct {
    panic("web: Context needs to be a struct type")
  }

  if parentCtxType != nil && parentCtxType != ctxType {
    if ctxType.NumField() == 0 {
      panic("web: Context needs to have first field be a pointer to parent context")
    }

    fldType := ctxType.Field(0).Type

    // Ensure fld is a pointer to parentCtxType
    if fldType != reflect.PtrTo(parentCtxType) {
      panic("web: Context needs to have first field be a pointer to parent context")
    }
  }
}

// Panics unless fn is a proper handler wrt ctxType
// eg, func(ctx *ctxType, writer, request)
func validateHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req)) {
    panic("web: handler be a function with signature TODO")
  }
}

func validateErrorHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  var interfaceType func() interface{}  // This is weird. I need to get an interface{} reflect.Type; var x interface{}; TypeOf(x) doesn't work, because it returns nil
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req), reflect.TypeOf(interfaceType).Out(0)) {
    panic("web: error handler be a function with signature TODO")
  }
}

func validateNotFoundHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req)) {
    panic("web: not found handler be a function with signature TODO")
  }
}

// Either of:
//    f(*context, *web.ResponseWriter, *web.Request, NextMiddlewareFunc)
//    f(*web.ResponseWriter, *web.Request, NextMiddlewareFunc)
func validateMiddleware(fnv reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  var n NextMiddlewareFunc
  if !isValidHandler(fnv, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req), reflect.TypeOf(n)) {
    panic("web: middlware must be a function with signature TODO")
  }
}

// Ensures fnv is a function, that optionally takes a *ctxType as the first argument, followed by the specified types. Handlers have no return value.
// Returns true if valid, false otherwise.
func isValidHandler(fnv reflect.Value, ctxType reflect.Type, types ...reflect.Type) bool {
  fnType := fnv.Type()

  if fnType.Kind() != reflect.Func {
    return false
  }

  typesStartIdx := 0
  typesLen := len(types)
  numIn := fnType.NumIn()
  numOut := fnType.NumOut()

  if numOut != 0 {
    return false
  }

  if numIn == typesLen {
    // No context
  } else if numIn == (typesLen + 1) {
    // context, types
    if fnType.In(0) != reflect.PtrTo(ctxType) {
      return false
    }
    typesStartIdx = 1
  } else {
    return false
  }

  for _, typeArg := range types {
    if fnType.In(typesStartIdx) != typeArg {
      return false
    }
    typesStartIdx += 1
  }

  return true
}

// Both rootPath/childPath are like "/" and "/users"
// Assumption is that both are well-formed paths.
func appendPath(rootPath, childPath string) string {
  if rootPath == "/" {
    return childPath
  }

  return rootPath + childPath
}

<commit_msg>Instructive error messaging when adding handlers<commit_after>package web

import (
  "reflect"
  "strings"
)

type HttpMethod string
const (
  HttpMethodGet = HttpMethod("GET")
  HttpMethodPost = HttpMethod("POST")
  HttpMethodPut = HttpMethod("PUT")
  HttpMethodDelete = HttpMethod("DELETE")
  HttpMethodPatch = HttpMethod("PATCH")
)
var HttpMethods = []HttpMethod{HttpMethodGet, HttpMethodPost, HttpMethodPut, HttpMethodDelete, HttpMethodPatch}

type Router struct {
  // Hierarchy:
  parent *Router      // nil if root router.
  children []*Router

  // For each request we'll create one of these objects
  contextType reflect.Type

  // Eg, "/" or "/admin". Any routes added to this router will be prefixed with this.
  pathPrefix string

  // Routeset contents:
  middleware []reflect.Value
  routes []*Route

  // The root pathnode is the same for a tree of Routers
  root map[HttpMethod]*PathNode

  // This can can be set on any router. The target's ErrorHandler will be invoked if it exists
  errorHandler reflect.Value

  // This can only be set on the root handler, since by virtue of not finding a route, we don't have a target.
  // (That being said, in the future we could investigate namespace matches)
  notFoundHandler reflect.Value
}

type Route struct {
  Router *Router
  Method HttpMethod
  Path string
  Handler reflect.Value // Dynamic method sig.
}

type NextMiddlewareFunc func(ResponseWriter, *Request)

func New(ctx interface{}) *Router {
  validateContext(ctx, nil)

  r := &Router{}
  r.contextType = reflect.TypeOf(ctx)
  r.pathPrefix = "/"
  r.root = make(map[HttpMethod]*PathNode)
  for _, method := range HttpMethods {
    r.root[method] = newPathNode()
  }
  return r
}

func (r *Router) Subrouter(ctx interface{}, pathPrefix string) *Router {
  validateContext(ctx, r.contextType)

  // Create new router, link up hierarchy
  newRouter := &Router{parent: r}
  r.children = append(r.children, newRouter)

  newRouter.contextType = reflect.TypeOf(ctx)
  newRouter.pathPrefix = appendPath(r.pathPrefix, pathPrefix)
  newRouter.root = r.root

  return newRouter
}

func (r *Router) Middleware(fn interface{}) *Router {
  vfn := reflect.ValueOf(fn)
  validateMiddleware(vfn, r.contextType)
  r.middleware = append(r.middleware, vfn)
  return r
}

func (r *Router) ErrorHandler(fn interface{}) {
  vfn := reflect.ValueOf(fn)
  validateErrorHandler(vfn, r.contextType)
  r.errorHandler = vfn
}

func (r *Router) NotFoundHandler(fn interface{}) {
  if r.parent != nil {
    panic("You can only set a NotFoundHandler on the root router.")
  }
  vfn := reflect.ValueOf(fn)
  validateNotFoundHandler(vfn, r.contextType)
  r.notFoundHandler = vfn
}

func (r *Router) Get(path string, fn interface{}) {
  r.addRoute(HttpMethodGet, path, fn)
}

func (r *Router) Post(path string, fn interface{}) {
  r.addRoute(HttpMethodPost, path, fn)
}

func (r *Router) Put(path string, fn interface{}) {
  r.addRoute(HttpMethodPut, path, fn)
}

func (r *Router) Delete(path string, fn interface{}) {
  r.addRoute(HttpMethodDelete, path, fn)
}

func (r *Router) Patch(path string, fn interface{}) {
  r.addRoute(HttpMethodPatch, path, fn)
}

//
//
//
func (r *Router) addRoute(method HttpMethod, path string, fn interface{}) {
  fnv := reflect.ValueOf(fn)
  validateHandler(fnv, r.contextType)
  fullPath := appendPath(r.pathPrefix, path)
  route := &Route{Method: method, Path: fullPath, Handler: fnv, Router: r}
  r.routes = append(r.routes, route)
  r.root[method].add(fullPath, route)
}

//
// Private methods:
//

// Panics unless validation is correct
func validateContext(ctx interface{}, parentCtxType reflect.Type) {
  ctxType := reflect.TypeOf(ctx)

  if ctxType.Kind() != reflect.Struct {
    panic("web: Context needs to be a struct type")
  }

  if parentCtxType != nil && parentCtxType != ctxType {
    if ctxType.NumField() == 0 {
      panic("web: Context needs to have first field be a pointer to parent context")
    }

    fldType := ctxType.Field(0).Type

    // Ensure fld is a pointer to parentCtxType
    if fldType != reflect.PtrTo(parentCtxType) {
      panic("web: Context needs to have first field be a pointer to parent context")
    }
  }
}

// Panics unless fn is a proper handler wrt ctxType
// eg, func(ctx *ctxType, writer, request)
func validateHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req)) {
    panic(instructiveMessage(vfn, "a handler", "handler", "rw web.ResponseWriter, req *web.Request", ctxType))
  }
}

func validateErrorHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  var interfaceType func() interface{}  // This is weird. I need to get an interface{} reflect.Type; var x interface{}; TypeOf(x) doesn't work, because it returns nil
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req), reflect.TypeOf(interfaceType).Out(0)) {
    panic(instructiveMessage(vfn, "an error handler", "error handler", "rw web.ResponseWriter, req *web.Request, err interface{}", ctxType))
  }
}

func validateNotFoundHandler(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req)) {
    panic(instructiveMessage(vfn, "a 'not found' handler", "not found handler", "rw web.ResponseWriter, req *web.Request", ctxType))
  }
}

func validateMiddleware(vfn reflect.Value, ctxType reflect.Type) {
  var req *Request
  var resp func() ResponseWriter
  var n NextMiddlewareFunc
  if !isValidHandler(vfn, ctxType, reflect.TypeOf(resp).Out(0), reflect.TypeOf(req), reflect.TypeOf(n)) {
    panic(instructiveMessage(vfn, "middleware", "middleware", "rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc", ctxType))
  }
}

// Ensures vfn is a function, that optionally takes a *ctxType as the first argument, followed by the specified types. Handlers have no return value.
// Returns true if valid, false otherwise.
func isValidHandler(vfn reflect.Value, ctxType reflect.Type, types ...reflect.Type) bool {
  fnType := vfn.Type()

  if fnType.Kind() != reflect.Func {
    return false
  }

  typesStartIdx := 0
  typesLen := len(types)
  numIn := fnType.NumIn()
  numOut := fnType.NumOut()

  if numOut != 0 {
    return false
  }

  if numIn == typesLen {
    // No context
  } else if numIn == (typesLen + 1) {
    // context, types
    if fnType.In(0) != reflect.PtrTo(ctxType) {
      return false
    }
    typesStartIdx = 1
  } else {
    return false
  }

  for _, typeArg := range types {
    if fnType.In(typesStartIdx) != typeArg {
      return false
    }
    typesStartIdx += 1
  }

  return true
}

// Since it's easy to pass the wrong method to a middleware/handler route, and since the user can't rely on static type checking since we use reflection,
// lets be super helpful about what they did and what they need to do.
// Arguments:
//  - vfn is the failed method
//  - addingType is for "You are adding {addingType} to a router...". Eg, "middleware" or "a handler" or "an error handler"
//  - yourType is for "Your {yourType} function can have...". Eg, "middleware" or "handler" or "error handler"
//  - args is like "rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc"
//    - NOTE: args can be calculated if you pass in each type. BUT, it doesn't have example argument name, so it has less copy/paste value.
func instructiveMessage(vfn reflect.Value, addingType string, yourType string, args string, ctxType reflect.Type) string {
  // Get context type without package.
  ctxString := ctxType.String()
  splitted := strings.Split(ctxString, ".")
  if len(splitted) <= 1 {
    ctxString = splitted[0]
  } else {
    ctxString = splitted[1]
  }

  str := "\n" + strings.Repeat("*", 120) + "\n"
  str += "* You are adding " + addingType + " to a router with context type '" + ctxString + "'\n"
  str += "*\n*\n"
  str += "* Your " + yourType + " function can have one of these signatures:\n"
  str += "*\n"
  str += "* // If you don't need context:\n"
  str += "* func YourFunctionName(" + args + ")\n"
  str += "*\n"
  str += "* // If you want your " + yourType + " to accept a context:\n"
  str += "* func (c *" + ctxString + ") YourFunctionName(" + args + ")  // or,\n"
  str += "* func YourFunctionName(c *" + ctxString + ", " + args + ")\n"
  str += "*\n"
  str += "* Unfortunately, your function has this signature: " + vfn.Type().String() + "\n"
  str += "*\n"
  str += strings.Repeat("*", 120) + "\n"

  return str
}

// Both rootPath/childPath are like "/" and "/users"
// Assumption is that both are well-formed paths.
func appendPath(rootPath, childPath string) string {
  if rootPath == "/" {
    return childPath
  }

  return rootPath + childPath
}

<|endoftext|>


<commit_before>//go:generate mapstructure-to-hcl2 -type Config

// NB this code was based on https://github.com/hashicorp/packer/blob/81522dced0b25084a824e79efda02483b12dc7cd/provisioner/windows-restart/provisioner.go

package update

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"
	"time"
	"unicode/utf16"

	"github.com/hashicorp/hcl/v2/hcldec"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/common/retry"
	"github.com/hashicorp/packer/common/uuid"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

const (
	elevatedPath                 = "C:/Windows/Temp/packer-windows-update-elevated.ps1"
	elevatedCommand              = "PowerShell -ExecutionPolicy Bypass -OutputFormat Text -File C:/Windows/Temp/packer-windows-update-elevated.ps1"
	windowsUpdatePath            = "C:/Windows/Temp/packer-windows-update.ps1"
	pendingRebootElevatedPath    = "C:/Windows/Temp/packer-windows-update-pending-reboot-elevated.ps1"
	pendingRebootElevatedCommand = "PowerShell -ExecutionPolicy Bypass -OutputFormat Text -File C:/Windows/Temp/packer-windows-update-pending-reboot-elevated.ps1"
	restartCommand               = "shutdown.exe -f -r -t 0 -c \"packer restart\""
	testRestartCommand           = "shutdown.exe -f -r -t 60 -c \"packer restart test\""
	abortTestRestartCommand      = "shutdown.exe -a"
	retryableDelay               = 5 * time.Second
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	// Instructs the communicator to run the remote script as a
	// Windows scheduled task, effectively elevating the remote
	// user by impersonating a logged-in user.
	Username string `mapstructure:"username"`
	Password string `mapstructure:"password"`

	// The updates search criteria.
	// See the IUpdateSearcher::Search method at https://docs.microsoft.com/en-us/windows/desktop/api/wuapi/nf-wuapi-iupdatesearcher-search.
	SearchCriteria string `mapstructure:"search_criteria"`

	// Filters the installed Windows updates. If no filter is
	// matched the update is NOT installed.
	Filters []string `mapstructure:"filters"`

	// Adds a limit to how many updates are installed at a time
	UpdateLimit int `mapstructure:"update_limit"`

	ctx interpolate.Context
}

type Provisioner struct {
	config Config
}

func (b *Provisioner) ConfigSpec() hcldec.ObjectSpec {
	return b.config.FlatMapstructure().HCL2Spec()
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 4 * time.Hour
	}

	if p.config.Username == "" {
		p.config.Username = "SYSTEM"
	}

	var errs error

	if p.config.Username == "" {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Must supply an 'username'"))
	}

	if p.config.UpdateLimit == 0 {
		p.config.UpdateLimit = 1000
	}

	return errs
}

func (p *Provisioner) Provision(ctx context.Context, ui packer.Ui, comm packer.Communicator, _ map[string]interface{}) error {
	ui.Say("Uploading the Windows update elevated script...")
	var buffer bytes.Buffer
	err := elevatedTemplate.Execute(&buffer, elevatedOptions{
		Username:        p.config.Username,
		Password:        p.config.Password,
		TaskDescription: "Packer Windows update elevated task",
		TaskName:        fmt.Sprintf("packer-windows-update-%s", uuid.TimeOrderedUUID()),
		Command:         p.windowsUpdateCommand(),
	})
	if err != nil {
		fmt.Printf("Error creating elevated template: %s", err)
		return err
	}
	err = comm.Upload(
		elevatedPath,
		bytes.NewReader(buffer.Bytes()),
		nil)
	if err != nil {
		return err
	}

	ui.Say("Uploading the Windows update check for reboot required elevated script...")
	buffer.Reset()
	err = elevatedTemplate.Execute(&buffer, elevatedOptions{
		Username:        p.config.Username,
		Password:        p.config.Password,
		TaskDescription: "Packer Windows update pending reboot elevated task",
		TaskName:        fmt.Sprintf("packer-windows-update-pending-reboot-%s", uuid.TimeOrderedUUID()),
		Command:         p.windowsUpdateCheckForRebootRequiredCommand(),
	})
	if err != nil {
		fmt.Printf("Error creating elevated template: %s", err)
		return err
	}
	err = comm.Upload(
		pendingRebootElevatedPath,
		bytes.NewReader(buffer.Bytes()),
		nil)
	if err != nil {
		return err
	}

	ui.Say("Uploading the Windows update script...")
	err = comm.Upload(
		windowsUpdatePath,
		bytes.NewReader(MustAsset("windows-update.ps1")),
		nil)
	if err != nil {
		return err
	}

	for {
		restartPending, err := p.update(ctx, ui, comm)
		if err != nil {
			return err
		}

		if !restartPending {
			return nil
		}

		err = p.restart(ctx, ui, comm)
		if err != nil {
			return err
		}
	}
}

func (p *Provisioner) update(ctx context.Context, ui packer.Ui, comm packer.Communicator) (bool, error) {
	ui.Say("Running Windows update...")
	cmd := &packer.RemoteCmd{Command: elevatedCommand}
	err := cmd.RunWithUi(ctx, comm, ui)
	if err != nil {
		return false, err
	}
	var exitStatus = cmd.ExitStatus()
	switch exitStatus {
	case 0:
		return false, nil
	case 101:
		return true, nil
	default:
		return false, fmt.Errorf("Windows update script exited with non-zero exit status: %d", exitStatus)
	}
}

func (p *Provisioner) restart(ctx context.Context, ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Restarting the machine...")
	err := p.retryable(ctx, func(ctx context.Context) error {
		cmd := &packer.RemoteCmd{Command: restartCommand}
		err := cmd.RunWithUi(ctx, comm, ui)
		if err != nil {
			return err
		}
		exitStatus := cmd.ExitStatus()
		if exitStatus != 0 {
			return fmt.Errorf("Failed to restart the machine with exit status: %d", exitStatus)
		}
		return nil
	})
	if err != nil {
		return err
	}

	ui.Say("Waiting for machine to become available...")
	err = p.retryable(ctx, func(ctx context.Context) error {
		// wait for the machine to reboot.
		cmd := &packer.RemoteCmd{Command: testRestartCommand}
		err := cmd.RunWithUi(ctx, comm, ui)
		if err != nil {
			return err
		}
		exitStatus := cmd.ExitStatus()
		if exitStatus != 0 {
			return fmt.Errorf("Machine not yet available (exit status %d)", exitStatus)
		}
		cmd = &packer.RemoteCmd{Command: abortTestRestartCommand}
		err = cmd.RunWithUi(ctx, comm, ui)
		if err != nil {
			return err
		}

		// wait for pending tasks to finish.
		cmd = &packer.RemoteCmd{Command: pendingRebootElevatedCommand}
		err = cmd.RunWithUi(ctx, comm, ui)
		if err != nil {
			return err
		}
		exitStatus = cmd.ExitStatus()
		if exitStatus != 0 {
			return fmt.Errorf("Machine not yet available (exit status %d)", exitStatus)
		}

		return nil
	})
	return err
}

// retryable will retry the given function over and over until a
// non-error is returned, RestartTimeout expires, or ctx is
// cancelled.
func (p *Provisioner) retryable(ctx context.Context, f func(ctx context.Context) error) error {
	return retry.Config{
		RetryDelay:   func() time.Duration { return retryableDelay },
		StartTimeout: p.config.RestartTimeout,
	}.Run(ctx, f)
}

func (p *Provisioner) windowsUpdateCommand() string {
	return fmt.Sprintf(
		"PowerShell -ExecutionPolicy Bypass -OutputFormat Text -EncodedCommand %s",
		base64.StdEncoding.EncodeToString(
			encodeUtf16Le(fmt.Sprintf(
				"%s%s%s -UpdateLimit %d",
				windowsUpdatePath,
				searchCriteriaArgument(p.config.SearchCriteria),
				filtersArgument(p.config.Filters),
				p.config.UpdateLimit))))
}

func (p *Provisioner) windowsUpdateCheckForRebootRequiredCommand() string {
	return fmt.Sprintf(
		"PowerShell -ExecutionPolicy Bypass -OutputFormat Text -EncodedCommand %s",
		base64.StdEncoding.EncodeToString(
			encodeUtf16Le(fmt.Sprintf(
				"%s -OnlyCheckForRebootRequired",
				windowsUpdatePath))))
}

func encodeUtf16Le(s string) []byte {
	d := utf16.Encode([]rune(s))
	b := make([]byte, len(d)*2)
	for i, r := range d {
		b[i*2] = byte(r)
		b[i*2+1] = byte(r >> 8)
	}
	return b
}

func searchCriteriaArgument(searchCriteria string) string {
	if searchCriteria == "" {
		return ""
	}

	var buffer bytes.Buffer

	buffer.WriteString(" -SearchCriteria ")
	buffer.WriteString(escapePowerShellString(searchCriteria))

	return buffer.String()
}

func filtersArgument(filters []string) string {
	if filters == nil {
		return ""
	}

	var buffer bytes.Buffer

	buffer.WriteString(" -Filters ")

	for i, filter := range filters {
		if i > 0 {
			buffer.WriteString(",")
		}
		buffer.WriteString(escapePowerShellString(filter))
	}

	return buffer.String()
}

func escapePowerShellString(value string) string {
	return fmt.Sprintf(
		"'%s'",
		// escape single quotes with another single quote.
		strings.Replace(value, "'", "''", -1))
}
<commit_msg>Fix issue #61 - Spurious issue observed where the restartCommand silently fails - Enhanced logic in 'restart' function to be robust against this failure - Added additional ui output in 'restart' function<commit_after>//go:generate mapstructure-to-hcl2 -type Config

// NB this code was based on https://github.com/hashicorp/packer/blob/81522dced0b25084a824e79efda02483b12dc7cd/provisioner/windows-restart/provisioner.go

package update

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"
	"time"
	"unicode/utf16"

	"github.com/hashicorp/hcl/v2/hcldec"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/common/retry"
	"github.com/hashicorp/packer/common/uuid"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

const (
	elevatedPath                 = "C:/Windows/Temp/packer-windows-update-elevated.ps1"
	elevatedCommand              = "PowerShell -ExecutionPolicy Bypass -OutputFormat Text -File C:/Windows/Temp/packer-windows-update-elevated.ps1"
	windowsUpdatePath            = "C:/Windows/Temp/packer-windows-update.ps1"
	pendingRebootElevatedPath    = "C:/Windows/Temp/packer-windows-update-pending-reboot-elevated.ps1"
	pendingRebootElevatedCommand = "PowerShell -ExecutionPolicy Bypass -OutputFormat Text -File C:/Windows/Temp/packer-windows-update-pending-reboot-elevated.ps1"
	restartCommand               = "shutdown.exe -f -r -t 0 -c \"packer restart\""
	testRestartCommand           = "shutdown.exe -f -r -t 60 -c \"packer restart test\""
	abortTestRestartCommand      = "shutdown.exe -a"
	retryableDelay               = 5 * time.Second
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	// Instructs the communicator to run the remote script as a
	// Windows scheduled task, effectively elevating the remote
	// user by impersonating a logged-in user.
	Username string `mapstructure:"username"`
	Password string `mapstructure:"password"`

	// The updates search criteria.
	// See the IUpdateSearcher::Search method at https://docs.microsoft.com/en-us/windows/desktop/api/wuapi/nf-wuapi-iupdatesearcher-search.
	SearchCriteria string `mapstructure:"search_criteria"`

	// Filters the installed Windows updates. If no filter is
	// matched the update is NOT installed.
	Filters []string `mapstructure:"filters"`

	// Adds a limit to how many updates are installed at a time
	UpdateLimit int `mapstructure:"update_limit"`

	ctx interpolate.Context
}

type Provisioner struct {
	config Config
}

func (b *Provisioner) ConfigSpec() hcldec.ObjectSpec {
	return b.config.FlatMapstructure().HCL2Spec()
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 4 * time.Hour
	}

	if p.config.Username == "" {
		p.config.Username = "SYSTEM"
	}

	var errs error

	if p.config.Username == "" {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Must supply an 'username'"))
	}

	if p.config.UpdateLimit == 0 {
		p.config.UpdateLimit = 1000
	}

	return errs
}

func (p *Provisioner) Provision(ctx context.Context, ui packer.Ui, comm packer.Communicator, _ map[string]interface{}) error {
	ui.Say("Uploading the Windows update elevated script...")
	var buffer bytes.Buffer
	err := elevatedTemplate.Execute(&buffer, elevatedOptions{
		Username:        p.config.Username,
		Password:        p.config.Password,
		TaskDescription: "Packer Windows update elevated task",
		TaskName:        fmt.Sprintf("packer-windows-update-%s", uuid.TimeOrderedUUID()),
		Command:         p.windowsUpdateCommand(),
	})
	if err != nil {
		fmt.Printf("Error creating elevated template: %s", err)
		return err
	}
	err = comm.Upload(
		elevatedPath,
		bytes.NewReader(buffer.Bytes()),
		nil)
	if err != nil {
		return err
	}

	ui.Say("Uploading the Windows update check for reboot required elevated script...")
	buffer.Reset()
	err = elevatedTemplate.Execute(&buffer, elevatedOptions{
		Username:        p.config.Username,
		Password:        p.config.Password,
		TaskDescription: "Packer Windows update pending reboot elevated task",
		TaskName:        fmt.Sprintf("packer-windows-update-pending-reboot-%s", uuid.TimeOrderedUUID()),
		Command:         p.windowsUpdateCheckForRebootRequiredCommand(),
	})
	if err != nil {
		fmt.Printf("Error creating elevated template: %s", err)
		return err
	}
	err = comm.Upload(
		pendingRebootElevatedPath,
		bytes.NewReader(buffer.Bytes()),
		nil)
	if err != nil {
		return err
	}

	ui.Say("Uploading the Windows update script...")
	err = comm.Upload(
		windowsUpdatePath,
		bytes.NewReader(MustAsset("windows-update.ps1")),
		nil)
	if err != nil {
		return err
	}

	for {
		restartPending, err := p.update(ctx, ui, comm)
		if err != nil {
			return err
		}

		if !restartPending {
			return nil
		}

		err = p.restart(ctx, ui, comm)
		if err != nil {
			return err
		}
	}
}

func (p *Provisioner) update(ctx context.Context, ui packer.Ui, comm packer.Communicator) (bool, error) {
	ui.Say("Running Windows update...")
	cmd := &packer.RemoteCmd{Command: elevatedCommand}
	err := cmd.RunWithUi(ctx, comm, ui)
	if err != nil {
		return false, err
	}
	var exitStatus = cmd.ExitStatus()
	switch exitStatus {
	case 0:
		return false, nil
	case 101:
		return true, nil
	default:
		return false, fmt.Errorf("Windows update script exited with non-zero exit status: %d", exitStatus)
	}
}

func (p *Provisioner) restart(ctx context.Context, ui packer.Ui, comm packer.Communicator) error {
	restartPending := true
	for restartPending {
		ui.Say("Restarting the machine...")
		err := p.retryable(ctx, func(ctx context.Context) error {
			cmd := &packer.RemoteCmd{Command: restartCommand}
			err := cmd.RunWithUi(ctx, comm, ui)
			if err != nil {
				return err
			}
			exitStatus := cmd.ExitStatus()
			if exitStatus != 0 {
				return fmt.Errorf("Failed to restart the machine with exit status: %d", exitStatus)
			}
			return err
		})
		if err != nil {
			return err
		}

		ui.Say("Waiting for machine to become available...")
		err = p.retryable(ctx, func(ctx context.Context) error {
			// wait for the machine to reboot.
			cmd := &packer.RemoteCmd{Command: testRestartCommand}
			err := cmd.RunWithUi(ctx, comm, ui)
			if err != nil {
				return err
			}
			exitStatus := cmd.ExitStatus()
			if exitStatus != 0 {
				return fmt.Errorf("Machine not yet available (exit status %d)", exitStatus)
			}
			cmd = &packer.RemoteCmd{Command: abortTestRestartCommand}
			err = cmd.RunWithUi(ctx, comm, ui)

			return err
		})
		if err != nil {
			return err
		}

		ui.Say("Checking for pending restart...")
		err = p.retryable(ctx, func(ctx context.Context) error {
			cmd := &packer.RemoteCmd{Command: pendingRebootElevatedCommand}
			err = cmd.RunWithUi(ctx, comm, ui)
			if err != nil {
				return err
			}

			exitStatus := cmd.ExitStatus()
			switch {
			case exitStatus == 0:
				restartPending = false
			case exitStatus == 101:
				restartPending = true
			default:
				return fmt.Errorf("Machine not yet available (exit status %d)", exitStatus)
			}

			return err
		})
		if err != nil {
			return err
		}

		if (restartPending) {
			ui.Say("Restart is still pending...")
		} else {
			ui.Say("Restart complete")
		}
	}

	return nil
}

// retryable will retry the given function over and over until a
// non-error is returned, RestartTimeout expires, or ctx is
// cancelled.
func (p *Provisioner) retryable(ctx context.Context, f func(ctx context.Context) error) error {
	return retry.Config{
		RetryDelay:   func() time.Duration { return retryableDelay },
		StartTimeout: p.config.RestartTimeout,
	}.Run(ctx, f)
}

func (p *Provisioner) windowsUpdateCommand() string {
	return fmt.Sprintf(
		"PowerShell -ExecutionPolicy Bypass -OutputFormat Text -EncodedCommand %s",
		base64.StdEncoding.EncodeToString(
			encodeUtf16Le(fmt.Sprintf(
				"%s%s%s -UpdateLimit %d",
				windowsUpdatePath,
				searchCriteriaArgument(p.config.SearchCriteria),
				filtersArgument(p.config.Filters),
				p.config.UpdateLimit))))
}

func (p *Provisioner) windowsUpdateCheckForRebootRequiredCommand() string {
	return fmt.Sprintf(
		"PowerShell -ExecutionPolicy Bypass -OutputFormat Text -EncodedCommand %s",
		base64.StdEncoding.EncodeToString(
			encodeUtf16Le(fmt.Sprintf(
				"%s -OnlyCheckForRebootRequired",
				windowsUpdatePath))))
}

func encodeUtf16Le(s string) []byte {
	d := utf16.Encode([]rune(s))
	b := make([]byte, len(d)*2)
	for i, r := range d {
		b[i*2] = byte(r)
		b[i*2+1] = byte(r >> 8)
	}
	return b
}

func searchCriteriaArgument(searchCriteria string) string {
	if searchCriteria == "" {
		return ""
	}

	var buffer bytes.Buffer

	buffer.WriteString(" -SearchCriteria ")
	buffer.WriteString(escapePowerShellString(searchCriteria))

	return buffer.String()
}

func filtersArgument(filters []string) string {
	if filters == nil {
		return ""
	}

	var buffer bytes.Buffer

	buffer.WriteString(" -Filters ")

	for i, filter := range filters {
		if i > 0 {
			buffer.WriteString(",")
		}
		buffer.WriteString(escapePowerShellString(filter))
	}

	return buffer.String()
}

func escapePowerShellString(value string) string {
	return fmt.Sprintf(
		"'%s'",
		// escape single quotes with another single quote.
		strings.Replace(value, "'", "''", -1))
}
<|endoftext|>


<commit_before><commit_msg>Implement moveFile to move the generated index<commit_after><|endoftext|>


<commit_before><commit_msg>cleanup spacing<commit_after><|endoftext|>


<commit_before>package server

import (
	"context"
	"expvar"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	afex "github.com/afex/hystrix-go/hystrix"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
	"github.com/justinas/alice"
	gsh "github.com/mchudgins/go-service-helper/handlers"
	"github.com/mchudgins/playground/pkg/healthz"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type Config struct {
	Insecure          bool
	Compress          bool // if true, add compression handling to messages
	UseZipkin         bool // if true, add zipkin tracing
	CertFilename      string
	KeyFilename       string
	HTTPListenPort    int
	RPCListenPort     int
	MetricsListenPort int
	Handler           http.Handler
	Hostname          string // if present, enforce canonical hostnames
	RPCRegister       RPCRegistration
	logger            *zap.Logger
	rpcServer         *grpc.Server
	httpServer        *http.Server
	metricsServer     *http.Server
	serviceName       string
}

type Option func(*Config) error

type RPCRegistration func(*grpc.Server) error

const (
	zipkinHTTPEndpoint = "http://localhost:9411/api/v1/spans"
)

func WithCanonicalHost(hostname string) Option {
	return func(cfg *Config) error {
		cfg.Hostname = hostname

		return nil
	}
}

func WithCertificate(certFilename, keyFilename string) Option {
	return func(cfg *Config) error {
		cfg.CertFilename = certFilename
		cfg.KeyFilename = keyFilename
		cfg.Insecure = false
		return nil
	}
}

func WithHTTPListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.HTTPListenPort = port
		return nil
	}
}

func WithHTTPServer(h http.Handler) Option {
	return func(cfg *Config) error {
		cfg.Handler = h
		return nil
	}
}

func WithLogger(l *zap.Logger) Option {
	return func(cfg *Config) error {
		cfg.logger = l
		return nil
	}
}

func WithMetricsListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.MetricsListenPort = port
		return nil
	}
}

func WithRPCListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.RPCListenPort = port
		return nil
	}
}

func WithRPCServer(fn RPCRegistration) Option {
	return func(cfg *Config) error {
		cfg.RPCRegister = fn

		return nil
	}
}

func WithServiceName(serviceName string) Option {
	return func(cfg *Config) error {
		cfg.serviceName = serviceName
		return nil
	}
}

func WithZipkinTracer() Option {
	return func(cfg *Config) error {
		cfg.UseZipkin = true
		return nil
	}
}

func Run(ctx context.Context, opts ...Option) {

	// default config
	cfg := &Config{
		Insecure:          true,
		HTTPListenPort:    8443,
		MetricsListenPort: 8080,
		RPCListenPort:     50050,
	}

	// process the Run() options
	for _, o := range opts {
		o(cfg)
	}

	// make a channel to listen on events,
	// then launch the servers.

	errc := make(chan eventSource)
	defer close(errc)

	// interrupt handler
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
		errc <- eventSource{
			source: interrupt,
			err:    fmt.Errorf("%s", <-c),
		}
	}()

	// gRPC server
	if cfg.RPCRegister != nil {
		go func() {
			rpcListenPort := ":" + strconv.Itoa(cfg.RPCListenPort)
			lis, err := net.Listen("tcp", rpcListenPort)
			if err != nil {
				errc <- eventSource{
					err:    err,
					source: rpcServer,
				}
				return
			}

			// configure the RPC server
			var grpcMiddleware grpc.ServerOption
			if cfg.UseZipkin {
				grpcMiddleware = grpc_middleware.WithUnaryServerChain(
					grpc_prometheus.UnaryServerInterceptor,
					otgrpc.OpenTracingServerInterceptor(opentracing.GlobalTracer(), otgrpc.LogPayloads()),
					grpcEndpointLog(cfg.logger, cfg.serviceName))
			} else {
				grpcMiddleware = grpc_middleware.WithUnaryServerChain(
					grpc_prometheus.UnaryServerInterceptor,
					grpcEndpointLog(cfg.logger, cfg.serviceName))
			}

			if cfg.Insecure {
				cfg.rpcServer = grpc.NewServer(
					grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
					grpcMiddleware)
			} else {
				tlsCreds, err := credentials.NewServerTLSFromFile(cfg.CertFilename, cfg.KeyFilename)
				if err != nil {
					cfg.logger.Fatal("Failed to generate grpc TLS credentials", zap.Error(err))
				}
				cfg.rpcServer = grpc.NewServer(
					grpc.Creds(tlsCreds),
					grpc.RPCCompressor(grpc.NewGZIPCompressor()),
					grpc.RPCDecompressor(grpc.NewGZIPDecompressor()),
					grpcMiddleware)
			}

			cfg.RPCRegister(cfg.rpcServer)

			// register w. prometheus
			grpc_prometheus.Register(cfg.rpcServer)
			grpc_prometheus.EnableHandlingTimeHistogram()

			// run the server & send an event upon termination
			errc <- eventSource{
				err:    cfg.rpcServer.Serve(lis),
				source: rpcServer,
			}
		}()
	}

	// http/https server
	if cfg.Handler != nil {
		go func() {
			rootMux := mux.NewRouter()

			hc, err := healthz.NewConfig()
			healthzHandler, err := healthz.Handler(hc)
			if err != nil {
				cfg.logger.Panic("Constructing healthz.Handler", zap.Error(err))
			}

			// TODO: move these three handlers to the metrics listener
			// set up handlers for THIS instance
			// (these are not expected to be proxied)
			rootMux.Handle("/debug/vars", expvar.Handler())
			rootMux.Handle("/healthz", healthzHandler)
			rootMux.Handle("/metrics", prometheus.Handler())

			rootMux.PathPrefix("/").Handler(cfg.Handler)

			var tracer func(http.Handler) http.Handler
			tracer = gsh.TracerFromHTTPRequest(gsh.NewTracer(cfg.serviceName), "proxy")

			chain := alice.New(tracer,
				gsh.HTTPMetricsCollector,
				gsh.HTTPLogrusLogger)

			if cfg.UseZipkin {
				var tracer func(http.Handler) http.Handler
				tracer = gsh.TracerFromHTTPRequest(gsh.NewTracer("commandName"), "proxy")
				chain.Append(tracer)
			}

			if len(cfg.Hostname) > 0 {
				canonical := handlers.CanonicalHost(cfg.Hostname, http.StatusPermanentRedirect)
				chain = chain.Append(canonical)
			}

			if cfg.Compress {
				chain = chain.Append(handlers.CompressHandler)
			}

			httpListenAddress := ":" + strconv.Itoa(cfg.HTTPListenPort)
			cfg.httpServer = &http.Server{
				Addr:              httpListenAddress,
				Handler:           chain.Then(rootMux),
				ReadTimeout:       time.Duration(5) * time.Second,
				ReadHeaderTimeout: time.Duration(2) * time.Second,
				TLSConfig:         tlsConfigFactory(),
			}

			if cfg.Insecure {
				err = cfg.httpServer.ListenAndServe()
			} else {
				err = cfg.httpServer.ListenAndServeTLS(cfg.CertFilename, cfg.KeyFilename)
			}

			errc <- eventSource{
				err:    err,
				source: httpServer,
			}
		}()
	}

	// start the hystrix stream provider
	go func() {
		hystrixStreamHandler := afex.NewStreamHandler()
		hystrixStreamHandler.Start()
		listenPort := ":" + strconv.Itoa(cfg.MetricsListenPort)
		cfg.metricsServer = &http.Server{
			Addr:    listenPort,
			Handler: hystrixStreamHandler,
		}

		errc <- eventSource{
			err:    cfg.metricsServer.ListenAndServe(),
			source: metricsServer,
		}
	}()

	cfg.logLaunch()

	// wait for somthin'
	rc := <-errc

	// somethin happened, now shut everything down gracefully, if possible
	cfg.performGracefulShutdown(ctx, rc)
}

func (cfg *Config) logLaunch() {
	serverList := make([]zapcore.Field, 0, 10)

	if cfg.RPCRegister != nil {
		serverList = append(serverList, zap.Int("gRPC port", cfg.RPCListenPort))
	}
	if cfg.Handler != nil {
		var key = "HTTPS port"
		if cfg.Insecure {
			key = "HTTP port"
		}
		serverList = append(serverList, zap.Int(key, cfg.HTTPListenPort))
	}
	serverList = append(serverList, zap.Int("metrics port", cfg.MetricsListenPort))

	if cfg.Insecure {
		cfg.logger.Info("Server listening insecurely on one or more ports", serverList...)
	} else {
		cfg.logger.Info("Server", serverList...)
	}
}
<commit_msg>Fix bad merge<commit_after>package server

import (
	"context"
	"expvar"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	afex "github.com/afex/hystrix-go/hystrix"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
	"github.com/justinas/alice"
	gsh "github.com/mchudgins/go-service-helper/handlers"
	"github.com/mchudgins/playground/pkg/healthz"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type Config struct {
	Insecure          bool
	Compress          bool // if true, add compression handling to messages
	UseZipkin         bool // if true, add zipkin tracing
	CertFilename      string
	KeyFilename       string
	HTTPListenPort    int
	RPCListenPort     int
	MetricsListenPort int
	Handler           http.Handler
	Hostname          string // if present, enforce canonical hostnames
	RPCRegister       RPCRegistration
	logger            *zap.Logger
	rpcServer         *grpc.Server
	httpServer        *http.Server
	metricsServer     *http.Server
	serviceName       string
}

type Option func(*Config) error

type RPCRegistration func(*grpc.Server) error

const (
	zipkinHTTPEndpoint = "http://localhost:9411/api/v1/spans"
)

func WithCanonicalHost(hostname string) Option {
	return func(cfg *Config) error {
		cfg.Hostname = hostname

		return nil
	}
}

func WithCertificate(certFilename, keyFilename string) Option {
	return func(cfg *Config) error {
		cfg.CertFilename = certFilename
		cfg.KeyFilename = keyFilename
		cfg.Insecure = false
		return nil
	}
}

func WithHTTPListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.HTTPListenPort = port
		return nil
	}
}

func WithHTTPServer(h http.Handler) Option {
	return func(cfg *Config) error {
		cfg.Handler = h
		return nil
	}
}

func WithLogger(l *zap.Logger) Option {
	return func(cfg *Config) error {
		cfg.logger = l
		return nil
	}
}

func WithMetricsListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.MetricsListenPort = port
		return nil
	}
}

func WithRPCListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.RPCListenPort = port
		return nil
	}
}

func WithRPCServer(fn RPCRegistration) Option {
	return func(cfg *Config) error {
		cfg.RPCRegister = fn

		return nil
	}
}

func WithServiceName(serviceName string) Option {
	return func(cfg *Config) error {
		cfg.serviceName = serviceName
		return nil
	}
}

func WithZipkinTracer() Option {
	return func(cfg *Config) error {
		cfg.UseZipkin = true
		return nil
	}
}

func Run(ctx context.Context, opts ...Option) {

	// default config
	cfg := &Config{
		Insecure:          true,
		HTTPListenPort:    8443,
		MetricsListenPort: 8080,
		RPCListenPort:     50050,
	}

	// process the Run() options
	for _, o := range opts {
		o(cfg)
	}

	// make a channel to listen on events,
	// then launch the servers.

	errc := make(chan eventSource)
	defer close(errc)

	// interrupt handler
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
		errc <- eventSource{
			source: interrupt,
			err:    fmt.Errorf("%s", <-c),
		}
	}()

	// gRPC server
	if cfg.RPCRegister != nil {
		go func() {
			rpcListenPort := ":" + strconv.Itoa(cfg.RPCListenPort)
			lis, err := net.Listen("tcp", rpcListenPort)
			if err != nil {
				errc <- eventSource{
					err:    err,
					source: rpcServer,
				}
				return
			}

			// configure the RPC server
			var grpcMiddleware grpc.ServerOption
			if cfg.UseZipkin {
				grpcMiddleware = grpc_middleware.WithUnaryServerChain(
					grpc_prometheus.UnaryServerInterceptor,
					otgrpc.OpenTracingServerInterceptor(opentracing.GlobalTracer(), otgrpc.LogPayloads()),
					grpcEndpointLog(cfg.logger, cfg.serviceName))
			} else {
				grpcMiddleware = grpc_middleware.WithUnaryServerChain(
					grpc_prometheus.UnaryServerInterceptor,
					grpcEndpointLog(cfg.logger, cfg.serviceName))
			}

			if cfg.Insecure {
				cfg.rpcServer = grpc.NewServer(
					grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
					grpcMiddleware)
			} else {
				tlsCreds, err := credentials.NewServerTLSFromFile(cfg.CertFilename, cfg.KeyFilename)
				if err != nil {
					cfg.logger.Fatal("Failed to generate grpc TLS credentials", zap.Error(err))
				}
				cfg.rpcServer = grpc.NewServer(
					grpc.Creds(tlsCreds),
					grpc.RPCCompressor(grpc.NewGZIPCompressor()),
					grpc.RPCDecompressor(grpc.NewGZIPDecompressor()),
					grpcMiddleware)
			}

			cfg.RPCRegister(cfg.rpcServer)

			// register w. prometheus
			grpc_prometheus.Register(cfg.rpcServer)
			grpc_prometheus.EnableHandlingTimeHistogram()

			// run the server & send an event upon termination
			errc <- eventSource{
				err:    cfg.rpcServer.Serve(lis),
				source: rpcServer,
			}
		}()
	}

	// http/https server
	if cfg.Handler != nil {
		go func() {
			rootMux := mux.NewRouter()

			hc, err := healthz.NewConfig()
			healthzHandler, err := healthz.Handler(hc)
			if err != nil {
				cfg.logger.Panic("Constructing healthz.Handler", zap.Error(err))
			}

			// TODO: move these three handlers to the metrics listener
			// set up handlers for THIS instance
			// (these are not expected to be proxied)
			rootMux.Handle("/debug/vars", expvar.Handler())
			rootMux.Handle("/healthz", healthzHandler)
			rootMux.Handle("/metrics", prometheus.Handler())

			rootMux.PathPrefix("/").Handler(cfg.Handler)

			chain := alice.New(gsh.HTTPMetricsCollector, gsh.HTTPLogrusLogger)

			if cfg.UseZipkin {
				var tracer func(http.Handler) http.Handler
				tracer = gsh.TracerFromHTTPRequest(gsh.NewTracer("commandName"), "proxy")
				chain.Append(tracer)
			}

			if len(cfg.Hostname) > 0 {
				canonical := handlers.CanonicalHost(cfg.Hostname, http.StatusPermanentRedirect)
				chain = chain.Append(canonical)
			}

			if cfg.Compress {
				chain = chain.Append(handlers.CompressHandler)
			}

			httpListenAddress := ":" + strconv.Itoa(cfg.HTTPListenPort)
			cfg.httpServer = &http.Server{
				Addr:              httpListenAddress,
				Handler:           chain.Then(rootMux),
				ReadTimeout:       time.Duration(5) * time.Second,
				ReadHeaderTimeout: time.Duration(2) * time.Second,
				TLSConfig:         tlsConfigFactory(),
			}

			if cfg.Insecure {
				err = cfg.httpServer.ListenAndServe()
			} else {
				err = cfg.httpServer.ListenAndServeTLS(cfg.CertFilename, cfg.KeyFilename)
			}

			errc <- eventSource{
				err:    err,
				source: httpServer,
			}
		}()
	}

	// start the hystrix stream provider
	go func() {
		hystrixStreamHandler := afex.NewStreamHandler()
		hystrixStreamHandler.Start()
		listenPort := ":" + strconv.Itoa(cfg.MetricsListenPort)
		cfg.metricsServer = &http.Server{
			Addr:    listenPort,
			Handler: hystrixStreamHandler,
		}

		errc <- eventSource{
			err:    cfg.metricsServer.ListenAndServe(),
			source: metricsServer,
		}
	}()

	cfg.logLaunch()

	// wait for somthin'
	rc := <-errc

	// somethin happened, now shut everything down gracefully, if possible
	cfg.performGracefulShutdown(ctx, rc)
}

func (cfg *Config) logLaunch() {
	serverList := make([]zapcore.Field, 0, 10)

	if cfg.RPCRegister != nil {
		serverList = append(serverList, zap.Int("gRPC port", cfg.RPCListenPort))
	}
	if cfg.Handler != nil {
		var key = "HTTPS port"
		if cfg.Insecure {
			key = "HTTP port"
		}
		serverList = append(serverList, zap.Int(key, cfg.HTTPListenPort))
	}
	serverList = append(serverList, zap.Int("metrics port", cfg.MetricsListenPort))

	if cfg.Insecure {
		cfg.logger.Info("Server listening insecurely on one or more ports", serverList...)
	} else {
		cfg.logger.Info("Server", serverList...)
	}
}
<|endoftext|>


<commit_before>package main

func main() {}

/*
import (
	"container/heap"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/boltdb/bolt"
	"github.com/gojp/goreportcard/check"
	"github.com/gojp/goreportcard/handlers"

	"gopkg.in/mgo.v2"
)

const (
	dbPath     string = "goreportcard.db"
	repoBucket string = "repos"
	metaBucket string = "meta"

	mongoURL        = "mongodb://127.0.0.1:27017"
	mongoDatabase   = "goreportcard"
	mongoCollection = "reports"
)

type Grade string

type score struct {
	Name          string              `json:"name"`
	Description   string              `json:"description"`
	FileSummaries []check.FileSummary `json:"file_summaries"`
	Weight        float64             `json:"weight"`
	Percentage    float64             `json:"percentage"`
}

type checksResp struct {
	Checks      []score   `json:"checks"`
	Average     float64   `json:"average"`
	Grade       Grade     `json:"grade"`
	Files       int       `json:"files"`
	Issues      int       `json:"issues"`
	Repo        string    `json:"repo"`
	LastRefresh time.Time `json:"last_refresh"`
}

// initDB opens the bolt database file (or creates it if it does not exist), and creates
// a bucket for saving the repos, also only if it does not exist.
func initDB() error {
	db, err := bolt.Open(handlers.DBPath, 0600, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return err
	}
	defer db.Close()

	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(repoBucket))
		if err != nil {
			return err
		}
		_, err = tx.CreateBucketIfNotExists([]byte(metaBucket))
		return err
	})
	return err
}

func main() {
	// initialize bolt database
	if err := initDB(); err != nil {
		log.Fatal("ERROR: could not open bolt db: ", err)
	}

	session, err := mgo.Dial(mongoURL)
	if err != nil {
		log.Fatal("ERROR: could not get collection:", err)
	}
	defer session.Close()
	coll := session.DB(mongoDatabase).C(mongoCollection)

	db, err := bolt.Open(dbPath, 0755, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		log.Println("Failed to open bolt database: ", err)
		return
	}
	defer db.Close()

	var repos []checksResp
	coll.Find(nil).All(&repos)

	for _, repo := range repos {
		fmt.Printf("inserting %q into bolt...\n", repo.Repo)
		err = db.Update(func(tx *bolt.Tx) error {
			bkt := tx.Bucket([]byte(repoBucket))
			if bkt == nil {
				return fmt.Errorf("repo bucket not found")
			}
			b, err := json.Marshal(repo)
			if err != nil {
				return err
			}

			mb := tx.Bucket([]byte(metaBucket))
			if mb == nil {
				return fmt.Errorf("repo bucket not found")
			}
			updateHighScores(mb, repo, repo.Repo)

			return bkt.Put([]byte(repo.Repo), b)
		})
		if err != nil {
			log.Println("Bolt writing error:", err)
		}
	}

	err = db.Update(func(tx *bolt.Tx) error {
		mb := tx.Bucket([]byte(metaBucket))
		if mb == nil {
			return fmt.Errorf("repo bucket not found")
		}
		totalInt := len(repos)
		total, err := json.Marshal(totalInt)
		if err != nil {
			return fmt.Errorf("could not marshal total repos count: %v", err)
		}
		return mb.Put([]byte("total_repos"), total)
	})
	if err != nil {
		log.Fatal(err)
	}
}

func updateHighScores(mb *bolt.Bucket, resp checksResp, repo string) error {
	// check if we need to update the high score list
	if resp.Files < 100 {
		// only repos with >= 100 files are considered for the high score list
		return nil
	}

	// start updating high score list
	scoreBytes := mb.Get([]byte("scores"))
	if scoreBytes == nil {
		scoreBytes, _ = json.Marshal([]scoreHeap{})
	}
	scores := &scoreHeap{}
	json.Unmarshal(scoreBytes, scores)

	heap.Init(scores)
	if len(*scores) > 0 && (*scores)[0].Score > resp.Average*100.0 && len(*scores) == 50 {
		// lowest score on list is higher than this repo's score, so no need to add, unless
		// we do not have 50 high scores yet
		return nil
	}
	// if this repo is already in the list, remove the original entry:
	for i := range *scores {
		if (*scores)[i].Repo == repo {
			heap.Remove(scores, i)
			break
		}
	}
	// now we can safely push it onto the heap
	heap.Push(scores, scoreItem{
		Repo:  repo,
		Score: resp.Average * 100.0,
		Files: resp.Files,
	})
	if len(*scores) > 50 {
		// trim heap if it's grown to over 50
		*scores = (*scores)[1:51]
	}
	scoreBytes, err := json.Marshal(&scores)
	if err != nil {
		return err
	}
	err = mb.Put([]byte("scores"), scoreBytes)
	if err != nil {
		return err
	}

	return nil
}

type scoreItem struct {
	Repo  string  `json:"repo"`
	Score float64 `json:"score"`
	Files int     `json:"files"`
}

// An scoreHeap is a min-heap of ints.
type scoreHeap []scoreItem

func (h scoreHeap) Len() int           { return len(h) }
func (h scoreHeap) Less(i, j int) bool { return h[i].Score < h[j].Score }
func (h scoreHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *scoreHeap) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(scoreItem))
}

func (h *scoreHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}
*/
<commit_msg>remove tools/mongotobolt.go<commit_after><|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"log"
	"net"
)

type TflowHeader struct {
	Version          uint16
	FlowRecords      uint16
	Uptime           uint32
	UnixSec          uint32
	UnixNsec         uint32
	FlowSeqNum       uint32
	EngineType       uint8
	EngineId         uint8
	SamplingInterval uint16
}

type TflowRecord struct {
	Ipv4SrcAddr   uint32
	Ipv4DstAddr   uint32
	Ipv4NextHop   uint32
	InputSnmp     uint16
	OutputSnmp    uint16
	InPkts        uint32
	InBytes       uint32
	FirstSwitched uint32
	LastSwitched  uint32
	L4SrcPort     uint16
	L4DstPort     uint16
	_             uint8
	TcpFlags      uint8
	Protocol      uint8
	SrcTos        uint8
	SrcAs         uint16
	DstAs         uint16
	SrcMask       uint8
	DstMask       uint8
	_             uint16
}

func intToNetIp(intAddr uint32) net.IP {
	return net.IPv4(
		byte(intAddr>>24),
		byte(intAddr>>16),
		byte(intAddr>>8),
		byte(intAddr))
}

func netIpToInt(ip net.IP) uint32 {
	p := ip.To4()
	vip := uint32(p[0]) << 24
	vip |= uint32(p[1]) << 16
	vip |= uint32(p[2]) << 8
	vip |= uint32(p[3])
	return vip
}

func intToStrIP(intAddr uint32) string {
	return intToNetIp(intAddr).String()
}

func handleNetFlowPacket(buf *bytes.Buffer, remoteAddr *net.UDPAddr) {
	var id uint32

	header := TflowHeader{}
	err := binary.Read(buf, binary.BigEndian, &header)
	if err != nil {
		log.Fatalf("Error:", err)
	}

	for i := 0; i < int(header.FlowRecords); i++ {
		record := TflowRecord{}
		err := binary.Read(buf, binary.BigEndian, &record)
		if err != nil {
			log.Fatalf("binary.Read failed: %v\n", err)
		}

		if (record.Ipv4SrcAddr >= minCheckIP) && (record.Ipv4SrcAddr <= maxCheckIP) {
			id = record.Ipv4SrcAddr
		} else {
			if (record.Ipv4DstAddr >= minCheckIP) && (record.Ipv4DstAddr <= maxCheckIP) {
				id = record.Ipv4DstAddr
			}
		}
		// record.Ipv4SrcAddrInt
		// record.Ipv4DstAddrInt
		// record.InBytes
		if id > 0 {
			//log.Printf("Start web listening on %v\n", SccCfg.WebPort)
			//buf, err := json.Marshal(record)
			//if err != nil {
			//	log.Fatalf("json.Marshal failed: %v\n", err)
			//}

			//fmt.Printf("%v\n", string(buf))
			AddTraffic(id, TIpTraffic(record.InBytes))
		}

	}
}

func ListenNetflow(inSource string, receiveBufferSizeBytes int) {
	/* Start listerning on the specified port */
	// log.Printf("Start netflow listening on %v\n", inSource)
	addr, err := net.ResolveUDPAddr("udp", inSource)
	if err != nil {
		log.Fatalf("Error: %v\n", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatalln(err)
	}

	err = conn.SetReadBuffer(receiveBufferSizeBytes)
	if err != nil {
		log.Fatalln(err)
	}

	defer conn.Close()

	for {
		buf := make([]byte, 4096)
		rlen, remote, err := conn.ReadFromUDP(buf)

		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		stream := bytes.NewBuffer(buf[:rlen])

		go handleNetFlowPacket(stream, remote)
	}
}
<commit_msg>- возможна ошибка если много записей в одном пакете<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"log"
	"net"
)

type TflowHeader struct {
	Version          uint16
	FlowRecords      uint16
	Uptime           uint32
	UnixSec          uint32
	UnixNsec         uint32
	FlowSeqNum       uint32
	EngineType       uint8
	EngineId         uint8
	SamplingInterval uint16
}

type TflowRecord struct {
	Ipv4SrcAddr   uint32
	Ipv4DstAddr   uint32
	Ipv4NextHop   uint32
	InputSnmp     uint16
	OutputSnmp    uint16
	InPkts        uint32
	InBytes       uint32
	FirstSwitched uint32
	LastSwitched  uint32
	L4SrcPort     uint16
	L4DstPort     uint16
	_             uint8
	TcpFlags      uint8
	Protocol      uint8
	SrcTos        uint8
	SrcAs         uint16
	DstAs         uint16
	SrcMask       uint8
	DstMask       uint8
	_             uint16
}

func intToNetIp(intAddr uint32) net.IP {
	return net.IPv4(
		byte(intAddr>>24),
		byte(intAddr>>16),
		byte(intAddr>>8),
		byte(intAddr))
}

func netIpToInt(ip net.IP) uint32 {
	p := ip.To4()
	vip := uint32(p[0]) << 24
	vip |= uint32(p[1]) << 16
	vip |= uint32(p[2]) << 8
	vip |= uint32(p[3])
	return vip
}

func intToStrIP(intAddr uint32) string {
	return intToNetIp(intAddr).String()
}

func handleNetFlowPacket(buf *bytes.Buffer, remoteAddr *net.UDPAddr) {
	var id uint32

	header := TflowHeader{}
	err := binary.Read(buf, binary.BigEndian, &header)
	if err != nil {
		log.Fatalf("Error:", err)
	}

	for i := 0; i < int(header.FlowRecords); i++ {
		record := TflowRecord{}
		id = 0
		err := binary.Read(buf, binary.BigEndian, &record)
		if err != nil {
			log.Fatalf("binary.Read failed: %v\n", err)
		}

		if (record.Ipv4SrcAddr >= minCheckIP) && (record.Ipv4SrcAddr <= maxCheckIP) {
			id = record.Ipv4SrcAddr
		} else {
			if (record.Ipv4DstAddr >= minCheckIP) && (record.Ipv4DstAddr <= maxCheckIP) {
				id = record.Ipv4DstAddr
			}
		}

		if id > 0 {
			AddTraffic(id, TIpTraffic(record.InBytes))
		}
	}
}

func ListenNetflow(inSource string, receiveBufferSizeBytes int) {
	/* Start listerning on the specified port */
	// log.Printf("Start netflow listening on %v\n", inSource)
	addr, err := net.ResolveUDPAddr("udp", inSource)
	if err != nil {
		log.Fatalf("Error: %v\n", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatalln(err)
	}

	err = conn.SetReadBuffer(receiveBufferSizeBytes)
	if err != nil {
		log.Fatalln(err)
	}

	defer conn.Close()

	for {
		buf := make([]byte, 4096)
		rlen, remote, err := conn.ReadFromUDP(buf)

		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		stream := bytes.NewBuffer(buf[:rlen])

		go handleNetFlowPacket(stream, remote)
	}
}
<|endoftext|>


<commit_before>package docker

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"net"
	"os/exec"
	"strconv"
	"strings"
)

const (
	networkBridgeIface = "lxcbr0"
	portRangeStart     = 49153
	portRangeEnd       = 65535
)

// Calculates the first and last IP addresses in an IPNet
func networkRange(network *net.IPNet) (net.IP, net.IP) {
	netIP := network.IP.To4()
	firstIP := netIP.Mask(network.Mask)
	lastIP := net.IPv4(0, 0, 0, 0).To4()
	for i := 0; i < len(lastIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return firstIP, lastIP
}

// Converts a 4 bytes IP into a 32 bit integer
func ipToInt(ip net.IP) (int32, error) {
	buf := bytes.NewBuffer(ip.To4())
	var n int32
	if err := binary.Read(buf, binary.BigEndian, &n); err != nil {
		return 0, err
	}
	return n, nil
}

// Converts 32 bit integer into a 4 bytes IP address
func intToIp(n int32) (net.IP, error) {
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, &n); err != nil {
		return net.IP{}, err
	}
	ip := net.IPv4(0, 0, 0, 0).To4()
	for i := 0; i < net.IPv4len; i++ {
		ip[i] = buf.Bytes()[i]
	}
	return ip, nil
}

// Given a netmask, calculates the number of available hosts
func networkSize(mask net.IPMask) (int32, error) {
	m := net.IPv4Mask(0, 0, 0, 0)
	for i := 0; i < net.IPv4len; i++ {
		m[i] = ^mask[i]
	}
	buf := bytes.NewBuffer(m)
	var n int32
	if err := binary.Read(buf, binary.BigEndian, &n); err != nil {
		return 0, err
	}
	return n + 1, nil
}

// Wrapper around the iptables command
func iptables(args ...string) error {
	if err := exec.Command("/sbin/iptables", args...).Run(); err != nil {
		return fmt.Errorf("iptables failed: iptables %v", strings.Join(args, " "))
	}
	return nil
}

// Return the IPv4 address of a network interface
func getIfaceAddr(name string) (net.Addr, error) {
	iface, err := net.InterfaceByName(name)
	if err != nil {
		return nil, err
	}
	addrs, err := iface.Addrs()
	if err != nil {
		return nil, err
	}
	var addrs4 []net.Addr
	for _, addr := range addrs {
		ip := (addr.(*net.IPNet)).IP
		if ip4 := ip.To4(); len(ip4) == net.IPv4len {
			addrs4 = append(addrs4, addr)
		}
	}
	switch {
	case len(addrs4) == 0:
		return nil, fmt.Errorf("Interface %v has no IP addresses", name)
	case len(addrs4) > 1:
		fmt.Printf("Interface %v has more than 1 IPv4 address. Defaulting to using %v\n",
			name, (addrs4[0].(*net.IPNet)).IP)
	}
	return addrs4[0], nil
}

// Port mapper takes care of mapping external ports to containers by setting
// up iptables rules.
// It keeps track of all mappings and is able to unmap at will
type PortMapper struct {
	mapping map[int]net.TCPAddr
}

func (mapper *PortMapper) cleanup() error {
	// Ignore errors - This could mean the chains were never set up
	iptables("-t", "nat", "-D", "PREROUTING", "-j", "DOCKER")
	iptables("-t", "nat", "-F", "DOCKER")
	iptables("-t", "nat", "-X", "DOCKER")
	mapper.mapping = make(map[int]net.TCPAddr)
	return nil
}

func (mapper *PortMapper) setup() error {
	if err := iptables("-t", "nat", "-N", "DOCKER"); err != nil {
		return errors.New("Unable to setup port networking: Failed to create DOCKER chain")
	}
	if err := iptables("-t", "nat", "-A", "PREROUTING", "-j", "DOCKER"); err != nil {
		return errors.New("Unable to setup port networking: Failed to inject docker in PREROUTING chain")
	}
	return nil
}

func (mapper *PortMapper) iptablesForward(rule string, port int, dest net.TCPAddr) error {
	return iptables("-t", "nat", rule, "DOCKER", "-p", "tcp", "--dport", strconv.Itoa(port),
		"-j", "DNAT", "--to-destination", net.JoinHostPort(dest.IP.String(), strconv.Itoa(dest.Port)))
}

func (mapper *PortMapper) Map(port int, dest net.TCPAddr) error {
	if err := mapper.iptablesForward("-A", port, dest); err != nil {
		return err
	}
	mapper.mapping[port] = dest
	return nil
}

func (mapper *PortMapper) Unmap(port int) error {
	dest, ok := mapper.mapping[port]
	if !ok {
		return errors.New("Port is not mapped")
	}
	if err := mapper.iptablesForward("-D", port, dest); err != nil {
		return err
	}
	delete(mapper.mapping, port)
	return nil
}

func newPortMapper() (*PortMapper, error) {
	mapper := &PortMapper{}
	if err := mapper.cleanup(); err != nil {
		return nil, err
	}
	if err := mapper.setup(); err != nil {
		return nil, err
	}
	return mapper, nil
}

// Port allocator: Atomatically allocate and release networking ports
type PortAllocator struct {
	ports chan (int)
}

func (alloc *PortAllocator) populate(start, end int) {
	alloc.ports = make(chan int, end-start)
	for port := start; port < end; port++ {
		alloc.ports <- port
	}
}

func (alloc *PortAllocator) Acquire() (int, error) {
	select {
	case port := <-alloc.ports:
		return port, nil
	default:
		return -1, errors.New("No more ports available")
	}
	return -1, nil
}

func (alloc *PortAllocator) Release(port int) error {
	select {
	case alloc.ports <- port:
		return nil
	default:
		return errors.New("Too many ports have been released")
	}
	return nil
}

func newPortAllocator(start, end int) (*PortAllocator, error) {
	allocator := &PortAllocator{}
	allocator.populate(start, end)
	return allocator, nil
}

// IP allocator: Atomatically allocate and release networking ports
type IPAllocator struct {
	network *net.IPNet
	queue   chan (net.IP)
}

func (alloc *IPAllocator) populate() error {
	firstIP, _ := networkRange(alloc.network)
	size, err := networkSize(alloc.network.Mask)
	if err != nil {
		return err
	}
	// The queue size should be the network size - 3
	// -1 for the network address, -1 for the broadcast address and
	// -1 for the gateway address
	alloc.queue = make(chan net.IP, size-3)
	for i := int32(1); i < size-1; i++ {
		ipNum, err := ipToInt(firstIP)
		if err != nil {
			return err
		}
		ip, err := intToIp(ipNum + int32(i))
		if err != nil {
			return err
		}
		// Discard the network IP (that's the host IP address)
		if ip.Equal(alloc.network.IP) {
			continue
		}
		alloc.queue <- ip
	}
	return nil
}

func (alloc *IPAllocator) Acquire() (net.IP, error) {
	select {
	case ip := <-alloc.queue:
		return ip, nil
	default:
		return net.IP{}, errors.New("No more IP addresses available")
	}
	return net.IP{}, nil
}

func (alloc *IPAllocator) Release(ip net.IP) error {
	select {
	case alloc.queue <- ip:
		return nil
	default:
		return errors.New("Too many IP addresses have been released")
	}
	return nil
}

func newIPAllocator(network *net.IPNet) (*IPAllocator, error) {
	alloc := &IPAllocator{
		network: network,
	}
	if err := alloc.populate(); err != nil {
		return nil, err
	}
	return alloc, nil
}

// Network interface represents the networking stack of a container
type NetworkInterface struct {
	IPNet   net.IPNet
	Gateway net.IP

	manager  *NetworkManager
	extPorts []int
}

// Allocate an external TCP port and map it to the interface
func (iface *NetworkInterface) AllocatePort(port int) (int, error) {
	extPort, err := iface.manager.portAllocator.Acquire()
	if err != nil {
		return -1, err
	}
	if err := iface.manager.portMapper.Map(extPort, net.TCPAddr{IP: iface.IPNet.IP, Port: port}); err != nil {
		iface.manager.portAllocator.Release(extPort)
		return -1, err
	}
	iface.extPorts = append(iface.extPorts, extPort)
	return extPort, nil
}

// Release: Network cleanup - release all resources
func (iface *NetworkInterface) Release() error {
	for _, port := range iface.extPorts {
		if err := iface.manager.portMapper.Unmap(port); err != nil {
			log.Printf("Unable to unmap port %v: %v", port, err)
		}
		if err := iface.manager.portAllocator.Release(port); err != nil {
			log.Printf("Unable to release port %v: %v", port, err)
		}

	}
	return iface.manager.ipAllocator.Release(iface.IPNet.IP)
}

// Network Manager manages a set of network interfaces
// Only *one* manager per host machine should be used
type NetworkManager struct {
	bridgeIface   string
	bridgeNetwork *net.IPNet

	ipAllocator   *IPAllocator
	portAllocator *PortAllocator
	portMapper    *PortMapper
}

// Allocate a network interface
func (manager *NetworkManager) Allocate() (*NetworkInterface, error) {
	ip, err := manager.ipAllocator.Acquire()
	if err != nil {
		return nil, err
	}
	iface := &NetworkInterface{
		IPNet:   net.IPNet{IP: ip, Mask: manager.bridgeNetwork.Mask},
		Gateway: manager.bridgeNetwork.IP,
		manager: manager,
	}
	return iface, nil
}

func newNetworkManager(bridgeIface string) (*NetworkManager, error) {
	addr, err := getIfaceAddr(bridgeIface)
	if err != nil {
		return nil, err
	}
	network := addr.(*net.IPNet)

	ipAllocator, err := newIPAllocator(network)
	if err != nil {
		return nil, err
	}

	portAllocator, err := newPortAllocator(portRangeStart, portRangeEnd)
	if err != nil {
		return nil, err
	}

	portMapper, err := newPortMapper()

	manager := &NetworkManager{
		bridgeIface:   bridgeIface,
		bridgeNetwork: network,
		ipAllocator:   ipAllocator,
		portAllocator: portAllocator,
		portMapper:    portMapper,
	}
	return manager, nil
}
<commit_msg>Fixing Issue #98: Adding DOCKER to output chain during iptables setup<commit_after>package docker

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"net"
	"os/exec"
	"strconv"
	"strings"
)

const (
	networkBridgeIface = "lxcbr0"
	portRangeStart     = 49153
	portRangeEnd       = 65535
)

// Calculates the first and last IP addresses in an IPNet
func networkRange(network *net.IPNet) (net.IP, net.IP) {
	netIP := network.IP.To4()
	firstIP := netIP.Mask(network.Mask)
	lastIP := net.IPv4(0, 0, 0, 0).To4()
	for i := 0; i < len(lastIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return firstIP, lastIP
}

// Converts a 4 bytes IP into a 32 bit integer
func ipToInt(ip net.IP) (int32, error) {
	buf := bytes.NewBuffer(ip.To4())
	var n int32
	if err := binary.Read(buf, binary.BigEndian, &n); err != nil {
		return 0, err
	}
	return n, nil
}

// Converts 32 bit integer into a 4 bytes IP address
func intToIp(n int32) (net.IP, error) {
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, &n); err != nil {
		return net.IP{}, err
	}
	ip := net.IPv4(0, 0, 0, 0).To4()
	for i := 0; i < net.IPv4len; i++ {
		ip[i] = buf.Bytes()[i]
	}
	return ip, nil
}

// Given a netmask, calculates the number of available hosts
func networkSize(mask net.IPMask) (int32, error) {
	m := net.IPv4Mask(0, 0, 0, 0)
	for i := 0; i < net.IPv4len; i++ {
		m[i] = ^mask[i]
	}
	buf := bytes.NewBuffer(m)
	var n int32
	if err := binary.Read(buf, binary.BigEndian, &n); err != nil {
		return 0, err
	}
	return n + 1, nil
}

// Wrapper around the iptables command
func iptables(args ...string) error {
	if err := exec.Command("/sbin/iptables", args...).Run(); err != nil {
		return fmt.Errorf("iptables failed: iptables %v", strings.Join(args, " "))
	}
	return nil
}

// Return the IPv4 address of a network interface
func getIfaceAddr(name string) (net.Addr, error) {
	iface, err := net.InterfaceByName(name)
	if err != nil {
		return nil, err
	}
	addrs, err := iface.Addrs()
	if err != nil {
		return nil, err
	}
	var addrs4 []net.Addr
	for _, addr := range addrs {
		ip := (addr.(*net.IPNet)).IP
		if ip4 := ip.To4(); len(ip4) == net.IPv4len {
			addrs4 = append(addrs4, addr)
		}
	}
	switch {
	case len(addrs4) == 0:
		return nil, fmt.Errorf("Interface %v has no IP addresses", name)
	case len(addrs4) > 1:
		fmt.Printf("Interface %v has more than 1 IPv4 address. Defaulting to using %v\n",
			name, (addrs4[0].(*net.IPNet)).IP)
	}
	return addrs4[0], nil
}

// Port mapper takes care of mapping external ports to containers by setting
// up iptables rules.
// It keeps track of all mappings and is able to unmap at will
type PortMapper struct {
	mapping map[int]net.TCPAddr
}

func (mapper *PortMapper) cleanup() error {
	// Ignore errors - This could mean the chains were never set up
	iptables("-t", "nat", "-D", "PREROUTING", "-j", "DOCKER")
	iptables("-t", "nat", "-F", "DOCKER")
	iptables("-t", "nat", "-X", "DOCKER")
	mapper.mapping = make(map[int]net.TCPAddr)
	return nil
}

func (mapper *PortMapper) setup() error {
	if err := iptables("-t", "nat", "-N", "DOCKER"); err != nil {
		return errors.New("Unable to setup port networking: Failed to create DOCKER chain")
	}
	if err := iptables("-t", "nat", "-A", "PREROUTING", "-j", "DOCKER"); err != nil {
		return errors.New("Unable to setup port networking: Failed to inject docker in PREROUTING chain")
	}
	if err := iptables("-t", "nat", "-A", "OUTPUT", "-j", "DOCKER"); err != nil {
		return errors.New("Unable to setup port networking: Failed to inject docker in OUTPUT chain")
	}
	return nil
}

func (mapper *PortMapper) iptablesForward(rule string, port int, dest net.TCPAddr) error {
	return iptables("-t", "nat", rule, "DOCKER", "-p", "tcp", "--dport", strconv.Itoa(port),
		"-j", "DNAT", "--to-destination", net.JoinHostPort(dest.IP.String(), strconv.Itoa(dest.Port)))
}

func (mapper *PortMapper) Map(port int, dest net.TCPAddr) error {
	if err := mapper.iptablesForward("-A", port, dest); err != nil {
		return err
	}
	mapper.mapping[port] = dest
	return nil
}

func (mapper *PortMapper) Unmap(port int) error {
	dest, ok := mapper.mapping[port]
	if !ok {
		return errors.New("Port is not mapped")
	}
	if err := mapper.iptablesForward("-D", port, dest); err != nil {
		return err
	}
	delete(mapper.mapping, port)
	return nil
}

func newPortMapper() (*PortMapper, error) {
	mapper := &PortMapper{}
	if err := mapper.cleanup(); err != nil {
		return nil, err
	}
	if err := mapper.setup(); err != nil {
		return nil, err
	}
	return mapper, nil
}

// Port allocator: Atomatically allocate and release networking ports
type PortAllocator struct {
	ports chan (int)
}

func (alloc *PortAllocator) populate(start, end int) {
	alloc.ports = make(chan int, end-start)
	for port := start; port < end; port++ {
		alloc.ports <- port
	}
}

func (alloc *PortAllocator) Acquire() (int, error) {
	select {
	case port := <-alloc.ports:
		return port, nil
	default:
		return -1, errors.New("No more ports available")
	}
	return -1, nil
}

func (alloc *PortAllocator) Release(port int) error {
	select {
	case alloc.ports <- port:
		return nil
	default:
		return errors.New("Too many ports have been released")
	}
	return nil
}

func newPortAllocator(start, end int) (*PortAllocator, error) {
	allocator := &PortAllocator{}
	allocator.populate(start, end)
	return allocator, nil
}

// IP allocator: Atomatically allocate and release networking ports
type IPAllocator struct {
	network *net.IPNet
	queue   chan (net.IP)
}

func (alloc *IPAllocator) populate() error {
	firstIP, _ := networkRange(alloc.network)
	size, err := networkSize(alloc.network.Mask)
	if err != nil {
		return err
	}
	// The queue size should be the network size - 3
	// -1 for the network address, -1 for the broadcast address and
	// -1 for the gateway address
	alloc.queue = make(chan net.IP, size-3)
	for i := int32(1); i < size-1; i++ {
		ipNum, err := ipToInt(firstIP)
		if err != nil {
			return err
		}
		ip, err := intToIp(ipNum + int32(i))
		if err != nil {
			return err
		}
		// Discard the network IP (that's the host IP address)
		if ip.Equal(alloc.network.IP) {
			continue
		}
		alloc.queue <- ip
	}
	return nil
}

func (alloc *IPAllocator) Acquire() (net.IP, error) {
	select {
	case ip := <-alloc.queue:
		return ip, nil
	default:
		return net.IP{}, errors.New("No more IP addresses available")
	}
	return net.IP{}, nil
}

func (alloc *IPAllocator) Release(ip net.IP) error {
	select {
	case alloc.queue <- ip:
		return nil
	default:
		return errors.New("Too many IP addresses have been released")
	}
	return nil
}

func newIPAllocator(network *net.IPNet) (*IPAllocator, error) {
	alloc := &IPAllocator{
		network: network,
	}
	if err := alloc.populate(); err != nil {
		return nil, err
	}
	return alloc, nil
}

// Network interface represents the networking stack of a container
type NetworkInterface struct {
	IPNet   net.IPNet
	Gateway net.IP

	manager  *NetworkManager
	extPorts []int
}

// Allocate an external TCP port and map it to the interface
func (iface *NetworkInterface) AllocatePort(port int) (int, error) {
	extPort, err := iface.manager.portAllocator.Acquire()
	if err != nil {
		return -1, err
	}
	if err := iface.manager.portMapper.Map(extPort, net.TCPAddr{IP: iface.IPNet.IP, Port: port}); err != nil {
		iface.manager.portAllocator.Release(extPort)
		return -1, err
	}
	iface.extPorts = append(iface.extPorts, extPort)
	return extPort, nil
}

// Release: Network cleanup - release all resources
func (iface *NetworkInterface) Release() error {
	for _, port := range iface.extPorts {
		if err := iface.manager.portMapper.Unmap(port); err != nil {
			log.Printf("Unable to unmap port %v: %v", port, err)
		}
		if err := iface.manager.portAllocator.Release(port); err != nil {
			log.Printf("Unable to release port %v: %v", port, err)
		}

	}
	return iface.manager.ipAllocator.Release(iface.IPNet.IP)
}

// Network Manager manages a set of network interfaces
// Only *one* manager per host machine should be used
type NetworkManager struct {
	bridgeIface   string
	bridgeNetwork *net.IPNet

	ipAllocator   *IPAllocator
	portAllocator *PortAllocator
	portMapper    *PortMapper
}

// Allocate a network interface
func (manager *NetworkManager) Allocate() (*NetworkInterface, error) {
	ip, err := manager.ipAllocator.Acquire()
	if err != nil {
		return nil, err
	}
	iface := &NetworkInterface{
		IPNet:   net.IPNet{IP: ip, Mask: manager.bridgeNetwork.Mask},
		Gateway: manager.bridgeNetwork.IP,
		manager: manager,
	}
	return iface, nil
}

func newNetworkManager(bridgeIface string) (*NetworkManager, error) {
	addr, err := getIfaceAddr(bridgeIface)
	if err != nil {
		return nil, err
	}
	network := addr.(*net.IPNet)

	ipAllocator, err := newIPAllocator(network)
	if err != nil {
		return nil, err
	}

	portAllocator, err := newPortAllocator(portRangeStart, portRangeEnd)
	if err != nil {
		return nil, err
	}

	portMapper, err := newPortMapper()

	manager := &NetworkManager{
		bridgeIface:   bridgeIface,
		bridgeNetwork: network,
		ipAllocator:   ipAllocator,
		portAllocator: portAllocator,
		portMapper:    portMapper,
	}
	return manager, nil
}
<|endoftext|>