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

//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris

package syscall_test

import (
	"internal/testenv"
	"io"
	"math/rand"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"syscall"
	"testing"
	"time"
	"unsafe"
)

type command struct {
	pipe io.WriteCloser
	proc *exec.Cmd
	test *testing.T
}

func (c *command) Info() (pid, pgrp int) {
	pid = c.proc.Process.Pid

	pgrp, err := syscall.Getpgid(pid)
	if err != nil {
		c.test.Fatal(err)
	}

	return
}

func (c *command) Start() {
	if err := c.proc.Start(); err != nil {
		c.test.Fatal(err)
	}
}

func (c *command) Stop() {
	c.pipe.Close()
	if err := c.proc.Wait(); err != nil {
		c.test.Fatal(err)
	}
}

func create(t *testing.T) *command {
	testenv.MustHaveExec(t)

	proc := exec.Command("cat")
	stdin, err := proc.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}

	return &command{stdin, proc, t}
}

func parent() (pid, pgrp int) {
	return syscall.Getpid(), syscall.Getpgrp()
}

func TestZeroSysProcAttr(t *testing.T) {
	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.Start()
	defer cmd.Stop()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp != ppgrp {
		t.Fatalf("Child is not in parent's process group")
	}
}

func TestSetpgid(t *testing.T) {
	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Start()
	defer cmd.Stop()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp == ppgrp {
		t.Fatalf("Parent and child are in the same process group")
	}

	if cpid != cpgrp {
		t.Fatalf("Child's process group is not the child's process ID")
	}
}

func TestPgid(t *testing.T) {
	ppid, ppgrp := parent()

	cmd1 := create(t)

	cmd1.proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd1.Start()
	defer cmd1.Stop()

	cpid1, cpgrp1 := cmd1.Info()

	if cpid1 == ppid {
		t.Fatalf("Parent and child 1 have the same process ID")
	}

	if cpgrp1 == ppgrp {
		t.Fatalf("Parent and child 1 are in the same process group")
	}

	if cpid1 != cpgrp1 {
		t.Fatalf("Child 1's process group is not its process ID")
	}

	cmd2 := create(t)

	cmd2.proc.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
		Pgid:    cpgrp1,
	}
	cmd2.Start()
	defer cmd2.Stop()

	cpid2, cpgrp2 := cmd2.Info()

	if cpid2 == ppid {
		t.Fatalf("Parent and child 2 have the same process ID")
	}

	if cpgrp2 == ppgrp {
		t.Fatalf("Parent and child 2 are in the same process group")
	}

	if cpid2 == cpgrp2 {
		t.Fatalf("Child 2's process group is its process ID")
	}

	if cpid1 == cpid2 {
		t.Fatalf("Child 1 and 2 have the same process ID")
	}

	if cpgrp1 != cpgrp2 {
		t.Fatalf("Child 1 and 2 are not in the same process group")
	}
}

func TestForeground(t *testing.T) {
	signal.Ignore(syscall.SIGTTIN, syscall.SIGTTOU)
	defer signal.Reset()

	tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0)
	if err != nil {
		t.Skipf("Can't test Foreground. Couldn't open /dev/tty: %s", err)
	}
	defer tty.Close()

	// This should really be pid_t, however _C_int (aka int32) is generally
	// equivalent.
	fpgrp := int32(0)

	errno := syscall.Ioctl(tty.Fd(), syscall.TIOCGPGRP, uintptr(unsafe.Pointer(&fpgrp)))
	if errno != 0 {
		t.Fatalf("TIOCGPGRP failed with error code: %s", errno)
	}

	if fpgrp == 0 {
		t.Fatalf("Foreground process group is zero")
	}

	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.proc.SysProcAttr = &syscall.SysProcAttr{
		Ctty:       int(tty.Fd()),
		Foreground: true,
	}
	cmd.Start()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp == ppgrp {
		t.Fatalf("Parent and child are in the same process group")
	}

	if cpid != cpgrp {
		t.Fatalf("Child's process group is not the child's process ID")
	}

	cmd.Stop()

	// This call fails on darwin/arm64. The failure doesn't matter, though.
	// This is just best effort.
	syscall.Ioctl(tty.Fd(), syscall.TIOCSPGRP, uintptr(unsafe.Pointer(&fpgrp)))
}

func TestForegroundSignal(t *testing.T) {
	tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0)
	if err != nil {
		t.Skipf("couldn't open /dev/tty: %s", err)
	}
	defer tty.Close()

	// This should really be pid_t, however _C_int (aka int32) is generally
	// equivalent.
	fpgrp := int32(0)

	errno := syscall.Ioctl(tty.Fd(), syscall.TIOCGPGRP, uintptr(unsafe.Pointer(&fpgrp)))
	if errno != 0 {
		t.Fatalf("TIOCGPGRP failed with error code: %s", errno)
	}

	if fpgrp == 0 {
		t.Fatalf("Foreground process group is zero")
	}

	defer func() {
		signal.Ignore(syscall.SIGTTIN, syscall.SIGTTOU)
		syscall.Ioctl(tty.Fd(), syscall.TIOCSPGRP, uintptr(unsafe.Pointer(&fpgrp)))
		signal.Reset()
	}()

	ch1 := make(chan os.Signal, 1)
	ch2 := make(chan bool)

	signal.Notify(ch1, syscall.SIGTTIN, syscall.SIGTTOU)
	defer signal.Stop(ch1)

	cmd := create(t)

	go func() {
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Ctty:       int(tty.Fd()),
			Foreground: true,
		}
		cmd.Start()
		cmd.Stop()
		close(ch2)
	}()

	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()
	for {
		select {
		case sig := <-ch1:
			t.Errorf("unexpected signal %v", sig)
		case <-ch2:
			// Success.
			return
		case <-timer.C:
			t.Fatal("timed out waiting for child process")
		}
	}
}

// Test a couple of cases that SysProcAttr can't handle. Issue 29458.
func TestInvalidExec(t *testing.T) {
	t.Parallel()
	t.Run("SetCtty-Foreground", func(t *testing.T) {
		t.Parallel()
		cmd := create(t)
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Setctty:    true,
			Foreground: true,
			Ctty:       0,
		}
		if err := cmd.proc.Start(); err == nil {
			t.Error("expected error setting both SetCtty and Foreground")
		}
	})
	t.Run("invalid-Ctty", func(t *testing.T) {
		t.Parallel()
		cmd := create(t)
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Setctty: true,
			Ctty:    3,
		}
		if err := cmd.proc.Start(); err == nil {
			t.Error("expected error with invalid Ctty value")
		}
	})
}

// TestExec is for issue #41702.
func TestExec(t *testing.T) {
	testenv.MustHaveExec(t)
	cmd := exec.Command(os.Args[0], "-test.run=TestExecHelper")
	cmd.Env = append(os.Environ(), "GO_WANT_HELPER_PROCESS=2")
	o, err := cmd.CombinedOutput()
	if err != nil {
		t.Errorf("%s\n%v", o, err)
	}
}

// TestExecHelper is used by TestExec. It does nothing by itself.
// In testing on macOS 10.14, this used to fail with
// "signal: illegal instruction" more than half the time.
func TestExecHelper(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "2" {
		return
	}

	// We don't have to worry about restoring these values.
	// We are in a child process that only runs this test,
	// and we are going to call syscall.Exec anyhow.
	runtime.GOMAXPROCS(50)
	os.Setenv("GO_WANT_HELPER_PROCESS", "3")

	stop := time.Now().Add(time.Second)
	for i := 0; i < 100; i++ {
		go func(i int) {
			r := rand.New(rand.NewSource(int64(i)))
			for time.Now().Before(stop) {
				r.Uint64()
			}
		}(i)
	}

	time.Sleep(10 * time.Millisecond)

	argv := []string{os.Args[0], "-test.run=TestExecHelper"}
	syscall.Exec(os.Args[0], argv, os.Environ())

	t.Error("syscall.Exec returned")
}
<commit_msg>syscall: remove GOMAXPROCS change in TestExecHelper<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris

package syscall_test

import (
	"internal/testenv"
	"io"
	"math/rand"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
	"testing"
	"time"
	"unsafe"
)

type command struct {
	pipe io.WriteCloser
	proc *exec.Cmd
	test *testing.T
}

func (c *command) Info() (pid, pgrp int) {
	pid = c.proc.Process.Pid

	pgrp, err := syscall.Getpgid(pid)
	if err != nil {
		c.test.Fatal(err)
	}

	return
}

func (c *command) Start() {
	if err := c.proc.Start(); err != nil {
		c.test.Fatal(err)
	}
}

func (c *command) Stop() {
	c.pipe.Close()
	if err := c.proc.Wait(); err != nil {
		c.test.Fatal(err)
	}
}

func create(t *testing.T) *command {
	testenv.MustHaveExec(t)

	proc := exec.Command("cat")
	stdin, err := proc.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}

	return &command{stdin, proc, t}
}

func parent() (pid, pgrp int) {
	return syscall.Getpid(), syscall.Getpgrp()
}

func TestZeroSysProcAttr(t *testing.T) {
	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.Start()
	defer cmd.Stop()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp != ppgrp {
		t.Fatalf("Child is not in parent's process group")
	}
}

func TestSetpgid(t *testing.T) {
	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Start()
	defer cmd.Stop()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp == ppgrp {
		t.Fatalf("Parent and child are in the same process group")
	}

	if cpid != cpgrp {
		t.Fatalf("Child's process group is not the child's process ID")
	}
}

func TestPgid(t *testing.T) {
	ppid, ppgrp := parent()

	cmd1 := create(t)

	cmd1.proc.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd1.Start()
	defer cmd1.Stop()

	cpid1, cpgrp1 := cmd1.Info()

	if cpid1 == ppid {
		t.Fatalf("Parent and child 1 have the same process ID")
	}

	if cpgrp1 == ppgrp {
		t.Fatalf("Parent and child 1 are in the same process group")
	}

	if cpid1 != cpgrp1 {
		t.Fatalf("Child 1's process group is not its process ID")
	}

	cmd2 := create(t)

	cmd2.proc.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
		Pgid:    cpgrp1,
	}
	cmd2.Start()
	defer cmd2.Stop()

	cpid2, cpgrp2 := cmd2.Info()

	if cpid2 == ppid {
		t.Fatalf("Parent and child 2 have the same process ID")
	}

	if cpgrp2 == ppgrp {
		t.Fatalf("Parent and child 2 are in the same process group")
	}

	if cpid2 == cpgrp2 {
		t.Fatalf("Child 2's process group is its process ID")
	}

	if cpid1 == cpid2 {
		t.Fatalf("Child 1 and 2 have the same process ID")
	}

	if cpgrp1 != cpgrp2 {
		t.Fatalf("Child 1 and 2 are not in the same process group")
	}
}

func TestForeground(t *testing.T) {
	signal.Ignore(syscall.SIGTTIN, syscall.SIGTTOU)
	defer signal.Reset()

	tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0)
	if err != nil {
		t.Skipf("Can't test Foreground. Couldn't open /dev/tty: %s", err)
	}
	defer tty.Close()

	// This should really be pid_t, however _C_int (aka int32) is generally
	// equivalent.
	fpgrp := int32(0)

	errno := syscall.Ioctl(tty.Fd(), syscall.TIOCGPGRP, uintptr(unsafe.Pointer(&fpgrp)))
	if errno != 0 {
		t.Fatalf("TIOCGPGRP failed with error code: %s", errno)
	}

	if fpgrp == 0 {
		t.Fatalf("Foreground process group is zero")
	}

	ppid, ppgrp := parent()

	cmd := create(t)

	cmd.proc.SysProcAttr = &syscall.SysProcAttr{
		Ctty:       int(tty.Fd()),
		Foreground: true,
	}
	cmd.Start()

	cpid, cpgrp := cmd.Info()

	if cpid == ppid {
		t.Fatalf("Parent and child have the same process ID")
	}

	if cpgrp == ppgrp {
		t.Fatalf("Parent and child are in the same process group")
	}

	if cpid != cpgrp {
		t.Fatalf("Child's process group is not the child's process ID")
	}

	cmd.Stop()

	// This call fails on darwin/arm64. The failure doesn't matter, though.
	// This is just best effort.
	syscall.Ioctl(tty.Fd(), syscall.TIOCSPGRP, uintptr(unsafe.Pointer(&fpgrp)))
}

func TestForegroundSignal(t *testing.T) {
	tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0)
	if err != nil {
		t.Skipf("couldn't open /dev/tty: %s", err)
	}
	defer tty.Close()

	// This should really be pid_t, however _C_int (aka int32) is generally
	// equivalent.
	fpgrp := int32(0)

	errno := syscall.Ioctl(tty.Fd(), syscall.TIOCGPGRP, uintptr(unsafe.Pointer(&fpgrp)))
	if errno != 0 {
		t.Fatalf("TIOCGPGRP failed with error code: %s", errno)
	}

	if fpgrp == 0 {
		t.Fatalf("Foreground process group is zero")
	}

	defer func() {
		signal.Ignore(syscall.SIGTTIN, syscall.SIGTTOU)
		syscall.Ioctl(tty.Fd(), syscall.TIOCSPGRP, uintptr(unsafe.Pointer(&fpgrp)))
		signal.Reset()
	}()

	ch1 := make(chan os.Signal, 1)
	ch2 := make(chan bool)

	signal.Notify(ch1, syscall.SIGTTIN, syscall.SIGTTOU)
	defer signal.Stop(ch1)

	cmd := create(t)

	go func() {
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Ctty:       int(tty.Fd()),
			Foreground: true,
		}
		cmd.Start()
		cmd.Stop()
		close(ch2)
	}()

	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()
	for {
		select {
		case sig := <-ch1:
			t.Errorf("unexpected signal %v", sig)
		case <-ch2:
			// Success.
			return
		case <-timer.C:
			t.Fatal("timed out waiting for child process")
		}
	}
}

// Test a couple of cases that SysProcAttr can't handle. Issue 29458.
func TestInvalidExec(t *testing.T) {
	t.Parallel()
	t.Run("SetCtty-Foreground", func(t *testing.T) {
		t.Parallel()
		cmd := create(t)
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Setctty:    true,
			Foreground: true,
			Ctty:       0,
		}
		if err := cmd.proc.Start(); err == nil {
			t.Error("expected error setting both SetCtty and Foreground")
		}
	})
	t.Run("invalid-Ctty", func(t *testing.T) {
		t.Parallel()
		cmd := create(t)
		cmd.proc.SysProcAttr = &syscall.SysProcAttr{
			Setctty: true,
			Ctty:    3,
		}
		if err := cmd.proc.Start(); err == nil {
			t.Error("expected error with invalid Ctty value")
		}
	})
}

// TestExec is for issue #41702.
func TestExec(t *testing.T) {
	testenv.MustHaveExec(t)
	cmd := exec.Command(os.Args[0], "-test.run=TestExecHelper")
	cmd.Env = append(os.Environ(), "GO_WANT_HELPER_PROCESS=2")
	o, err := cmd.CombinedOutput()
	if err != nil {
		t.Errorf("%s\n%v", o, err)
	}
}

// TestExecHelper is used by TestExec. It does nothing by itself.
// In testing on macOS 10.14, this used to fail with
// "signal: illegal instruction" more than half the time.
func TestExecHelper(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "2" {
		return
	}

	// We don't have to worry about restoring these values.
	// We are in a child process that only runs this test,
	// and we are going to call syscall.Exec anyhow.
	os.Setenv("GO_WANT_HELPER_PROCESS", "3")

	stop := time.Now().Add(time.Second)
	for i := 0; i < 100; i++ {
		go func(i int) {
			r := rand.New(rand.NewSource(int64(i)))
			for time.Now().Before(stop) {
				r.Uint64()
			}
		}(i)
	}

	time.Sleep(10 * time.Millisecond)

	argv := []string{os.Args[0], "-test.run=TestExecHelper"}
	syscall.Exec(os.Args[0], argv, os.Environ())

	t.Error("syscall.Exec returned")
}
<|endoftext|>


<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
	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
)

var createSchemaSQL = `
CREATE TABLE IF NOT EXISTS sets (
	id integer NOT NULL PRIMARY KEY,
	thumb_photo_id integer UNIQUE REFERENCES photos,
	name varchar(4096) NOT NULL UNIQUE,
	photos_count integer,
	taken_at datetime
);

CREATE INDEX IF NOT EXISTS sets_thumb_photo_id_index ON sets (thumb_photo_id);

CREATE TABLE IF NOT EXISTS photos (
	id integer NOT NULL PRIMARY KEY,
	set_id integer NOT NULL REFERENCES sets,
	prev_photo_id integer UNIQUE REFERENCES photos,
	next_photo_id integer UNIQUE REFERENCES photos,
	path varchar(4096) NOT NULL UNIQUE,
	size integer NOT NULL,
	width integer NOT NULL,
	height integer NOT NULL,
	aperture decimal(2, 1),
	camera varchar(1000),
	exposure_comp integer,
	exposure_time decimal(9, 5),
	flash varchar(51),
	focal_length decimal(3, 1),
	focal_length_35 integer,
	iso integer,
	lat decimal(9, 6),
	lens varchar(1000),
	lng decimal(9, 6),
	taken_at datetime
);

CREATE INDEX IF NOT EXISTS photos_set_id_index ON photos (set_id);

CREATE INDEX IF NOT EXISTS photos_prev_photo_id_index ON photos (prev_photo_id);

CREATE INDEX IF NOT EXISTS photos_next_photo_id_index ON photos (next_photo_id);
`

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

	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

	setName := filepath.Base(filepath.Dir(photo.Path))
	row := selectSetStmt.QueryRow(setName)
	if err := row.Scan(&setId); err == sql.ErrNoRows { // set does not exist
		result, err := insertSetStmt.Exec(setName) // 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.Printf("photos id=%d path=%s\n", photoId, photo.Path)
	}

	return nil
}

func walkPath(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.Printf("photos id=%d prev_photo_id=%d\n", id, prevId)
			updateNextPhotoStmt.Exec(id, prevId)
			fmt.Printf("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.Printf("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()

	_, err = db.Exec(createSchemaSQL)
	if err != nil {
		log.Fatal(err)
	}

	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], walkPath)
	}

	updatePhotoSiblings()
	updateSets()
}
<commit_msg>Break long var line<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
	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   *sql.Stmt
	selectPhotoStmt *sql.Stmt
	insertSetStmt   *sql.Stmt
	insertPhotoStmt *sql.Stmt
)

var createSchemaSQL = `
CREATE TABLE IF NOT EXISTS sets (
	id integer NOT NULL PRIMARY KEY,
	thumb_photo_id integer UNIQUE REFERENCES photos,
	name varchar(4096) NOT NULL UNIQUE,
	photos_count integer,
	taken_at datetime
);

CREATE INDEX IF NOT EXISTS sets_thumb_photo_id_index ON sets (thumb_photo_id);

CREATE TABLE IF NOT EXISTS photos (
	id integer NOT NULL PRIMARY KEY,
	set_id integer NOT NULL REFERENCES sets,
	prev_photo_id integer UNIQUE REFERENCES photos,
	next_photo_id integer UNIQUE REFERENCES photos,
	path varchar(4096) NOT NULL UNIQUE,
	size integer NOT NULL,
	width integer NOT NULL,
	height integer NOT NULL,
	aperture decimal(2, 1),
	camera varchar(1000),
	exposure_comp integer,
	exposure_time decimal(9, 5),
	flash varchar(51),
	focal_length decimal(3, 1),
	focal_length_35 integer,
	iso integer,
	lat decimal(9, 6),
	lens varchar(1000),
	lng decimal(9, 6),
	taken_at datetime
);

CREATE INDEX IF NOT EXISTS photos_set_id_index ON photos (set_id);

CREATE INDEX IF NOT EXISTS photos_prev_photo_id_index ON photos (prev_photo_id);

CREATE INDEX IF NOT EXISTS photos_next_photo_id_index ON photos (next_photo_id);
`

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

	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

	setName := filepath.Base(filepath.Dir(photo.Path))
	row := selectSetStmt.QueryRow(setName)
	if err := row.Scan(&setId); err == sql.ErrNoRows { // set does not exist
		result, err := insertSetStmt.Exec(setName) // 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.Printf("photos id=%d path=%s\n", photoId, photo.Path)
	}

	return nil
}

func walkPath(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.Printf("photos id=%d prev_photo_id=%d\n", id, prevId)
			updateNextPhotoStmt.Exec(id, prevId)
			fmt.Printf("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.Printf("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()

	_, err = db.Exec(createSchemaSQL)
	if err != nil {
		log.Fatal(err)
	}

	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], walkPath)
	}

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


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// A webserver that only serves a 404 page. Used as a default backend.
package main

import (
	"flag"
	"fmt"
	"net/http"
)

var port = flag.Int("port", 8080, "Port number to serve default backend 404 page.")

func main() {
	flag.Parse()
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "default backend - 404")
	})
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, "ok")
	})
	http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)
}
<commit_msg>Catch error from listenandserve<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// A webserver that only serves a 404 page. Used as a default backend.
package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
)

var port = flag.Int("port", 8080, "Port number to serve default backend 404 page.")

func main() {
	flag.Parse()
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "default backend - 404")
	})
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, "ok")
	})
	http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)
	// TODO: Use .Shutdown from Go 1.8
	err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)
	if err != nil {
		fmt.Fprintf(os.Stderr, "could not start http server: %s\n", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestGetAllRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rr := recordedRequests{}
	err = json.Unmarshal(body, &rr)

	expect(t, len(rr.Data), 0)
}

func TestGetAllRecordsWRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}
	// performing query
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rr := recordedRequests{}
	err = json.Unmarshal(body, &rr)

	expect(t, len(rr.Data), 5)
}

func TestGetRecordsCount(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/count", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rc := recordsCount{}
	err = json.Unmarshal(body, &rc)

	expect(t, rc.Count, 0)
}

func TestGetRecordsCountWRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}
	// performing query
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/count", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rc := recordsCount{}
	err = json.Unmarshal(body, &rc)

	expect(t, rc.Count, 5)
}

func TestExportImportRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	// deleting records
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	expect(t, err, nil)

	// using body to import records again
	importReq, err := http.NewRequest("POST", "/records", ioutil.NopCloser(bytes.NewBuffer(body)))
	//The response recorder used to record HTTP responses
	importRec := httptest.NewRecorder()

	m.ServeHTTP(importRec, importReq)
	expect(t, importRec.Code, http.StatusOK)

	// records should be there
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 5)

}

func TestDeleteHandler(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	// checking whether we have records
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 5)

	// deleting through handler
	deleteReq, err := http.NewRequest("DELETE", "/records", nil)
	//The response recorder used to record HTTP responses
	rec := httptest.NewRecorder()

	m.ServeHTTP(rec, deleteReq)
	expect(t, rec.Code, http.StatusOK)
}

func TestDeleteHandlerNoBucket(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// deleting through handler
	importReq, err := http.NewRequest("DELETE", "/records", nil)
	expect(t, err, nil)
	//The response recorder used to record HTTP responses
	importRec := httptest.NewRecorder()

	m.ServeHTTP(importRec, importReq)
	expect(t, importRec.Code, http.StatusOK)
}

<commit_msg>test get state<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestGetAllRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rr := recordedRequests{}
	err = json.Unmarshal(body, &rr)

	expect(t, len(rr.Data), 0)
}

func TestGetAllRecordsWRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}
	// performing query
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rr := recordedRequests{}
	err = json.Unmarshal(body, &rr)

	expect(t, len(rr.Data), 5)
}

func TestGetRecordsCount(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/count", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rc := recordsCount{}
	err = json.Unmarshal(body, &rc)

	expect(t, rc.Count, 0)
}

func TestGetRecordsCountWRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}
	// performing query
	m := getBoneRouter(*dbClient)

	req, err := http.NewRequest("GET", "/count", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	rc := recordsCount{}
	err = json.Unmarshal(body, &rc)

	expect(t, rc.Count, 5)
}

func TestExportImportRecords(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	req, err := http.NewRequest("GET", "/records", nil)
	expect(t, err, nil)

	//The response recorder used to record HTTP responses
	respRec := httptest.NewRecorder()

	m.ServeHTTP(respRec, req)

	expect(t, respRec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(respRec.Body)

	// deleting records
	err = dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	expect(t, err, nil)

	// using body to import records again
	importReq, err := http.NewRequest("POST", "/records", ioutil.NopCloser(bytes.NewBuffer(body)))
	//The response recorder used to record HTTP responses
	importRec := httptest.NewRecorder()

	m.ServeHTTP(importRec, importReq)
	expect(t, importRec.Code, http.StatusOK)

	// records should be there
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 5)

}

func TestDeleteHandler(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// inserting some payloads
	for i := 0; i < 5; i++ {
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com/q=%d", i), nil)
		expect(t, err, nil)
		dbClient.captureRequest(req)
	}

	// checking whether we have records
	payloads, err := dbClient.cache.GetAllRequests()
	expect(t, err, nil)
	expect(t, len(payloads), 5)

	// deleting through handler
	deleteReq, err := http.NewRequest("DELETE", "/records", nil)
	//The response recorder used to record HTTP responses
	rec := httptest.NewRecorder()

	m.ServeHTTP(rec, deleteReq)
	expect(t, rec.Code, http.StatusOK)
}

func TestDeleteHandlerNoBucket(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// deleting through handler
	importReq, err := http.NewRequest("DELETE", "/records", nil)
	expect(t, err, nil)
	//The response recorder used to record HTTP responses
	importRec := httptest.NewRecorder()

	m.ServeHTTP(importRec, importReq)
	expect(t, importRec.Code, http.StatusOK)
}

func TestGetState(t *testing.T) {
	server, dbClient := testTools(200, `{'message': 'here'}`)
	defer server.Close()
	defer dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
	m := getBoneRouter(*dbClient)

	// setting initial mode
	dbClient.cfg.SetMode("virtualize")

	// deleting through handler
	req, err := http.NewRequest("GET", "/state", nil)
	expect(t, err, nil)
	//The response recorder used to record HTTP responses
	rec := httptest.NewRecorder()

	m.ServeHTTP(rec, req)
	expect(t, rec.Code, http.StatusOK)

	body, err := ioutil.ReadAll(rec.Body)

	sr := stateRequest{}
	err = json.Unmarshal(body, &sr)

	expect(t, sr.Mode, "virtualize")

}
<|endoftext|>


<commit_before>package vault

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"

	"github.com/hashicorp/vault/helper/pgpkeys"
	"github.com/hashicorp/vault/shamir"
)

// InitParams keeps the init function from being littered with too many
// params, that's it!
type InitParams struct {
	BarrierConfig   *SealConfig
	RecoveryConfig  *SealConfig
	RootTokenPGPKey string
}

// InitResult is used to provide the key parts back after
// they are generated as part of the initialization.
type InitResult struct {
	SecretShares   [][]byte
	RecoveryShares [][]byte
	RootToken      string
}

// Initialized checks if the Vault is already initialized
func (c *Core) Initialized() (bool, error) {
	// Check the barrier first
	init, err := c.barrier.Initialized()
	if err != nil {
		c.logger.Error("core: barrier init check failed", "error", err)
		return false, err
	}
	if !init {
		c.logger.Info("core: security barrier not initialized")
		return false, nil
	}

	// Verify the seal configuration
	sealConf, err := c.seal.BarrierConfig()
	if err != nil {
		return false, err
	}
	if sealConf == nil {
		return false, fmt.Errorf("core: barrier reports initialized but no seal configuration found")
	}

	return true, nil
}

func (c *Core) generateShares(sc *SealConfig) ([]byte, [][]byte, error) {
	// Generate a master key
	masterKey, err := c.barrier.GenerateKey()
	if err != nil {
		return nil, nil, fmt.Errorf("key generation failed: %v", err)
	}

	// Return the master key if only a single key part is used
	var unsealKeys [][]byte
	if sc.SecretShares == 1 {
		unsealKeys = append(unsealKeys, masterKey)
	} else {
		// Split the master key using the Shamir algorithm
		shares, err := shamir.Split(masterKey, sc.SecretShares, sc.SecretThreshold)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to generate barrier shares: %v", err)
		}
		unsealKeys = shares
	}

	// If we have PGP keys, perform the encryption
	if len(sc.PGPKeys) > 0 {
		hexEncodedShares := make([][]byte, len(unsealKeys))
		for i, _ := range unsealKeys {
			hexEncodedShares[i] = []byte(hex.EncodeToString(unsealKeys[i]))
		}
		_, encryptedShares, err := pgpkeys.EncryptShares(hexEncodedShares, sc.PGPKeys)
		if err != nil {
			return nil, nil, err
		}
		unsealKeys = encryptedShares
	}

	return masterKey, unsealKeys, nil
}

// Initialize is used to initialize the Vault with the given
// configurations.
func (c *Core) Initialize(initParams *InitParams) (*InitResult, error) {
	barrierConfig := initParams.BarrierConfig
	recoveryConfig := initParams.RecoveryConfig

	if c.seal.RecoveryKeySupported() {
		if recoveryConfig == nil {
			return nil, fmt.Errorf("recovery configuration must be supplied")
		}

		if recoveryConfig.SecretShares < 1 {
			return nil, fmt.Errorf("recovery configuration must specify a positive number of shares")
		}

		// Check if the seal configuration is valid
		if err := recoveryConfig.Validate(); err != nil {
			c.logger.Error("core: invalid recovery configuration", "error", err)
			return nil, fmt.Errorf("invalid recovery configuration: %v", err)
		}
	}

	// Check if the seal configuration is valid
	if err := barrierConfig.Validate(); err != nil {
		c.logger.Error("core: invalid seal configuration", "error", err)
		return nil, fmt.Errorf("invalid seal configuration: %v", err)
	}

	// Avoid an initialization race
	c.stateLock.Lock()
	defer c.stateLock.Unlock()

	// Check if we are initialized
	init, err := c.Initialized()
	if err != nil {
		return nil, err
	}
	if init {
		return nil, ErrAlreadyInit
	}

	err = c.seal.Init()
	if err != nil {
		c.logger.Error("core: failed to initialize seal", "error", err)
		return nil, fmt.Errorf("error initializing seal: %v", err)
	}

	err = c.seal.SetBarrierConfig(barrierConfig)
	if err != nil {
		c.logger.Error("core: failed to save barrier configuration", "error", err)
		return nil, fmt.Errorf("barrier configuration saving failed: %v", err)
	}

	barrierKey, barrierUnsealKeys, err := c.generateShares(barrierConfig)
	if err != nil {
		c.logger.Error("core: error generating shares", "error", err)
		return nil, err
	}

	// If we are storing shares, pop them out of the returned results and push
	// them through the seal
	if barrierConfig.StoredShares > 0 {
		var keysToStore [][]byte
		for i := 0; i < barrierConfig.StoredShares; i++ {
			keysToStore = append(keysToStore, barrierUnsealKeys[0])
			barrierUnsealKeys = barrierUnsealKeys[1:]
		}
		if err := c.seal.SetStoredKeys(keysToStore); err != nil {
			c.logger.Error("core: failed to store keys", "error", err)
			return nil, fmt.Errorf("failed to store keys: %v", err)
		}
	}

	results := &InitResult{
		SecretShares: barrierUnsealKeys,
	}

	// Initialize the barrier
	if err := c.barrier.Initialize(barrierKey); err != nil {
		c.logger.Error("core: failed to initialize barrier", "error", err)
		return nil, fmt.Errorf("failed to initialize barrier: %v", err)
	}
	if c.logger.IsInfo() {
		c.logger.Info("core: security barrier initialized", "shares", barrierConfig.SecretShares, "threshold", barrierConfig.SecretThreshold)
	}

	// Unseal the barrier
	if err := c.barrier.Unseal(barrierKey); err != nil {
		c.logger.Error("core: failed to unseal barrier", "error", err)
		return nil, fmt.Errorf("failed to unseal barrier: %v", err)
	}

	// Ensure the barrier is re-sealed
	defer func() {
		if err := c.barrier.Seal(); err != nil {
			c.logger.Error("core: failed to seal barrier", "error", err)
		}
	}()

	// Perform initial setup
	if err := c.setupCluster(); err != nil {
		c.stateLock.Unlock()
		c.logger.Error("core: cluster setup failed during init", "error", err)
		return nil, err
	}
	if err := c.postUnseal(); err != nil {
		c.logger.Error("core: post-unseal setup failed during init", "error", err)
		return nil, err
	}

	// Save the configuration regardless, but only generate a key if it's not
	// disabled. When using recovery keys they are stored in the barrier, so
	// this must happen post-unseal.
	if c.seal.RecoveryKeySupported() {
		err = c.seal.SetRecoveryConfig(recoveryConfig)
		if err != nil {
			c.logger.Error("core: failed to save recovery configuration", "error", err)
			return nil, fmt.Errorf("recovery configuration saving failed: %v", err)
		}

		if recoveryConfig.SecretShares > 0 {
			recoveryKey, recoveryUnsealKeys, err := c.generateShares(recoveryConfig)
			if err != nil {
				c.logger.Error("core: failed to generate recovery shares", "error", err)
				return nil, err
			}

			err = c.seal.SetRecoveryKey(recoveryKey)
			if err != nil {
				return nil, err
			}

			results.RecoveryShares = recoveryUnsealKeys
		}
	}

	// Generate a new root token
	rootToken, err := c.tokenStore.rootToken()
	if err != nil {
		c.logger.Error("core: root token generation failed", "error", err)
		return nil, err
	}
	results.RootToken = rootToken.ID
	c.logger.Info("core: root token generated")

	if initParams.RootTokenPGPKey != "" {
		_, encryptedVals, err := pgpkeys.EncryptShares([][]byte{[]byte(results.RootToken)}, []string{initParams.RootTokenPGPKey})
		if err != nil {
			c.logger.Error("core: root token encryption failed", "error", err)
			return nil, err
		}
		results.RootToken = base64.StdEncoding.EncodeToString(encryptedVals[0])
	}

	// Prepare to re-seal
	if err := c.preSeal(); err != nil {
		c.logger.Error("core: pre-seal teardown failed", "error", err)
		return nil, err
	}

	return results, nil
}

func (c *Core) UnsealWithStoredKeys() error {
	if !c.seal.StoredKeysSupported() {
		return nil
	}

	sealed, err := c.Sealed()
	if err != nil {
		c.logger.Error("core: error checking sealed status in auto-unseal", "error", err)
		return fmt.Errorf("error checking sealed status in auto-unseal: %s", err)
	}
	if !sealed {
		return nil
	}

	c.logger.Info("core: stored unseal keys supported, attempting fetch")
	keys, err := c.seal.GetStoredKeys()
	if err != nil {
		c.logger.Error("core: fetching stored unseal keys failed", "error", err)
		return &NonFatalError{Err: fmt.Errorf("fetching stored unseal keys failed: %v", err)}
	}
	if len(keys) == 0 {
		c.logger.Warn("core: stored unseal key(s) supported but none found")
	} else {
		unsealed := false
		keysUsed := 0
		for _, key := range keys {
			unsealed, err = c.Unseal(key)
			if err != nil {
				c.logger.Error("core: unseal with stored unseal key failed", "error", err)
				return &NonFatalError{Err: fmt.Errorf("unseal with stored key failed: %v", err)}
			}
			keysUsed += 1
			if unsealed {
				break
			}
		}
		if !unsealed {
			if c.logger.IsWarn() {
				c.logger.Warn("core: stored unseal key(s) used but Vault not unsealed yet", "stored_keys_used", keysUsed)
			}
		} else {
			if c.logger.IsInfo() {
				c.logger.Info("core: successfully unsealed with stored key(s)", "stored_keys_used", keysUsed)
			}
		}
	}

	return nil
}
<commit_msg>Remove errant unlock of state lock<commit_after>package vault

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"

	"github.com/hashicorp/vault/helper/pgpkeys"
	"github.com/hashicorp/vault/shamir"
)

// InitParams keeps the init function from being littered with too many
// params, that's it!
type InitParams struct {
	BarrierConfig   *SealConfig
	RecoveryConfig  *SealConfig
	RootTokenPGPKey string
}

// InitResult is used to provide the key parts back after
// they are generated as part of the initialization.
type InitResult struct {
	SecretShares   [][]byte
	RecoveryShares [][]byte
	RootToken      string
}

// Initialized checks if the Vault is already initialized
func (c *Core) Initialized() (bool, error) {
	// Check the barrier first
	init, err := c.barrier.Initialized()
	if err != nil {
		c.logger.Error("core: barrier init check failed", "error", err)
		return false, err
	}
	if !init {
		c.logger.Info("core: security barrier not initialized")
		return false, nil
	}

	// Verify the seal configuration
	sealConf, err := c.seal.BarrierConfig()
	if err != nil {
		return false, err
	}
	if sealConf == nil {
		return false, fmt.Errorf("core: barrier reports initialized but no seal configuration found")
	}

	return true, nil
}

func (c *Core) generateShares(sc *SealConfig) ([]byte, [][]byte, error) {
	// Generate a master key
	masterKey, err := c.barrier.GenerateKey()
	if err != nil {
		return nil, nil, fmt.Errorf("key generation failed: %v", err)
	}

	// Return the master key if only a single key part is used
	var unsealKeys [][]byte
	if sc.SecretShares == 1 {
		unsealKeys = append(unsealKeys, masterKey)
	} else {
		// Split the master key using the Shamir algorithm
		shares, err := shamir.Split(masterKey, sc.SecretShares, sc.SecretThreshold)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to generate barrier shares: %v", err)
		}
		unsealKeys = shares
	}

	// If we have PGP keys, perform the encryption
	if len(sc.PGPKeys) > 0 {
		hexEncodedShares := make([][]byte, len(unsealKeys))
		for i, _ := range unsealKeys {
			hexEncodedShares[i] = []byte(hex.EncodeToString(unsealKeys[i]))
		}
		_, encryptedShares, err := pgpkeys.EncryptShares(hexEncodedShares, sc.PGPKeys)
		if err != nil {
			return nil, nil, err
		}
		unsealKeys = encryptedShares
	}

	return masterKey, unsealKeys, nil
}

// Initialize is used to initialize the Vault with the given
// configurations.
func (c *Core) Initialize(initParams *InitParams) (*InitResult, error) {
	barrierConfig := initParams.BarrierConfig
	recoveryConfig := initParams.RecoveryConfig

	if c.seal.RecoveryKeySupported() {
		if recoveryConfig == nil {
			return nil, fmt.Errorf("recovery configuration must be supplied")
		}

		if recoveryConfig.SecretShares < 1 {
			return nil, fmt.Errorf("recovery configuration must specify a positive number of shares")
		}

		// Check if the seal configuration is valid
		if err := recoveryConfig.Validate(); err != nil {
			c.logger.Error("core: invalid recovery configuration", "error", err)
			return nil, fmt.Errorf("invalid recovery configuration: %v", err)
		}
	}

	// Check if the seal configuration is valid
	if err := barrierConfig.Validate(); err != nil {
		c.logger.Error("core: invalid seal configuration", "error", err)
		return nil, fmt.Errorf("invalid seal configuration: %v", err)
	}

	// Avoid an initialization race
	c.stateLock.Lock()
	defer c.stateLock.Unlock()

	// Check if we are initialized
	init, err := c.Initialized()
	if err != nil {
		return nil, err
	}
	if init {
		return nil, ErrAlreadyInit
	}

	err = c.seal.Init()
	if err != nil {
		c.logger.Error("core: failed to initialize seal", "error", err)
		return nil, fmt.Errorf("error initializing seal: %v", err)
	}

	err = c.seal.SetBarrierConfig(barrierConfig)
	if err != nil {
		c.logger.Error("core: failed to save barrier configuration", "error", err)
		return nil, fmt.Errorf("barrier configuration saving failed: %v", err)
	}

	barrierKey, barrierUnsealKeys, err := c.generateShares(barrierConfig)
	if err != nil {
		c.logger.Error("core: error generating shares", "error", err)
		return nil, err
	}

	// If we are storing shares, pop them out of the returned results and push
	// them through the seal
	if barrierConfig.StoredShares > 0 {
		var keysToStore [][]byte
		for i := 0; i < barrierConfig.StoredShares; i++ {
			keysToStore = append(keysToStore, barrierUnsealKeys[0])
			barrierUnsealKeys = barrierUnsealKeys[1:]
		}
		if err := c.seal.SetStoredKeys(keysToStore); err != nil {
			c.logger.Error("core: failed to store keys", "error", err)
			return nil, fmt.Errorf("failed to store keys: %v", err)
		}
	}

	results := &InitResult{
		SecretShares: barrierUnsealKeys,
	}

	// Initialize the barrier
	if err := c.barrier.Initialize(barrierKey); err != nil {
		c.logger.Error("core: failed to initialize barrier", "error", err)
		return nil, fmt.Errorf("failed to initialize barrier: %v", err)
	}
	if c.logger.IsInfo() {
		c.logger.Info("core: security barrier initialized", "shares", barrierConfig.SecretShares, "threshold", barrierConfig.SecretThreshold)
	}

	// Unseal the barrier
	if err := c.barrier.Unseal(barrierKey); err != nil {
		c.logger.Error("core: failed to unseal barrier", "error", err)
		return nil, fmt.Errorf("failed to unseal barrier: %v", err)
	}

	// Ensure the barrier is re-sealed
	defer func() {
		if err := c.barrier.Seal(); err != nil {
			c.logger.Error("core: failed to seal barrier", "error", err)
		}
	}()

	// Perform initial setup
	if err := c.setupCluster(); err != nil {
		c.logger.Error("core: cluster setup failed during init", "error", err)
		return nil, err
	}
	if err := c.postUnseal(); err != nil {
		c.logger.Error("core: post-unseal setup failed during init", "error", err)
		return nil, err
	}

	// Save the configuration regardless, but only generate a key if it's not
	// disabled. When using recovery keys they are stored in the barrier, so
	// this must happen post-unseal.
	if c.seal.RecoveryKeySupported() {
		err = c.seal.SetRecoveryConfig(recoveryConfig)
		if err != nil {
			c.logger.Error("core: failed to save recovery configuration", "error", err)
			return nil, fmt.Errorf("recovery configuration saving failed: %v", err)
		}

		if recoveryConfig.SecretShares > 0 {
			recoveryKey, recoveryUnsealKeys, err := c.generateShares(recoveryConfig)
			if err != nil {
				c.logger.Error("core: failed to generate recovery shares", "error", err)
				return nil, err
			}

			err = c.seal.SetRecoveryKey(recoveryKey)
			if err != nil {
				return nil, err
			}

			results.RecoveryShares = recoveryUnsealKeys
		}
	}

	// Generate a new root token
	rootToken, err := c.tokenStore.rootToken()
	if err != nil {
		c.logger.Error("core: root token generation failed", "error", err)
		return nil, err
	}
	results.RootToken = rootToken.ID
	c.logger.Info("core: root token generated")

	if initParams.RootTokenPGPKey != "" {
		_, encryptedVals, err := pgpkeys.EncryptShares([][]byte{[]byte(results.RootToken)}, []string{initParams.RootTokenPGPKey})
		if err != nil {
			c.logger.Error("core: root token encryption failed", "error", err)
			return nil, err
		}
		results.RootToken = base64.StdEncoding.EncodeToString(encryptedVals[0])
	}

	// Prepare to re-seal
	if err := c.preSeal(); err != nil {
		c.logger.Error("core: pre-seal teardown failed", "error", err)
		return nil, err
	}

	return results, nil
}

func (c *Core) UnsealWithStoredKeys() error {
	if !c.seal.StoredKeysSupported() {
		return nil
	}

	sealed, err := c.Sealed()
	if err != nil {
		c.logger.Error("core: error checking sealed status in auto-unseal", "error", err)
		return fmt.Errorf("error checking sealed status in auto-unseal: %s", err)
	}
	if !sealed {
		return nil
	}

	c.logger.Info("core: stored unseal keys supported, attempting fetch")
	keys, err := c.seal.GetStoredKeys()
	if err != nil {
		c.logger.Error("core: fetching stored unseal keys failed", "error", err)
		return &NonFatalError{Err: fmt.Errorf("fetching stored unseal keys failed: %v", err)}
	}
	if len(keys) == 0 {
		c.logger.Warn("core: stored unseal key(s) supported but none found")
	} else {
		unsealed := false
		keysUsed := 0
		for _, key := range keys {
			unsealed, err = c.Unseal(key)
			if err != nil {
				c.logger.Error("core: unseal with stored unseal key failed", "error", err)
				return &NonFatalError{Err: fmt.Errorf("unseal with stored key failed: %v", err)}
			}
			keysUsed += 1
			if unsealed {
				break
			}
		}
		if !unsealed {
			if c.logger.IsWarn() {
				c.logger.Warn("core: stored unseal key(s) used but Vault not unsealed yet", "stored_keys_used", keysUsed)
			}
		} else {
			if c.logger.IsInfo() {
				c.logger.Info("core: successfully unsealed with stored key(s)", "stored_keys_used", keysUsed)
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "flag"
    "encoding/json"
    "os"
    "sync/atomic"
    "log"
	"time"
    "github.com/APTrust/bagman"
    "github.com/APTrust/bagman/fluctus/client"
    "github.com/APTrust/bagman/fluctus/models"
    "github.com/bitly/go-nsq"
	"github.com/nu7hatch/gouuid"
	"github.com/crowdmob/goamz/aws"
)

type Channels struct {
	FedoraChannel    chan *bagman.ProcessResult
    CleanUpChannel   chan *bagman.ProcessResult
    ResultsChannel   chan *bagman.ProcessResult
}

// Global vars.
var channels *Channels
var config bagman.Config
var jsonLog *log.Logger
var messageLog *log.Logger
var succeeded = int64(0)
var failed = int64(0)
var s3Client *bagman.S3Client
var fluctusClient *client.Client

func main() {

    loadConfig()
	err := config.EnsureFluctusConfig()
    if err != nil {
        messageLog.Fatalf("Required Fluctus config vars are missing: %v", err)
    }

	fluctusClient, err = client.New(
		config.FluctusURL,
		config.FluctusAPIVersion,
		os.Getenv("FLUCTUS_API_USER"),
		os.Getenv("FLUCTUS_API_KEY"),
		messageLog)
	if err != nil {
		messageLog.Fatalf("Cannot initialize Fluctus Client: %v", err)
	}

    initChannels()
    initGoRoutines()

	err = initS3Client()
    if err != nil {
        messageLog.Fatalf("Cannot initialize S3Client: %v", err)
    }

	nsqConfig := nsq.NewConfig()
	nsqConfig.Set("max_in_flight", 20)
	nsqConfig.Set("heartbeat_interval", "10s")
	nsqConfig.Set("max_attempts", uint16(config.MaxMetadataAttempts))
	nsqConfig.Set("read_timeout", "60s")
	nsqConfig.Set("write_timeout", "10s")
	nsqConfig.Set("msg_timeout", "60m")
    consumer, err := nsq.NewConsumer(config.MetadataTopic,
		config.MetadataChannel, nsqConfig)
    if err != nil {
        messageLog.Fatalf(err.Error())
    }

    handler := &RecordProcessor{}
    consumer.SetHandler(handler)
    consumer.ConnectToNSQLookupd(config.NsqLookupd)

    // This reader blocks until we get an interrupt, so our program does not exit.
    <-consumer.StopChan
}

func loadConfig() {
    // Load the config or die.
    requestedConfig := flag.String("config", "", "configuration to run")
    flag.Parse()
    config = bagman.LoadRequestedConfig(requestedConfig)
    jsonLog, messageLog = bagman.InitLoggers(config.LogDirectory, "metarecord")
}

// Set up the channels.
func initChannels() {
    workerBufferSize := config.Workers * 10
    channels = &Channels{}
	channels.FedoraChannel = make(chan *bagman.ProcessResult, workerBufferSize)
    channels.CleanUpChannel = make(chan *bagman.ProcessResult, workerBufferSize)
    channels.ResultsChannel = make(chan *bagman.ProcessResult, workerBufferSize)
}

// Set up our go routines. We want to limit the number of
// go routines so we do not have 1000+ simultaneous connections
// to Fluctus. That would just cause Fluctus to crash.
func initGoRoutines() {
    for i := 0; i < config.Workers; i++ {
 		go recordInFedora()
        go logResult()
        go doCleanUp()
    }
}

// Initialize the reusable S3 client.
func initS3Client() (err error) {
    s3Client, err = bagman.NewS3Client(aws.USEast)
	return err
}

type RecordProcessor struct {

}

// MessageHandler handles messages from the queue, putting each
// item into the pipleline.
func (*RecordProcessor) HandleMessage(message *nsq.Message) (error) {
	message.DisableAutoResponse()
    var result bagman.ProcessResult
    err := json.Unmarshal(message.Body, &result)
    if err != nil {
		detailedError := fmt.Errorf(
			"[ERROR] Could not unmarshal JSON data from nsq: %v. JSON: %s",
            err, string(message.Body))
        messageLog.Println("[ERROR]", detailedError)
        message.Finish()
        return detailedError
    }
	result.NsqMessage = message
    channels.FedoraChannel <- &result
    messageLog.Println("[INFO]", "Put", result.S3File.Key.Key, "into Fluctus channel")
    return nil
}

func recordInFedora() {
    for result := range channels.FedoraChannel {
		messageLog.Println("[INFO] Recording Fedora metadata for",
			result.S3File.Key.Key)
		result.NsqMessage.Touch()
		result.Stage = "Record"
		err := recordAllFedoraData(result)
		if err != nil {
			result.ErrorMessage += fmt.Sprintf(" %s",err.Error())
		}
		if result.FedoraResult.AllRecordsSucceeded() == false {
			result.ErrorMessage += " When recording IntellectualObject, GenericFiles and " +
				"PremisEvents, one or more calls to Fluctus failed."
		}
		if result.ErrorMessage == "" {
			messageLog.Println("[INFO] Successfully recorded Fedora metadata for",
				result.S3File.Key.Key)
		} else {
			// If any errors in occur while talking to Fluctus,
			// we'll want to requeue and try again. Just leave
			// the result.Retry flag alone, and that will happen.
			messageLog.Println("[ERROR]", result.ErrorMessage)
		}
        channels.ResultsChannel <- result
	}
}

// TODO: This code is duplicated in bag_processor.go
func logResult() {
    for result := range channels.ResultsChannel {
        // Log full results to the JSON log
        json, err := json.Marshal(result)
        if err != nil {
            messageLog.Println("[ERROR]", err)
        }
        jsonLog.Println(string(json))

        // Add a message to the message log
        if(result.ErrorMessage != "") {
            atomic.AddInt64(&failed, 1)
            messageLog.Println("[ERROR]",
                result.S3File.BucketName,
                result.S3File.Key.Key,
                "->", result.ErrorMessage)
        } else {
            atomic.AddInt64(&succeeded, 1)
            messageLog.Println("[INFO]", result.S3File.Key.Key, "-> finished OK")
        }

        // Add some stats to the message log
        messageLog.Printf("[STATS] Succeeded: %d, Failed: %d\n", succeeded, failed)

		if result.NsqMessage.Attempts >= uint16(config.MaxMetadataAttempts) && result.ErrorMessage != "" {
			result.Retry = false
			result.ErrorMessage += fmt.Sprintf("Failure is due to a technical error " +
				"in Fedora. Giving up after %d failed attempts. This item has been " +
				"queued for administrative review. ",
				result.NsqMessage.Attempts)
			err = bagman.Enqueue(config.NsqdHttpAddress, config.TroubleTopic, result)
			if err != nil {
				messageLog.Printf("[ERROR] Could not send '%s' to trouble queue: %v\n",
					result.S3File.Key.Key, err)
			} else {
				messageLog.Printf("[WARN] Sent '%s' to trouble queue\n", result.S3File.Key.Key)
			}
		}

		// Tell Fluctus what happened
		go func() {
			err := fluctusClient.SendProcessedItem(result.IngestStatus())
			if err != nil {
				result.ErrorMessage += fmt.Sprintf("Attempt to record processed " +
					"item status returned error %v. ", err)
				messageLog.Println("[ERROR] Error sending ProcessedItem to Fluctus:",
					err)
			}
		}()

        // Clean up the bag/tar files
        channels.CleanUpChannel <- result
    }
}

func doCleanUp() {
    for result := range channels.CleanUpChannel {
        messageLog.Println("[INFO]", "Cleaning up", result.S3File.Key.Key)
        // Build and send message back to NSQ, indicating whether
        // processing succeeded.
        if result.ErrorMessage != "" && result.Retry == true {
			messageLog.Printf("[INFO] Requeueing %s", result.S3File.Key.Key)
            result.NsqMessage.Requeue(1 * time.Minute)
        } else {
            result.NsqMessage.Finish()
        }
    }
}

// Send all metadata about the bag to Fluctus/Fedora. This includes
// the IntellectualObject, the GenericFiles, and all PremisEvents
// related to the object and the files.
func recordAllFedoraData(result *bagman.ProcessResult) (err error) {
	intellectualObject, err := result.IntellectualObject()
    if err != nil {
        return err
    }
	result.FedoraResult = bagman.NewFedoraResult(
		intellectualObject.Identifier,
		result.TarResult.GenericFilePaths())
	existingObj, err := fluctusClient.IntellectualObjectGet(
		intellectualObject.Identifier, false)
	if err != nil {
		result.FedoraResult.ErrorMessage = fmt.Sprintf(
			"[ERROR] Error checking Fluctus for existing IntellectualObject '%s': %v",
			intellectualObject.Identifier, err)
		return err
	}
	if existingObj != nil {
		result.FedoraResult.IsNewObject = false
		fedoraUpdateIntellectualObject(result, intellectualObject)
		for i := range(result.TarResult.GenericFiles) {
			genericFile := result.TarResult.GenericFiles[i]
			fedoraRecordGenericFile(result, intellectualObject.Identifier, genericFile)
		}
	} else {
		result.FedoraResult.IsNewObject = true
		newObj, err := fluctusClient.IntellectualObjectCreate(intellectualObject)
		if err != nil {
			result.FedoraResult.ErrorMessage = fmt.Sprintf(
				"[ERROR] Error creating new IntellectualObject '%s' in Fluctus: %v",
				intellectualObject.Identifier, err)
			return err
		} else {
			intellectualObject.Id = newObj.Id
		}
	}
	return nil
}

func fedoraRecordGenericFile(result *bagman.ProcessResult, objId string, gf *bagman.GenericFile) (error) {
	// Save the GenericFile metadata in Fedora, and add a metadata
	// record so we know whether the call to Fluctus succeeded or failed.
	fluctusGenericFile, err := gf.ToFluctusModel()
	if err != nil {
		return fmt.Errorf("Error converting GenericFile to Fluctus model: %v", err)
	}
	_, err = fluctusClient.GenericFileSave(objId, fluctusGenericFile)
	if err != nil {
		handleFedoraError(result,
			fmt.Sprintf("Error saving generic file '%s' to Fedora", objId),
			err)
		return err
	}
	addMetadataRecord(result, "GenericFile", "file_registered", gf.Path, err)

	for _, event := range(fluctusGenericFile.Events) {
		_, err = fluctusClient.PremisEventSave(fluctusGenericFile.Identifier, "GenericFile", event)
		if err != nil {
			message := fmt.Sprintf("Error saving event '%s' for generic file " +
				"'%s' to Fedora", event, objId)
			handleFedoraError(result, message, err)
			return err
		}
		addMetadataRecord(result, "PremisEvent", event.EventType, gf.Path, err)
	}

	return nil
}

// Creates/Updates an IntellectualObject in Fedora, and sends the
// Ingest PremisEvent to Fedora.
func fedoraUpdateIntellectualObject(result *bagman.ProcessResult, intellectualObject *models.IntellectualObject) (error) {
	// Create/Update the IntellectualObject
	savedObj, err := fluctusClient.IntellectualObjectUpdate(intellectualObject)
	if err != nil {
		message := fmt.Sprintf("Error saving intellectual object '%s' to Fedora",
			intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "IntellectualObject", "object_registered", intellectualObject.Identifier, err)
	if savedObj != nil {
		intellectualObject.Id = savedObj.Id
	}

	// Add PremisEvents for the ingest
	eventId, err := uuid.NewV4()
	if err != nil {
		return fmt.Errorf("Error generating UUID for ingest event: %v", err)
	}
	ingestEvent := &models.PremisEvent{
		Identifier: eventId.String(),
		EventType: "ingest",
		DateTime: time.Now(),
		Detail: "Copied all files to perservation bucket",
		Outcome: "Success",
		OutcomeDetail: fmt.Sprintf("%d files copied", len(result.FedoraResult.GenericFilePaths)),
		Object: "goamz S3 client",
		Agent: "https://launchpad.net/goamz",
		OutcomeInformation: "Multipart put using md5 checksum",
	}
	_, err = fluctusClient.PremisEventSave(intellectualObject.Identifier, "IntellectualObject", ingestEvent)
	if err != nil {
		message := fmt.Sprintf("Error saving ingest event for intellectual " +
			"object '%s' to Fedora", intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "PremisEvent", "ingest", intellectualObject.Identifier, err)

	idEvent := &models.PremisEvent{
		Identifier: eventId.String(),
		EventType: "identifier_assignment",
		DateTime: time.Now(),
		Detail: "Assigned bag identifier",
		Outcome: "Success",
		OutcomeDetail: intellectualObject.Identifier,
		Object: "APTrust bagman",
		Agent: "https://github.com/APTrust/bagman",
		OutcomeInformation: "Institution domain + tar file name",
	}
	_, err = fluctusClient.PremisEventSave(intellectualObject.Identifier, "IntellectualObject", idEvent)
	if err != nil {
		message := fmt.Sprintf("Error saving identifier_assignment event for " +
			"intellectual object '%s' to Fedora", intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "PremisEvent", "identifier_assignment", intellectualObject.Identifier, err)

	return nil
}

func addMetadataRecord(result *bagman.ProcessResult, eventType, action, eventObject string, err error) {
	errMsg := ""
	if err != nil {
		errMsg = err.Error()
	}
	// Die on bad records. This is entirely within the developer's control
	// and should never happen.
	recError := result.FedoraResult.AddRecord(eventType, action, eventObject, errMsg)
	if recError != nil {
		messageLog.Fatal(recError)
	}
}

func handleFedoraError(result *bagman.ProcessResult, message string, err error) {
	result.FedoraResult.ErrorMessage = fmt.Sprintf("%s: %v", message, err)
	result.ErrorMessage = result.FedoraResult.ErrorMessage
}
<commit_msg>Added some temp code for debugging<commit_after>package main

import (
    "fmt"
    "flag"
    "encoding/json"
    "os"
    "sync/atomic"
    "log"
	"time"
    "github.com/APTrust/bagman"
    "github.com/APTrust/bagman/fluctus/client"
    "github.com/APTrust/bagman/fluctus/models"
    "github.com/bitly/go-nsq"
	"github.com/nu7hatch/gouuid"
	"github.com/crowdmob/goamz/aws"
)

type Channels struct {
	FedoraChannel    chan *bagman.ProcessResult
    CleanUpChannel   chan *bagman.ProcessResult
    ResultsChannel   chan *bagman.ProcessResult
}

// Global vars.
var channels *Channels
var config bagman.Config
var jsonLog *log.Logger
var messageLog *log.Logger
var succeeded = int64(0)
var failed = int64(0)
var s3Client *bagman.S3Client
var fluctusClient *client.Client

func main() {

    loadConfig()
	err := config.EnsureFluctusConfig()
    if err != nil {
        messageLog.Fatalf("Required Fluctus config vars are missing: %v", err)
    }

	fluctusClient, err = client.New(
		config.FluctusURL,
		config.FluctusAPIVersion,
		os.Getenv("FLUCTUS_API_USER"),
		os.Getenv("FLUCTUS_API_KEY"),
		messageLog)
	if err != nil {
		messageLog.Fatalf("Cannot initialize Fluctus Client: %v", err)
	}

    initChannels()
    initGoRoutines()

	err = initS3Client()
    if err != nil {
        messageLog.Fatalf("Cannot initialize S3Client: %v", err)
    }

	nsqConfig := nsq.NewConfig()
	nsqConfig.Set("max_in_flight", 20)
	nsqConfig.Set("heartbeat_interval", "10s")
	nsqConfig.Set("max_attempts", uint16(config.MaxMetadataAttempts))
	nsqConfig.Set("read_timeout", "60s")
	nsqConfig.Set("write_timeout", "10s")
	nsqConfig.Set("msg_timeout", "60m")
    consumer, err := nsq.NewConsumer(config.MetadataTopic,
		config.MetadataChannel, nsqConfig)
    if err != nil {
        messageLog.Fatalf(err.Error())
    }

    handler := &RecordProcessor{}
    consumer.SetHandler(handler)
    consumer.ConnectToNSQLookupd(config.NsqLookupd)

    // This reader blocks until we get an interrupt, so our program does not exit.
    <-consumer.StopChan
}

func loadConfig() {
    // Load the config or die.
    requestedConfig := flag.String("config", "", "configuration to run")
    flag.Parse()
    config = bagman.LoadRequestedConfig(requestedConfig)
    jsonLog, messageLog = bagman.InitLoggers(config.LogDirectory, "metarecord")
}

// Set up the channels.
func initChannels() {
    workerBufferSize := config.Workers * 10
    channels = &Channels{}
	channels.FedoraChannel = make(chan *bagman.ProcessResult, workerBufferSize)
    channels.CleanUpChannel = make(chan *bagman.ProcessResult, workerBufferSize)
    channels.ResultsChannel = make(chan *bagman.ProcessResult, workerBufferSize)
}

// Set up our go routines. We want to limit the number of
// go routines so we do not have 1000+ simultaneous connections
// to Fluctus. That would just cause Fluctus to crash.
func initGoRoutines() {
    for i := 0; i < config.Workers; i++ {
 		go recordInFedora()
        go logResult()
        go doCleanUp()
    }
}

// Initialize the reusable S3 client.
func initS3Client() (err error) {
    s3Client, err = bagman.NewS3Client(aws.USEast)
	return err
}

type RecordProcessor struct {

}

// MessageHandler handles messages from the queue, putting each
// item into the pipleline.
func (*RecordProcessor) HandleMessage(message *nsq.Message) (error) {
	message.DisableAutoResponse()
    var result bagman.ProcessResult
    err := json.Unmarshal(message.Body, &result)
    if err != nil {
		detailedError := fmt.Errorf(
			"[ERROR] Could not unmarshal JSON data from nsq: %v. JSON: %s",
            err, string(message.Body))
        messageLog.Println("[ERROR]", detailedError)
        message.Finish()
        return detailedError
    }
	result.NsqMessage = message
    channels.FedoraChannel <- &result
    messageLog.Println("[INFO]", "Put", result.S3File.Key.Key, "into Fluctus channel")
    return nil
}

func recordInFedora() {
    for result := range channels.FedoraChannel {
		messageLog.Println("[INFO] Recording Fedora metadata for",
			result.S3File.Key.Key)
		result.NsqMessage.Touch()
		result.Stage = "Record"
		err := recordAllFedoraData(result)
		if err != nil {
			result.ErrorMessage += fmt.Sprintf(" %s",err.Error())
		}
		if result.FedoraResult.AllRecordsSucceeded() == false {
			result.ErrorMessage += " When recording IntellectualObject, GenericFiles and " +
				"PremisEvents, one or more calls to Fluctus failed."
		}
		if result.ErrorMessage == "" {
			messageLog.Println("[INFO] Successfully recorded Fedora metadata for",
				result.S3File.Key.Key)
		} else {
			// If any errors in occur while talking to Fluctus,
			// we'll want to requeue and try again. Just leave
			// the result.Retry flag alone, and that will happen.
			messageLog.Println("[ERROR]", result.ErrorMessage)
		}
        channels.ResultsChannel <- result
	}
}

// TODO: This code is duplicated in bag_processor.go
func logResult() {
    for result := range channels.ResultsChannel {
        // Log full results to the JSON log
        json, err := json.Marshal(result)
        if err != nil {
            messageLog.Println("[ERROR]", err)
        }
        jsonLog.Println(string(json))

        // Add a message to the message log
        if(result.ErrorMessage != "") {
            atomic.AddInt64(&failed, 1)
            messageLog.Println("[ERROR]",
                result.S3File.BucketName,
                result.S3File.Key.Key,
                "->", result.ErrorMessage)
        } else {
            atomic.AddInt64(&succeeded, 1)
            messageLog.Println("[INFO]", result.S3File.Key.Key, "-> finished OK")
        }

        // Add some stats to the message log
        messageLog.Printf("[STATS] Succeeded: %d, Failed: %d\n", succeeded, failed)

		if result.NsqMessage.Attempts >= uint16(config.MaxMetadataAttempts) && result.ErrorMessage != "" {
			result.Retry = false
			result.ErrorMessage += fmt.Sprintf("Failure is due to a technical error " +
				"in Fedora. Giving up after %d failed attempts. This item has been " +
				"queued for administrative review. ",
				result.NsqMessage.Attempts)
			err = bagman.Enqueue(config.NsqdHttpAddress, config.TroubleTopic, result)
			if err != nil {
				messageLog.Printf("[ERROR] Could not send '%s' to trouble queue: %v\n",
					result.S3File.Key.Key, err)
			} else {
				messageLog.Printf("[WARN] Sent '%s' to trouble queue\n", result.S3File.Key.Key)
			}
		}

		// Tell Fluctus what happened
		go func() {
			err := fluctusClient.SendProcessedItem(result.IngestStatus())
			if err != nil {
				result.ErrorMessage += fmt.Sprintf("Attempt to record processed " +
					"item status returned error %v. ", err)
				messageLog.Println("[ERROR] Error sending ProcessedItem to Fluctus:",
					err)
			}
		}()

        // Clean up the bag/tar files
        channels.CleanUpChannel <- result
    }
}

func doCleanUp() {
    for result := range channels.CleanUpChannel {
        messageLog.Println("[INFO]", "Cleaning up", result.S3File.Key.Key)
        // Build and send message back to NSQ, indicating whether
        // processing succeeded.
        if result.ErrorMessage != "" && result.Retry == true {
			messageLog.Printf("[INFO] Requeueing %s", result.S3File.Key.Key)
            result.NsqMessage.Requeue(1 * time.Minute)
        } else {
            result.NsqMessage.Finish()
        }
    }
}

// Send all metadata about the bag to Fluctus/Fedora. This includes
// the IntellectualObject, the GenericFiles, and all PremisEvents
// related to the object and the files.
func recordAllFedoraData(result *bagman.ProcessResult) (err error) {
	intellectualObject, err := result.IntellectualObject()
    if err != nil {
        return err
    }
	result.FedoraResult = bagman.NewFedoraResult(
		intellectualObject.Identifier,
		result.TarResult.GenericFilePaths())
	existingObj, err := fluctusClient.IntellectualObjectGet(
		intellectualObject.Identifier, false)
	if err != nil {
		result.FedoraResult.ErrorMessage = fmt.Sprintf(
			"[ERROR] Error checking Fluctus for existing IntellectualObject '%s': %v",
			intellectualObject.Identifier, err)
		return err
	}
	if existingObj != nil {
		result.FedoraResult.IsNewObject = false
		fedoraUpdateIntellectualObject(result, intellectualObject)
		for i := range(result.TarResult.GenericFiles) {
			genericFile := result.TarResult.GenericFiles[i]
			// -------------------------------------------------------------------------
			// TEMP - For debugging a specific error
			// -------------------------------------------------------------------------
			if genericFile.MimeType == "" {
				messageLog.Printf("[WARN] Generic file %s of object %s has no mime type",
					genericFile.Path, intellectualObject.Identifier)
			}
			// -------------------------------------------------------------------------
			// END OF TEMP CODE
			// -------------------------------------------------------------------------

			fedoraRecordGenericFile(result, intellectualObject.Identifier, genericFile)
		}
	} else {
		result.FedoraResult.IsNewObject = true

		// -------------------------------------------------------------------------
		// TEMP - For debugging a specific error
		// -------------------------------------------------------------------------
		for i := range(result.TarResult.GenericFiles) {
			genericFile := result.TarResult.GenericFiles[i]
			if genericFile.MimeType == "" {
				messageLog.Printf("[WARN] Generic file %s of object %s has no mime type",
					genericFile.Path, intellectualObject.Identifier)
			}
		}
		// -------------------------------------------------------------------------
		// END OF TEMP CODE
		// -------------------------------------------------------------------------

		newObj, err := fluctusClient.IntellectualObjectCreate(intellectualObject)
		if err != nil {
			result.FedoraResult.ErrorMessage = fmt.Sprintf(
				"[ERROR] Error creating new IntellectualObject '%s' in Fluctus: %v",
				intellectualObject.Identifier, err)
			return err
		} else {
			intellectualObject.Id = newObj.Id
		}
	}
	return nil
}

func fedoraRecordGenericFile(result *bagman.ProcessResult, objId string, gf *bagman.GenericFile) (error) {
	// Save the GenericFile metadata in Fedora, and add a metadata
	// record so we know whether the call to Fluctus succeeded or failed.
	fluctusGenericFile, err := gf.ToFluctusModel()
	if err != nil {
		return fmt.Errorf("Error converting GenericFile to Fluctus model: %v", err)
	}
	_, err = fluctusClient.GenericFileSave(objId, fluctusGenericFile)
	if err != nil {
		handleFedoraError(result,
			fmt.Sprintf("Error saving generic file '%s' to Fedora", objId),
			err)
		return err
	}
	addMetadataRecord(result, "GenericFile", "file_registered", gf.Path, err)

	for _, event := range(fluctusGenericFile.Events) {
		_, err = fluctusClient.PremisEventSave(fluctusGenericFile.Identifier, "GenericFile", event)
		if err != nil {
			message := fmt.Sprintf("Error saving event '%s' for generic file " +
				"'%s' to Fedora", event, objId)
			handleFedoraError(result, message, err)
			return err
		}
		addMetadataRecord(result, "PremisEvent", event.EventType, gf.Path, err)
	}

	return nil
}

// Creates/Updates an IntellectualObject in Fedora, and sends the
// Ingest PremisEvent to Fedora.
func fedoraUpdateIntellectualObject(result *bagman.ProcessResult, intellectualObject *models.IntellectualObject) (error) {
	// Create/Update the IntellectualObject
	savedObj, err := fluctusClient.IntellectualObjectUpdate(intellectualObject)
	if err != nil {
		message := fmt.Sprintf("Error saving intellectual object '%s' to Fedora",
			intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "IntellectualObject", "object_registered", intellectualObject.Identifier, err)
	if savedObj != nil {
		intellectualObject.Id = savedObj.Id
	}

	// Add PremisEvents for the ingest
	eventId, err := uuid.NewV4()
	if err != nil {
		return fmt.Errorf("Error generating UUID for ingest event: %v", err)
	}
	ingestEvent := &models.PremisEvent{
		Identifier: eventId.String(),
		EventType: "ingest",
		DateTime: time.Now(),
		Detail: "Copied all files to perservation bucket",
		Outcome: "Success",
		OutcomeDetail: fmt.Sprintf("%d files copied", len(result.FedoraResult.GenericFilePaths)),
		Object: "goamz S3 client",
		Agent: "https://launchpad.net/goamz",
		OutcomeInformation: "Multipart put using md5 checksum",
	}
	_, err = fluctusClient.PremisEventSave(intellectualObject.Identifier, "IntellectualObject", ingestEvent)
	if err != nil {
		message := fmt.Sprintf("Error saving ingest event for intellectual " +
			"object '%s' to Fedora", intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "PremisEvent", "ingest", intellectualObject.Identifier, err)

	idEvent := &models.PremisEvent{
		Identifier: eventId.String(),
		EventType: "identifier_assignment",
		DateTime: time.Now(),
		Detail: "Assigned bag identifier",
		Outcome: "Success",
		OutcomeDetail: intellectualObject.Identifier,
		Object: "APTrust bagman",
		Agent: "https://github.com/APTrust/bagman",
		OutcomeInformation: "Institution domain + tar file name",
	}
	_, err = fluctusClient.PremisEventSave(intellectualObject.Identifier, "IntellectualObject", idEvent)
	if err != nil {
		message := fmt.Sprintf("Error saving identifier_assignment event for " +
			"intellectual object '%s' to Fedora", intellectualObject.Identifier)
		handleFedoraError(result, message, err)
		return err
	}
	addMetadataRecord(result, "PremisEvent", "identifier_assignment", intellectualObject.Identifier, err)

	return nil
}

func addMetadataRecord(result *bagman.ProcessResult, eventType, action, eventObject string, err error) {
	errMsg := ""
	if err != nil {
		errMsg = err.Error()
	}
	// Die on bad records. This is entirely within the developer's control
	// and should never happen.
	recError := result.FedoraResult.AddRecord(eventType, action, eventObject, errMsg)
	if recError != nil {
		messageLog.Fatal(recError)
	}
}

func handleFedoraError(result *bagman.ProcessResult, message string, err error) {
	result.FedoraResult.ErrorMessage = fmt.Sprintf("%s: %v", message, err)
	result.ErrorMessage = result.FedoraResult.ErrorMessage
}
<|endoftext|>


<commit_before>// Package image contains methods and helpers for parsing the docker image
// format. Unfortunately, Docker's codebase does not provide a simple library
// for encoding/decoding the string representation of an image into a consistent
// structure.
//
// The general format is:
//
//	NAME[:TAG|@DIGEST]
//
// Example:
//
//	ubuntu:14.04
//	localhost.localdomain:5000/samalba/hipache:latest
//	localhost:5000/foo/bar@sha256:bc8813ea7b3603864987522f02a76101c17ad122e1c46d790efc0fca78ca7bfb
package image

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// DefaultTag is used when json decoding an image. If there is no tag part
// present, this will be used as the tag.
const DefaultTag = "latest"

var (
	// ErrInvalidImage is returned when the image does not specify a repo.
	ErrInvalidImage = errors.New("image: invalid")

	// ErrInvalidRepo is returned by Split when the repo is not a valid repo.
	ErrInvalidRepo = errors.New("image: not a valid repo")
)

// Image represents all the information about an image.
type Image struct {
	// Registry is the registry that the image belongs to.
	Registry string

	// Repository is the repository part of the image.
	Repository string

	// If provided a tag.
	Tag string

	// If provided, a digest for the image.
	Digest string
}

func (i Image) String() string {
	return Encode(i)
}

// Scan implements the sql.Scanner interface.
func (i *Image) Scan(src interface{}) error {
	if src, ok := src.([]byte); ok {
		image, err := Decode(string(src))
		if err != nil {
			return err
		}

		*i = image
	}

	return nil
}

// Value implements the driver.Value interface.
func (i Image) Value() (driver.Value, error) {
	return driver.Value(i.String()), nil
}

func (i *Image) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}

	image, err := Decode(s)
	if err != nil {
		return err
	}

	*i = image

	return nil
}

// Decode decodes the string representation of an image into an Image structure.
func Decode(in string) (image Image, err error) {
	p := strings.Split(in, ":")

	if len(p) == 0 {
		err = ErrInvalidImage
		return
	}

	image.Registry, image.Repository = splitRepository(p[0])

	if image.Repository == "" {
		err = ErrInvalidImage
		return
	}

	if len(p) == 1 {
		image.Tag = DefaultTag
		return
	}

	image.Tag = p[1]
	return
}

// Encode encodes an Image to it's string representation.
func Encode(image Image) string {
	repo := image.Repository
	if image.Registry != "" {
		repo = fmt.Sprintf("%s/%s", image.Registry, repo)
	}
	return fmt.Sprintf("%s:%s", repo, image.Tag)
}

// splitRepository splits a full docker repo into registry and path segments.
func splitRepository(fullRepo string) (registry string, path string) {
	parts := strings.Split(fullRepo, "/")

	if len(parts) < 2 {
		return "", parts[0]
	}

	if len(parts) == 2 {
		return "", strings.Join(parts, "/")
	}

	return parts[0], strings.Join(parts[1:], "/")
}
<commit_msg>Cleanup.<commit_after>// Package image contains methods and helpers for parsing the docker image
// format. Unfortunately, Docker's codebase does not provide a simple library
// for encoding/decoding the string representation of an image into a consistent
// structure.
//
// The general format is:
//
//	NAME[:TAG|@DIGEST]
//
// Example:
//
//	ubuntu:14.04
//	localhost.localdomain:5000/samalba/hipache:latest
//	localhost:5000/foo/bar@sha256:bc8813ea7b3603864987522f02a76101c17ad122e1c46d790efc0fca78ca7bfb
package image

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// DefaultTag is used when json decoding an image. If there is no tag part
// present, this will be used as the tag.
const DefaultTag = "latest"

// ErrInvalidImage is returned when the image does not specify a repo.
var ErrInvalidImage = errors.New("image: invalid")

// Image represents all the information about an image.
type Image struct {
	// Registry is the registry that the image belongs to.
	Registry string

	// Repository is the repository part of the image.
	Repository string

	// If provided a tag.
	Tag string

	// If provided, a digest for the image.
	Digest string
}

func (i Image) String() string {
	return Encode(i)
}

// Scan implements the sql.Scanner interface.
func (i *Image) Scan(src interface{}) error {
	if src, ok := src.([]byte); ok {
		image, err := Decode(string(src))
		if err != nil {
			return err
		}

		*i = image
	}

	return nil
}

// Value implements the driver.Value interface.
func (i Image) Value() (driver.Value, error) {
	return driver.Value(i.String()), nil
}

func (i *Image) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}

	image, err := Decode(s)
	if err != nil {
		return err
	}

	*i = image

	return nil
}

// Decode decodes the string representation of an image into an Image structure.
func Decode(in string) (image Image, err error) {
	p := strings.Split(in, ":")

	if len(p) == 0 {
		err = ErrInvalidImage
		return
	}

	image.Registry, image.Repository = splitRepository(p[0])

	if image.Repository == "" {
		err = ErrInvalidImage
		return
	}

	if len(p) == 1 {
		image.Tag = DefaultTag
		return
	}

	image.Tag = p[1]
	return
}

// Encode encodes an Image to it's string representation.
func Encode(image Image) string {
	repo := image.Repository
	if image.Registry != "" {
		repo = fmt.Sprintf("%s/%s", image.Registry, repo)
	}
	return fmt.Sprintf("%s:%s", repo, image.Tag)
}

// splitRepository splits a full docker repo into registry and path segments.
func splitRepository(fullRepo string) (registry string, path string) {
	parts := strings.Split(fullRepo, "/")

	if len(parts) < 2 {
		return "", parts[0]
	}

	if len(parts) == 2 {
		return "", strings.Join(parts, "/")
	}

	return parts[0], strings.Join(parts[1:], "/")
}
<|endoftext|>