<commit_before>package main

import (
	"errors"
	"flag"
	"github.com/dchest/captcha"
	"html/template"
	"log"
	"net/http"
	"math/rand"
	"time"
)

var port = flag.String("port", "8080", "Listening HTTP port")

var rtmpl = template.Must(
	template.New("register.html").ParseFiles("templates/register.html"))
var ltmpl = template.Must(
	template.New("login.html").ParseFiles("templates/login.html"))
var stmpl = template.Must(
	template.New("settings.html").Funcs( template.FuncMap{
		"GetService": func(key string) *Service {
			return services[key]
		},
		}).ParseFiles("templates/settings.html"))
var atmpl = template.Must(
	template.New("admin.html").Funcs( template.FuncMap{
		"GetUser": func(id int32) *User {
			return db.GetUser(id)
		},
		"GetService": func(key string) *Service {
			return services[key]
		},
		}).ParseFiles("templates/admin.html"))

func index(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		ko(w); return
	}

	writeFiles(w, "templates/header.html", GetNavbar(r),
		"templates/index.html", "templates/footer.html")
}

func getRegister(w http.ResponseWriter, r *http.Request) {
	writeFiles(w, "templates/header.html", GetNavbar(r))

	d := struct { CaptchaId string }{ captcha.New() }

	if err := rtmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postRegister(w http.ResponseWriter, r *http.Request) {
/*
	if !captcha.VerifyString(r.FormValue("captchaId"), r.FormValue("captchaRes")) {
		w.Write([]byte("
Bad captcha; try again.
"))
		return
	}
*/

	if err := Register(r.FormValue("name"), r.FormValue("email")); err != nil {
		LogHttp(w, err); return
	}

	w.Write([]byte("
Check your email account, "+
		`and <a href="/login">login</a>!</p>`))
}

func register(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getRegister(w, r)
	case "POST":
		postRegister(w, r)
	default:
		ko(w)
	}
}

func getLogin(w http.ResponseWriter, r *http.Request) {
	writeFiles(w, "templates/header.html", "templates/navbar.html")

	d := struct { CaptchaId string }{ captcha.New() }

	if err := ltmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postLogin(w http.ResponseWriter, r *http.Request) {
/*
	if !captcha.VerifyString(r.FormValue("captchaId"), r.FormValue("captchaRes")) {
		w.Write([]byte("
Bad captcha; try again.
"))
		return
	}
*/

	token, err := Login(r.FormValue("login"))
	if err != nil {
		LogHttp(w, err)
		return
	}
	if token == "" {
		w.Write([]byte("
Check your email account, "+
			`and <a href="/login">login</a>!</p>`))
		return
	}

	err = SetToken(w, token)
	if err != nil { LogHttp(w, err); return }

	http.Redirect(w, r, "/settings", http.StatusFound)
}

func login(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getLogin(w, r)
	case "POST":
		postLogin(w, r)
	default:
		ko(w)
	}
}

func logout(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	token, err := GetToken(r)
	if err != nil { LogHttp(w, err); return }

	Logout(token)
	UnsetToken(w)

	http.Redirect(w, r, "/", http.StatusFound)
}

func admin(w http.ResponseWriter, r *http.Request) {
	token, err := GetToken(r)
	if err != nil || !ACheckToken(token) || !IsAdmin(token) {
		ko(w); return
	}

	writeFiles(w, "templates/header.html", "templates/navbar3.html")

	d := struct {
		Users		map[int32][]*Token
		Services	map[string]*Service
	}{ utokens, services }

	if err := atmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func getSettings(w http.ResponseWriter, r *http.Request) {
	var tokens []*Token
	token, err := GetToken(r)
	if err == nil {
		tokens = GetTokens(token)
	}

	if err != nil { LogHttp(w, err); return }

	writeFiles(w, "templates/header.html", GetNavbar(r))
	d := struct { Tokens []*Token }{ tokens }
	if err := stmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postSettings(w http.ResponseWriter, r *http.Request) {
	// TODO
}

func settings(w http.ResponseWriter, r *http.Request) {
	token, err := GetToken(r)
	if err == nil {
		if !ACheckToken(token) {
			err = errors.New("Wrong Token")
		}
	}
	if err != nil { LogHttp(w, err); return }

	ntoken := RandomString(LenToken)
	err = SetToken(w, ntoken)
	if err != nil { LogHttp(w, err); return }

	switch r.Method {
	case "GET":
		getSettings(w, r)
	case "POST":
		postSettings(w, r)
	default:
		ko(w)
	}

	UpdateToken(token, ntoken)
}

func discover(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { ko(w); return }

	name, url := r.FormValue("name"), r.FormValue("url")

	key, err := AddService(name, url)
	if err != nil { ko(w); return }

	w.Write([]byte(key))
}

func update(w http.ResponseWriter, r *http.Request) {
	ko(w)
}

func info(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	u := db.GetUser2(r.FormValue("login"))
	if u == nil { ko(w); return }

	w.Write([]byte(u.Name+"\n"+u.Email))
}

func alogin(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	login, key := r.FormValue("login"), r.FormValue("key")

	if isToken(login) {
		if CheckToken(&Token{ key, login }) {
			ok(w)
		} else {
			ko(w)
		}
	} else {
		u := db.GetUser2(login)
		s := services[key]
		if s == nil { ko(w); return }
		token := NewToken(s.Key)
		StoreToken(u.Id, token)
		w.Write([]byte("new"))
	}
}

func generate(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	u := db.GetUser2(r.FormValue("login"))

	s := services[r.FormValue("key")]
	if s == nil { ko(w); return }

	// XXX check the ip address

	token := NewToken(s.Key)
	StoreToken(u.Id, token)

	ok(w)
}

func check(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	if CheckToken(&Token{ r.FormValue("key"), r.FormValue("token") }) {
		ok(w)
	} else {
		ko(w)
	}
}

func chain(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	token := Token{ r.FormValue("key"), r.FormValue("token") }

	ntoken := ChainToken(&token)
	if ntoken != nil {
		w.Write([]byte(ntoken.Token))
	} else {
		ko(w)
	}
}

func main() {
	// Data init
	services = map[string]*Service{}
	utokens = map[int32][]*Token{}
	tokens = map[string]int32{}
	// db requires services to be  created
	db = NewDatabase()
//	services[Auth.Key] = &Auth
	rand.Seed(time.Now().UnixNano())

	// XXX load services

	// Auth website
	http.HandleFunc("/", index)
	http.HandleFunc("/register/", register)
//	http.HandleFunc("/unregister/", unregister)
	http.HandleFunc("/login/", login)
	http.HandleFunc("/logout/", logout)
	http.HandleFunc("/admin/", admin)
	http.HandleFunc("/settings/", settings)

	// API
	http.HandleFunc("/api/discover", discover)
	http.HandleFunc("/api/update", update)
	http.HandleFunc("/api/info", info)
//	http.HandleFunc("/api/generate", generate)
//	http.HandleFunc("/api/check", check)
	http.HandleFunc("/api/login", alogin)
	http.HandleFunc("/api/chain", chain)

	// Captchas
	http.Handle("/captcha/",
		captcha.Server(captcha.StdWidth, captcha.StdHeight))

	// Static files
	http.Handle("/static/",
		http.StripPrefix("/static/",
			http.FileServer(http.Dir("./static/"))))

	log.Println("Launching on http://localhost:"+*port)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<commit_msg>bug when attempting to connect with dummy user name<commit_after>package main

import (
	"errors"
	"flag"
	"github.com/dchest/captcha"
	"html/template"
	"log"
	"net/http"
	"math/rand"
	"time"
)

var port = flag.String("port", "8080", "Listening HTTP port")

var rtmpl = template.Must(
	template.New("register.html").ParseFiles("templates/register.html"))
var ltmpl = template.Must(
	template.New("login.html").ParseFiles("templates/login.html"))
var stmpl = template.Must(
	template.New("settings.html").Funcs( template.FuncMap{
		"GetService": func(key string) *Service {
			return services[key]
		},
		}).ParseFiles("templates/settings.html"))
var atmpl = template.Must(
	template.New("admin.html").Funcs( template.FuncMap{
		"GetUser": func(id int32) *User {
			return db.GetUser(id)
		},
		"GetService": func(key string) *Service {
			return services[key]
		},
		}).ParseFiles("templates/admin.html"))

func index(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		ko(w); return
	}

	writeFiles(w, "templates/header.html", GetNavbar(r),
		"templates/index.html", "templates/footer.html")
}

func getRegister(w http.ResponseWriter, r *http.Request) {
	writeFiles(w, "templates/header.html", GetNavbar(r))

	d := struct { CaptchaId string }{ captcha.New() }

	if err := rtmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postRegister(w http.ResponseWriter, r *http.Request) {
/*
	if !captcha.VerifyString(r.FormValue("captchaId"), r.FormValue("captchaRes")) {
		w.Write([]byte("
Bad captcha; try again.
"))
		return
	}
*/

	if err := Register(r.FormValue("name"), r.FormValue("email")); err != nil {
		LogHttp(w, err); return
	}

	w.Write([]byte("
Check your email account, "+
		`and <a href="/login">login</a>!</p>`))
}

func register(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getRegister(w, r)
	case "POST":
		postRegister(w, r)
	default:
		ko(w)
	}
}

func getLogin(w http.ResponseWriter, r *http.Request) {
	writeFiles(w, "templates/header.html", "templates/navbar.html")

	d := struct { CaptchaId string }{ captcha.New() }

	if err := ltmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postLogin(w http.ResponseWriter, r *http.Request) {
/*
	if !captcha.VerifyString(r.FormValue("captchaId"), r.FormValue("captchaRes")) {
		w.Write([]byte("
Bad captcha; try again.
"))
		return
	}
*/

	token, err := Login(r.FormValue("login"))
	if err != nil {
		LogHttp(w, err)
		return
	}
	if token == "" {
		w.Write([]byte("
Check your email account, "+
			`and <a href="/login">login</a>!</p>`))
		return
	}

	err = SetToken(w, token)
	if err != nil { LogHttp(w, err); return }

	http.Redirect(w, r, "/settings", http.StatusFound)
}

func login(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		getLogin(w, r)
	case "POST":
		postLogin(w, r)
	default:
		ko(w)
	}
}

func logout(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	token, err := GetToken(r)
	if err != nil { LogHttp(w, err); return }

	Logout(token)
	UnsetToken(w)

	http.Redirect(w, r, "/", http.StatusFound)
}

func admin(w http.ResponseWriter, r *http.Request) {
	token, err := GetToken(r)
	if err != nil || !ACheckToken(token) || !IsAdmin(token) {
		ko(w); return
	}

	writeFiles(w, "templates/header.html", "templates/navbar3.html")

	d := struct {
		Users		map[int32][]*Token
		Services	map[string]*Service
	}{ utokens, services }

	if err := atmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func getSettings(w http.ResponseWriter, r *http.Request) {
	var tokens []*Token
	token, err := GetToken(r)
	if err == nil {
		tokens = GetTokens(token)
	}

	if err != nil { LogHttp(w, err); return }

	writeFiles(w, "templates/header.html", GetNavbar(r))
	d := struct { Tokens []*Token }{ tokens }
	if err := stmpl.Execute(w, &d); err != nil {
		LogHttp(w, err); return
	}

	writeFiles(w, "templates/footer.html")
}

func postSettings(w http.ResponseWriter, r *http.Request) {
	// TODO
}

func settings(w http.ResponseWriter, r *http.Request) {
	token, err := GetToken(r)
	if err == nil {
		if !ACheckToken(token) {
			err = errors.New("Wrong Token")
		}
	}
	if err != nil { LogHttp(w, err); return }

	ntoken := RandomString(LenToken)
	err = SetToken(w, ntoken)
	if err != nil { LogHttp(w, err); return }

	switch r.Method {
	case "GET":
		getSettings(w, r)
	case "POST":
		postSettings(w, r)
	default:
		ko(w)
	}

	UpdateToken(token, ntoken)
}

func discover(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { ko(w); return }

	name, url := r.FormValue("name"), r.FormValue("url")

	key, err := AddService(name, url)
	if err != nil { ko(w); return }

	w.Write([]byte(key))
}

func update(w http.ResponseWriter, r *http.Request) {
	ko(w)
}

func info(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	id := tokens[r.FormValue("token")]

	u := db.GetUser(id)
	if u == nil { ko(w); return }

	w.Write([]byte(u.Name+"\n"+u.Email))
}

func alogin(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	login, key := r.FormValue("login"), r.FormValue("key")

	if isToken(login) {
		if CheckToken(&Token{ key, login }) {
			ok(w)
		} else {
			ko(w)
		}
	} else {
		u := db.GetUser2(login)
		if u == nil { ko(w) }
		s := services[key]
		if s == nil { ko(w); return }
		token := NewToken(s.Key)
		StoreToken(u.Id, token)
		w.Write([]byte("new"))
	}
}

/*
func generate(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	u := db.GetUser2(r.FormValue("login"))

	s := services[r.FormValue("key")]
	if s == nil { ko(w); return }

	// XXX check the ip address

	token := NewToken(s.Key)
	StoreToken(u.Id, token)

	ok(w)
}

func check(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	if CheckToken(&Token{ r.FormValue("key"), r.FormValue("token") }) {
		ok(w)
	} else {
		ko(w)
	}
}
*/

func chain(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" { ko(w); return }

	token := Token{ r.FormValue("key"), r.FormValue("token") }

	ntoken := ChainToken(&token)
	if ntoken != nil {
		w.Write([]byte(ntoken.Token))
	} else {
		ko(w)
	}
}

func main() {
	// Data init
	services = map[string]*Service{}
	utokens = map[int32][]*Token{}
	tokens = map[string]int32{}
	// db requires services to be  created
	db = NewDatabase()
//	services[Auth.Key] = &Auth
	rand.Seed(time.Now().UnixNano())

	// XXX load services

	// Auth website
	http.HandleFunc("/", index)
	http.HandleFunc("/register/", register)
//	http.HandleFunc("/unregister/", unregister)
	http.HandleFunc("/login/", login)
	http.HandleFunc("/logout/", logout)
	http.HandleFunc("/admin/", admin)
	http.HandleFunc("/settings/", settings)

	// API
	http.HandleFunc("/api/discover", discover)
	http.HandleFunc("/api/update", update)
	http.HandleFunc("/api/info", info)
//	http.HandleFunc("/api/generate", generate)
//	http.HandleFunc("/api/check", check)
	http.HandleFunc("/api/login", alogin)
	http.HandleFunc("/api/chain", chain)

	// Captchas
	http.Handle("/captcha/",
		captcha.Server(captcha.StdWidth, captcha.StdHeight))

	// Static files
	http.Handle("/static/",
		http.StripPrefix("/static/",
			http.FileServer(http.Dir("./static/"))))

	log.Println("Launching on http://localhost:"+*port)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net/http"
)

func main() {

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "hello world")
	})

	log.Fatal(http.ListenAndServe(":80", nil))
}
<commit_msg>test: Make hello.go use all cores<commit_after>package main

import (
	"io"
	"log"
	"net/http"
	"runtime"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "hello world")
	})

	log.Fatal(http.ListenAndServe(":80", nil))
}
<|endoftext|>


<commit_before>// Copyright 2013 Travis Keep. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or
// at http://opensource.org/licenses/BSD-3-Clause.

// Package weblogs provides access logs for webservers written in go.
package weblogs

import (
  "bytes"
  "fmt"
  "github.com/gorilla/context"
  "github.com/keep94/weblogs/loggers"
  "io"
  "net/http"
  "os"
  "runtime/debug"
  "sync"
  "time"
)

type contextKeyType int

const (
  kBufferKey contextKeyType = iota
  kValuesKey
)

// Snapshot represents a snapshot of an HTTP request.
type Snapshot interface{}

// Capture captures a server response. Implementations delegate to an
// underlying ResponseWriter.
type Capture interface {
  http.ResponseWriter
  // HasStatus returns true if server has sent a status. False means that
  // server failed to send a response.
  HasStatus() bool
}

// LogRecord represents a single entry in the access logs.
type LogRecord struct {
  // The time request was received.
  T time.Time
  // The request snapshot
  R Snapshot
  // The capture of the response
  W Capture
  // Time spent processing the request
  Duration time.Duration
  // Additional information added with the Writer method.
  Extra string
  // Key-value pairs to be logged.
  Values map[interface{}]interface{}
}

// Logger represents an access log format. Clients are free to provide their
// own implementations.
type Logger interface {
  // NewSnapshot creates a new snapshot of a request.
  NewSnapshot(r *http.Request) Snapshot
  // NewCapture creates a new capture for capturing a response. w is the
  // original ResponseWriter.
  NewCapture(w http.ResponseWriter) Capture
  // Log writes the log record.
  Log(w io.Writer, record *LogRecord)
}

// Options specifies options for writing to access logs.
type Options struct {
  // Where to write the web logs. nil means write to stderr,
  Writer io.Writer
  // How to write the web logs. nil means SimpleLogger().
  Logger Logger
  // How to get current time. nil means use time.Now(). This field is used
  // for testing purposes.
  Now func() time.Time
}

func (o *Options) writer() io.Writer {
  if o.Writer == nil {
    return os.Stderr
  }
  return o.Writer
}

func (o *Options) logger() Logger {
  if o.Logger == nil {
    return simpleLogger{}
  }
  return o.Logger
}

func (o *Options) now() func() time.Time {
  if o.Now == nil {
    return time.Now
  }
  return o.Now
}

// Handler wraps a handler creating access logs. Access logs are written to
// stderr using SimpleLogger(). Returned handler must be wrapped by
// context.ClearHandler.
func Handler(handler http.Handler) http.Handler {
  return HandlerWithOptions(handler, nil)
}

// HandlerWithOptions wraps a handler creating access logs and allows caller to
// configure how access logs are written. Returned handler must be
// wrapped by context.ClearHandler.
func HandlerWithOptions(
    handler http.Handler, options *Options) http.Handler {
  if options == nil {
    options = &Options{}
  }
  return &logHandler{
      handler: handler,
      w: options.writer(),
      logger: options.logger(),
      now: options.now()}
}

// Writer returns a writer whereby the caller can add additional information
// to the current log entry. If the handler calling this is not wrapped by
// the Handler() method, then writing to the returned io.Writer does
// nothing.
func Writer(r *http.Request) io.Writer {
  value := context.Get(r, kBufferKey)
  if value == nil {
    return nilWriter{}
  }
  return value.(*bytes.Buffer)
}

// Values returns the current key-value pairs to be logged.
// If the handler calling this is not wrapped by the Handler() method,
// then this method returns nil.
func Values(r *http.Request) map[interface{}]interface{} {
  instance := context.Get(r, kValuesKey)
  if instance == nil {
    return nil
  }
  return instance.(map[interface{}]interface{})
}

type logHandler struct {
  // mutex protects the w field.
  mutex sync.Mutex
  handler http.Handler
  w io.Writer
  logger Logger
  now func() time.Time
}

func (h *logHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  snapshot :=  h.logger.NewSnapshot(r)
  capture := h.logger.NewCapture(w)
  additional := &bytes.Buffer{}
  values := make(map[interface{}]interface{})
  context.Set(r, kBufferKey, additional)
  context.Set(r, kValuesKey, values)
  startTime := h.now()
  defer func() {
    endTime := h.now()
    err := recover()
    maybeSend500(capture)
    h.writeLogRecord(
        &LogRecord{
            T: startTime,
            R: snapshot,
            W: capture,
            Duration: endTime.Sub(startTime),
            Extra: additional.String(),
            Values: values})
    if err != nil {
      h.writePanic(err, debug.Stack());
    }
  }()
  h.handler.ServeHTTP(capture, r)
}

func (h *logHandler) writeLogRecord(logRecord *LogRecord) {
  h.mutex.Lock()
  defer h.mutex.Unlock()
  h.logger.Log(h.w, logRecord)
}

func (h *logHandler) writePanic(
    panicError interface{}, debugStack []byte) {
  h.mutex.Lock()
  defer h.mutex.Unlock()
  fmt.Fprintf(h.w, "Panic: %v\n", panicError)
  h.w.Write(debugStack)
}

// SimpleLogger provides access logs with the following columns:
// date, remote address, method, URI, status, time elapsed milliseconds,
// followed by any additional information provided via the Writer method.
func SimpleLogger() Logger {
  return simpleLogger{}
}

// ApacheCommonLogger provides access logs in apache common log format.
func ApacheCommonLogger() Logger {
  return apacheCommonLogger{}
}

// ApacheCombinedLogger provides access logs in apache combined log format.
func ApacheCombinedLogger() Logger {
  return apacheCombinedLogger{}
}

type loggerBase struct {
}

func (l loggerBase) NewSnapshot(r *http.Request) Snapshot {
  return loggers.NewSnapshot(r)
}

func (l loggerBase) NewCapture(w http.ResponseWriter) Capture {
  return &loggers.Capture{ResponseWriter: w}
}

type simpleLogger struct {
  loggerBase
}

func (l simpleLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s %s %s %s %d %d%s\n",
      log.T.Format("01/02/2006 15:04:05"),
      loggers.StripPort(s.RemoteAddr),
      s.Method,
      s.URL,
      c.Status(),
      log.Duration / time.Millisecond,
      log.Extra)
}

type apacheCommonLogger struct {
  loggerBase
}

func (l apacheCommonLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s - %s [%s] \"%s %s %s\" %d %d\n",
        loggers.StripPort(s.RemoteAddr),
        loggers.ApacheUser(s.URL.User),
        log.T.Format("02/Jan/2006:15:04:05 -0700"),
        s.Method,
        s.URL.RequestURI(),
        s.Proto,
        c.Status(),
        c.Size())
}

type apacheCombinedLogger struct {
  loggerBase
}

func (l apacheCombinedLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s - %s [%s] \"%s %s %s\" %d %d \"%s\" \"%s\"\n",
        loggers.StripPort(s.RemoteAddr),
        loggers.ApacheUser(s.URL.User),
        log.T.Format("02/Jan/2006:15:04:05 -0700"),
        s.Method,
        s.URL.RequestURI(),
        s.Proto,
        c.Status(),
        c.Size(),
        s.Referer,
        s.UserAgent)
}

func maybeSend500(c Capture) {
  if !c.HasStatus() {
    sendError(c, http.StatusInternalServerError)
  }
}

func sendError(w http.ResponseWriter, status int) {
  http.Error(w, fmt.Sprintf("%d %s", status, http.StatusText(status)), status)
}

type nilWriter struct {
}

func (w nilWriter) Write(p []byte) (n int, err error) {
  return len(p), nil
}
<commit_msg>Performance: Store empty options in a variable.<commit_after>// Copyright 2013 Travis Keep. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or
// at http://opensource.org/licenses/BSD-3-Clause.

// Package weblogs provides access logs for webservers written in go.
package weblogs

import (
  "bytes"
  "fmt"
  "github.com/gorilla/context"
  "github.com/keep94/weblogs/loggers"
  "io"
  "net/http"
  "os"
  "runtime/debug"
  "sync"
  "time"
)

type contextKeyType int

const (
  kBufferKey contextKeyType = iota
  kValuesKey
)

var (
  kNoOptions = &Options{}
)

// Snapshot represents a snapshot of an HTTP request.
type Snapshot interface{}

// Capture captures a server response. Implementations delegate to an
// underlying ResponseWriter.
type Capture interface {
  http.ResponseWriter
  // HasStatus returns true if server has sent a status. False means that
  // server failed to send a response.
  HasStatus() bool
}

// LogRecord represents a single entry in the access logs.
type LogRecord struct {
  // The time request was received.
  T time.Time
  // The request snapshot
  R Snapshot
  // The capture of the response
  W Capture
  // Time spent processing the request
  Duration time.Duration
  // Additional information added with the Writer method.
  Extra string
  // Key-value pairs to be logged.
  Values map[interface{}]interface{}
}

// Logger represents an access log format. Clients are free to provide their
// own implementations.
type Logger interface {
  // NewSnapshot creates a new snapshot of a request.
  NewSnapshot(r *http.Request) Snapshot
  // NewCapture creates a new capture for capturing a response. w is the
  // original ResponseWriter.
  NewCapture(w http.ResponseWriter) Capture
  // Log writes the log record.
  Log(w io.Writer, record *LogRecord)
}

// Options specifies options for writing to access logs.
type Options struct {
  // Where to write the web logs. nil means write to stderr,
  Writer io.Writer
  // How to write the web logs. nil means SimpleLogger().
  Logger Logger
  // How to get current time. nil means use time.Now(). This field is used
  // for testing purposes.
  Now func() time.Time
}

func (o *Options) writer() io.Writer {
  if o.Writer == nil {
    return os.Stderr
  }
  return o.Writer
}

func (o *Options) logger() Logger {
  if o.Logger == nil {
    return simpleLogger{}
  }
  return o.Logger
}

func (o *Options) now() func() time.Time {
  if o.Now == nil {
    return time.Now
  }
  return o.Now
}

// Handler wraps a handler creating access logs. Access logs are written to
// stderr using SimpleLogger(). Returned handler must be wrapped by
// context.ClearHandler.
func Handler(handler http.Handler) http.Handler {
  return HandlerWithOptions(handler, nil)
}

// HandlerWithOptions wraps a handler creating access logs and allows caller to
// configure how access logs are written. Returned handler must be
// wrapped by context.ClearHandler.
func HandlerWithOptions(
    handler http.Handler, options *Options) http.Handler {
  if options == nil {
    options = kNoOptions
  }
  return &logHandler{
      handler: handler,
      w: options.writer(),
      logger: options.logger(),
      now: options.now()}
}

// Writer returns a writer whereby the caller can add additional information
// to the current log entry. If the handler calling this is not wrapped by
// the Handler() method, then writing to the returned io.Writer does
// nothing.
func Writer(r *http.Request) io.Writer {
  value := context.Get(r, kBufferKey)
  if value == nil {
    return nilWriter{}
  }
  return value.(*bytes.Buffer)
}

// Values returns the current key-value pairs to be logged.
// If the handler calling this is not wrapped by the Handler() method,
// then this method returns nil.
func Values(r *http.Request) map[interface{}]interface{} {
  instance := context.Get(r, kValuesKey)
  if instance == nil {
    return nil
  }
  return instance.(map[interface{}]interface{})
}

type logHandler struct {
  // mutex protects the w field.
  mutex sync.Mutex
  handler http.Handler
  w io.Writer
  logger Logger
  now func() time.Time
}

func (h *logHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  snapshot :=  h.logger.NewSnapshot(r)
  capture := h.logger.NewCapture(w)
  additional := &bytes.Buffer{}
  values := make(map[interface{}]interface{})
  context.Set(r, kBufferKey, additional)
  context.Set(r, kValuesKey, values)
  startTime := h.now()
  defer func() {
    endTime := h.now()
    err := recover()
    maybeSend500(capture)
    h.writeLogRecord(
        &LogRecord{
            T: startTime,
            R: snapshot,
            W: capture,
            Duration: endTime.Sub(startTime),
            Extra: additional.String(),
            Values: values})
    if err != nil {
      h.writePanic(err, debug.Stack());
    }
  }()
  h.handler.ServeHTTP(capture, r)
}

func (h *logHandler) writeLogRecord(logRecord *LogRecord) {
  h.mutex.Lock()
  defer h.mutex.Unlock()
  h.logger.Log(h.w, logRecord)
}

func (h *logHandler) writePanic(
    panicError interface{}, debugStack []byte) {
  h.mutex.Lock()
  defer h.mutex.Unlock()
  fmt.Fprintf(h.w, "Panic: %v\n", panicError)
  h.w.Write(debugStack)
}

// SimpleLogger provides access logs with the following columns:
// date, remote address, method, URI, status, time elapsed milliseconds,
// followed by any additional information provided via the Writer method.
func SimpleLogger() Logger {
  return simpleLogger{}
}

// ApacheCommonLogger provides access logs in apache common log format.
func ApacheCommonLogger() Logger {
  return apacheCommonLogger{}
}

// ApacheCombinedLogger provides access logs in apache combined log format.
func ApacheCombinedLogger() Logger {
  return apacheCombinedLogger{}
}

type loggerBase struct {
}

func (l loggerBase) NewSnapshot(r *http.Request) Snapshot {
  return loggers.NewSnapshot(r)
}

func (l loggerBase) NewCapture(w http.ResponseWriter) Capture {
  return &loggers.Capture{ResponseWriter: w}
}

type simpleLogger struct {
  loggerBase
}

func (l simpleLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s %s %s %s %d %d%s\n",
      log.T.Format("01/02/2006 15:04:05"),
      loggers.StripPort(s.RemoteAddr),
      s.Method,
      s.URL,
      c.Status(),
      log.Duration / time.Millisecond,
      log.Extra)
}

type apacheCommonLogger struct {
  loggerBase
}

func (l apacheCommonLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s - %s [%s] \"%s %s %s\" %d %d\n",
        loggers.StripPort(s.RemoteAddr),
        loggers.ApacheUser(s.URL.User),
        log.T.Format("02/Jan/2006:15:04:05 -0700"),
        s.Method,
        s.URL.RequestURI(),
        s.Proto,
        c.Status(),
        c.Size())
}

type apacheCombinedLogger struct {
  loggerBase
}

func (l apacheCombinedLogger) Log(w io.Writer, log *LogRecord) {
  s := log.R.(*loggers.Snapshot)
  c := log.W.(*loggers.Capture)
  fmt.Fprintf(w, "%s - %s [%s] \"%s %s %s\" %d %d \"%s\" \"%s\"\n",
        loggers.StripPort(s.RemoteAddr),
        loggers.ApacheUser(s.URL.User),
        log.T.Format("02/Jan/2006:15:04:05 -0700"),
        s.Method,
        s.URL.RequestURI(),
        s.Proto,
        c.Status(),
        c.Size(),
        s.Referer,
        s.UserAgent)
}

func maybeSend500(c Capture) {
  if !c.HasStatus() {
    sendError(c, http.StatusInternalServerError)
  }
}

func sendError(w http.ResponseWriter, status int) {
  http.Error(w, fmt.Sprintf("%d %s", status, http.StatusText(status)), status)
}

type nilWriter struct {
}

func (w nilWriter) Write(p []byte) (n int, err error) {
  return len(p), nil
}
<|endoftext|>


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

package flate

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

type deflateTest struct {
	in    []byte
	level int
	out   []byte
}

type deflateInflateTest struct {
	in []byte
}

type reverseBitsTest struct {
	in       uint16
	bitCount uint8
	out      uint16
}

var deflateTests = []*deflateTest{
	&deflateTest{[]byte{}, 0, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, -1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, DefaultCompression, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 4, []byte{18, 4, 4, 0, 0, 255, 255}},

	&deflateTest{[]byte{0x11}, 0, []byte{0, 1, 0, 254, 255, 17, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 0, []byte{0, 2, 0, 253, 255, 17, 18, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 0,
		[]byte{0, 8, 0, 247, 255, 17, 17, 17, 17, 17, 17, 17, 17, 1, 0, 0, 255, 255},
	},
	&deflateTest{[]byte{}, 1, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 1, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 1, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
	&deflateTest{[]byte{}, 9, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 9, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 9, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 9, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
}

var deflateInflateTests = []*deflateInflateTest{
	&deflateInflateTest{[]byte{}},
	&deflateInflateTest{[]byte{0x11}},
	&deflateInflateTest{[]byte{0x11, 0x12}},
	&deflateInflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}},
	&deflateInflateTest{[]byte{0x11, 0x10, 0x13, 0x41, 0x21, 0x21, 0x41, 0x13, 0x87, 0x78, 0x13}},
	&deflateInflateTest{getLargeDataChunk()},
}

var reverseBitsTests = []*reverseBitsTest{
	&reverseBitsTest{1, 1, 1},
	&reverseBitsTest{1, 2, 2},
	&reverseBitsTest{1, 3, 4},
	&reverseBitsTest{1, 4, 8},
	&reverseBitsTest{1, 5, 16},
	&reverseBitsTest{17, 5, 17},
	&reverseBitsTest{257, 9, 257},
	&reverseBitsTest{29, 5, 23},
}

func getLargeDataChunk() []byte {
	result := make([]byte, 100000)
	for i := range result {
		result[i] = byte(int64(i) * int64(i) & 0xFF)
	}
	return result
}

func TestDeflate(t *testing.T) {
	for _, h := range deflateTests {
		buffer := bytes.NewBuffer(nil)
		w := NewWriter(buffer, h.level)
		w.Write(h.in)
		w.Close()
		if bytes.Compare(buffer.Bytes(), h.out) != 0 {
			t.Errorf("buffer is wrong; level = %v, buffer.Bytes() = %v, expected output = %v",
				h.level, buffer.Bytes(), h.out)
		}
	}
}

type syncBuffer struct {
	buf    bytes.Buffer
	mu     sync.RWMutex
	closed bool
	ready  chan bool
}

func newSyncBuffer() *syncBuffer {
	return &syncBuffer{ready: make(chan bool, 1)}
}

func (b *syncBuffer) Read(p []byte) (n int, err os.Error) {
	for {
		b.mu.RLock()
		n, err = b.buf.Read(p)
		b.mu.RUnlock()
		if n > 0 || b.closed {
			return
		}
		<-b.ready
	}
	panic("unreachable")
}

func (b *syncBuffer) signal() {
	select {
	case b.ready <- true:
	default:
	}
}

func (b *syncBuffer) Write(p []byte) (n int, err os.Error) {
	n, err = b.buf.Write(p)
	b.signal()
	return
}

func (b *syncBuffer) WriteMode() {
	b.mu.Lock()
}

func (b *syncBuffer) ReadMode() {
	b.mu.Unlock()
	b.signal()
}

func (b *syncBuffer) Close() os.Error {
	b.closed = true
	b.signal()
	return nil
}

func testSync(t *testing.T, level int, input []byte, name string) {
	if len(input) == 0 {
		return
	}

	t.Logf("--testSync %d, %d, %s", level, len(input), name)
	buf := newSyncBuffer()
	buf1 := new(bytes.Buffer)
	buf.WriteMode()
	w := NewWriter(io.MultiWriter(buf, buf1), level)
	r := NewReader(buf)

	// Write half the input and read back.
	for i := 0; i < 2; i++ {
		var lo, hi int
		if i == 0 {
			lo, hi = 0, (len(input)+1)/2
		} else {
			lo, hi = (len(input)+1)/2, len(input)
		}
		t.Logf("#%d: write %d-%d", i, lo, hi)
		if _, err := w.Write(input[lo:hi]); err != nil {
			t.Errorf("testSync: write: %v", err)
			return
		}
		if i == 0 {
			if err := w.Flush(); err != nil {
				t.Errorf("testSync: flush: %v", err)
				return
			}
		} else {
			if err := w.Close(); err != nil {
				t.Errorf("testSync: close: %v", err)
			}
		}
		buf.ReadMode()
		out := make([]byte, hi-lo+1)
		m, err := io.ReadAtLeast(r, out, hi-lo)
		t.Logf("#%d: read %d", i, m)
		if m != hi-lo || err != nil {
			t.Errorf("testSync/%d (%d, %d, %s): read %d: %d, %v (%d left)", i, level, len(input), name, hi-lo, m, err, buf.buf.Len())
			return
		}
		if !bytes.Equal(input[lo:hi], out[:hi-lo]) {
			t.Errorf("testSync/%d: read wrong bytes: %x vs %x", i, input[lo:hi], out[:hi-lo])
			return
		}
		if i == 0 && buf.buf.Len() != 0 {
			t.Errorf("testSync/%d (%d, %d, %s): extra data after %d", i, level, len(input), name, hi-lo)
		}
		buf.WriteMode()
	}
	buf.ReadMode()
	out := make([]byte, 10)
	if n, err := r.Read(out); n > 0 || err != os.EOF {
		t.Errorf("testSync (%d, %d, %s): final Read: %d, %v (hex: %x)", level, len(input), name, n, err, out[0:n])
	}
	if buf.buf.Len() != 0 {
		t.Errorf("testSync (%d, %d, %s): extra data at end", level, len(input), name)
	}
	r.Close()

	// stream should work for ordinary reader too
	r = NewReader(buf1)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("testSync: read: %s", err)
		return
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("testSync: decompress(compress(data)) != data: level=%d input=%s", level, name)
	}
}

func testToFromWithLevel(t *testing.T, level int, input []byte, name string) os.Error {
	buffer := bytes.NewBuffer(nil)
	w := NewWriter(buffer, level)
	w.Write(input)
	w.Close()
	r := NewReader(buffer)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("read: %s", err)
		return err
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("decompress(compress(data)) != data: level=%d input=%s", level, name)
	}

	testSync(t, level, input, name)
	return nil
}

func testToFrom(t *testing.T, input []byte, name string) {
	for i := 0; i < 10; i++ {
		testToFromWithLevel(t, i, input, name)
	}
}

func TestDeflateInflate(t *testing.T) {
	for i, h := range deflateInflateTests {
		testToFrom(t, h.in, fmt.Sprintf("#%d", i))
	}
}

func TestReverseBits(t *testing.T) {
	for _, h := range reverseBitsTests {
		if v := reverseBits(h.in, h.bitCount); v != h.out {
			t.Errorf("reverseBits(%v,%v) = %v, want %v",
				h.in, h.bitCount, v, h.out)
		}
	}
}

func TestDeflateInflateString(t *testing.T) {
	gold, err := ioutil.ReadFile("../testdata/e.txt")
	if err != nil {
		t.Error(err)
	}
	testToFromWithLevel(t, 1, gold, "2.718281828...")
}
<commit_msg>compress/flate: fix test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package flate

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

type deflateTest struct {
	in    []byte
	level int
	out   []byte
}

type deflateInflateTest struct {
	in []byte
}

type reverseBitsTest struct {
	in       uint16
	bitCount uint8
	out      uint16
}

var deflateTests = []*deflateTest{
	&deflateTest{[]byte{}, 0, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, -1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, DefaultCompression, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 4, []byte{18, 4, 4, 0, 0, 255, 255}},

	&deflateTest{[]byte{0x11}, 0, []byte{0, 1, 0, 254, 255, 17, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 0, []byte{0, 2, 0, 253, 255, 17, 18, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 0,
		[]byte{0, 8, 0, 247, 255, 17, 17, 17, 17, 17, 17, 17, 17, 1, 0, 0, 255, 255},
	},
	&deflateTest{[]byte{}, 1, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 1, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 1, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
	&deflateTest{[]byte{}, 9, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 9, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 9, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 9, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
}

var deflateInflateTests = []*deflateInflateTest{
	&deflateInflateTest{[]byte{}},
	&deflateInflateTest{[]byte{0x11}},
	&deflateInflateTest{[]byte{0x11, 0x12}},
	&deflateInflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}},
	&deflateInflateTest{[]byte{0x11, 0x10, 0x13, 0x41, 0x21, 0x21, 0x41, 0x13, 0x87, 0x78, 0x13}},
	&deflateInflateTest{getLargeDataChunk()},
}

var reverseBitsTests = []*reverseBitsTest{
	&reverseBitsTest{1, 1, 1},
	&reverseBitsTest{1, 2, 2},
	&reverseBitsTest{1, 3, 4},
	&reverseBitsTest{1, 4, 8},
	&reverseBitsTest{1, 5, 16},
	&reverseBitsTest{17, 5, 17},
	&reverseBitsTest{257, 9, 257},
	&reverseBitsTest{29, 5, 23},
}

func getLargeDataChunk() []byte {
	result := make([]byte, 100000)
	for i := range result {
		result[i] = byte(int64(i) * int64(i) & 0xFF)
	}
	return result
}

func TestDeflate(t *testing.T) {
	for _, h := range deflateTests {
		buffer := bytes.NewBuffer(nil)
		w := NewWriter(buffer, h.level)
		w.Write(h.in)
		w.Close()
		if bytes.Compare(buffer.Bytes(), h.out) != 0 {
			t.Errorf("buffer is wrong; level = %v, buffer.Bytes() = %v, expected output = %v",
				h.level, buffer.Bytes(), h.out)
		}
	}
}

type syncBuffer struct {
	buf    bytes.Buffer
	mu     sync.RWMutex
	closed bool
	ready  chan bool
}

func newSyncBuffer() *syncBuffer {
	return &syncBuffer{ready: make(chan bool, 1)}
}

func (b *syncBuffer) Read(p []byte) (n int, err os.Error) {
	for {
		b.mu.RLock()
		n, err = b.buf.Read(p)
		b.mu.RUnlock()
		if n > 0 || b.closed {
			return
		}
		<-b.ready
	}
	panic("unreachable")
}

func (b *syncBuffer) signal() {
	select {
	case b.ready <- true:
	default:
	}
}

func (b *syncBuffer) Write(p []byte) (n int, err os.Error) {
	n, err = b.buf.Write(p)
	b.signal()
	return
}

func (b *syncBuffer) WriteMode() {
	b.mu.Lock()
}

func (b *syncBuffer) ReadMode() {
	b.mu.Unlock()
	b.signal()
}

func (b *syncBuffer) Close() os.Error {
	b.closed = true
	b.signal()
	return nil
}

func testSync(t *testing.T, level int, input []byte, name string) {
	if len(input) == 0 {
		return
	}

	t.Logf("--testSync %d, %d, %s", level, len(input), name)
	buf := newSyncBuffer()
	buf1 := new(bytes.Buffer)
	buf.WriteMode()
	w := NewWriter(io.MultiWriter(buf, buf1), level)
	r := NewReader(buf)

	// Write half the input and read back.
	for i := 0; i < 2; i++ {
		var lo, hi int
		if i == 0 {
			lo, hi = 0, (len(input)+1)/2
		} else {
			lo, hi = (len(input)+1)/2, len(input)
		}
		t.Logf("#%d: write %d-%d", i, lo, hi)
		if _, err := w.Write(input[lo:hi]); err != nil {
			t.Errorf("testSync: write: %v", err)
			return
		}
		if i == 0 {
			if err := w.Flush(); err != nil {
				t.Errorf("testSync: flush: %v", err)
				return
			}
		} else {
			if err := w.Close(); err != nil {
				t.Errorf("testSync: close: %v", err)
			}
		}
		buf.ReadMode()
		out := make([]byte, hi-lo+1)
		m, err := io.ReadAtLeast(r, out, hi-lo)
		t.Logf("#%d: read %d", i, m)
		if m != hi-lo || err != nil {
			t.Errorf("testSync/%d (%d, %d, %s): read %d: %d, %v (%d left)", i, level, len(input), name, hi-lo, m, err, buf.buf.Len())
			return
		}
		if !bytes.Equal(input[lo:hi], out[:hi-lo]) {
			t.Errorf("testSync/%d: read wrong bytes: %x vs %x", i, input[lo:hi], out[:hi-lo])
			return
		}
		// This test originally checked that after reading
		// the first half of the input, there was nothing left
		// in the read buffer (buf.buf.Len() != 0) but that is
		// not necessarily the case: the write Flush may emit
		// some extra framing bits that are not necessary
		// to process to obtain the first half of the uncompressed
		// data.  The test ran correctly most of the time, because
		// the background goroutine had usually read even
		// those extra bits by now, but it's not a useful thing to
		// check.
		buf.WriteMode()
	}
	buf.ReadMode()
	out := make([]byte, 10)
	if n, err := r.Read(out); n > 0 || err != os.EOF {
		t.Errorf("testSync (%d, %d, %s): final Read: %d, %v (hex: %x)", level, len(input), name, n, err, out[0:n])
	}
	if buf.buf.Len() != 0 {
		t.Errorf("testSync (%d, %d, %s): extra data at end", level, len(input), name)
	}
	r.Close()

	// stream should work for ordinary reader too
	r = NewReader(buf1)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("testSync: read: %s", err)
		return
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("testSync: decompress(compress(data)) != data: level=%d input=%s", level, name)
	}
}

func testToFromWithLevel(t *testing.T, level int, input []byte, name string) os.Error {
	buffer := bytes.NewBuffer(nil)
	w := NewWriter(buffer, level)
	w.Write(input)
	w.Close()
	r := NewReader(buffer)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("read: %s", err)
		return err
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("decompress(compress(data)) != data: level=%d input=%s", level, name)
	}

	testSync(t, level, input, name)
	return nil
}

func testToFrom(t *testing.T, input []byte, name string) {
	for i := 0; i < 10; i++ {
		testToFromWithLevel(t, i, input, name)
	}
}

func TestDeflateInflate(t *testing.T) {
	for i, h := range deflateInflateTests {
		testToFrom(t, h.in, fmt.Sprintf("#%d", i))
	}
}

func TestReverseBits(t *testing.T) {
	for _, h := range reverseBitsTests {
		if v := reverseBits(h.in, h.bitCount); v != h.out {
			t.Errorf("reverseBits(%v,%v) = %v, want %v",
				h.in, h.bitCount, v, h.out)
		}
	}
}

func TestDeflateInflateString(t *testing.T) {
	gold, err := ioutil.ReadFile("../testdata/e.txt")
	if err != nil {
		t.Error(err)
	}
	testToFromWithLevel(t, 1, gold, "2.718281828...")
}
<|endoftext|>


<commit_before>package wkb

import (
	"encoding/binary"
	"fmt"
	"github.com/twpayne/gogeom/geom"
	"reflect"
)

const (
	wkbXDR = 0
	wkbNDR = 1
)

const (
	wkbPoint                = 1
	wkbPointM               = 2001
	wkbPointZ               = 1001
	wkbPointZM              = 3001
	wkbLineString           = 2
	wkbLineStringM          = 2002
	wkbLineStringZ          = 1002
	wkbLineStringZM         = 3002
	wkbPolygon              = 3
	wkbPolygonM             = 2003
	wkbPolygonZ             = 1003
	wkbPolygonZM            = 3003
	wkbMultiPoint           = 4
	wkbMultiPointM          = 2004
	wkbMultiPointZ          = 1004
	wkbMultiPointZM         = 3004
	wkbMultiLineString      = 5
	wkbMultiLineStringM     = 2005
	wkbMultiLineStringZ     = 1005
	wkbMultiLineStringZM    = 3005
	wkbMultiPolygon         = 6
	wkbMultiPolygonM        = 2006
	wkbMultiPolygonZ        = 1006
	wkbMultiPolygonZM       = 3006
	wkbGeometryCollection   = 7
	wkbGeometryCollectionM  = 2007
	wkbGeometryCollectionZ  = 1007
	wkbGeometryCollectionZM = 3007
	wkbPolyhedralSurface    = 15
	wkbPolyhedralSurfaceM   = 2015
	wkbPolyhedralSurfaceZ   = 1015
	wkbPolyhedralSurfaceZM  = 3015
	wkbTIN                  = 16
	wkbTINM                 = 2016
	wkbTINZ                 = 1016
	wkbTINZM                = 3016
	wkbTriangle             = 17
	wkbTriangleM            = 2017
	wkbTriangleZ            = 1017
	wkbTriangleZM           = 3017
)

var (
	XDR = binary.BigEndian
	NDR = binary.LittleEndian
)

type UnexpectedGeometryError struct {
	Geom geom.T
}

func (e UnexpectedGeometryError) Error() string {
	return fmt.Sprintf("wkb: unexpected geometry %v", e.Geom)
}

type UnsupportedGeometryError struct {
	Type reflect.Type
}

func (e UnsupportedGeometryError) Error() string {
	return "wkb: unsupported type: " + e.Type.String()
}
<commit_msg>Sort values more numerically<commit_after>package wkb

import (
	"encoding/binary"
	"fmt"
	"github.com/twpayne/gogeom/geom"
	"reflect"
)

const (
	wkbXDR = 0
	wkbNDR = 1
)

const (
	wkbPoint                = 1
	wkbPointZ               = 1001
	wkbPointM               = 2001
	wkbPointZM              = 3001
	wkbLineString           = 2
	wkbLineStringZ          = 1002
	wkbLineStringM          = 2002
	wkbLineStringZM         = 3002
	wkbPolygon              = 3
	wkbPolygonZ             = 1003
	wkbPolygonM             = 2003
	wkbPolygonZM            = 3003
	wkbMultiPoint           = 4
	wkbMultiPointZ          = 1004
	wkbMultiPointM          = 2004
	wkbMultiPointZM         = 3004
	wkbMultiLineString      = 5
	wkbMultiLineStringZ     = 1005
	wkbMultiLineStringM     = 2005
	wkbMultiLineStringZM    = 3005
	wkbMultiPolygon         = 6
	wkbMultiPolygonZ        = 1006
	wkbMultiPolygonM        = 2006
	wkbMultiPolygonZM       = 3006
	wkbGeometryCollection   = 7
	wkbGeometryCollectionZ  = 1007
	wkbGeometryCollectionM  = 2007
	wkbGeometryCollectionZM = 3007
	wkbPolyhedralSurface    = 15
	wkbPolyhedralSurfaceZ   = 1015
	wkbPolyhedralSurfaceM   = 2015
	wkbPolyhedralSurfaceZM  = 3015
	wkbTIN                  = 16
	wkbTINZ                 = 1016
	wkbTINM                 = 2016
	wkbTINZM                = 3016
	wkbTriangle             = 17
	wkbTriangleZ            = 1017
	wkbTriangleM            = 2017
	wkbTriangleZM           = 3017
)

var (
	XDR = binary.BigEndian
	NDR = binary.LittleEndian
)

type UnexpectedGeometryError struct {
	Geom geom.T
}

func (e UnexpectedGeometryError) Error() string {
	return fmt.Sprintf("wkb: unexpected geometry %v", e.Geom)
}

type UnsupportedGeometryError struct {
	Type reflect.Type
}

func (e UnsupportedGeometryError) Error() string {
	return "wkb: unsupported type: " + e.Type.String()
}
<|endoftext|>


<commit_before>// Copyright 2013, 2015 Canonical Ltd.
// Copyright 2015 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package containerinit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"path/filepath"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/utils/proxy"
	"github.com/juju/utils/set"

	"github.com/juju/juju/cloudconfig"
	"github.com/juju/juju/cloudconfig/cloudinit"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/container"
	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/service/common"
)

var (
	logger = loggo.GetLogger("juju.cloudconfig.containerinit")
)

// WriteUserData generates the cloud-init user-data using the
// specified machine and network config for a container, and writes
// the serialized form out to a cloud-init file in the directory
// specified.
func WriteUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
	directory string,
) (string, error) {
	userData, err := CloudInitUserData(instanceConfig, networkConfig)
	if err != nil {
		logger.Errorf("failed to create user data: %v", err)
		return "", err
	}
	return WriteCloudInitFile(directory, userData)
}

// WriteCloudInitFile writes the data out to a cloud-init file in the
// directory specified, and returns the filename.
func WriteCloudInitFile(directory string, userData []byte) (string, error) {
	userDataFilename := filepath.Join(directory, "cloud-init")
	if err := ioutil.WriteFile(userDataFilename, userData, 0644); err != nil {
		logger.Errorf("failed to write user data: %v", err)
		return "", err
	}
	return userDataFilename, nil
}

var (
	systemNetworkInterfacesFile = "/etc/network/interfaces"
	networkInterfacesFile       = systemNetworkInterfacesFile + "-juju"
)

// GenerateNetworkConfig renders a network config for one or more network
// interfaces, using the given non-nil networkConfig containing a non-empty
// Interfaces field.
func GenerateNetworkConfig(networkConfig *container.NetworkConfig) (string, error) {
	if networkConfig == nil || len(networkConfig.Interfaces) == 0 {
		return "", errors.Errorf("missing container network config")
	}
	logger.Debugf("generating network config from %#v", *networkConfig)

	prepared := PrepareNetworkConfigFromInterfaces(networkConfig.Interfaces)

	var output bytes.Buffer
	gatewayWritten := false
	for _, name := range prepared.InterfaceNames {
		output.WriteString("\n")
		if name == "lo" {
			output.WriteString("auto ")
			autoStarted := strings.Join(prepared.AutoStarted, " ")
			output.WriteString(autoStarted + "\n\n")
			output.WriteString("iface lo inet loopback\n")

			dnsServers := strings.Join(prepared.DNSServers, " ")
			if dnsServers != "" {
				output.WriteString("  dns-nameservers ")
				output.WriteString(dnsServers + "\n")
			}

			dnsSearchDomains := strings.Join(prepared.DNSSearchDomains, " ")
			if dnsSearchDomains != "" {
				output.WriteString("  dns-search ")
				output.WriteString(dnsSearchDomains + "\n")
			}
			continue
		}

		address, hasAddress := prepared.NameToAddress[name]
		if !hasAddress {
			output.WriteString("iface " + name + " inet manual\n")
			continue
		} else if address == string(network.ConfigDHCP) {
			output.WriteString("iface " + name + " inet dhcp\n")
			// We're expecting to get a default gateway
			// from the DHCP lease.
			gatewayWritten = true
			continue
		}

		output.WriteString("iface " + name + " inet static\n")
		output.WriteString("  address " + address + "\n")
		if !gatewayWritten && prepared.GatewayAddress != "" {
			_, network, err := net.ParseCIDR(address)

			if err != nil {
				return "", errors.Trace(err)
			}

			gatewayIP := net.ParseIP(prepared.GatewayAddress)
			if network.Contains(gatewayIP) {
				output.WriteString("  gateway " + prepared.GatewayAddress + "\n")
				gatewayWritten = true // write it only once
			}
		}
	}

	generatedConfig := output.String()
	logger.Debugf("generated network config:\n%s", generatedConfig)

	if !gatewayWritten {
		logger.Infof("generated network config has no gateway")
	}

	return generatedConfig, nil
}

// PreparedConfig holds all the necessary information to render a persistent
// network config to a file.
type PreparedConfig struct {
	InterfaceNames   []string
	AutoStarted      []string
	DNSServers       []string
	DNSSearchDomains []string
	NameToAddress    map[string]string
	GatewayAddress   string
}

// PrepareNetworkConfigFromInterfaces collects the necessary information to
// render a persistent network config from the given slice of
// network.InterfaceInfo. The result always includes the loopback interface.
func PrepareNetworkConfigFromInterfaces(interfaces []network.InterfaceInfo) *PreparedConfig {
	dnsServers := set.NewStrings()
	dnsSearchDomains := set.NewStrings()
	gatewayAddress := ""
	namesInOrder := make([]string, 1, len(interfaces)+1)
	nameToAddress := make(map[string]string)

	// Always include the loopback.
	namesInOrder[0] = "lo"
	autoStarted := set.NewStrings("lo")

	for _, info := range interfaces {
		if !info.NoAutoStart {
			autoStarted.Add(info.InterfaceName)
		}

		if cidr := info.CIDRAddress(); cidr != "" {
			nameToAddress[info.InterfaceName] = cidr
		} else if info.ConfigType == network.ConfigDHCP {
			nameToAddress[info.InterfaceName] = string(network.ConfigDHCP)
		}

		for _, dns := range info.DNSServers {
			dnsServers.Add(dns.Value)
		}

		dnsSearchDomains = dnsSearchDomains.Union(set.NewStrings(info.DNSSearchDomains...))

		if gatewayAddress == "" && info.GatewayAddress.Value != "" {
			gatewayAddress = info.GatewayAddress.Value
		}

		namesInOrder = append(namesInOrder, info.InterfaceName)
	}

	prepared := &PreparedConfig{
		InterfaceNames:   namesInOrder,
		NameToAddress:    nameToAddress,
		AutoStarted:      autoStarted.SortedValues(),
		DNSServers:       dnsServers.SortedValues(),
		DNSSearchDomains: dnsSearchDomains.SortedValues(),
		GatewayAddress:   gatewayAddress,
	}

	logger.Debugf("prepared network config for rendering: %+v", prepared)
	return prepared
}

// newCloudInitConfigWithNetworks creates a cloud-init config which
// might include per-interface networking config if both networkConfig
// is not nil and its Interfaces field is not empty.
func newCloudInitConfigWithNetworks(series string, networkConfig *container.NetworkConfig) (cloudinit.CloudConfig, error) {
	config, err := GenerateNetworkConfig(networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cloudConfig, err := cloudinit.New(series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cloudConfig.AddBootTextFile(networkInterfacesFile, config, 0644)
	cloudConfig.AddRunCmd(raiseJujuNetworkInterfacesScript(systemNetworkInterfacesFile, networkInterfacesFile))

	return cloudConfig, nil
}

func CloudInitUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	cloudConfig, err := newCloudInitConfigWithNetworks(instanceConfig.Series, networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	udata, err := cloudconfig.NewUserdataConfig(instanceConfig, cloudConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if err = udata.Configure(); err != nil {
		return nil, errors.Trace(err)
	}
	// Run ifconfig to get the addresses of the internal container at least
	// logged in the host.
	cloudConfig.AddRunCmd("ifconfig")

	if instanceConfig.MachineContainerHostname != "" {
		cloudConfig.SetAttr("hostname", instanceConfig.MachineContainerHostname)
	}

	data, err := cloudConfig.RenderYAML()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return data, nil
}

// TemplateUserData returns a minimal user data necessary for the template.
// This should have the authorized keys, base packages, the cloud archive if
// necessary,  initial apt proxy config, and it should do the apt-get
// update/upgrade initially.
func TemplateUserData(
	series string,
	authorizedKeys string,
	aptProxy proxy.Settings,
	aptMirror string,
	enablePackageUpdates bool,
	enableOSUpgrades bool,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	var config cloudinit.CloudConfig
	var err error
	if networkConfig != nil {
		config, err = newCloudInitConfigWithNetworks(series, networkConfig)
		if err != nil {
			return nil, errors.Trace(err)
		}
	} else {
		config, err = cloudinit.New(series)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	cloudconfig.SetUbuntuUser(config, authorizedKeys)
	config.AddScripts(
		"set -xe", // ensure we run all the scripts or abort.
	)
	// For LTS series which need support for the cloud-tools archive,
	// we need to enable apt-get update regardless of the environ
	// setting, otherwise provisioning will fail.
	if series == "precise" && !enablePackageUpdates {
		logger.Infof("series %q requires cloud-tools archive: enabling updates", series)
		enablePackageUpdates = true
	}

	if enablePackageUpdates && config.RequiresCloudArchiveCloudTools() {
		config.AddCloudArchiveCloudTools()
	}
	config.AddPackageCommands(aptProxy, aptMirror, enablePackageUpdates, enableOSUpgrades)

	initSystem, err := service.VersionInitSystem(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds, err := shutdownInitCommands(initSystem, series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config.AddScripts(strings.Join(cmds, "\n"))

	data, err := config.RenderYAML()
	if err != nil {
		return nil, err
	}
	return data, nil
}

func shutdownInitCommands(initSystem, series string) ([]string, error) {
	shutdownCmd := "/sbin/shutdown -h now"
	name := "juju-template-restart"
	desc := "juju shutdown job"

	execStart := shutdownCmd

	conf := common.Conf{
		Desc:         desc,
		Transient:    true,
		AfterStopped: "cloud-final",
		ExecStart:    execStart,
	}
	// systemd uses targets for synchronization of services
	if initSystem == service.InitSystemSystemd {
		conf.AfterStopped = "cloud-config.target"
	}

	svc, err := service.NewService(name, conf, series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cmds, err := svc.InstallCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}

	startCommands, err := svc.StartCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds = append(cmds, startCommands...)

	return cmds, nil
}

// raiseJujuNetworkInterfacesScript returns a cloud-init script to
// raise Juju's network interfaces supplied via cloud-init.
//
// Note: we sleep to mitigate against LP #1337873 and LP #1269921.
func raiseJujuNetworkInterfacesScript(oldInterfacesFile, newInterfacesFile string) string {
	return fmt.Sprintf(`
if [ -f %[2]s ]; then
    ifdown -a
    sleep 1.5
    if ifup -a --interfaces=%[2]s; then
        cp %[1]s %[1]s-orig
        cp %[2]s %[1]s
    else
        ifup -a
    fi
fi`[1:],
		oldInterfacesFile, newInterfacesFile)
}
<commit_msg>Couple of minor cleanups<commit_after>// Copyright 2013, 2015 Canonical Ltd.
// Copyright 2015 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package containerinit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"path/filepath"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/utils/proxy"
	"github.com/juju/utils/set"

	"github.com/juju/juju/cloudconfig"
	"github.com/juju/juju/cloudconfig/cloudinit"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/container"
	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/service/common"
)

var (
	logger = loggo.GetLogger("juju.cloudconfig.containerinit")
)

// WriteUserData generates the cloud-init user-data using the
// specified machine and network config for a container, and writes
// the serialized form out to a cloud-init file in the directory
// specified.
func WriteUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
	directory string,
) (string, error) {
	userData, err := CloudInitUserData(instanceConfig, networkConfig)
	if err != nil {
		logger.Errorf("failed to create user data: %v", err)
		return "", err
	}
	return WriteCloudInitFile(directory, userData)
}

// WriteCloudInitFile writes the data out to a cloud-init file in the
// directory specified, and returns the filename.
func WriteCloudInitFile(directory string, userData []byte) (string, error) {
	userDataFilename := filepath.Join(directory, "cloud-init")
	if err := ioutil.WriteFile(userDataFilename, userData, 0644); err != nil {
		logger.Errorf("failed to write user data: %v", err)
		return "", err
	}
	return userDataFilename, nil
}

var (
	systemNetworkInterfacesFile = "/etc/network/interfaces"
	networkInterfacesFile       = systemNetworkInterfacesFile + "-juju"
)

// GenerateNetworkConfig renders a network config for one or more network
// interfaces, using the given non-nil networkConfig containing a non-empty
// Interfaces field.
func GenerateNetworkConfig(networkConfig *container.NetworkConfig) (string, error) {
	if networkConfig == nil || len(networkConfig.Interfaces) == 0 {
		return "", errors.Errorf("missing container network config")
	}
	logger.Debugf("generating network config from %#v", *networkConfig)

	prepared := PrepareNetworkConfigFromInterfaces(networkConfig.Interfaces)

	var output bytes.Buffer
	gatewayWritten := false
	for _, name := range prepared.InterfaceNames {
		output.WriteString("\n")
		if name == "lo" {
			output.WriteString("auto ")
			autoStarted := strings.Join(prepared.AutoStarted, " ")
			output.WriteString(autoStarted + "\n\n")
			output.WriteString("iface lo inet loopback\n")

			dnsServers := strings.Join(prepared.DNSServers, " ")
			if dnsServers != "" {
				output.WriteString("  dns-nameservers ")
				output.WriteString(dnsServers + "\n")
			}

			dnsSearchDomains := strings.Join(prepared.DNSSearchDomains, " ")
			if dnsSearchDomains != "" {
				output.WriteString("  dns-search ")
				output.WriteString(dnsSearchDomains + "\n")
			}
			continue
		}

		address, hasAddress := prepared.NameToAddress[name]
		if !hasAddress {
			output.WriteString("iface " + name + " inet manual\n")
			continue
		} else if address == string(network.ConfigDHCP) {
			output.WriteString("iface " + name + " inet dhcp\n")
			// We're expecting to get a default gateway
			// from the DHCP lease.
			gatewayWritten = true
			continue
		}

		output.WriteString("iface " + name + " inet static\n")
		output.WriteString("  address " + address + "\n")
		if !gatewayWritten && prepared.GatewayAddress != "" {
			_, network, err := net.ParseCIDR(address)
			if err != nil {
				return "", errors.Annotatef(err, "invalid gateway for interface %q with address %q", name, address)
			}

			gatewayIP := net.ParseIP(prepared.GatewayAddress)
			if network.Contains(gatewayIP) {
				output.WriteString("  gateway " + prepared.GatewayAddress + "\n")
				gatewayWritten = true // write it only once
			}
		}
	}

	generatedConfig := output.String()
	logger.Debugf("generated network config:\n%s", generatedConfig)

	if !gatewayWritten {
		logger.Infof("generated network config has no gateway")
	}

	return generatedConfig, nil
}

// PreparedConfig holds all the necessary information to render a persistent
// network config to a file.
type PreparedConfig struct {
	InterfaceNames   []string
	AutoStarted      []string
	DNSServers       []string
	DNSSearchDomains []string
	NameToAddress    map[string]string
	GatewayAddress   string
}

// PrepareNetworkConfigFromInterfaces collects the necessary information to
// render a persistent network config from the given slice of
// network.InterfaceInfo. The result always includes the loopback interface.
func PrepareNetworkConfigFromInterfaces(interfaces []network.InterfaceInfo) *PreparedConfig {
	dnsServers := set.NewStrings()
	dnsSearchDomains := set.NewStrings()
	gatewayAddress := ""
	namesInOrder := make([]string, 1, len(interfaces)+1)
	nameToAddress := make(map[string]string)

	// Always include the loopback.
	namesInOrder[0] = "lo"
	autoStarted := set.NewStrings("lo")

	for _, info := range interfaces {
		if !info.NoAutoStart {
			autoStarted.Add(info.InterfaceName)
		}

		if cidr := info.CIDRAddress(); cidr != "" {
			nameToAddress[info.InterfaceName] = cidr
		} else if info.ConfigType == network.ConfigDHCP {
			nameToAddress[info.InterfaceName] = string(network.ConfigDHCP)
		}

		for _, dns := range info.DNSServers {
			dnsServers.Add(dns.Value)
		}

		dnsSearchDomains = dnsSearchDomains.Union(set.NewStrings(info.DNSSearchDomains...))

		if gatewayAddress == "" && info.GatewayAddress.Value != "" {
			gatewayAddress = info.GatewayAddress.Value
		}

		namesInOrder = append(namesInOrder, info.InterfaceName)
	}

	prepared := &PreparedConfig{
		InterfaceNames:   namesInOrder,
		NameToAddress:    nameToAddress,
		AutoStarted:      autoStarted.SortedValues(),
		DNSServers:       dnsServers.SortedValues(),
		DNSSearchDomains: dnsSearchDomains.SortedValues(),
		GatewayAddress:   gatewayAddress,
	}

	logger.Debugf("prepared network config for rendering: %+v", prepared)
	return prepared
}

// newCloudInitConfigWithNetworks creates a cloud-init config which
// might include per-interface networking config if both networkConfig
// is not nil and its Interfaces field is not empty.
func newCloudInitConfigWithNetworks(series string, networkConfig *container.NetworkConfig) (cloudinit.CloudConfig, error) {
	config, err := GenerateNetworkConfig(networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cloudConfig, err := cloudinit.New(series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cloudConfig.AddBootTextFile(networkInterfacesFile, config, 0644)
	cloudConfig.AddRunCmd(raiseJujuNetworkInterfacesScript(systemNetworkInterfacesFile, networkInterfacesFile))

	return cloudConfig, nil
}

func CloudInitUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	cloudConfig, err := newCloudInitConfigWithNetworks(instanceConfig.Series, networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	udata, err := cloudconfig.NewUserdataConfig(instanceConfig, cloudConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if err = udata.Configure(); err != nil {
		return nil, errors.Trace(err)
	}
	// Run ifconfig to get the addresses of the internal container at least
	// logged in the host.
	cloudConfig.AddRunCmd("ifconfig")

	if instanceConfig.MachineContainerHostname != "" {
		cloudConfig.SetAttr("hostname", instanceConfig.MachineContainerHostname)
	}

	data, err := cloudConfig.RenderYAML()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return data, nil
}

// TemplateUserData returns a minimal user data necessary for the template.
// This should have the authorized keys, base packages, the cloud archive if
// necessary,  initial apt proxy config, and it should do the apt-get
// update/upgrade initially.
func TemplateUserData(
	series string,
	authorizedKeys string,
	aptProxy proxy.Settings,
	aptMirror string,
	enablePackageUpdates bool,
	enableOSUpgrades bool,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	var config cloudinit.CloudConfig
	var err error
	if networkConfig != nil {
		config, err = newCloudInitConfigWithNetworks(series, networkConfig)
		if err != nil {
			return nil, errors.Trace(err)
		}
	} else {
		config, err = cloudinit.New(series)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	cloudconfig.SetUbuntuUser(config, authorizedKeys)
	config.AddScripts(
		"set -xe", // ensure we run all the scripts or abort.
	)
	// For LTS series which need support for the cloud-tools archive,
	// we need to enable apt-get update regardless of the environ
	// setting, otherwise provisioning will fail.
	if series == "precise" && !enablePackageUpdates {
		logger.Infof("series %q requires cloud-tools archive: enabling updates", series)
		enablePackageUpdates = true
	}

	if enablePackageUpdates && config.RequiresCloudArchiveCloudTools() {
		config.AddCloudArchiveCloudTools()
	}
	config.AddPackageCommands(aptProxy, aptMirror, enablePackageUpdates, enableOSUpgrades)

	initSystem, err := service.VersionInitSystem(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds, err := shutdownInitCommands(initSystem, series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config.AddScripts(strings.Join(cmds, "\n"))

	data, err := config.RenderYAML()
	if err != nil {
		return nil, err
	}
	return data, nil
}

func shutdownInitCommands(initSystem, series string) ([]string, error) {
	shutdownCmd := "/sbin/shutdown -h now"
	name := "juju-template-restart"
	desc := "juju shutdown job"

	execStart := shutdownCmd

	conf := common.Conf{
		Desc:         desc,
		Transient:    true,
		AfterStopped: "cloud-final",
		ExecStart:    execStart,
	}
	// systemd uses targets for synchronization of services
	if initSystem == service.InitSystemSystemd {
		conf.AfterStopped = "cloud-config.target"
	}

	svc, err := service.NewService(name, conf, series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cmds, err := svc.InstallCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}

	startCommands, err := svc.StartCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds = append(cmds, startCommands...)

	return cmds, nil
}

// raiseJujuNetworkInterfacesScript returns a cloud-init script to
// raise Juju's network interfaces supplied via cloud-init.
//
// Note: we sleep to mitigate against LP #1337873 and LP #1269921.
func raiseJujuNetworkInterfacesScript(oldInterfacesFile, newInterfacesFile string) string {
	return fmt.Sprintf(`
if [ -f %[2]s ]; then
    ifdown -a
    sleep 1.5
    if ifup -a --interfaces=%[2]s; then
        cp %[1]s %[1]s-orig
        cp %[2]s %[1]s
    else
        ifup -a
    fi
fi`[1:],
		oldInterfacesFile, newInterfacesFile)
}
<|endoftext|>


<commit_before>package etcdclient

import (
	"fmt"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/etcd/client"
)

// EtcdClient interface lets your Get/Set from Etcd
type EtcdClient interface {
	// Del deletes a key from Etcd
	Del(key string) error

	// DelDir deletes a dir from Etcd
	DelDir(key string) error

	// Get gets a value in Etcd
	Get(key string) (string, error)

	// Set sets a value in Etcd
	Set(key, value string) error

	// UpdateDirWithTTL updates a directory with a ttl value
	UpdateDirWithTTL(key string, ttl time.Duration) error

	// Ls returns all the keys available in the directory
	Ls(directory string) ([]string, error)

	// LsRecursive returns all the keys available in the directory, recursively
	LsRecursive(directory string) ([]string, error)

	// MkDir creates an empty etcd directory
	MkDir(directory string) error

	// WatchRecursive watches a directory and calls the callback everytime something changes.
	// The callback is called with the key of the thing that changed along with the value
	// that the thing was changed to.
	// This method only returns if there is an error
	WatchRecursive(directory string, onChangeCallback OnChangeCallback) error
}

// OnChangeCallback is used for passing callbacks to
// WatchRecursive
type OnChangeCallback func(key, newValue string)

// SimpleEtcdClient implements EtcdClient
type SimpleEtcdClient struct {
	etcd client.Client
}

// Dial constructs a new EtcdClient
func Dial(etcdURI string) (EtcdClient, error) {
	etcd, err := client.New(client.Config{
		Endpoints: []string{etcdURI},
	})
	if err != nil {
		return nil, err
	}
	return &SimpleEtcdClient{etcd}, nil
}

// Del deletes a key from Etcd
func (etcdClient *SimpleEtcdClient) Del(key string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Delete(context.Background(), key, nil)
	if err != nil {
		if client.IsKeyNotFound(err) {
			return nil
		}
	}
	return err
}

// DelDir deletes a dir from Etcd
func (etcdClient *SimpleEtcdClient) DelDir(key string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Delete(context.Background(), key, &client.DeleteOptions{Dir: true, Recursive: true})
	if err != nil {
		if client.IsKeyNotFound(err) {
			return nil
		}
	}
	return err
}

// Get gets a value in Etcd
func (etcdClient *SimpleEtcdClient) Get(key string) (string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	response, err := api.Get(context.Background(), key, nil)
	if err != nil {
		if client.IsKeyNotFound(err) {
			return "", nil
		}
		return "", err
	}
	return response.Node.Value, nil
}

// Set sets a value in Etcd
func (etcdClient *SimpleEtcdClient) Set(key, value string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Set(context.Background(), key, value, nil)
	return err
}

// UpdateDirWithTTL updates a directory with a ttl value
func (etcdClient *SimpleEtcdClient) UpdateDirWithTTL(key string, ttl time.Duration) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Set(context.Background(), key, "", &client.SetOptions{TTL: ttl, Dir: true, PrevExist: client.PrevExist})
	return err
}

// Ls returns all the keys available in the directory
func (etcdClient *SimpleEtcdClient) Ls(directory string) ([]string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	options := &client.GetOptions{Sort: true, Recursive: false}
	response, err := api.Get(context.Background(), directory, options)

	if err != nil {
		if client.IsKeyNotFound(err) {
			return make([]string, 0), nil
		}
		return make([]string, 0), err
	}

	return nodesToStringSlice(response.Node.Nodes), nil
}

// LsRecursive returns all the keys available in the directory, recursively
func (etcdClient *SimpleEtcdClient) LsRecursive(directory string) ([]string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	options := &client.GetOptions{Sort: true, Recursive: true}
	response, err := api.Get(context.Background(), directory, options)

	if err != nil {
		if client.IsKeyNotFound(err) {
			return make([]string, 0), nil
		}
		return make([]string, 0), err
	}

	return nodesToStringSlice(response.Node.Nodes), nil
}

// MkDir creates an empty etcd directory
func (etcdClient *SimpleEtcdClient) MkDir(directory string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	results, err := api.Get(context.Background(), directory, nil)

	if err != nil && !client.IsKeyNotFound(err) {
		return err
	}

	if err != nil && client.IsKeyNotFound(err) {
		_, err = api.Set(context.Background(), directory, "", &client.SetOptions{Dir: true, PrevExist: client.PrevIgnore})
		return err
	}

	if !results.Node.Dir {
		return fmt.Errorf("Refusing to overwrite key/value with a directory: %v", directory)
	}
	return nil
}

// WatchRecursive watches a directory and calls the callback everytime something changes.
// The callback is called with the key of the thing that changed along with the value
// that the thing was changed to.
// This method only returns if there is an error
func (etcdClient *SimpleEtcdClient) WatchRecursive(directory string, onChange OnChangeCallback) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	afterIndex := uint64(0)

	for {
		watcher := api.Watcher(directory, &client.WatcherOptions{Recursive: true, AfterIndex: afterIndex})
		response, err := watcher.Next(context.Background())
		if err != nil {
			return err
		}

		afterIndex = response.Index
		onChange(response.Node.Key, response.Node.Value)
	}
}

func nodesToStringSlice(nodes client.Nodes) []string {
	var keys []string

	for _, node := range nodes {
		keys = append(keys, node.Key)

		for _, key := range nodesToStringSlice(node.Nodes) {
			keys = append(keys, key)
		}
	}

	return keys
}
<commit_msg>Recover from etcd 401 errors, v2.7.1<commit_after>package etcdclient

import (
	"fmt"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/coreos/etcd/client"
)

// EtcdClient interface lets your Get/Set from Etcd
type EtcdClient interface {
	// Del deletes a key from Etcd
	Del(key string) error

	// DelDir deletes a dir from Etcd
	DelDir(key string) error

	// Get gets a value in Etcd
	Get(key string) (string, error)

	// Set sets a value in Etcd
	Set(key, value string) error

	// UpdateDirWithTTL updates a directory with a ttl value
	UpdateDirWithTTL(key string, ttl time.Duration) error

	// Ls returns all the keys available in the directory
	Ls(directory string) ([]string, error)

	// LsRecursive returns all the keys available in the directory, recursively
	LsRecursive(directory string) ([]string, error)

	// MkDir creates an empty etcd directory
	MkDir(directory string) error

	// WatchRecursive watches a directory and calls the callback everytime something changes.
	// The callback is called with the key of the thing that changed along with the value
	// that the thing was changed to.
	// This method only returns if there is an error
	WatchRecursive(directory string, onChangeCallback OnChangeCallback) error
}

// OnChangeCallback is used for passing callbacks to
// WatchRecursive
type OnChangeCallback func(key, newValue string)

// SimpleEtcdClient implements EtcdClient
type SimpleEtcdClient struct {
	etcd client.Client
}

// Dial constructs a new EtcdClient
func Dial(etcdURI string) (EtcdClient, error) {
	etcd, err := client.New(client.Config{
		Endpoints: []string{etcdURI},
	})
	if err != nil {
		return nil, err
	}
	return &SimpleEtcdClient{etcd}, nil
}

// Del deletes a key from Etcd
func (etcdClient *SimpleEtcdClient) Del(key string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Delete(context.Background(), key, nil)
	if err != nil {
		if client.IsKeyNotFound(err) {
			return nil
		}
	}
	return err
}

// DelDir deletes a dir from Etcd
func (etcdClient *SimpleEtcdClient) DelDir(key string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Delete(context.Background(), key, &client.DeleteOptions{Dir: true, Recursive: true})
	if err != nil {
		if client.IsKeyNotFound(err) {
			return nil
		}
	}
	return err
}

// Get gets a value in Etcd
func (etcdClient *SimpleEtcdClient) Get(key string) (string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	response, err := api.Get(context.Background(), key, nil)
	if err != nil {
		if client.IsKeyNotFound(err) {
			return "", nil
		}
		return "", err
	}
	return response.Node.Value, nil
}

// Set sets a value in Etcd
func (etcdClient *SimpleEtcdClient) Set(key, value string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Set(context.Background(), key, value, nil)
	return err
}

// UpdateDirWithTTL updates a directory with a ttl value
func (etcdClient *SimpleEtcdClient) UpdateDirWithTTL(key string, ttl time.Duration) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	_, err := api.Set(context.Background(), key, "", &client.SetOptions{TTL: ttl, Dir: true, PrevExist: client.PrevExist})
	return err
}

// Ls returns all the keys available in the directory
func (etcdClient *SimpleEtcdClient) Ls(directory string) ([]string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	options := &client.GetOptions{Sort: true, Recursive: false}
	response, err := api.Get(context.Background(), directory, options)

	if err != nil {
		if client.IsKeyNotFound(err) {
			return make([]string, 0), nil
		}
		return make([]string, 0), err
	}

	return nodesToStringSlice(response.Node.Nodes), nil
}

// LsRecursive returns all the keys available in the directory, recursively
func (etcdClient *SimpleEtcdClient) LsRecursive(directory string) ([]string, error) {
	api := client.NewKeysAPI(etcdClient.etcd)
	options := &client.GetOptions{Sort: true, Recursive: true}
	response, err := api.Get(context.Background(), directory, options)

	if err != nil {
		if client.IsKeyNotFound(err) {
			return make([]string, 0), nil
		}
		return make([]string, 0), err
	}

	return nodesToStringSlice(response.Node.Nodes), nil
}

// MkDir creates an empty etcd directory
func (etcdClient *SimpleEtcdClient) MkDir(directory string) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	results, err := api.Get(context.Background(), directory, nil)

	if err != nil && !client.IsKeyNotFound(err) {
		return err
	}

	if err != nil && client.IsKeyNotFound(err) {
		_, err = api.Set(context.Background(), directory, "", &client.SetOptions{Dir: true, PrevExist: client.PrevIgnore})
		return err
	}

	if !results.Node.Dir {
		return fmt.Errorf("Refusing to overwrite key/value with a directory: %v", directory)
	}
	return nil
}

// WatchRecursive watches a directory and calls the callback everytime something changes.
// The callback is called with the key of the thing that changed along with the value
// that the thing was changed to.
// This method only returns if there is an error
func (etcdClient *SimpleEtcdClient) WatchRecursive(directory string, onChange OnChangeCallback) error {
	api := client.NewKeysAPI(etcdClient.etcd)
	afterIndex := uint64(0)

	for {
		watcher := api.Watcher(directory, &client.WatcherOptions{Recursive: true, AfterIndex: afterIndex})
		response, err := watcher.Next(context.Background())
		if err != nil {
			if err.(client.Error).Code == client.ErrorCodeEventIndexCleared {
				continue
			}
			return err
		}

		afterIndex = response.Index
		onChange(response.Node.Key, response.Node.Value)
	}
}

func nodesToStringSlice(nodes client.Nodes) []string {
	var keys []string

	for _, node := range nodes {
		keys = append(keys, node.Key)

		for _, key := range nodesToStringSlice(node.Nodes) {
			keys = append(keys, key)
		}
	}

	return keys
}
<|endoftext|>


<commit_before>package raft

import (
	"fmt"
	"github.com/armon/go-metrics"
	"net"
	"time"
)

const (
	maxFailureScale = 12
	failureWait     = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm     uint64
	matchIndex      uint64
	nextIndex       uint64
	lastCommitIndex uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
	var start time.Time
START:
	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Skip the RPC call if there is nothing to do
	if len(req.Entries) == 0 && req.LeaderCommitIndex == s.lastCommitIndex {
		return
	}

	// Make the RPC call
	start = time.Now()
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		select {
		case <-time.After(backoff(failureWait, s.failures, maxFailureScale)):
		case <-r.shutdownCh:
		}
		return
	}
	s.failures = 0
	metrics.MeasureSince([]string{"raft", "replication", "appendEntries", "rpc", s.peer.String()}, start)
	metrics.IncrCounter([]string{"raft", "replication", "appendEntries", "logs", s.peer.String()}, float32(len(req.Entries)))

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
		s.lastCommitIndex = req.LeaderCommitIndex
	} else {
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs (next: %d)", s.peer, s.nextIndex)
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	start := time.Now()
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		s.failures++
		select {
		case <-time.After(backoff(failureWait, s.failures, maxFailureScale)):
		case <-r.shutdownCh:
		}
		return false, err
	}
	s.failures = 0
	metrics.MeasureSince([]string{"raft", "replication", "installSnapshot", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	var failures uint64
	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout(r.conf.HeartbeatTimeout / 8):
			start := time.Now()
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				failures++
				select {
				case <-time.After(backoff(failureWait, failures, maxFailureScale)):
				case <-stopCh:
				}
			} else {
				s.lastContact = time.Now()
				failures = 0
				metrics.MeasureSince([]string{"raft", "replication", "heartbeat", s.peer.String()}, start)
			}

		case <-stopCh:
			return
		}
	}
}
<commit_msg>Improve the exponential backoff of replication<commit_after>package raft

import (
	"fmt"
	"github.com/armon/go-metrics"
	"net"
	"time"
)

const (
	maxFailureScale = 12
	failureWait     = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm     uint64
	matchIndex      uint64
	nextIndex       uint64
	lastCommitIndex uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
	var start time.Time
START:
	// Prevent an excessive retry rate on errors
	if s.failures > 0 {
		select {
		case <-time.After(backoff(failureWait, s.failures, maxFailureScale)):
		case <-r.shutdownCh:
		}
	}

	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Skip the RPC call if there is nothing to do
	if len(req.Entries) == 0 && req.LeaderCommitIndex == s.lastCommitIndex {
		return
	}

	// Make the RPC call
	start = time.Now()
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		return
	}
	metrics.MeasureSince([]string{"raft", "replication", "appendEntries", "rpc", s.peer.String()}, start)
	metrics.IncrCounter([]string{"raft", "replication", "appendEntries", "logs", s.peer.String()}, float32(len(req.Entries)))

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
		s.lastCommitIndex = req.LeaderCommitIndex

		// Clear any failures
		s.failures = 0
	} else {
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
		s.failures++
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs (next: %d)", s.peer, s.nextIndex)
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	start := time.Now()
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		s.failures++
		return false, err
	}
	metrics.MeasureSince([]string{"raft", "replication", "installSnapshot", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1

		// Clear any failures
		s.failures = 0
	} else {
		s.failures++
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	var failures uint64
	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout(r.conf.HeartbeatTimeout / 8):
			start := time.Now()
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				failures++
				select {
				case <-time.After(backoff(failureWait, failures, maxFailureScale)):
				case <-stopCh:
				}
			} else {
				s.lastContact = time.Now()
				failures = 0
				metrics.MeasureSince([]string{"raft", "replication", "heartbeat", s.peer.String()}, start)
			}

		case <-stopCh:
			return
		}
	}
}
<|endoftext|>