<commit_before>package main

import (
	"database/sql"
	"fmt"
	"image"
	_ "image/jpeg"
	"log"
	"mime"
	"os"
	"path/filepath"

	"github.com/agorf/goexif/exif"
	_ "github.com/mattn/go-sqlite3"
)

type Photo struct {
	Aperture      sql.NullFloat64
	Camera        sql.NullString
	ExposureComp  sql.NullInt64
	ExposureTime  sql.NullFloat64
	Flash         sql.NullString
	FocalLength   sql.NullFloat64
	FocalLength35 sql.NullInt64
	Folder        string
	Height        int
	ISO           sql.NullInt64
	Lat           sql.NullFloat64
	Lens          sql.NullString
	Lng           sql.NullFloat64
	Path          string
	Size          int64
	Taken         sql.NullString
	Width         int
}

var (
	db                                                             *sql.DB
	selectSetStmt, selectPhotoStmt, insertSetStmt, insertPhotoStmt *sql.Stmt
)

func decodePhotoExif(photo *Photo, x *exif.Exif) {
	taken, err := x.DateTime()
	if err == nil {
		photo.Taken.String = taken.UTC().Format("2006-01-02 15:04:05")
		photo.Taken.Valid = true
	}

	lat, lng, err := x.LatLong()
	if err == nil {
		photo.Lat.Float64 = lat
		photo.Lat.Valid = true
		photo.Lng.Float64 = lng
		photo.Lng.Valid = true
	}

	orientTag, err := x.Get(exif.Orientation)
	if err == nil {
		switch orient, _ := orientTag.Int(0); orient {
		case 5, 6, 7, 8: // rotated
			photo.Width, photo.Height = photo.Height, photo.Width // swap
		}
	}

	camMakeTag, err := x.Get(exif.Make)
	if err == nil {
		photo.Camera.String, _ = camMakeTag.StringVal()
		photo.Camera.Valid = true
	}

	camModelTag, err := x.Get(exif.Model)
	if err == nil {
		cameraModel, _ := camModelTag.StringVal()

		if photo.Camera.Valid {
			photo.Camera.String = fmt.Sprintf("%s %s", photo.Camera.String, cameraModel)
		} else {
			photo.Camera.String = cameraModel
			photo.Camera.Valid = true
		}
	}

	lensMakeTag, err := x.Get(exif.LensMake)
	if err == nil {
		photo.Lens.String, _ = lensMakeTag.StringVal()
		photo.Lens.Valid = true
	}

	lensModelTag, err := x.Get(exif.LensModel)
	if err == nil {
		lensModel, _ := lensModelTag.StringVal()

		if photo.Lens.Valid {
			photo.Lens.String = fmt.Sprintf("%s %s", photo.Lens.String, lensModel)
		} else {
			photo.Lens.String = lensModel
			photo.Lens.Valid = true
		}
	}

	focalLenTag, err := x.Get(exif.FocalLength)
	if err == nil {
		focalLen, _ := focalLenTag.Rat(0)
		photo.FocalLength.Float64, _ = focalLen.Float64()
		photo.FocalLength.Valid = true
	}

	focalLen35Tag, err := x.Get(exif.FocalLengthIn35mmFilm)
	if err == nil {
		photo.FocalLength35.Int64, _ = focalLen35Tag.Int64(0)
		photo.FocalLength35.Valid = true
	}

	apertureTag, err := x.Get(exif.FNumber)
	if err == nil {
		aperture, _ := apertureTag.Rat(0)
		photo.Aperture.Float64, _ = aperture.Float64()
		photo.Aperture.Valid = true
	}

	expTimeTag, err := x.Get(exif.ExposureTime)
	if err == nil {
		expTime, _ := expTimeTag.Rat(0)
		photo.ExposureTime.Float64, _ = expTime.Float64()
		photo.ExposureTime.Valid = true
	}

	isoTag, err := x.Get(exif.ISOSpeedRatings)
	if err == nil {
		photo.ISO.Int64, _ = isoTag.Int64(0)
		photo.ISO.Valid = true
	}

	expBiasTag, err := x.Get(exif.ExposureBiasValue)
	if err == nil {
		photo.ExposureComp.Int64, _ = expBiasTag.Int64(0)
		photo.ExposureComp.Valid = true
	}

	flash, err := x.Flash()
	if err == nil {
		photo.Flash.String = flash
		photo.Flash.Valid = true
	}
}

func decodePhoto(path string) (*Photo, error) {
	var photo Photo

	photo.Path = path
	photo.Folder = filepath.Base(filepath.Dir(path))

	f, err := os.Open(path)
	if err != nil {
		return &photo, err
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		return &photo, err
	}
	photo.Size = fi.Size()

	img, _, err := image.DecodeConfig(f)
	if err != nil {
		return &photo, err
	}
	photo.Width, photo.Height = img.Width, img.Height

	f.Seek(0, 0) // rewind

	x, err := exif.Decode(f)
	if err == nil { // EXIF data exists
		decodePhotoExif(&photo, x)
	}

	return &photo, nil
}

func storePhoto(photo *Photo) error {
	var setId, photoId int64

	row := selectSetStmt.QueryRow(photo.Folder)
	if err := row.Scan(&setId); err == sql.ErrNoRows { // set does not exist
		result, err := insertSetStmt.Exec(photo.Folder) // create it
		if err != nil {
			return err
		}

		setId, err = result.LastInsertId()
		if err != nil {
			return err
		}
	}

	row = selectPhotoStmt.QueryRow(photo.Path)
	if err := row.Scan(&photoId); err == sql.ErrNoRows { // photo does not exist
		result, err := insertPhotoStmt.Exec(photo.Aperture, photo.Camera,
			photo.ExposureComp, photo.ExposureTime, photo.Flash, photo.FocalLength,
			photo.FocalLength35, photo.Height, photo.ISO, photo.Lat, photo.Lens,
			photo.Lng, photo.Path, setId, photo.Size, photo.Taken, photo.Width) // create it
		if err != nil {
			return err
		}

		photoId, err = result.LastInsertId()
		if err != nil {
			return err
		}

		fmt.Fprintf(os.Stderr, "photos id=%d path=%s\n", photoId, photo.Path)
	}

	return nil
}

func walk(path string, info os.FileInfo, err error) error {
	if err != nil { // error walking "path"
		return nil // skip
	}

	if info.IsDir() {
		return nil // skip
	}

	if mime.TypeByExtension(filepath.Ext(path)) != "image/jpeg" { // not JPEG
		return nil // skip
	}

	photo, err := decodePhoto(path)
	if err == nil {
		storePhoto(photo)
	}

	return nil // next
}

func updatePhotoSiblings() error {
	var prevId, prevSetId int

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	updatePrevPhotoStmt, err := tx.Prepare(`
	UPDATE photos SET prev_photo_id = ? WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updatePrevPhotoStmt.Close()

	updateNextPhotoStmt, err := tx.Prepare(`
	UPDATE photos SET next_photo_id = ? WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updateNextPhotoStmt.Close()

	rows, err := tx.Query(`
	SELECT id, set_id FROM photos ORDER BY set_id, taken_at
	`)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id, setId int
		rows.Scan(&id, &setId)

		if setId == prevSetId && prevId > 0 {
			updatePrevPhotoStmt.Exec(prevId, id)
			fmt.Fprintf(os.Stderr, "photos id=%d prev_photo_id=%d\n", id, prevId)
			updateNextPhotoStmt.Exec(id, prevId)
			fmt.Fprintf(os.Stderr, "photos id=%d next_photo_id=%d\n", prevId, id)
		}

		prevId = id
		prevSetId = setId
	}

	if err := rows.Err(); err != nil {
		return err
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func updateSets() error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	photosCountStmt, err := tx.Prepare(`
	SELECT COUNT(*) FROM photos WHERE set_id = ?
	`)
	if err != nil {
		return err
	}
	defer photosCountStmt.Close()

	updateSetStmt, err := tx.Prepare(`
	UPDATE sets
	SET photos_count = ?, taken_at = ?, thumb_photo_id = ?
	WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updateSetStmt.Close()

	rows, err := tx.Query(`
	SELECT id, set_id, MIN(taken_at) FROM photos GROUP BY set_id
	`)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id, setId, photosCount int
		var taken sql.NullString

		rows.Scan(&id, &setId, &taken)

		row := photosCountStmt.QueryRow(setId)
		row.Scan(&photosCount)

		updateSetStmt.Exec(photosCount, taken, id, setId)
		fmt.Fprintf(os.Stderr, "sets id=%d photos_count=%d taken=%q thumb_photo_id=%d\n", setId, photosCount, taken.String, id)
	}

	if err := rows.Err(); err != nil {
		return err
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func main() {
	var err error

	if len(os.Args) == 1 {
		fmt.Fprintf(os.Stderr, "usage: %v path [path...]\n", os.Args[0])
		return
	}

	db, err = sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	selectSetStmt, err = db.Prepare("SELECT id FROM sets WHERE name = ?")
	if err != nil {
		log.Fatal(err)
	}
	defer selectSetStmt.Close()

	selectPhotoStmt, err = db.Prepare("SELECT id FROM photos WHERE path = ?")
	if err != nil {
		log.Fatal(err)
	}
	defer selectPhotoStmt.Close()

	insertSetStmt, err = db.Prepare("INSERT INTO sets (name) VALUES (?)")
	if err != nil {
		log.Fatal(err)
	}
	defer insertSetStmt.Close()

	insertPhotoStmt, err = db.Prepare(`
	INSERT INTO photos (
	aperture, camera, exposure_comp, exposure_time, flash, focal_length,
	focal_length_35, height, iso, lat, lens, lng, path, set_id, size, taken_at,
	width
	)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`)
	if err != nil {
		log.Fatal(err)
	}
	defer insertPhotoStmt.Close()

	for i := 1; i < len(os.Args); i++ {
		filepath.Walk(os.Args[i], walk)
	}

	updatePhotoSiblings()
	updateSets()
}
<commit_msg>Use fmt.Sprint instead of Sprintf to join strings<commit_after>package main

import (
	"database/sql"
	"fmt"
	"image"
	_ "image/jpeg"
	"log"
	"mime"
	"os"
	"path/filepath"

	"github.com/agorf/goexif/exif"
	_ "github.com/mattn/go-sqlite3"
)

type Photo struct {
	Aperture      sql.NullFloat64
	Camera        sql.NullString
	ExposureComp  sql.NullInt64
	ExposureTime  sql.NullFloat64
	Flash         sql.NullString
	FocalLength   sql.NullFloat64
	FocalLength35 sql.NullInt64
	Folder        string
	Height        int
	ISO           sql.NullInt64
	Lat           sql.NullFloat64
	Lens          sql.NullString
	Lng           sql.NullFloat64
	Path          string
	Size          int64
	Taken         sql.NullString
	Width         int
}

var (
	db                                                             *sql.DB
	selectSetStmt, selectPhotoStmt, insertSetStmt, insertPhotoStmt *sql.Stmt
)

func decodePhotoExif(photo *Photo, x *exif.Exif) {
	taken, err := x.DateTime()
	if err == nil {
		photo.Taken.String = taken.UTC().Format("2006-01-02 15:04:05")
		photo.Taken.Valid = true
	}

	lat, lng, err := x.LatLong()
	if err == nil {
		photo.Lat.Float64 = lat
		photo.Lat.Valid = true
		photo.Lng.Float64 = lng
		photo.Lng.Valid = true
	}

	orientTag, err := x.Get(exif.Orientation)
	if err == nil {
		switch orient, _ := orientTag.Int(0); orient {
		case 5, 6, 7, 8: // rotated
			photo.Width, photo.Height = photo.Height, photo.Width // swap
		}
	}

	camMakeTag, err := x.Get(exif.Make)
	if err == nil {
		photo.Camera.String, _ = camMakeTag.StringVal()
		photo.Camera.Valid = true
	}

	camModelTag, err := x.Get(exif.Model)
	if err == nil {
		cameraModel, _ := camModelTag.StringVal()

		if photo.Camera.Valid {
			photo.Camera.String = fmt.Sprint(photo.Camera.String, " ", cameraModel)
		} else {
			photo.Camera.String = cameraModel
			photo.Camera.Valid = true
		}
	}

	lensMakeTag, err := x.Get(exif.LensMake)
	if err == nil {
		photo.Lens.String, _ = lensMakeTag.StringVal()
		photo.Lens.Valid = true
	}

	lensModelTag, err := x.Get(exif.LensModel)
	if err == nil {
		lensModel, _ := lensModelTag.StringVal()

		if photo.Lens.Valid {
			photo.Lens.String = fmt.Sprint(photo.Lens.String, " ", lensModel)
		} else {
			photo.Lens.String = lensModel
			photo.Lens.Valid = true
		}
	}

	focalLenTag, err := x.Get(exif.FocalLength)
	if err == nil {
		focalLen, _ := focalLenTag.Rat(0)
		photo.FocalLength.Float64, _ = focalLen.Float64()
		photo.FocalLength.Valid = true
	}

	focalLen35Tag, err := x.Get(exif.FocalLengthIn35mmFilm)
	if err == nil {
		photo.FocalLength35.Int64, _ = focalLen35Tag.Int64(0)
		photo.FocalLength35.Valid = true
	}

	apertureTag, err := x.Get(exif.FNumber)
	if err == nil {
		aperture, _ := apertureTag.Rat(0)
		photo.Aperture.Float64, _ = aperture.Float64()
		photo.Aperture.Valid = true
	}

	expTimeTag, err := x.Get(exif.ExposureTime)
	if err == nil {
		expTime, _ := expTimeTag.Rat(0)
		photo.ExposureTime.Float64, _ = expTime.Float64()
		photo.ExposureTime.Valid = true
	}

	isoTag, err := x.Get(exif.ISOSpeedRatings)
	if err == nil {
		photo.ISO.Int64, _ = isoTag.Int64(0)
		photo.ISO.Valid = true
	}

	expBiasTag, err := x.Get(exif.ExposureBiasValue)
	if err == nil {
		photo.ExposureComp.Int64, _ = expBiasTag.Int64(0)
		photo.ExposureComp.Valid = true
	}

	flash, err := x.Flash()
	if err == nil {
		photo.Flash.String = flash
		photo.Flash.Valid = true
	}
}

func decodePhoto(path string) (*Photo, error) {
	var photo Photo

	photo.Path = path
	photo.Folder = filepath.Base(filepath.Dir(path))

	f, err := os.Open(path)
	if err != nil {
		return &photo, err
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		return &photo, err
	}
	photo.Size = fi.Size()

	img, _, err := image.DecodeConfig(f)
	if err != nil {
		return &photo, err
	}
	photo.Width, photo.Height = img.Width, img.Height

	f.Seek(0, 0) // rewind

	x, err := exif.Decode(f)
	if err == nil { // EXIF data exists
		decodePhotoExif(&photo, x)
	}

	return &photo, nil
}

func storePhoto(photo *Photo) error {
	var setId, photoId int64

	row := selectSetStmt.QueryRow(photo.Folder)
	if err := row.Scan(&setId); err == sql.ErrNoRows { // set does not exist
		result, err := insertSetStmt.Exec(photo.Folder) // create it
		if err != nil {
			return err
		}

		setId, err = result.LastInsertId()
		if err != nil {
			return err
		}
	}

	row = selectPhotoStmt.QueryRow(photo.Path)
	if err := row.Scan(&photoId); err == sql.ErrNoRows { // photo does not exist
		result, err := insertPhotoStmt.Exec(photo.Aperture, photo.Camera,
			photo.ExposureComp, photo.ExposureTime, photo.Flash, photo.FocalLength,
			photo.FocalLength35, photo.Height, photo.ISO, photo.Lat, photo.Lens,
			photo.Lng, photo.Path, setId, photo.Size, photo.Taken, photo.Width) // create it
		if err != nil {
			return err
		}

		photoId, err = result.LastInsertId()
		if err != nil {
			return err
		}

		fmt.Fprintf(os.Stderr, "photos id=%d path=%s\n", photoId, photo.Path)
	}

	return nil
}

func walk(path string, info os.FileInfo, err error) error {
	if err != nil { // error walking "path"
		return nil // skip
	}

	if info.IsDir() {
		return nil // skip
	}

	if mime.TypeByExtension(filepath.Ext(path)) != "image/jpeg" { // not JPEG
		return nil // skip
	}

	photo, err := decodePhoto(path)
	if err == nil {
		storePhoto(photo)
	}

	return nil // next
}

func updatePhotoSiblings() error {
	var prevId, prevSetId int

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	updatePrevPhotoStmt, err := tx.Prepare(`
	UPDATE photos SET prev_photo_id = ? WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updatePrevPhotoStmt.Close()

	updateNextPhotoStmt, err := tx.Prepare(`
	UPDATE photos SET next_photo_id = ? WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updateNextPhotoStmt.Close()

	rows, err := tx.Query(`
	SELECT id, set_id FROM photos ORDER BY set_id, taken_at
	`)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id, setId int
		rows.Scan(&id, &setId)

		if setId == prevSetId && prevId > 0 {
			updatePrevPhotoStmt.Exec(prevId, id)
			fmt.Fprintf(os.Stderr, "photos id=%d prev_photo_id=%d\n", id, prevId)
			updateNextPhotoStmt.Exec(id, prevId)
			fmt.Fprintf(os.Stderr, "photos id=%d next_photo_id=%d\n", prevId, id)
		}

		prevId = id
		prevSetId = setId
	}

	if err := rows.Err(); err != nil {
		return err
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func updateSets() error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	photosCountStmt, err := tx.Prepare(`
	SELECT COUNT(*) FROM photos WHERE set_id = ?
	`)
	if err != nil {
		return err
	}
	defer photosCountStmt.Close()

	updateSetStmt, err := tx.Prepare(`
	UPDATE sets
	SET photos_count = ?, taken_at = ?, thumb_photo_id = ?
	WHERE id = ?
	`)
	if err != nil {
		return err
	}
	defer updateSetStmt.Close()

	rows, err := tx.Query(`
	SELECT id, set_id, MIN(taken_at) FROM photos GROUP BY set_id
	`)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id, setId, photosCount int
		var taken sql.NullString

		rows.Scan(&id, &setId, &taken)

		row := photosCountStmt.QueryRow(setId)
		row.Scan(&photosCount)

		updateSetStmt.Exec(photosCount, taken, id, setId)
		fmt.Fprintf(os.Stderr, "sets id=%d photos_count=%d taken=%q thumb_photo_id=%d\n", setId, photosCount, taken.String, id)
	}

	if err := rows.Err(); err != nil {
		return err
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func main() {
	var err error

	if len(os.Args) == 1 {
		fmt.Fprintf(os.Stderr, "usage: %v path [path...]\n", os.Args[0])
		return
	}

	db, err = sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	selectSetStmt, err = db.Prepare("SELECT id FROM sets WHERE name = ?")
	if err != nil {
		log.Fatal(err)
	}
	defer selectSetStmt.Close()

	selectPhotoStmt, err = db.Prepare("SELECT id FROM photos WHERE path = ?")
	if err != nil {
		log.Fatal(err)
	}
	defer selectPhotoStmt.Close()

	insertSetStmt, err = db.Prepare("INSERT INTO sets (name) VALUES (?)")
	if err != nil {
		log.Fatal(err)
	}
	defer insertSetStmt.Close()

	insertPhotoStmt, err = db.Prepare(`
	INSERT INTO photos (
	aperture, camera, exposure_comp, exposure_time, flash, focal_length,
	focal_length_35, height, iso, lat, lens, lng, path, set_id, size, taken_at,
	width
	)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`)
	if err != nil {
		log.Fatal(err)
	}
	defer insertPhotoStmt.Close()

	for i := 1; i < len(os.Args); i++ {
		filepath.Walk(os.Args[i], walk)
	}

	updatePhotoSiblings()
	updateSets()
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/pajlada/pajbot2/common"
)

// Format xD
type Format struct {
	re    *regexp.Regexp
	cmdRe *regexp.Regexp
}

// command is not a good name, but idk what else to call it
type command struct {
	c       string
	subC    []string
	rawCmd  string
	outcome string
}

// InitFormatter compiles format regexes
func (bot *Bot) InitFormatter() *Format {
	return &Format{
		re:    regexp.MustCompile(`\$\([a-z\.]+\)`),
		cmdRe: regexp.MustCompile(`[a-z]+`),
	}
}

// Format formats the given line xD
func (bot *Bot) Format(line string, msg *common.Msg) string {
	// catch all errors until we have proper error handling
	defer func() {
		if r := recover(); r != nil {
			log.Error(r)
		}
	}()
	f := bot.Fmt
	fmtline, rawCommands := f.parseLine(line)
	for i := range rawCommands {
		bot.execCommand(&rawCommands[i], msg)
	}
	return f.formatCommands(fmtline, rawCommands)
}

func (bot *Bot) execCommand(cmd *command, msg *common.Msg) {
	switch cmd.c {
	case "source", "sender":
		cmd.outcome = bot.formatUser(&msg.User, cmd.subC)
	case "user":
		if msg.Args != nil {
			if bot.Redis.IsValidUser(msg.Channel, msg.Args[0]) {
				user := bot.Redis.LoadUser(msg.Channel, msg.Args[0])
				cmd.outcome = bot.formatUser(&user, cmd.subC)
				return
			}
		}
		cmd.outcome = bot.formatUser(&msg.User, cmd.subC)
	}
}

func (bot *Bot) formatUser(user *common.User, cmds []string) string {
	if cmds == nil {
		return user.DisplayName
	}
	switch cmds[0] {
	case "name":
		return user.Name
	case "points":
		return fmt.Sprintf("%d", user.Points)
	case "level":
		return fmt.Sprintf("%d", user.Level)
	default:
		return user.DisplayName
	}
}

func (f *Format) formatCommands(line string, cmds []command) string {
	log.Debug(line)
	log.Debug(cmds)
	for _, c := range cmds {
		line = strings.Replace(line, c.rawCmd, c.outcome, 1)
	}
	log.Debug(line)
	return line
}

func (f *Format) parseLine(line string) (string, []command) {
	log.Debug(line)
	matches := f.re.FindAllString(line, -1)
	log.Debug(matches)
	var cmds []command
	for _, match := range matches {
		cmdlist := f.cmdRe.FindAllString(match, -1)
		c := command{
			c: cmdlist[0],
		}
		if len(cmdlist) > 1 {
			c.subC = cmdlist[1:]
		}
		c.rawCmd = match
		cmds = append(cmds, c)
	}
	log.Debug(cmds, line)
	return line, cmds
}
<commit_msg>added placeholder for line counts<commit_after>package bot

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/pajlada/pajbot2/common"
)

// Format xD
type Format struct {
	re    *regexp.Regexp
	cmdRe *regexp.Regexp
}

// command is not a good name, but idk what else to call it
type command struct {
	c       string
	subC    []string
	rawCmd  string
	outcome string
}

// InitFormatter compiles format regexes
func (bot *Bot) InitFormatter() *Format {
	return &Format{
		re:    regexp.MustCompile(`\$\([a-z\.]+\)`),
		cmdRe: regexp.MustCompile(`[a-z]+`),
	}
}

// Format formats the given line xD
func (bot *Bot) Format(line string, msg *common.Msg) string {
	// catch all errors until we have proper error handling
	defer func() {
		if r := recover(); r != nil {
			log.Error(r)
		}
	}()
	f := bot.Fmt
	fmtline, rawCommands := f.parseLine(line)
	for i := range rawCommands {
		bot.execCommand(&rawCommands[i], msg)
	}
	return f.formatCommands(fmtline, rawCommands)
}

func (bot *Bot) execCommand(cmd *command, msg *common.Msg) {
	switch cmd.c {
	case "source", "sender":
		cmd.outcome = bot.formatUser(&msg.User, cmd.subC)
	case "user":
		if msg.Args != nil {
			if bot.Redis.IsValidUser(msg.Channel, msg.Args[0]) {
				user := bot.Redis.LoadUser(msg.Channel, msg.Args[0])
				cmd.outcome = bot.formatUser(&user, cmd.subC)
				return
			}
		}
		cmd.outcome = bot.formatUser(&msg.User, cmd.subC)
	}
}

func (bot *Bot) formatUser(user *common.User, cmds []string) string {
	if cmds == nil {
		return user.DisplayName
	}
	switch cmds[0] {
	case "name":
		return user.Name
	case "points":
		return fmt.Sprintf("%d", user.Points)
	case "level":
		return fmt.Sprintf("%d", user.Level)
	case "lines":
		return bot.Fmt.lines(user, cmds[1])
	default:
		return user.DisplayName
	}
}

func (f *Format) lines(user *common.User, arg string) string {
	switch arg {
	case "online":
		return fmt.Sprintf("%d", user.OnlineMessageCount)
	case "offline":
		return fmt.Sprintf("%d", user.OfflineMessageCount)
	default:
		return fmt.Sprintf("%d", user.TotalMessageCount)
	}
}

func (f *Format) formatCommands(line string, cmds []command) string {
	log.Debug(line)
	log.Debug(cmds)
	for _, c := range cmds {
		line = strings.Replace(line, c.rawCmd, c.outcome, 1)
	}
	log.Debug(line)
	return line
}

func (f *Format) parseLine(line string) (string, []command) {
	log.Debug(line)
	matches := f.re.FindAllString(line, -1)
	log.Debug(matches)
	var cmds []command
	for _, match := range matches {
		cmdlist := f.cmdRe.FindAllString(match, -1)
		c := command{
			c: cmdlist[0],
		}
		if len(cmdlist) > 1 {
			c.subC = cmdlist[1:]
		}
		c.rawCmd = match
		// lazy fix to avoid out of range error
		c.subC = append(c.subC, "")
		c.subC = append(c.subC, "")
		cmds = append(cmds, c)
	}
	log.Debug(cmds, line)
	return line, cmds
}
<|endoftext|>


<commit_before>package packer

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"

	"github.com/mitchellh/packer/builder/digitalocean"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/provisioner/file"
	"github.com/mitchellh/packer/provisioner/shell"
)

type Provider struct {
	BuildName    string
	TemplatePath string
	Data         []byte
	Vars         map[string]string
	EnableDebug  bool
}

func (p *Provider) Build() error {
	var template *packer.Template
	var err error

	if p.Data != nil {
		template, err = p.NewTemplate()
		if err != nil {
			return err
		}
	} else if p.TemplatePath != "" {
		template, err = p.NewTemplateFile()
		if err != nil {
			return err
		}
	} else {
		return errors.New("Can't find Template source, neither p.Data or p.TemplatePath is defined")
	}

	if len(template.Builders) == 0 {
		return errors.New("No builder is available")
	}

	if len(template.Builders) != 1 {
		return errors.New("Only on builder is supported currently")
	}

	if _, ok := template.Builders[p.BuildName]; !ok {
		return fmt.Errorf("Build '%s' does not exist", p.BuildName)
	}

	build, err := template.Build(p.BuildName, newComponentFinder())
	if err != nil {
		return err
	}

	if !p.EnableDebug {
		build.SetDebug(false)
		log.SetOutput(ioutil.Discard)
	}

	_, err = build.Prepare()
	if err != nil {
		return err
	}

	// Handle interrupts for this build
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)
	defer signal.Stop(sig)
	go func(b packer.Build) {
		<-sig

		log.Printf("Stopping build: %s", b.Name())
		b.Cancel()
		log.Printf("Build cancelled: %s", b.Name())
	}(build)

	artifacts, err := build.Run(p.coloredUi(), p.cache())
	if err != nil {
		return err
	}

	for _, a := range artifacts {
		fmt.Println(a.Files(), a.BuilderId(), a.Id(), a.String())
	}

	return nil
}

func (p *Provider) coloredUi() packer.Ui {
	return &packer.ColoredUi{
		Color: packer.UiColorGreen,
		Ui:    p.basicUI(),
	}
}

func (p *Provider) cache() packer.Cache {
	return &packer.FileCache{
		CacheDir: os.TempDir(),
	}
}

func (p *Provider) basicUI() packer.Ui {
	return &packer.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stdout,
	}
}

func (p *Provider) NewTemplate() (*packer.Template, error) {
	template, err := packer.ParseTemplate(p.Data, p.Vars)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse template: %s", err)
	}

	return template, nil
}

func (p *Provider) NewTemplateFile() (*packer.Template, error) {
	template, err := packer.ParseTemplateFile(p.TemplatePath, p.Vars)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse template: %s", err)
	}

	return template, nil
}

func newComponentFinder() *packer.ComponentFinder {
	return &packer.ComponentFinder{
		Builder:       builderFunc,
		Command:       commandFunc,
		Provisioner:   provisionerFunc,
		Hook:          hookFunc,
		PostProcessor: postProcessorFunc,
	}
}

func builderFunc(name string) (packer.Builder, error) {
	switch name {
	case "digitalocean":
		return &digitalocean.Builder{}, nil
	}

	return nil, errors.New("no suitable build found")
}

func provisionerFunc(name string) (packer.Provisioner, error) {
	switch name {
	case "file":
		return &file.Provisioner{}, nil
	case "shell":
		return &shell.Provisioner{}, nil
	}

	return nil, errors.New("no suitable provisioner found")
}

func hookFunc(name string) (packer.Hook, error) {
	return nil, errors.New("not supported")
}

func postProcessorFunc(name string) (packer.PostProcessor, error) {
	return nil, errors.New("not supported")
}

func commandFunc(name string) (packer.Command, error) {
	return nil, errors.New("not supported")
}
<commit_msg>kloud/packer: add amazon-ebs builder<commit_after>package packer

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"

	"github.com/mitchellh/packer/builder/amazon/ebs"
	"github.com/mitchellh/packer/builder/digitalocean"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/provisioner/file"
	"github.com/mitchellh/packer/provisioner/shell"
)

type Provider struct {
	BuildName    string
	TemplatePath string
	Data         []byte
	Vars         map[string]string
	EnableDebug  bool
}

func (p *Provider) Build() error {
	var template *packer.Template
	var err error

	if p.Data != nil {
		template, err = p.NewTemplate()
		if err != nil {
			return err
		}
	} else if p.TemplatePath != "" {
		template, err = p.NewTemplateFile()
		if err != nil {
			return err
		}
	} else {
		return errors.New("Can't find Template source, neither p.Data or p.TemplatePath is defined")
	}

	if len(template.Builders) == 0 {
		return errors.New("No builder is available")
	}

	if len(template.Builders) != 1 {
		return errors.New("Only on builder is supported currently")
	}

	if _, ok := template.Builders[p.BuildName]; !ok {
		return fmt.Errorf("Build '%s' does not exist", p.BuildName)
	}

	build, err := template.Build(p.BuildName, newComponentFinder())
	if err != nil {
		return err
	}

	if !p.EnableDebug {
		build.SetDebug(false)
		log.SetOutput(ioutil.Discard)
	}

	_, err = build.Prepare()
	if err != nil {
		return err
	}

	// Handle interrupts for this build
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt)
	defer signal.Stop(sig)
	go func(b packer.Build) {
		<-sig

		log.Printf("Stopping build: %s", b.Name())
		b.Cancel()
		log.Printf("Build cancelled: %s", b.Name())
	}(build)

	artifacts, err := build.Run(p.coloredUi(), p.cache())
	if err != nil {
		return err
	}

	for _, a := range artifacts {
		fmt.Println(a.Files(), a.BuilderId(), a.Id(), a.String())
	}

	return nil
}

func (p *Provider) coloredUi() packer.Ui {
	return &packer.ColoredUi{
		Color: packer.UiColorGreen,
		Ui:    p.basicUI(),
	}
}

func (p *Provider) cache() packer.Cache {
	return &packer.FileCache{
		CacheDir: os.TempDir(),
	}
}

func (p *Provider) basicUI() packer.Ui {
	return &packer.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stdout,
	}
}

func (p *Provider) NewTemplate() (*packer.Template, error) {
	template, err := packer.ParseTemplate(p.Data, p.Vars)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse template: %s", err)
	}

	return template, nil
}

func (p *Provider) NewTemplateFile() (*packer.Template, error) {
	template, err := packer.ParseTemplateFile(p.TemplatePath, p.Vars)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse template: %s", err)
	}

	return template, nil
}

func newComponentFinder() *packer.ComponentFinder {
	return &packer.ComponentFinder{
		Builder:       builderFunc,
		Command:       commandFunc,
		Provisioner:   provisionerFunc,
		Hook:          hookFunc,
		PostProcessor: postProcessorFunc,
	}
}

func builderFunc(name string) (packer.Builder, error) {
	switch name {
	case "digitalocean":
		return &digitalocean.Builder{}, nil
	case "amazon-ebs":
		return &ebs.Builder{}, nil
	}

	return nil, errors.New("no suitable build found")
}

func provisionerFunc(name string) (packer.Provisioner, error) {
	switch name {
	case "file":
		return &file.Provisioner{}, nil
	case "shell":
		return &shell.Provisioner{}, nil
	}

	return nil, errors.New("no suitable provisioner found")
}

func hookFunc(name string) (packer.Hook, error) {
	return nil, errors.New("not supported")
}

func postProcessorFunc(name string) (packer.PostProcessor, error) {
	return nil, errors.New("not supported")
}

func commandFunc(name string) (packer.Command, error) {
	return nil, errors.New("not supported")
}
<|endoftext|>


<commit_before>package kontrol

import (
	"errors"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"net"
	"net/url"

	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
)

// GetKodingKitesResult mirrors the Kite GetKitesResult protocol.
type GetKodingKitesResult struct {
	Kites []*KodingKiteWithToken `json:"kites"`
}

// KodingKiteWithToken mirrors the Kite library KiteWithToken with a handful of
// additional fields pertaining to Koding itself. Such as Machine Label, and
// Teams information.
type KodingKiteWithToken struct {
	// These fields match protocol.KiteWithToken
	Kite  protocol.Kite `json:"kite"`
	URL   string        `json:"url"`
	KeyID string        `json:"keyId,omitempty"`
	Token string        `json:"token"`

	// The machine label, as seen from the Koding UI
	MachineLabel string `json:"machineLabel"`

	// The team names, if any.
	Teams []string `json:"teams"`
}

// HandleGetKodingKites implements GetKites with additional Koding specific response
// data.
func HandleGetKodingKites(handleGetKites kite.HandlerFunc) kite.HandlerFunc {
	return func(req *kite.Request) (interface{}, error) {
		// getKitesResponse is an interface, not the actual GetKitesResult value.
		getKitesResponse, err := handleGetKites(req)
		if err != nil {
			req.LocalKite.Log.Error("GetKitesHandler returned an error: %s", err)
			return nil, err
		}

		getKitesResult, ok := getKitesResponse.(*protocol.GetKitesResult)
		if !ok {
			return nil, errors.New("GetKites returned unexpected protocol result")
		}

		// Get all machines for the given user.
		machines, err := modelhelper.GetMachinesByUsername(req.Username)
		if err != nil {
			return nil, err
		}

		// A map of GroupId a slice of all KodingKiteWithToken that contain that group.
		kkwtsByGroupID := map[string][]*KodingKiteWithToken{}

		// Create our slice of unique group ids that we will use to query all the
		// groups that we need.
		uniqueGroups := []string{}

		// Because we are using the Kite's IP as a means to identify the machine,
		// we are sorting the machines by IP so that we can easily look them up.
		machinesByIP := map[string]*models.Machine{}
		for _, m := range machines {
			if m.IpAddress != "" {
				machinesByIP[m.IpAddress] = m
			}
		}

		// Create our KodingKitesResult which will be returned after being populated.
		kodingKitesResult := &GetKodingKitesResult{
			Kites: make([]*KodingKiteWithToken, len(getKitesResult.Kites)),
		}

		// Populate the koding kites
		for i, kwt := range getKitesResult.Kites {
			kkwt := &KodingKiteWithToken{
				Kite:  kwt.Kite,
				URL:   kwt.URL,
				KeyID: kwt.KeyID,
				Token: kwt.Token,
			}
			kodingKitesResult.Kites[i] = kkwt

			host, err := hostFromURL(kwt.URL)
			if err != nil {
				req.LocalKite.Log.Warning(
					"Kite with badly formed URL, unable to extract host. Koding data will not be provided for this Kite. [username:%s, kiteId:%s, host:%s]",
					kwt.Kite.Username,
					kwt.Kite.ID,
					kwt.URL,
				)
				continue
			}

			// With a valid host, get the koding specific information such as Label
			// or Team names and apply it to each KodingKiteWithToken
			if machine, ok := machinesByIP[host]; ok {
				kkwt.MachineLabel = machine.Label

				for _, g := range machine.Groups {
					id := g.Id.Hex()
					// Get all of the kites (not including this one) which have
					// a reference to this kite, so we can append this kite to them.
					kkwts, ok := kkwtsByGroupID[id]

					// If the id is *not* in the kkwts map yet, it is unique.
					// Store the group id in the unique slice, so we can query for the groups.
					if !ok {
						uniqueGroups = append(uniqueGroups, id)
					}

					// kkwts could be a nil map, but that's okay as append will resolve
					// that issue.
					//
					// Note that if a Machine somehow has two entries for the same group,
					// this will be appended twice, but should not cause a problem.
					kkwtsByGroupID[id] = append(kkwts, kkwt)
				}
			}
		}

		// Now that we have all of our group ids that we want to query, get our Groups
		// from that.
		groups, err := modelhelper.GetGroupsByIds(uniqueGroups...)
		if err != nil {
			req.LocalKite.Log.Error(
				"Error returned when querying Team Names. [groupIds:%s, err:%s",
				uniqueGroups, err.Error(),
			)
		}

		// Now that we have the groups, go through our KodingKiteWithToken slice and
		// assign the team names.
		for _, group := range groups {
			id := group.Id.Hex()

			kkwts, ok := kkwtsByGroupID[id]

			// This shouldn't happen, so log a warning to aid in identifying a problem,
			// and continue.
			if !ok {
				req.LocalKite.Log.Warning(
					"Queried GroupId not present in KodingKitesWithToken by GroupId map. This is unexpected, and may represent a problem. [id:%s, ids:%s]",
					id, uniqueGroups,
				)
				continue
			}

			// Append this groups team title to the KodingKiteWithToken Teams slice.
			for _, kkwt := range kkwts {
				kkwt.Teams = append(kkwt.Teams, group.Title)
			}
		}

		return kodingKitesResult, nil
	}
}

// hostFromURL parses the given string, extracting the host (ip/domain) from the
// given string. Ignoring other data such as Port or Url Parameters.
func hostFromURL(s string) (string, error) {
	u, err := url.Parse(s)
	if err != nil {
		return "", err
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		return "", err
	}

	return host, nil
}
<commit_msg>styleguide: Added newline<commit_after>package kontrol

import (
	"errors"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"net"
	"net/url"

	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
)

// GetKodingKitesResult mirrors the Kite GetKitesResult protocol.
type GetKodingKitesResult struct {
	Kites []*KodingKiteWithToken `json:"kites"`
}

// KodingKiteWithToken mirrors the Kite library KiteWithToken with a handful of
// additional fields pertaining to Koding itself. Such as Machine Label, and
// Teams information.
type KodingKiteWithToken struct {
	// These fields match protocol.KiteWithToken
	Kite  protocol.Kite `json:"kite"`
	URL   string        `json:"url"`
	KeyID string        `json:"keyId,omitempty"`
	Token string        `json:"token"`

	// The machine label, as seen from the Koding UI
	MachineLabel string `json:"machineLabel"`

	// The team names, if any.
	Teams []string `json:"teams"`
}

// HandleGetKodingKites implements GetKites with additional Koding specific response
// data.
func HandleGetKodingKites(handleGetKites kite.HandlerFunc) kite.HandlerFunc {
	return func(req *kite.Request) (interface{}, error) {
		// getKitesResponse is an interface, not the actual GetKitesResult value.
		getKitesResponse, err := handleGetKites(req)
		if err != nil {
			req.LocalKite.Log.Error("GetKitesHandler returned an error: %s", err)
			return nil, err
		}

		getKitesResult, ok := getKitesResponse.(*protocol.GetKitesResult)
		if !ok {
			return nil, errors.New("GetKites returned unexpected protocol result")
		}

		// Get all machines for the given user.
		machines, err := modelhelper.GetMachinesByUsername(req.Username)
		if err != nil {
			return nil, err
		}

		// A map of GroupId a slice of all KodingKiteWithToken that contain that group.
		kkwtsByGroupID := map[string][]*KodingKiteWithToken{}

		// Create our slice of unique group ids that we will use to query all the
		// groups that we need.
		uniqueGroups := []string{}

		// Because we are using the Kite's IP as a means to identify the machine,
		// we are sorting the machines by IP so that we can easily look them up.
		machinesByIP := map[string]*models.Machine{}
		for _, m := range machines {
			if m.IpAddress != "" {
				machinesByIP[m.IpAddress] = m
			}
		}

		// Create our KodingKitesResult which will be returned after being populated.
		kodingKitesResult := &GetKodingKitesResult{
			Kites: make([]*KodingKiteWithToken, len(getKitesResult.Kites)),
		}

		// Populate the koding kites
		for i, kwt := range getKitesResult.Kites {
			kkwt := &KodingKiteWithToken{
				Kite:  kwt.Kite,
				URL:   kwt.URL,
				KeyID: kwt.KeyID,
				Token: kwt.Token,
			}
			kodingKitesResult.Kites[i] = kkwt

			host, err := hostFromURL(kwt.URL)
			if err != nil {
				req.LocalKite.Log.Warning(
					"Kite with badly formed URL, unable to extract host. Koding data will not be provided for this Kite. [username:%s, kiteId:%s, host:%s]",
					kwt.Kite.Username,
					kwt.Kite.ID,
					kwt.URL,
				)
				continue
			}

			// With a valid host, get the koding specific information such as Label
			// or Team names and apply it to each KodingKiteWithToken
			if machine, ok := machinesByIP[host]; ok {
				kkwt.MachineLabel = machine.Label

				for _, g := range machine.Groups {
					id := g.Id.Hex()

					// Get all of the kites (not including this one) which have
					// a reference to this kite, so we can append this kite to them.
					kkwts, ok := kkwtsByGroupID[id]

					// If the id is *not* in the kkwts map yet, it is unique.
					// Store the group id in the unique slice, so we can query for the groups.
					if !ok {
						uniqueGroups = append(uniqueGroups, id)
					}

					// kkwts could be a nil map, but that's okay as append will resolve
					// that issue.
					//
					// Note that if a Machine somehow has two entries for the same group,
					// this will be appended twice, but should not cause a problem.
					kkwtsByGroupID[id] = append(kkwts, kkwt)
				}
			}
		}

		// Now that we have all of our group ids that we want to query, get our Groups
		// from that.
		groups, err := modelhelper.GetGroupsByIds(uniqueGroups...)
		if err != nil {
			req.LocalKite.Log.Error(
				"Error returned when querying Team Names. [groupIds:%s, err:%s",
				uniqueGroups, err.Error(),
			)
		}

		// Now that we have the groups, go through our KodingKiteWithToken slice and
		// assign the team names.
		for _, group := range groups {
			id := group.Id.Hex()

			kkwts, ok := kkwtsByGroupID[id]

			// This shouldn't happen, so log a warning to aid in identifying a problem,
			// and continue.
			if !ok {
				req.LocalKite.Log.Warning(
					"Queried GroupId not present in KodingKitesWithToken by GroupId map. This is unexpected, and may represent a problem. [id:%s, ids:%s]",
					id, uniqueGroups,
				)
				continue
			}

			// Append this groups team title to the KodingKiteWithToken Teams slice.
			for _, kkwt := range kkwts {
				kkwt.Teams = append(kkwt.Teams, group.Title)
			}
		}

		return kodingKitesResult, nil
	}
}

// hostFromURL parses the given string, extracting the host (ip/domain) from the
// given string. Ignoring other data such as Port or Url Parameters.
func hostFromURL(s string) (string, error) {
	u, err := url.Parse(s)
	if err != nil {
		return "", err
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		return "", err
	}

	return host, nil
}
<|endoftext|>


<commit_before>package bufit

import (
	"bytes"
	"crypto/rand"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"
)

func ExampleBytes() {
	buf := newWriter(make([]byte, 0, 10))
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, io.NewSectionReader(*&buf, 0, 100))

	io.WriteString(buf, "Hello ")
	r := io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 5)
	io.WriteString(buf, "World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 6)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)

	io.WriteString(buf, "Hello World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 4)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)
	//Output:
	// HelloHello World WorldabcdefgHelloHello Worldabcdefg
}

type badBuffer []byte

func (b *badBuffer) Write(p []byte) (int, error) {
	*b = append(*b, p...)
	return len(p), nil
}

func (b *badBuffer) Read(p []byte) (n int, err error) {
	n = copy(p, *b)
	*b = (*b)[n:]
	if len(*b) == 0 {
		err = io.EOF
	}
	return n, err
}

func BenchmarkBuffer(b *testing.B) {
	for i := 0; i < b.N; i++ {
		benchBuffer(1)
	}
	b.ReportAllocs()
}

func BenchmarkBuffer100(b *testing.B) {
	for i := 0; i < b.N; i++ {
		benchBuffer(100)
	}
	b.ReportAllocs()
}

func BenchmarkBuffer1000(b *testing.B) {
	for i := 0; i < b.N; i++ {
		benchBuffer(1000)
	}
	b.ReportAllocs()
}

func benchBuffer(n int) {
	var grp sync.WaitGroup
	buf := New()
	go func() {
		io.CopyN(buf, rand.Reader, 32*1024*100)
		buf.Close()
	}()
	rs := []io.Reader{}
	for i := 0; i < n; i++ {
		rs = append(rs, buf.NextReader())
	}
	for _, rdr := range rs {
		grp.Add(1)
		go func(r io.Reader) {
			defer grp.Done()
			io.Copy(ioutil.Discard, r)
		}(rdr)
	}
	grp.Wait()
}

func BenchmarkStdBuffer(b *testing.B) {
	for i := 0; i < b.N; i++ {
		r, w := io.Pipe()
		go func() {
			io.CopyN(w, rand.Reader, 32*1024*100)
			w.Close()
		}()
		io.Copy(ioutil.Discard, r)
	}
	b.ReportAllocs()
}

func BenchmarkMyBytes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		tryBuffer(newWriter(nil))
	}
	b.ReportAllocs()
}

func BenchmarkStdBytes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		tryBuffer(bytes.NewBuffer(nil))
	}
	b.ReportAllocs()
}

func BenchmarkFwdMyBytes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		tryFwdBuffer(newWriter(nil))
	}
	b.ReportAllocs()
}

func BenchmarkStdFwdBytes(b *testing.B) {
	for i := 0; i < b.N; i++ {
		tryFwdBuffer(bytes.NewBuffer(nil))
	}
	b.ReportAllocs()
}

const tries = 1000

func tryFwdBuffer(buf io.ReadWriter) {
	for i := 0; i < tries; i++ {
		io.CopyN(buf, rand.Reader, 60*1024)
		io.Copy(ioutil.Discard, buf)
	}
}

func tryBuffer(buf io.ReadWriter) {
	for i := 0; i < tries; i++ {
		io.CopyN(buf, rand.Reader, 60*1024)
		io.CopyN(ioutil.Discard, buf, 32*1024)
		io.CopyN(buf, rand.Reader, 60*1024)
	}
	io.Copy(ioutil.Discard, buf)
}

func TestConcurrent(t *testing.T) {
	var grp sync.WaitGroup
	buf := New()

	var rs []io.ReadCloser
	for i := 0; i < 1000; i++ {
		rs = append(rs, buf.NextReader())
	}

	testData := bytes.NewBuffer(nil)
	io.CopyN(testData, rand.Reader, 32*1024*10)

	for _, r := range rs {
		grp.Add(1)
		go func(r io.ReadCloser) {
			defer grp.Done()
			defer r.Close()
			data, err := ioutil.ReadAll(r)
			if err != nil {
				t.Error(err)
			}
			if !bytes.Equal(testData.Bytes(), data) {
				t.Error("unexpected result...", testData.Len(), len(data))
			}
		}(r)
	}

	r := bytes.NewReader(testData.Bytes())
	for r.Len() > 0 {
		io.CopyN(buf, r, 32*1024*2)
		<-time.After(100 * time.Millisecond)
	}
	buf.Close()
	grp.Wait()
}

func TestQuit(t *testing.T) {
	buf := New()
	r := buf.NextReader()

	wait := make(chan struct{})
	go func() {
		io.Copy(ioutil.Discard, r)
		close(wait)
	}()

	r.Close()
	select {
	case <-wait:
	case <-time.After(100 * time.Millisecond):
		t.Error("timed out waiting for Reader to Close")
	}
}

func ExampleBuffer() {
	// Start a new buffer
	buf := New()

	// Create two readers
	r1, r2 := buf.NextReader(), buf.NextReader()

	// Broadcast a message
	io.WriteString(buf, "Hello World\n")

	// Wait
	var grp sync.WaitGroup
	grp.Add(4)

	// Read fast
	go func() {
		defer grp.Done()
		io.Copy(os.Stdout, r1) // "Hello World\n"
	}()

	// Read slow
	go func() {
		defer grp.Done()
		<-time.After(100 * time.Millisecond)
		io.CopyN(os.Stdout, r2, 5) // "Hello"
		<-time.After(time.Second)
		io.Copy(os.Stdout, r2) // "World\n"
	}()

	// Both readers will read the entire buffer! The slow reader
	// won't block the fast one from reading ahead either.

	// Late reader
	// Since this reader joins after all existing readers have Read "Hello"
	// "Hello" has already been cleared from the Buffer, this Reader will only see
	// "World\n" and beyond.
	go func() {
		defer grp.Done()
		<-time.After(500 * time.Millisecond)
		r3 := buf.NextReader()
		io.Copy(os.Stdout, r3) // "World\n"
	}()

	// Short Reader
	// **Important!** if your reader isn't going to read until the buffer is empty
	// you'll need to call Close() when you are done with it to tell the buffer
	// it's done reading data.
	go func() {
		defer grp.Done()
		<-time.After(100 * time.Millisecond)
		r4 := buf.NextReader()
		io.CopyN(os.Stdout, r4, 5) // "Hello"
		r4.Close()                 // tell the buffer you're done reading
	}()

	// **Important!** mark close so that readers can ret. io.EOF
	buf.Close()

	grp.Wait()
	// Output:
	// Hello World
	// HelloHelloHello World
	//  World
}
<commit_msg>changing benchmarks<commit_after>package bufit

import (
	"bytes"
	"crypto/rand"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"
)

func ExampleBytes() {
	buf := newWriter(make([]byte, 0, 10))
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, io.NewSectionReader(*&buf, 0, 100))

	io.WriteString(buf, "Hello ")
	r := io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 5)
	io.WriteString(buf, "World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 6)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)

	io.WriteString(buf, "Hello World")
	r = io.NewSectionReader(*&buf, 0, int64(buf.Len()))
	io.CopyN(os.Stdout, r, 5)
	io.CopyN(os.Stdout, buf, 4)

	io.WriteString(buf, "abcdefg")
	io.Copy(os.Stdout, buf)
	io.Copy(os.Stdout, buf)
	//Output:
	// HelloHello World WorldabcdefgHelloHello Worldabcdefg
}

type badBuffer []byte

func (b *badBuffer) Write(p []byte) (int, error) {
	*b = append(*b, p...)
	return len(p), nil
}

func (b *badBuffer) Read(p []byte) (n int, err error) {
	n = copy(p, *b)
	*b = (*b)[n:]
	if len(*b) == 0 {
		err = io.EOF
	}
	return n, err
}

func BenchmarkBuffer(b *testing.B) {
	buf := New()
	data, _ := ioutil.ReadAll(io.LimitReader(rand.Reader, 32*1024))

	go func() {
		for i := 0; i < 1000; i++ {
			buf.Write(data)
		}
		buf.Close()
	}()

	r := buf.NextReader()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			r := buf.NextReader()
			io.Copy(ioutil.Discard, r)
			r.Close()
		}
	})

	r.Close()
	b.ReportAllocs()
}

func BenchmarkReadWriter(b *testing.B) {
	buf := newWriter(nil)
	data, _ := ioutil.ReadAll(io.LimitReader(rand.Reader, 32*1024))
	temp := make([]byte, 32*1024)
	for i := 0; i < b.N; i++ {
		buf.Write(data)
		io.CopyBuffer(ioutil.Discard, buf, temp)
	}
	b.ReportAllocs()
}

func TestConcurrent(t *testing.T) {
	var grp sync.WaitGroup
	buf := New()

	var rs []io.ReadCloser
	for i := 0; i < 1000; i++ {
		rs = append(rs, buf.NextReader())
	}

	testData := bytes.NewBuffer(nil)
	io.CopyN(testData, rand.Reader, 32*1024*10)

	for _, r := range rs {
		grp.Add(1)
		go func(r io.ReadCloser) {
			defer grp.Done()
			defer r.Close()
			data, err := ioutil.ReadAll(r)
			if err != nil {
				t.Error(err)
			}
			if !bytes.Equal(testData.Bytes(), data) {
				t.Error("unexpected result...", testData.Len(), len(data))
			}
		}(r)
	}

	r := bytes.NewReader(testData.Bytes())
	for r.Len() > 0 {
		io.CopyN(buf, r, 32*1024*2)
		<-time.After(100 * time.Millisecond)
	}
	buf.Close()
	grp.Wait()
}

func TestQuit(t *testing.T) {
	buf := New()
	r := buf.NextReader()

	wait := make(chan struct{})
	go func() {
		io.Copy(ioutil.Discard, r)
		close(wait)
	}()

	r.Close()
	select {
	case <-wait:
	case <-time.After(100 * time.Millisecond):
		t.Error("timed out waiting for Reader to Close")
	}
}

func ExampleBuffer() {
	// Start a new buffer
	buf := New()

	// Create two readers
	r1, r2 := buf.NextReader(), buf.NextReader()

	// Broadcast a message
	io.WriteString(buf, "Hello World\n")

	// Wait
	var grp sync.WaitGroup
	grp.Add(4)

	// Read fast
	go func() {
		defer grp.Done()
		io.Copy(os.Stdout, r1) // "Hello World\n"
	}()

	// Read slow
	go func() {
		defer grp.Done()
		<-time.After(100 * time.Millisecond)
		io.CopyN(os.Stdout, r2, 5) // "Hello"
		<-time.After(time.Second)
		io.Copy(os.Stdout, r2) // "World\n"
	}()

	// Both readers will read the entire buffer! The slow reader
	// won't block the fast one from reading ahead either.

	// Late reader
	// Since this reader joins after all existing readers have Read "Hello"
	// "Hello" has already been cleared from the Buffer, this Reader will only see
	// "World\n" and beyond.
	go func() {
		defer grp.Done()
		<-time.After(500 * time.Millisecond)
		r3 := buf.NextReader()
		io.Copy(os.Stdout, r3) // "World\n"
	}()

	// Short Reader
	// **Important!** if your reader isn't going to read until the buffer is empty
	// you'll need to call Close() when you are done with it to tell the buffer
	// it's done reading data.
	go func() {
		defer grp.Done()
		<-time.After(100 * time.Millisecond)
		r4 := buf.NextReader()
		io.CopyN(os.Stdout, r4, 5) // "Hello"
		r4.Close()                 // tell the buffer you're done reading
	}()

	// **Important!** mark close so that readers can ret. io.EOF
	buf.Close()

	grp.Wait()
	// Output:
	// Hello World
	// HelloHelloHello World
	//  World
}
<|endoftext|>


<commit_before>package secio

import (
	"crypto/cipher"
	"errors"
	"fmt"
	"io"
	"sync"

	"crypto/hmac"

	proto "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/gogo/protobuf/proto"
	msgio "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
	mpool "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio/mpool"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
)

// ErrMACInvalid signals that a MAC verification failed
var ErrMACInvalid = errors.New("MAC verification failed")

// bufPool is a ByteSlicePool for messages. we need buffers because (sadly)
// we cannot encrypt in place-- the user needs their buffer back.
var bufPool = mpool.ByteSlicePool

type etmWriter struct {
	// params
	pool mpool.Pool        // for the buffers with encrypted data
	msg  msgio.WriteCloser // msgio for knowing where boundaries lie
	str  cipher.Stream     // the stream cipher to encrypt with
	mac  HMAC              // the mac to authenticate data with

	sync.Mutex
}

// NewETMWriter Encrypt-Then-MAC
func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac, pool: bufPool}
}

// Write writes passed in buffer as a single message.
func (w *etmWriter) Write(b []byte) (int, error) {
	if err := w.WriteMsg(b); err != nil {
		return 0, err
	}
	return len(b), nil
}

// WriteMsg writes the msg in the passed in buffer.
func (w *etmWriter) WriteMsg(b []byte) error {
	w.Lock()
	defer w.Unlock()

	// encrypt.
	data := w.pool.Get(uint32(len(b))).([]byte)
	data = data[:len(b)] // the pool's buffer may be larger
	w.str.XORKeyStream(data, b)

	// log.Debugf("ENC plaintext (%d): %s %v", len(b), b, b)
	// log.Debugf("ENC ciphertext (%d): %s %v", len(data), data, data)

	// then, mac.
	if _, err := w.mac.Write(data); err != nil {
		return err
	}

	// Sum appends.
	data = w.mac.Sum(data)
	w.mac.Reset()
	// it's sad to append here. our buffers are -- hopefully -- coming from
	// a shared buffer pool, so the append may not actually cause allocation
	// one can only hope. i guess we'll see.

	return w.msg.WriteMsg(data)
}

func (w *etmWriter) Close() error {
	return w.msg.Close()
}

type etmReader struct {
	msgio.Reader
	io.Closer

	// buffer
	buf []byte

	// params
	msg msgio.ReadCloser // msgio for knowing where boundaries lie
	str cipher.Stream    // the stream cipher to encrypt with
	mac HMAC             // the mac to authenticate data with

	sync.Mutex
}

// NewETMReader Encrypt-Then-MAC
func NewETMReader(r io.Reader, s cipher.Stream, mac HMAC) msgio.ReadCloser {
	return &etmReader{msg: msgio.NewReader(r), str: s, mac: mac}
}

func (r *etmReader) NextMsgLen() (int, error) {
	return r.msg.NextMsgLen()
}

func (r *etmReader) drainBuf(buf []byte) int {
	if r.buf == nil {
		return 0
	}

	n := copy(buf, r.buf)
	r.buf = r.buf[n:]
	return n
}

func (r *etmReader) Read(buf []byte) (int, error) {
	r.Lock()
	defer r.Unlock()

	// first, check if we have anything in the buffer
	copied := r.drainBuf(buf)
	buf = buf[copied:]
	if copied > 0 {
		return copied, nil
		// return here to avoid complicating the rest...
		// user can call io.ReadFull.
	}

	// check the buffer has enough space for the next msg
	fullLen, err := r.msg.NextMsgLen()
	if err != nil {
		return 0, err
	}

	buf2 := buf
	changed := false
	// if not enough space, allocate a new buffer.
	if cap(buf) < fullLen {
		buf2 = make([]byte, fullLen)
		changed = true
	}
	buf2 = buf2[:fullLen]

	n, err := io.ReadFull(r.msg, buf2)
	if err != nil {
		return n, err
	}

	m, err := r.macCheckThenDecrypt(buf2)
	if err != nil {
		return 0, err
	}
	buf2 = buf2[:m]
	if !changed {
		return m, nil
	}

	n = copy(buf, buf2)
	if len(buf2) > len(buf) {
		r.buf = buf2[len(buf):] // had some left over? save it.
	}
	return n, nil
}

func (r *etmReader) ReadMsg() ([]byte, error) {
	r.Lock()
	defer r.Unlock()

	msg, err := r.msg.ReadMsg()
	if err != nil {
		return nil, err
	}

	n, err := r.macCheckThenDecrypt(msg)
	if err != nil {
		return nil, err
	}
	return msg[:n], nil
}

func (r *etmReader) macCheckThenDecrypt(m []byte) (int, error) {
	l := len(m)
	if l < r.mac.size {
		return 0, fmt.Errorf("buffer (%d) shorter than MAC size (%d)", l, r.mac.size)
	}

	mark := l - r.mac.size
	data := m[:mark]
	macd := m[mark:]

	r.mac.Write(data)
	expected := r.mac.Sum(nil)
	r.mac.Reset()

	// check mac. if failed, return error.
	if !hmac.Equal(macd, expected) {
		log.Debug("MAC Invalid:", expected, "!=", macd)
		return 0, ErrMACInvalid
	}

	// ok seems good. decrypt. (can decrypt in place, yay!)
	// log.Debugf("DEC ciphertext (%d): %s %v", len(data), data, data)
	r.str.XORKeyStream(data, data)
	// log.Debugf("DEC plaintext (%d): %s %v", len(data), data, data)

	return mark, nil
}

func (w *etmReader) Close() error {
	return w.msg.Close()
}

// ReleaseMsg signals a buffer can be reused.
func (r *etmReader) ReleaseMsg(b []byte) {
	r.msg.ReleaseMsg(b)
}

// writeMsgCtx is used by the
func writeMsgCtx(ctx context.Context, w msgio.Writer, msg proto.Message) ([]byte, error) {
	enc, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}

	// write in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func(m []byte) {
		err := w.WriteMsg(m)
		select {
		case done <- err:
		case <-ctx.Done():
		}
	}(enc)

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		return enc, e
	}
}

func readMsgCtx(ctx context.Context, r msgio.Reader, p proto.Message) ([]byte, error) {
	var msg []byte

	// read in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func() {
		var err error
		msg, err = r.ReadMsg()
		select {
		case done <- err:
		case <-ctx.Done():
		}
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		if e != nil {
			return nil, e
		}
	}

	return msg, proto.Unmarshal(msg, p)
}
<commit_msg>quick fix for OOM panic that has been plaguing us<commit_after>package secio

import (
	"crypto/cipher"
	"errors"
	"fmt"
	"io"
	"sync"

	"crypto/hmac"

	proto "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/gogo/protobuf/proto"
	msgio "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
	mpool "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio/mpool"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
)

const MaxMsgSize = 8 * 1024 * 1024

var ErrMaxMessageSize = errors.New("attempted to read message larger than max size")

// ErrMACInvalid signals that a MAC verification failed
var ErrMACInvalid = errors.New("MAC verification failed")

// bufPool is a ByteSlicePool for messages. we need buffers because (sadly)
// we cannot encrypt in place-- the user needs their buffer back.
var bufPool = mpool.ByteSlicePool

type etmWriter struct {
	// params
	pool mpool.Pool        // for the buffers with encrypted data
	msg  msgio.WriteCloser // msgio for knowing where boundaries lie
	str  cipher.Stream     // the stream cipher to encrypt with
	mac  HMAC              // the mac to authenticate data with

	sync.Mutex
}

// NewETMWriter Encrypt-Then-MAC
func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac, pool: bufPool}
}

// Write writes passed in buffer as a single message.
func (w *etmWriter) Write(b []byte) (int, error) {
	if err := w.WriteMsg(b); err != nil {
		return 0, err
	}
	return len(b), nil
}

// WriteMsg writes the msg in the passed in buffer.
func (w *etmWriter) WriteMsg(b []byte) error {
	w.Lock()
	defer w.Unlock()

	// encrypt.
	data := w.pool.Get(uint32(len(b))).([]byte)
	data = data[:len(b)] // the pool's buffer may be larger
	w.str.XORKeyStream(data, b)

	// log.Debugf("ENC plaintext (%d): %s %v", len(b), b, b)
	// log.Debugf("ENC ciphertext (%d): %s %v", len(data), data, data)

	// then, mac.
	if _, err := w.mac.Write(data); err != nil {
		return err
	}

	// Sum appends.
	data = w.mac.Sum(data)
	w.mac.Reset()
	// it's sad to append here. our buffers are -- hopefully -- coming from
	// a shared buffer pool, so the append may not actually cause allocation
	// one can only hope. i guess we'll see.

	return w.msg.WriteMsg(data)
}

func (w *etmWriter) Close() error {
	return w.msg.Close()
}

type etmReader struct {
	msgio.Reader
	io.Closer

	// buffer
	buf []byte

	// params
	msg msgio.ReadCloser // msgio for knowing where boundaries lie
	str cipher.Stream    // the stream cipher to encrypt with
	mac HMAC             // the mac to authenticate data with

	sync.Mutex
}

// NewETMReader Encrypt-Then-MAC
func NewETMReader(r io.Reader, s cipher.Stream, mac HMAC) msgio.ReadCloser {
	return &etmReader{msg: msgio.NewReader(r), str: s, mac: mac}
}

func (r *etmReader) NextMsgLen() (int, error) {
	return r.msg.NextMsgLen()
}

func (r *etmReader) drainBuf(buf []byte) int {
	if r.buf == nil {
		return 0
	}

	n := copy(buf, r.buf)
	r.buf = r.buf[n:]
	return n
}

func (r *etmReader) Read(buf []byte) (int, error) {
	r.Lock()
	defer r.Unlock()

	// first, check if we have anything in the buffer
	copied := r.drainBuf(buf)
	buf = buf[copied:]
	if copied > 0 {
		return copied, nil
		// return here to avoid complicating the rest...
		// user can call io.ReadFull.
	}

	// check the buffer has enough space for the next msg
	fullLen, err := r.msg.NextMsgLen()
	if err != nil {
		return 0, err
	}

	if fullLen > MaxMsgSize {
		return 0, ErrMaxMessageSize
	}

	buf2 := buf
	changed := false
	// if not enough space, allocate a new buffer.
	if cap(buf) < fullLen {
		buf2 = make([]byte, fullLen)
		changed = true
	}
	buf2 = buf2[:fullLen]

	n, err := io.ReadFull(r.msg, buf2)
	if err != nil {
		return n, err
	}

	m, err := r.macCheckThenDecrypt(buf2)
	if err != nil {
		return 0, err
	}
	buf2 = buf2[:m]
	if !changed {
		return m, nil
	}

	n = copy(buf, buf2)
	if len(buf2) > len(buf) {
		r.buf = buf2[len(buf):] // had some left over? save it.
	}
	return n, nil
}

func (r *etmReader) ReadMsg() ([]byte, error) {
	r.Lock()
	defer r.Unlock()

	msg, err := r.msg.ReadMsg()
	if err != nil {
		return nil, err
	}

	n, err := r.macCheckThenDecrypt(msg)
	if err != nil {
		return nil, err
	}
	return msg[:n], nil
}

func (r *etmReader) macCheckThenDecrypt(m []byte) (int, error) {
	l := len(m)
	if l < r.mac.size {
		return 0, fmt.Errorf("buffer (%d) shorter than MAC size (%d)", l, r.mac.size)
	}

	mark := l - r.mac.size
	data := m[:mark]
	macd := m[mark:]

	r.mac.Write(data)
	expected := r.mac.Sum(nil)
	r.mac.Reset()

	// check mac. if failed, return error.
	if !hmac.Equal(macd, expected) {
		log.Debug("MAC Invalid:", expected, "!=", macd)
		return 0, ErrMACInvalid
	}

	// ok seems good. decrypt. (can decrypt in place, yay!)
	// log.Debugf("DEC ciphertext (%d): %s %v", len(data), data, data)
	r.str.XORKeyStream(data, data)
	// log.Debugf("DEC plaintext (%d): %s %v", len(data), data, data)

	return mark, nil
}

func (w *etmReader) Close() error {
	return w.msg.Close()
}

// ReleaseMsg signals a buffer can be reused.
func (r *etmReader) ReleaseMsg(b []byte) {
	r.msg.ReleaseMsg(b)
}

// writeMsgCtx is used by the
func writeMsgCtx(ctx context.Context, w msgio.Writer, msg proto.Message) ([]byte, error) {
	enc, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}

	// write in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func(m []byte) {
		err := w.WriteMsg(m)
		select {
		case done <- err:
		case <-ctx.Done():
		}
	}(enc)

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		return enc, e
	}
}

func readMsgCtx(ctx context.Context, r msgio.Reader, p proto.Message) ([]byte, error) {
	var msg []byte

	// read in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func() {
		var err error
		msg, err = r.ReadMsg()
		select {
		case done <- err:
		case <-ctx.Done():
		}
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		if e != nil {
			return nil, e
		}
	}

	return msg, proto.Unmarshal(msg, p)
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"strings"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/network"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
)

const nicRoutedIPv4GW = "169.254.0.1"
const nicRoutedIPv6GW = "fe80::1"

type nicRouted struct {
	deviceCommon
}

func (d *nicRouted) CanHotPlug() (bool, []string) {
	return false, []string{}
}

// validateConfig checks the supplied config for correctness.
func (d *nicRouted) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container) {
		return ErrUnsupportedDevType
	}

	requiredFields := []string{}
	optionalFields := []string{
		"name",
		"parent",
		"mtu",
		"hwaddr",
		"host_name",
		"vlan",
		"ipv4.gateway",
		"ipv6.gateway",
		"ipv4.host_address",
		"ipv6.host_address",
	}

	rules := nicValidationRules(requiredFields, optionalFields)
	rules["ipv4.address"] = func(value string) error {
		if value == "" {
			return nil
		}

		return NetworkValidAddressV4List(value)
	}
	rules["ipv6.address"] = func(value string) error {
		if value == "" {
			return nil
		}

		return NetworkValidAddressV6List(value)
	}

	err := d.config.Validate(rules)
	if err != nil {
		return err
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *nicRouted) validateEnvironment() error {
	if d.inst.Type() == instancetype.Container && d.config["name"] == "" {
		return fmt.Errorf("Requires name property to start")
	}

	if d.config["parent"] != "" && !shared.PathExists(fmt.Sprintf("/sys/class/net/%s", d.config["parent"])) {
		return fmt.Errorf("Parent device '%s' doesn't exist", d.config["parent"])
	}

	if d.config["parent"] == "" && d.config["vlan"] != "" {
		return fmt.Errorf("The vlan setting can only be used when combined with a parent interface")
	}

	extensions := d.state.OS.LXCFeatures
	if !extensions["network_veth_router"] || !extensions["network_l2proxy"] {
		return fmt.Errorf("Requires liblxc has following API extensions: network_veth_router, network_l2proxy")
	}

	// Check necessary sysctls are configured for use with l2proxy parent for routed mode.
	if d.config["parent"] != "" && d.config["ipv4.address"] != "" {
		ipv4FwdPath := fmt.Sprintf("net/ipv4/conf/%s/forwarding", d.config["parent"])
		sysctlVal, err := util.SysctlGet(ipv4FwdPath)
		if err != nil || sysctlVal != "1\n" {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv4FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv4.conf.%s.forwarding=1", d.config["parent"])
		}
	}

	// Check necessary sysctls are configured for use with l2proxy parent for routed mode.
	if d.config["parent"] != "" && d.config["ipv6.address"] != "" {
		// net.ipv6.conf.all.forwarding=1 is required to enable general packet forwarding for IPv6.
		ipv6FwdPath := fmt.Sprintf("net/ipv6/conf/%s/forwarding", "all")
		sysctlVal, err := util.SysctlGet(ipv6FwdPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.forwarding=1", "all")
		}

		ipv6FwdPath = fmt.Sprintf("net/ipv6/conf/%s/forwarding", d.config["parent"])
		sysctlVal, err = util.SysctlGet(ipv6FwdPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.forwarding=1", d.config["parent"])
		}

		// net.ipv6.conf.all.proxy_ndp=1 is needed otherwise unicast neighbour solicitations are rejected.
		// This causes periodic latency spikes every 15-20s as the neighbour has to resort to using
		// multicast NDP resolution and expires the previous neighbour entry.
		ipv6ProxyNdpPath := fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", "all")
		sysctlVal, err = util.SysctlGet(ipv6ProxyNdpPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.proxy_ndp=1", "all")
		}

		ipv6ProxyNdpPath = fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", d.config["parent"])
		sysctlVal, err = util.SysctlGet(ipv6ProxyNdpPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.proxy_ndp=1", d.config["parent"])
		}
	}

	return nil
}

// Start is run when the instance is starting up (Routed mode doesn't support hot plugging).
func (d *nicRouted) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	// Lock to avoid issues with containers starting in parallel.
	networkCreateSharedDeviceLock.Lock()
	defer networkCreateSharedDeviceLock.Unlock()

	saveData := make(map[string]string)

	// Decide which parent we should use based on VLAN setting.
	parentName := ""
	if d.config["parent"] != "" {
		parentName = network.GetHostDevice(d.config["parent"], d.config["vlan"])

		statusDev, err := networkCreateVlanDeviceIfNeeded(d.state, d.config["parent"], parentName, d.config["vlan"])
		if err != nil {
			return nil, err
		}

		// Record whether we created this device or not so it can be removed on stop.
		saveData["last_state.created"] = fmt.Sprintf("%t", statusDev != "existing")

		// If we created a VLAN interface, we need to setup the sysctls on that interface.
		if statusDev == "created" {
			err := d.setupParentSysctls(parentName)
			if err != nil {
				return nil, err
			}
		}
	}

	hostName := d.config["host_name"]
	if hostName == "" {
		hostName = networkRandomDevName("veth")
	}
	saveData["host_name"] = hostName

	err = d.volatileSet(saveData)
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}
	nic := []deviceConfig.RunConfigItem{
		{Key: "name", Value: d.config["name"]},
		{Key: "type", Value: "veth"},
		{Key: "flags", Value: "up"},
		{Key: "veth.mode", Value: "router"},
		{Key: "veth.pair", Value: saveData["host_name"]},
	}

	// If there is a designated parent interface, activate the layer2 proxy mode to advertise
	// the instance's IPs over that interface using proxy APR/NDP.
	if parentName != "" {
		nic = append(nic,
			deviceConfig.RunConfigItem{Key: "l2proxy", Value: "1"},
			deviceConfig.RunConfigItem{Key: "link", Value: parentName},
		)
	}

	if d.config["mtu"] != "" {
		nic = append(nic, deviceConfig.RunConfigItem{Key: "mtu", Value: d.config["mtu"]})
	}

	if d.config["ipv4.address"] != "" {
		for _, addr := range strings.Split(d.config["ipv4.address"], ",") {
			addr = strings.TrimSpace(addr)
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv4.address", Value: fmt.Sprintf("%s/32", addr)})
		}

		if nicHasAutoGateway(d.config["ipv4.gateway"]) {
			// Use a fixed link-local address as the next-hop default gateway.
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv4.gateway", Value: d.ipv4HostAddress()})
		}
	}

	if d.config["ipv6.address"] != "" {
		for _, addr := range strings.Split(d.config["ipv6.address"], ",") {
			addr = strings.TrimSpace(addr)
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv6.address", Value: fmt.Sprintf("%s/128", addr)})
		}

		if nicHasAutoGateway(d.config["ipv6.gateway"]) {
			// Use a fixed link-local address as the next-hop default gateway.
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv6.gateway", Value: d.ipv6HostAddress()})
		}
	}

	runConf.NetworkInterface = nic
	runConf.PostHooks = append(runConf.PostHooks, d.postStart)
	return &runConf, nil
}

// setupParentSysctls configures the required sysctls on the parent to allow l2proxy to work.
// Because of our policy not to modify sysctls on existing interfaces, this should only be called
// if we created the parent interface.
func (d *nicRouted) setupParentSysctls(parentName string) error {
	if d.config["ipv4.address"] != "" {
		// Set necessary sysctls for use with l2proxy parent in routed mode.
		ipv4FwdPath := fmt.Sprintf("net/ipv4/conf/%s/forwarding", parentName)
		err := util.SysctlSet(ipv4FwdPath, "1")
		if err != nil {
			return fmt.Errorf("Error setting net sysctl %s: %v", ipv4FwdPath, err)
		}
	}

	if d.config["ipv6.address"] != "" {
		// Set necessary sysctls use with l2proxy parent in routed mode.
		ipv6FwdPath := fmt.Sprintf("net/ipv6/conf/%s/forwarding", parentName)
		err := util.SysctlSet(ipv6FwdPath, "1")
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}

		ipv6ProxyNdpPath := fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", parentName)
		err = util.SysctlSet(ipv6ProxyNdpPath, "1")
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
	}

	return nil
}

// postStart is run after the instance is started.
func (d *nicRouted) postStart() error {
	v := d.volatileGet()

	// If host_name is defined (and it should be), then we add the dummy link-local gateway IPs
	// to the host end of the veth pair. This ensures that liveness detection of the gateways
	// inside the instance work and ensure that traffic doesn't periodically halt whilst ARP/NDP
	// is re-detected.
	if v["host_name"] != "" {
		if d.config["ipv4.address"] != "" {
			_, err := shared.RunCommand("ip", "-4", "addr", "add", fmt.Sprintf("%s/32", d.ipv4HostAddress()), "dev", v["host_name"])
			if err != nil {
				return err
			}
		}

		if d.config["ipv6.address"] != "" {
			_, err := shared.RunCommand("ip", "-6", "addr", "add", fmt.Sprintf("%s/128", d.ipv6HostAddress()), "dev", v["host_name"])
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Stop is run when the device is removed from the instance.
func (d *nicRouted) Stop() (*deviceConfig.RunConfig, error) {
	runConf := deviceConfig.RunConfig{
		PostHooks: []func() error{d.postStop},
	}

	return &runConf, nil
}

// postStop is run after the device is removed from the instance.
func (d *nicRouted) postStop() error {
	defer d.volatileSet(map[string]string{
		"last_state.created": "",
		"host_name":          "",
	})

	v := d.volatileGet()

	// This will delete the parent interface if we created it for VLAN parent.
	if shared.IsTrue(v["last_state.created"]) {
		parentName := network.GetHostDevice(d.config["parent"], d.config["vlan"])
		err := networkRemoveInterfaceIfNeeded(d.state, parentName, d.inst, d.config["parent"], d.config["vlan"])
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *nicRouted) ipv4HostAddress() string {
	if d.config["ipv4.host_address"] != "" {
		return d.config["ipv4.host_address"]
	}

	return nicRoutedIPv4GW
}

func (d *nicRouted) ipv6HostAddress() string {
	if d.config["ipv6.host_address"] != "" {
		return d.config["ipv6.host_address"]
	}

	return nicRoutedIPv6GW
}
<commit_msg>lxd/device/nic/routed: Sets accept_ra=0 on host interface<commit_after>package device

import (
	"fmt"
	"strings"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/network"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
)

const nicRoutedIPv4GW = "169.254.0.1"
const nicRoutedIPv6GW = "fe80::1"

type nicRouted struct {
	deviceCommon
}

func (d *nicRouted) CanHotPlug() (bool, []string) {
	return false, []string{}
}

// validateConfig checks the supplied config for correctness.
func (d *nicRouted) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container) {
		return ErrUnsupportedDevType
	}

	requiredFields := []string{}
	optionalFields := []string{
		"name",
		"parent",
		"mtu",
		"hwaddr",
		"host_name",
		"vlan",
		"ipv4.gateway",
		"ipv6.gateway",
		"ipv4.host_address",
		"ipv6.host_address",
	}

	rules := nicValidationRules(requiredFields, optionalFields)
	rules["ipv4.address"] = func(value string) error {
		if value == "" {
			return nil
		}

		return NetworkValidAddressV4List(value)
	}
	rules["ipv6.address"] = func(value string) error {
		if value == "" {
			return nil
		}

		return NetworkValidAddressV6List(value)
	}

	err := d.config.Validate(rules)
	if err != nil {
		return err
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *nicRouted) validateEnvironment() error {
	if d.inst.Type() == instancetype.Container && d.config["name"] == "" {
		return fmt.Errorf("Requires name property to start")
	}

	if d.config["parent"] != "" && !shared.PathExists(fmt.Sprintf("/sys/class/net/%s", d.config["parent"])) {
		return fmt.Errorf("Parent device '%s' doesn't exist", d.config["parent"])
	}

	if d.config["parent"] == "" && d.config["vlan"] != "" {
		return fmt.Errorf("The vlan setting can only be used when combined with a parent interface")
	}

	extensions := d.state.OS.LXCFeatures
	if !extensions["network_veth_router"] || !extensions["network_l2proxy"] {
		return fmt.Errorf("Requires liblxc has following API extensions: network_veth_router, network_l2proxy")
	}

	// Check necessary sysctls are configured for use with l2proxy parent for routed mode.
	if d.config["parent"] != "" && d.config["ipv4.address"] != "" {
		ipv4FwdPath := fmt.Sprintf("net/ipv4/conf/%s/forwarding", d.config["parent"])
		sysctlVal, err := util.SysctlGet(ipv4FwdPath)
		if err != nil || sysctlVal != "1\n" {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv4FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv4.conf.%s.forwarding=1", d.config["parent"])
		}
	}

	// Check necessary sysctls are configured for use with l2proxy parent for routed mode.
	if d.config["parent"] != "" && d.config["ipv6.address"] != "" {
		// net.ipv6.conf.all.forwarding=1 is required to enable general packet forwarding for IPv6.
		ipv6FwdPath := fmt.Sprintf("net/ipv6/conf/%s/forwarding", "all")
		sysctlVal, err := util.SysctlGet(ipv6FwdPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.forwarding=1", "all")
		}

		ipv6FwdPath = fmt.Sprintf("net/ipv6/conf/%s/forwarding", d.config["parent"])
		sysctlVal, err = util.SysctlGet(ipv6FwdPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.forwarding=1", d.config["parent"])
		}

		// net.ipv6.conf.all.proxy_ndp=1 is needed otherwise unicast neighbour solicitations are rejected.
		// This causes periodic latency spikes every 15-20s as the neighbour has to resort to using
		// multicast NDP resolution and expires the previous neighbour entry.
		ipv6ProxyNdpPath := fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", "all")
		sysctlVal, err = util.SysctlGet(ipv6ProxyNdpPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.proxy_ndp=1", "all")
		}

		ipv6ProxyNdpPath = fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", d.config["parent"])
		sysctlVal, err = util.SysctlGet(ipv6ProxyNdpPath)
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
		if sysctlVal != "1\n" {
			return fmt.Errorf("Routed mode requires sysctl net.ipv6.conf.%s.proxy_ndp=1", d.config["parent"])
		}
	}

	return nil
}

// Start is run when the instance is starting up (Routed mode doesn't support hot plugging).
func (d *nicRouted) Start() (*deviceConfig.RunConfig, error) {
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	// Lock to avoid issues with containers starting in parallel.
	networkCreateSharedDeviceLock.Lock()
	defer networkCreateSharedDeviceLock.Unlock()

	saveData := make(map[string]string)

	// Decide which parent we should use based on VLAN setting.
	parentName := ""
	if d.config["parent"] != "" {
		parentName = network.GetHostDevice(d.config["parent"], d.config["vlan"])

		statusDev, err := networkCreateVlanDeviceIfNeeded(d.state, d.config["parent"], parentName, d.config["vlan"])
		if err != nil {
			return nil, err
		}

		// Record whether we created this device or not so it can be removed on stop.
		saveData["last_state.created"] = fmt.Sprintf("%t", statusDev != "existing")

		// If we created a VLAN interface, we need to setup the sysctls on that interface.
		if statusDev == "created" {
			err := d.setupParentSysctls(parentName)
			if err != nil {
				return nil, err
			}
		}
	}

	hostName := d.config["host_name"]
	if hostName == "" {
		hostName = networkRandomDevName("veth")
	}
	saveData["host_name"] = hostName

	err = d.volatileSet(saveData)
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}
	nic := []deviceConfig.RunConfigItem{
		{Key: "name", Value: d.config["name"]},
		{Key: "type", Value: "veth"},
		{Key: "flags", Value: "up"},
		{Key: "veth.mode", Value: "router"},
		{Key: "veth.pair", Value: saveData["host_name"]},
	}

	// If there is a designated parent interface, activate the layer2 proxy mode to advertise
	// the instance's IPs over that interface using proxy APR/NDP.
	if parentName != "" {
		nic = append(nic,
			deviceConfig.RunConfigItem{Key: "l2proxy", Value: "1"},
			deviceConfig.RunConfigItem{Key: "link", Value: parentName},
		)
	}

	if d.config["mtu"] != "" {
		nic = append(nic, deviceConfig.RunConfigItem{Key: "mtu", Value: d.config["mtu"]})
	}

	if d.config["ipv4.address"] != "" {
		for _, addr := range strings.Split(d.config["ipv4.address"], ",") {
			addr = strings.TrimSpace(addr)
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv4.address", Value: fmt.Sprintf("%s/32", addr)})
		}

		if nicHasAutoGateway(d.config["ipv4.gateway"]) {
			// Use a fixed link-local address as the next-hop default gateway.
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv4.gateway", Value: d.ipv4HostAddress()})
		}
	}

	if d.config["ipv6.address"] != "" {
		for _, addr := range strings.Split(d.config["ipv6.address"], ",") {
			addr = strings.TrimSpace(addr)
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv6.address", Value: fmt.Sprintf("%s/128", addr)})
		}

		if nicHasAutoGateway(d.config["ipv6.gateway"]) {
			// Use a fixed link-local address as the next-hop default gateway.
			nic = append(nic, deviceConfig.RunConfigItem{Key: "ipv6.gateway", Value: d.ipv6HostAddress()})
		}
	}

	runConf.NetworkInterface = nic
	runConf.PostHooks = append(runConf.PostHooks, d.postStart)
	return &runConf, nil
}

// setupParentSysctls configures the required sysctls on the parent to allow l2proxy to work.
// Because of our policy not to modify sysctls on existing interfaces, this should only be called
// if we created the parent interface.
func (d *nicRouted) setupParentSysctls(parentName string) error {
	if d.config["ipv4.address"] != "" {
		// Set necessary sysctls for use with l2proxy parent in routed mode.
		ipv4FwdPath := fmt.Sprintf("net/ipv4/conf/%s/forwarding", parentName)
		err := util.SysctlSet(ipv4FwdPath, "1")
		if err != nil {
			return fmt.Errorf("Error setting net sysctl %s: %v", ipv4FwdPath, err)
		}
	}

	if d.config["ipv6.address"] != "" {
		// Set necessary sysctls use with l2proxy parent in routed mode.
		ipv6FwdPath := fmt.Sprintf("net/ipv6/conf/%s/forwarding", parentName)
		err := util.SysctlSet(ipv6FwdPath, "1")
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6FwdPath, err)
		}

		ipv6ProxyNdpPath := fmt.Sprintf("net/ipv6/conf/%s/proxy_ndp", parentName)
		err = util.SysctlSet(ipv6ProxyNdpPath, "1")
		if err != nil {
			return fmt.Errorf("Error reading net sysctl %s: %v", ipv6ProxyNdpPath, err)
		}
	}

	return nil
}

// postStart is run after the instance is started.
func (d *nicRouted) postStart() error {
	v := d.volatileGet()

	// If host_name is defined (and it should be), then we add the dummy link-local gateway IPs
	// to the host end of the veth pair. This ensures that liveness detection of the gateways
	// inside the instance work and ensure that traffic doesn't periodically halt whilst ARP/NDP
	// is re-detected.
	if v["host_name"] != "" {
		// Attempt to disable IPv6 router advertisement acceptance.
		err := util.SysctlSet(fmt.Sprintf("net/ipv6/conf/%s/accept_ra", v["host_name"]), "0")
		if err != nil {
			return err
		}

		if d.config["ipv4.address"] != "" {
			_, err := shared.RunCommand("ip", "-4", "addr", "add", fmt.Sprintf("%s/32", d.ipv4HostAddress()), "dev", v["host_name"])
			if err != nil {
				return err
			}
		}

		if d.config["ipv6.address"] != "" {
			_, err := shared.RunCommand("ip", "-6", "addr", "add", fmt.Sprintf("%s/128", d.ipv6HostAddress()), "dev", v["host_name"])
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Stop is run when the device is removed from the instance.
func (d *nicRouted) Stop() (*deviceConfig.RunConfig, error) {
	runConf := deviceConfig.RunConfig{
		PostHooks: []func() error{d.postStop},
	}

	return &runConf, nil
}

// postStop is run after the device is removed from the instance.
func (d *nicRouted) postStop() error {
	defer d.volatileSet(map[string]string{
		"last_state.created": "",
		"host_name":          "",
	})

	v := d.volatileGet()

	// This will delete the parent interface if we created it for VLAN parent.
	if shared.IsTrue(v["last_state.created"]) {
		parentName := network.GetHostDevice(d.config["parent"], d.config["vlan"])
		err := networkRemoveInterfaceIfNeeded(d.state, parentName, d.inst, d.config["parent"], d.config["vlan"])
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *nicRouted) ipv4HostAddress() string {
	if d.config["ipv4.host_address"] != "" {
		return d.config["ipv4.host_address"]
	}

	return nicRoutedIPv4GW
}

func (d *nicRouted) ipv6HostAddress() string {
	if d.config["ipv6.host_address"] != "" {
		return d.config["ipv6.host_address"]
	}

	return nicRoutedIPv6GW
}
<|endoftext|>


<commit_before>package warnings

import (
	"fmt"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/db"
)

// ResolveWarningsByLocalNodeAndType resolves warnings with the local node and type code.
// Returns error if no local node name.
func ResolveWarningsByLocalNodeAndType(cluster *db.Cluster, typeCode db.WarningType) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrapf(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return ResolveWarningsByNodeAndType(cluster, localName, typeCode)
}

// ResolveWarningsByNodeAndType resolves warnings with the given node and type code.
func ResolveWarningsByNodeAndType(cluster *db.Cluster, nodeName string, typeCode db.WarningType) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarningsByType(typeCode)
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName {
				continue
			}

			err = tx.UpdateWarningStatus(w.UUID, db.WarningStatusResolved)
			if err != nil {
				return err
			}
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// ResolveWarningsByNodeAndProjectAndType resolves warnings with the given node, project and type code.
func ResolveWarningsByNodeAndProjectAndType(cluster *db.Cluster, nodeName string, projectName string, typeCode db.WarningType) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarningsByType(typeCode)
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName || w.Project != projectName {
				continue
			}

			err = tx.UpdateWarningStatus(w.UUID, db.WarningStatusResolved)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// ResolveWarningsByLocalNodeAndProjectAndType resolves warnings with the given project and type code.
func ResolveWarningsByLocalNodeAndProjectAndType(cluster *db.Cluster, projectName string, typeCode db.WarningType) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return ResolveWarningsByNodeAndProjectAndType(cluster, localName, projectName, typeCode)
}
<commit_msg>lxd/warnings: Add resolve and delete functions<commit_after>package warnings

import (
	"fmt"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/db"
)

// ResolveWarningsByLocalNodeAndType resolves warnings with the local node and type code.
// Returns error if no local node name.
func ResolveWarningsByLocalNodeAndType(cluster *db.Cluster, typeCode db.WarningType) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrapf(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return ResolveWarningsByNodeAndType(cluster, localName, typeCode)
}

// ResolveWarningsByNodeAndType resolves warnings with the given node and type code.
func ResolveWarningsByNodeAndType(cluster *db.Cluster, nodeName string, typeCode db.WarningType) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarningsByType(typeCode)
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName {
				continue
			}

			err = tx.UpdateWarningStatus(w.UUID, db.WarningStatusResolved)
			if err != nil {
				return err
			}
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// ResolveWarningsByNodeAndProjectAndType resolves warnings with the given node, project and type code.
func ResolveWarningsByNodeAndProjectAndType(cluster *db.Cluster, nodeName string, projectName string, typeCode db.WarningType) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarningsByType(typeCode)
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName || w.Project != projectName {
				continue
			}

			err = tx.UpdateWarningStatus(w.UUID, db.WarningStatusResolved)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// ResolveWarningsByLocalNodeAndProjectAndType resolves warnings with the given project and type code.
func ResolveWarningsByLocalNodeAndProjectAndType(cluster *db.Cluster, projectName string, typeCode db.WarningType) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return ResolveWarningsByNodeAndProjectAndType(cluster, localName, projectName, typeCode)
}

// ResolveWarningsByNodeAndProjectAndTypeAndEntity resolves warnings with the given node, project, type code, and entity.
func ResolveWarningsByNodeAndProjectAndTypeAndEntity(cluster *db.Cluster, nodeName string, projectName string, typeCode db.WarningType, entityTypeCode int, entityID int) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarningsByType(typeCode)
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName || w.Project != projectName || w.EntityTypeCode != entityTypeCode || entityID != entityID {
				continue
			}

			err = tx.UpdateWarningStatus(w.UUID, db.WarningStatusResolved)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// ResolveWarningsByLocalNodeAndProjectAndTypeAndEntity resolves warnings with the given project, type code, and entity.
func ResolveWarningsByLocalNodeAndProjectAndTypeAndEntity(cluster *db.Cluster, projectName string, typeCode db.WarningType, entityTypeCode int, entityID int) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return ResolveWarningsByNodeAndProjectAndTypeAndEntity(cluster, localName, projectName, typeCode, entityTypeCode, entityID)
}

// DeleteWarningsByNodeAndProjectAndEntity deletes warnings with the given node, project, and entity.
func DeleteWarningsByNodeAndProjectAndEntity(cluster *db.Cluster, nodeName string, projectName string, entityTypeCode int, entityID int) error {
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		warnings, err := tx.GetWarnings()
		if err != nil {
			return err
		}

		for _, w := range warnings {
			if w.Node != nodeName || w.Project != projectName || w.EntityTypeCode != entityTypeCode || entityID != entityID {
				continue
			}

			err = tx.DeleteWarning(w.UUID)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed to resolve warnings")
	}

	return nil
}

// DeleteWarningsByLocalNodeAndProjectAndEntity deletes warnings with the given project, and entity.
func DeleteWarningsByLocalNodeAndProjectAndEntity(cluster *db.Cluster, projectName string, entityTypeCode int, entityID int) error {
	var err error
	var localName string

	err = cluster.Transaction(func(tx *db.ClusterTx) error {
		localName, err = tx.GetLocalNodeName()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return errors.Wrap(err, "Failed getting local member name")
	}

	if localName == "" {
		return fmt.Errorf("Local member name not available")
	}

	return DeleteWarningsByNodeAndProjectAndEntity(cluster, localName, projectName, entityTypeCode, entityID)
}
<|endoftext|>