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

// racebuild builds the race runtime (syso files) on all supported OSes using gomote.
// Usage:
//	$ racebuild -rev <llvm_git_revision> -goroot <path_to_go_repo>
package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/sync/errgroup"
)

var (
	flagGoroot    = flag.String("goroot", "", "path to Go repository to update (required)")
	flagRev       = flag.String("rev", "", "llvm compiler-rt git revision from http://llvm.org/git/compiler-rt.git (required)")
	flagPlatforms = flag.String("platforms", "all", `comma-separated platforms (such as "linux/amd64") to rebuild, or "all"`)
)

// TODO: use buildlet package instead of calling out to gomote.
var platforms = []*Platform{
	&Platform{
		OS:   "freebsd",
		Arch: "amd64",
		Type: "freebsd-amd64-race",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_freebsd_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "darwin",
		Arch: "amd64",
		Type: "darwin-amd64-10_12",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_darwin_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "linux",
		Arch: "amd64",
		Type: "linux-amd64-race",
		Script: `#!/usr/bin/env bash
set -e
apt-get update
apt-get install -y git g++
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_linux_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "linux",
		Arch: "ppc64le",
		Type: "linux-ppc64le-buildlet",
		Script: `#!/usr/bin/env bash
set -e
apt-get update
apt-get install -y git g++
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_linux_ppc64le.syso go/src/runtime/race
# TODO(#23731): Uncomment to test the syso file before accepting it.
# (cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "netbsd",
		Arch: "amd64",
		Type: "netbsd-amd64-8_0",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_netbsd_amd64.syso go/src/runtime/race
# TODO(#24322): Uncomment to test the syso file before accepting it.
# (cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "windows",
		Arch: "amd64",
		Type: "windows-amd64-race",
		Script: `
git clone https://go.googlesource.com/go
if %errorlevel% neq 0 exit /b %errorlevel%
git clone http://llvm.org/git/compiler-rt.git
if %errorlevel% neq 0 exit /b %errorlevel%
cd compiler-rt
git checkout %REV%
if %errorlevel% neq 0 exit /b %errorlevel%
cd ..
cd compiler-rt/lib/tsan/go
call build.bat
if %errorlevel% neq 0 exit /b %errorlevel%
cd ../../../..
xcopy compiler-rt\lib\tsan\go\race_windows_amd64.syso go\src\runtime\race\race_windows_amd64.syso /Y
if %errorlevel% neq 0 exit /b %errorlevel%
cd go/src
call race.bat
if %errorlevel% neq 0 exit /b %errorlevel%
			`,
	},
}

func init() {
	// Ensure that there are no duplicate platform entries.
	seen := make(map[string]bool)
	for _, p := range platforms {
		if seen[p.Name()] {
			log.Fatal("Duplicate platforms entry for %s.", p.Name())
		}
		seen[p.Name()] = true
	}
}

var platformEnabled = make(map[string]bool)

func parsePlatformsFlag() {
	if *flagPlatforms == "all" {
		for _, p := range platforms {
			platformEnabled[p.Name()] = true
		}
		return
	}

	var invalid []string
	for _, name := range strings.Split(*flagPlatforms, ",") {
		for _, p := range platforms {
			if name == p.Name() {
				platformEnabled[name] = true
				break
			}
		}
		if !platformEnabled[name] {
			invalid = append(invalid, name)
		}
	}

	if len(invalid) > 0 {
		var msg bytes.Buffer
		fmt.Fprintf(&msg, "Unrecognized platforms: %q. Supported platforms are:\n", invalid)
		for _, p := range platforms {
			fmt.Fprintf(&msg, "\t%s/%s\n", p.OS, p.Arch)
		}
		log.Fatal(&msg)
	}
}

func main() {
	flag.Parse()
	if *flagRev == "" || *flagGoroot == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}
	parsePlatformsFlag()

	// Update revision in the README file.
	// Do this early to check goroot correctness.
	readmeFile := filepath.Join(*flagGoroot, "src", "runtime", "race", "README")
	readme, err := ioutil.ReadFile(readmeFile)
	if err != nil {
		log.Fatalf("bad -goroot? %v", err)
	}
	readmeRev := regexp.MustCompile("Current runtime is built on rev ([0-9,a-z]+)\\.").FindSubmatchIndex(readme)
	if readmeRev == nil {
		log.Fatalf("failed to find current revision in src/runtime/race/README")
	}
	readme = bytes.Replace(readme, readme[readmeRev[2]:readmeRev[3]], []byte(*flagRev), -1)
	if err := ioutil.WriteFile(readmeFile, readme, 0640); err != nil {
		log.Fatalf("failed to write README file: %v", err)
	}

	// Start build on all platforms in parallel.
	// On interrupt, destroy any in-flight builders before exiting.
	ctx, cancel := context.WithCancel(context.Background())
	shutdown := make(chan os.Signal, 1)
	signal.Notify(shutdown, os.Interrupt)
	go func() {
		<-shutdown
		cancel()
	}()

	g, ctx := errgroup.WithContext(ctx)
	for _, p := range platforms {
		if !platformEnabled[p.Name()] {
			continue
		}

		p := p
		g.Go(func() error {
			if err := p.Build(ctx); err != nil {
				return fmt.Errorf("%v failed: %v", p.Name(), err)
			}
			return nil
		})
	}

	if err := g.Wait(); err != nil {
		log.Println(err)
		os.Exit(1)
	}
}

type Platform struct {
	OS     string
	Arch   string
	Type   string // gomote instance type
	Inst   string // actual gomote instance name
	Script string
}

func (p *Platform) Name() string {
	return fmt.Sprintf("%v/%v", p.OS, p.Arch)
}

func (p *Platform) Build(ctx context.Context) error {
	// Create gomote instance (or reuse an existing instance for debugging).
	var lastErr error
	for p.Inst == "" {
		inst, err := p.Gomote(ctx, "create", p.Type)
		if err != nil {
			select {
			case <-ctx.Done():
				if lastErr != nil {
					return lastErr
				}
				return err
			default:
				// Creation sometimes fails with transient errors like:
				// "buildlet didn't come up at http://10.240.0.13 in 3m0s".
				log.Printf("%v: instance creation failed, retrying", p.Name)
				lastErr = err
				continue
			}
		}
		p.Inst = strings.Trim(string(inst), " \t\n")
		defer p.Gomote(context.Background(), "destroy", p.Inst)
	}
	log.Printf("%s: using instance %v", p.Name(), p.Inst)

	// put14
	if _, err := p.Gomote(ctx, "put14", p.Inst); err != nil {
		return err
	}

	// Execute the script.
	script, err := ioutil.TempFile("", "racebuild")
	if err != nil {
		return fmt.Errorf("failed to create temp file: %v", err)
	}
	defer func() {
		script.Close()
		os.Remove(script.Name())
	}()
	if _, err := script.Write([]byte(p.Script)); err != nil {
		return fmt.Errorf("failed to write temp file: %v", err)
	}
	script.Close()
	targetName := "script.bash"
	if p.OS == "windows" {
		targetName = "script.bat"
	}
	if _, err := p.Gomote(ctx, "put", "-mode=0700", p.Inst, script.Name(), targetName); err != nil {
		return err
	}
	if _, err := p.Gomote(ctx, "run", "-e=REV="+*flagRev, p.Inst, targetName); err != nil {
		return err
	}

	// The script is supposed to leave updated runtime at that path. Copy it out.
	syso := fmt.Sprintf("race_%v_%s.syso", p.OS, p.Arch)
	targz, err := p.Gomote(ctx, "gettar", "-dir=go/src/runtime/race/"+syso, p.Inst)
	if err != nil {
		return err
	}

	// Untar the runtime and write it to goroot.
	if err := p.WriteSyso(filepath.Join(*flagGoroot, "src", "runtime", "race", syso), targz); err != nil {
		return fmt.Errorf("%v", err)
	}

	log.Printf("%v: build completed", p.Name())
	return nil
}

func (p *Platform) WriteSyso(sysof string, targz []byte) error {
	// Ungzip.
	gzipr, err := gzip.NewReader(bytes.NewReader(targz))
	if err != nil {
		return fmt.Errorf("failed to read gzip archive: %v", err)
	}
	defer gzipr.Close()
	tr := tar.NewReader(gzipr)
	if _, err := tr.Next(); err != nil {
		return fmt.Errorf("failed to read tar archive: %v", err)
	}

	// Copy the file.
	syso, err := os.Create(sysof)
	if err != nil {
		return fmt.Errorf("failed to open race runtime: %v", err)
	}
	defer syso.Close()
	if _, err := io.Copy(syso, tr); err != nil {
		return fmt.Errorf("failed to write race runtime: %v", err)
	}
	return nil
}

func (p *Platform) Gomote(ctx context.Context, args ...string) ([]byte, error) {
	log.Printf("%v: gomote %v", p.Name(), args)

	cmd := exec.CommandContext(ctx, "gomote", args...)
	outBuf := new(bytes.Buffer)

	// Combine stderr and stdout for everything except gettar: gettar's output is
	// huge, so we only want to log stderr for it.
	errBuf := outBuf
	if args[0] == "gettar" {
		errBuf = new(bytes.Buffer)
	}

	cmd.Stdout = outBuf
	cmd.Stderr = errBuf
	run := cmd.Run
	if len(platformEnabled) == 1 {
		// If building only one platform, stream gomote output to os.Stderr.
		r, w := io.Pipe()
		errTee := io.TeeReader(r, cmd.Stderr)
		if cmd.Stdout == cmd.Stderr {
			cmd.Stdout = w
		}
		cmd.Stderr = w

		run = func() (err error) {
			go func() {
				err = cmd.Run()
				w.Close()
			}()
			io.Copy(os.Stderr, errTee)
			return
		}
	}

	if err := run(); err != nil {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
		log.Printf("%v: gomote %v failed:\n%s", p.Name(), args, errBuf)
		return nil, err
	}

	if errBuf.Len() == 0 {
		log.Printf("%v: gomote %v succeeded: <no output>", p.Name(), args)
	} else {
		log.Printf("%v: gomote %v succeeded:\n%s", p.Name(), args, errBuf)
	}
	return outBuf.Bytes(), nil
}
<commit_msg>cmd/racebuild: install Git and GCC on the Windows builder<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.

// racebuild builds the race runtime (syso files) on all supported OSes using gomote.
// Usage:
//	$ racebuild -rev <llvm_git_revision> -goroot <path_to_go_repo>
package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/sync/errgroup"
)

var (
	flagGoroot    = flag.String("goroot", "", "path to Go repository to update (required)")
	flagRev       = flag.String("rev", "", "llvm compiler-rt git revision from http://llvm.org/git/compiler-rt.git (required)")
	flagPlatforms = flag.String("platforms", "all", `comma-separated platforms (such as "linux/amd64") to rebuild, or "all"`)
)

// TODO: use buildlet package instead of calling out to gomote.
var platforms = []*Platform{
	&Platform{
		OS:   "freebsd",
		Arch: "amd64",
		Type: "freebsd-amd64-race",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_freebsd_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "darwin",
		Arch: "amd64",
		Type: "darwin-amd64-10_12",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_darwin_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "linux",
		Arch: "amd64",
		Type: "linux-amd64-race",
		Script: `#!/usr/bin/env bash
set -e
apt-get update
apt-get install -y git g++
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_linux_amd64.syso go/src/runtime/race
(cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "linux",
		Arch: "ppc64le",
		Type: "linux-ppc64le-buildlet",
		Script: `#!/usr/bin/env bash
set -e
apt-get update
apt-get install -y git g++
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_linux_ppc64le.syso go/src/runtime/race
# TODO(#23731): Uncomment to test the syso file before accepting it.
# (cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "netbsd",
		Arch: "amd64",
		Type: "netbsd-amd64-8_0",
		Script: `#!/usr/bin/env bash
set -e
git clone https://go.googlesource.com/go
git clone http://llvm.org/git/compiler-rt.git
(cd compiler-rt && git checkout $REV)
(cd compiler-rt/lib/tsan/go && CC=clang ./buildgo.sh)
cp compiler-rt/lib/tsan/go/race_netbsd_amd64.syso go/src/runtime/race
# TODO(#24322): Uncomment to test the syso file before accepting it.
# (cd go/src && ./race.bash)
			`,
	},
	&Platform{
		OS:   "windows",
		Arch: "amd64",
		Type: "windows-amd64-race",
		Script: `
@"%SystemRoot%\System32\WindowsPowerShell\v1.0\powershell.exe" -NoProfile -InputFormat None -ExecutionPolicy Bypass -Command "iex ((New-Object System.Net.WebClient).DownloadString('https://chocolatey.org/install.ps1'))" && SET "PATH=%PATH%;%ALLUSERSPROFILE%\chocolatey\bin"
choco install git -y
if %errorlevel% neq 0 exit /b %errorlevel%
choco install mingw -y
if %errorlevel% neq 0 exit /b %errorlevel%
call refreshenv
git clone https://go.googlesource.com/go
if %errorlevel% neq 0 exit /b %errorlevel%
git clone http://llvm.org/git/compiler-rt.git
if %errorlevel% neq 0 exit /b %errorlevel%
cd compiler-rt
git checkout %REV%
if %errorlevel% neq 0 exit /b %errorlevel%
cd ..
cd compiler-rt/lib/tsan/go
call build.bat
if %errorlevel% neq 0 exit /b %errorlevel%
cd ../../../..
xcopy compiler-rt\lib\tsan\go\race_windows_amd64.syso go\src\runtime\race\race_windows_amd64.syso /Y
if %errorlevel% neq 0 exit /b %errorlevel%
cd go/src
call race.bat
if %errorlevel% neq 0 exit /b %errorlevel%
			`,
	},
}

func init() {
	// Ensure that there are no duplicate platform entries.
	seen := make(map[string]bool)
	for _, p := range platforms {
		if seen[p.Name()] {
			log.Fatal("Duplicate platforms entry for %s.", p.Name())
		}
		seen[p.Name()] = true
	}
}

var platformEnabled = make(map[string]bool)

func parsePlatformsFlag() {
	if *flagPlatforms == "all" {
		for _, p := range platforms {
			platformEnabled[p.Name()] = true
		}
		return
	}

	var invalid []string
	for _, name := range strings.Split(*flagPlatforms, ",") {
		for _, p := range platforms {
			if name == p.Name() {
				platformEnabled[name] = true
				break
			}
		}
		if !platformEnabled[name] {
			invalid = append(invalid, name)
		}
	}

	if len(invalid) > 0 {
		var msg bytes.Buffer
		fmt.Fprintf(&msg, "Unrecognized platforms: %q. Supported platforms are:\n", invalid)
		for _, p := range platforms {
			fmt.Fprintf(&msg, "\t%s/%s\n", p.OS, p.Arch)
		}
		log.Fatal(&msg)
	}
}

func main() {
	flag.Parse()
	if *flagRev == "" || *flagGoroot == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}
	parsePlatformsFlag()

	// Update revision in the README file.
	// Do this early to check goroot correctness.
	readmeFile := filepath.Join(*flagGoroot, "src", "runtime", "race", "README")
	readme, err := ioutil.ReadFile(readmeFile)
	if err != nil {
		log.Fatalf("bad -goroot? %v", err)
	}
	readmeRev := regexp.MustCompile("Current runtime is built on rev ([0-9,a-z]+)\\.").FindSubmatchIndex(readme)
	if readmeRev == nil {
		log.Fatalf("failed to find current revision in src/runtime/race/README")
	}
	readme = bytes.Replace(readme, readme[readmeRev[2]:readmeRev[3]], []byte(*flagRev), -1)
	if err := ioutil.WriteFile(readmeFile, readme, 0640); err != nil {
		log.Fatalf("failed to write README file: %v", err)
	}

	// Start build on all platforms in parallel.
	// On interrupt, destroy any in-flight builders before exiting.
	ctx, cancel := context.WithCancel(context.Background())
	shutdown := make(chan os.Signal, 1)
	signal.Notify(shutdown, os.Interrupt)
	go func() {
		<-shutdown
		cancel()
	}()

	g, ctx := errgroup.WithContext(ctx)
	for _, p := range platforms {
		if !platformEnabled[p.Name()] {
			continue
		}

		p := p
		g.Go(func() error {
			if err := p.Build(ctx); err != nil {
				return fmt.Errorf("%v failed: %v", p.Name(), err)
			}
			return nil
		})
	}

	if err := g.Wait(); err != nil {
		log.Println(err)
		os.Exit(1)
	}
}

type Platform struct {
	OS     string
	Arch   string
	Type   string // gomote instance type
	Inst   string // actual gomote instance name
	Script string
}

func (p *Platform) Name() string {
	return fmt.Sprintf("%v/%v", p.OS, p.Arch)
}

func (p *Platform) Build(ctx context.Context) error {
	// Create gomote instance (or reuse an existing instance for debugging).
	var lastErr error
	for p.Inst == "" {
		inst, err := p.Gomote(ctx, "create", p.Type)
		if err != nil {
			select {
			case <-ctx.Done():
				if lastErr != nil {
					return lastErr
				}
				return err
			default:
				// Creation sometimes fails with transient errors like:
				// "buildlet didn't come up at http://10.240.0.13 in 3m0s".
				log.Printf("%v: instance creation failed, retrying", p.Name)
				lastErr = err
				continue
			}
		}
		p.Inst = strings.Trim(string(inst), " \t\n")
		defer p.Gomote(context.Background(), "destroy", p.Inst)
	}
	log.Printf("%s: using instance %v", p.Name(), p.Inst)

	// put14
	if _, err := p.Gomote(ctx, "put14", p.Inst); err != nil {
		return err
	}

	// Execute the script.
	script, err := ioutil.TempFile("", "racebuild")
	if err != nil {
		return fmt.Errorf("failed to create temp file: %v", err)
	}
	defer func() {
		script.Close()
		os.Remove(script.Name())
	}()
	if _, err := script.Write([]byte(p.Script)); err != nil {
		return fmt.Errorf("failed to write temp file: %v", err)
	}
	script.Close()
	targetName := "script.bash"
	if p.OS == "windows" {
		targetName = "script.bat"
	}
	if _, err := p.Gomote(ctx, "put", "-mode=0700", p.Inst, script.Name(), targetName); err != nil {
		return err
	}
	if _, err := p.Gomote(ctx, "run", "-e=REV="+*flagRev, p.Inst, targetName); err != nil {
		return err
	}

	// The script is supposed to leave updated runtime at that path. Copy it out.
	syso := fmt.Sprintf("race_%v_%s.syso", p.OS, p.Arch)
	targz, err := p.Gomote(ctx, "gettar", "-dir=go/src/runtime/race/"+syso, p.Inst)
	if err != nil {
		return err
	}

	// Untar the runtime and write it to goroot.
	if err := p.WriteSyso(filepath.Join(*flagGoroot, "src", "runtime", "race", syso), targz); err != nil {
		return fmt.Errorf("%v", err)
	}

	log.Printf("%v: build completed", p.Name())
	return nil
}

func (p *Platform) WriteSyso(sysof string, targz []byte) error {
	// Ungzip.
	gzipr, err := gzip.NewReader(bytes.NewReader(targz))
	if err != nil {
		return fmt.Errorf("failed to read gzip archive: %v", err)
	}
	defer gzipr.Close()
	tr := tar.NewReader(gzipr)
	if _, err := tr.Next(); err != nil {
		return fmt.Errorf("failed to read tar archive: %v", err)
	}

	// Copy the file.
	syso, err := os.Create(sysof)
	if err != nil {
		return fmt.Errorf("failed to open race runtime: %v", err)
	}
	defer syso.Close()
	if _, err := io.Copy(syso, tr); err != nil {
		return fmt.Errorf("failed to write race runtime: %v", err)
	}
	return nil
}

func (p *Platform) Gomote(ctx context.Context, args ...string) ([]byte, error) {
	log.Printf("%v: gomote %v", p.Name(), args)

	cmd := exec.CommandContext(ctx, "gomote", args...)
	outBuf := new(bytes.Buffer)

	// Combine stderr and stdout for everything except gettar: gettar's output is
	// huge, so we only want to log stderr for it.
	errBuf := outBuf
	if args[0] == "gettar" {
		errBuf = new(bytes.Buffer)
	}

	cmd.Stdout = outBuf
	cmd.Stderr = errBuf
	run := cmd.Run
	if len(platformEnabled) == 1 {
		// If building only one platform, stream gomote output to os.Stderr.
		r, w := io.Pipe()
		errTee := io.TeeReader(r, cmd.Stderr)
		if cmd.Stdout == cmd.Stderr {
			cmd.Stdout = w
		}
		cmd.Stderr = w

		run = func() (err error) {
			go func() {
				err = cmd.Run()
				w.Close()
			}()
			io.Copy(os.Stderr, errTee)
			return
		}
	}

	if err := run(); err != nil {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
		log.Printf("%v: gomote %v failed:\n%s", p.Name(), args, errBuf)
		return nil, err
	}

	if errBuf.Len() == 0 {
		log.Printf("%v: gomote %v succeeded: <no output>", p.Name(), args)
	} else {
		log.Printf("%v: gomote %v succeeded:\n%s", p.Name(), args, errBuf)
	}
	return outBuf.Bytes(), nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/sha1"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"

	"github.com/bradfitz/iter"
	"github.com/edsrzf/mmap-go"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/mmap_span"
)

var (
	torrentPath = flag.String("torrent", "/path/to/the.torrent", "path of the torrent file")
	dataPath    = flag.String("path", "/torrent/data", "path of the torrent data")
)

func fileToMmap(filename string, length int64) mmap.MMap {
	osFile, err := os.Open(filename)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		log.Fatal(err)
	}
	goMMap, err := mmap.MapRegion(osFile, int(length), mmap.RDONLY, mmap.COPY, 0)
	if err != nil {
		log.Fatal(err)
	}
	if int64(len(goMMap)) != length {
		log.Printf("file mmap has wrong size: %#v", filename)
	}
	osFile.Close()

	return goMMap
}

func main() {
	log.SetFlags(log.Flags() | log.Lshortfile)
	flag.Parse()
	metaInfo, err := metainfo.LoadFromFile(*torrentPath)
	if err != nil {
		log.Fatal(err)
	}
	info, err := metaInfo.UnmarshalInfo()
	if err != nil {
		log.Fatalf("error unmarshalling info: %s", err)
	}
	mMapSpan := &mmap_span.MMapSpan{}
	if len(info.Files) > 0 {
		for _, file := range info.Files {
			filename := filepath.Join(append([]string{*dataPath, info.Name}, file.Path...)...)
			goMMap := fileToMmap(filename, file.Length)
			mMapSpan.Append(goMMap)
		}
		log.Println(len(info.Files))
	} else {
		goMMap := fileToMmap(*dataPath, info.Length)
		mMapSpan.Append(goMMap)
	}
	log.Println(mMapSpan.Size())
	log.Println(len(info.Pieces))
	for i := range iter.N(info.NumPieces()) {
		p := info.Piece(i)
		hash := sha1.New()
		_, err := io.Copy(hash, io.NewSectionReader(mMapSpan, p.Offset(), p.Length()))
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%d: %x: %v\n", i, p.Hash(), bytes.Equal(hash.Sum(nil), p.Hash().Bytes()))
	}
}
<commit_msg>Rewrite cmd/torrent-verify<commit_after>package main

import (
	"bytes"
	"crypto/sha1"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"

	"github.com/anacrolix/tagflag"
	"github.com/bradfitz/iter"
	"github.com/edsrzf/mmap-go"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/mmap_span"
)

func mmapFile(name string) (mm mmap.MMap, err error) {
	f, err := os.Open(name)
	if err != nil {
		return
	}
	defer f.Close()
	fi, err := f.Stat()
	if err != nil {
		return
	}
	if fi.Size() == 0 {
		return
	}
	return mmap.MapRegion(f, -1, mmap.RDONLY, mmap.COPY, 0)
}

func verifyTorrent(info *metainfo.Info, root string) error {
	span := new(mmap_span.MMapSpan)
	for _, file := range info.UpvertedFiles() {
		filename := filepath.Join(append([]string{root, info.Name}, file.Path...)...)
		mm, err := mmapFile(filename)
		if err != nil {
			return err
		}
		if int64(len(mm)) != file.Length {
			return fmt.Errorf("file %q has wrong length", filename)
		}
		span.Append(mm)
	}
	for i := range iter.N(info.NumPieces()) {
		p := info.Piece(i)
		hash := sha1.New()
		_, err := io.Copy(hash, io.NewSectionReader(span, p.Offset(), p.Length()))
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("%d: %x: %v\n", i, p.Hash(), bytes.Equal(hash.Sum(nil), p.Hash().Bytes()))
	}
	return nil
}

func main() {
	log.SetFlags(log.Flags() | log.Lshortfile)
	var flags = struct {
		DataDir string
		tagflag.StartPos
		TorrentFile string
	}{}
	tagflag.Parse(&flags)
	metaInfo, err := metainfo.LoadFromFile(flags.TorrentFile)
	if err != nil {
		log.Fatal(err)
	}
	info, err := metaInfo.UnmarshalInfo()
	if err != nil {
		log.Fatalf("error unmarshalling info: %s", err)
	}
	err = verifyTorrent(&info, flags.DataDir)
	if err != nil {
		log.Fatalf("torrent failed verification: %s", err)
	}
}
<|endoftext|>


<commit_before>// +build gofuzz

package ipfix

import (
	"bytes"
	"io"
)

func Fuzz(bs []byte) int {
	p := NewSession()
	r := bytes.NewReader(bs)

	_, err := p.ParseReader(r)
	for err == nil {
		_, err = p.ParseReader(r)
	}
	if err == io.EOF {
		return 1
	}

	return 0
}
<commit_msg>Also fuzz Interpreter<commit_after>// +build gofuzz

package ipfix

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"code.google.com/p/gcfg"
)

var extra []DictionaryEntry

func init() {
	if dictFile := os.Getenv("IPFIXDICT"); dictFile != "" {
		var dict UserDictionary
		err := gcfg.ReadFileInto(&dict, dictFile)
		if err != nil {
			panic(err)
		}

		for name, field := range dict.Field {
			extra = append(extra, field.DictionaryEntry(name))
		}
	}
}

func Fuzz(bs []byte) int {
	r := bytes.NewReader(bs)

	s := NewSession()
	i := NewInterpreter(s)
	for _, e := range extra {
		i.AddDictionaryEntry(e)
	}

	msg, err := s.ParseReader(r)
	for err == nil {
		for _, rec := range msg.DataRecords {
			i.Interpret(rec)
		}
		msg, err = s.ParseReader(r)
	}
	if err == io.EOF {
		return 1
	}

	return 0
}

type UserDictionary struct {
	Field map[string]*Field
}

type Field struct {
	ID         uint16
	Enterprise uint32
	Type       string
}

func (f Field) DictionaryEntry(name string) DictionaryEntry {
	return DictionaryEntry{
		Name:         name,
		EnterpriseID: f.Enterprise,
		FieldID:      f.ID,
		Type:         FieldTypes[f.Type],
	}
}
<|endoftext|>


<commit_before>package gosnow

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"testing"
)

const simpleSrc string = `
# My API
## GET /message
 + Response 200 (text/plain)

         Hello World
`

const namelessSrc string = `
FORMAT: 1A

# Group Messages

# Message [/messages/{id}]

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

`

const warningSrc string = `
FORMAT: 1A

# Group Messages

# Message [/messages/{id}]

+ Model (text/plain)

        Hello World

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

`

var (
	apibFile      = "test/fixtures/sample-api.apib"
	astFile       = "test/fixtures/sample-api-ast.json"
	sourcemapFile = "test/fixtures/sample-api-sourcemap.json"
)

// replace the variables with the contents of the file they point to
func init() {
	if c, err := ioutil.ReadFile(apibFile); err != nil {
		panic("apibFile not found")
	} else {
		apibFile = string(c)
	}

	if c, err := ioutil.ReadFile(astFile); err != nil {
		panic("astFile not found")
	} else {
		astFile = string(c)
	}

	if c, err := ioutil.ReadFile(sourcemapFile); err != nil {
		panic("sourcemapFile not found")
	} else {
		sourcemapFile = string(c)
	}
}

func TestSourceAnnotationOk(t *testing.T) {
	sa := new(SourceAnnotation)
	if !sa.Ok() {
		t.Error("empty source annotation should have zero value indicating ok")
	}
	sa.Code = 2
	if sa.Ok() {
		t.Error("source annotation should have non zero value indicating not ok")
	}
}

func TestNewPR(t *testing.T) {
	_, err := newPR([]byte(`{"unrelated": "json"}`))
	if err != nil {
		t.Errorf("newPR errored for valid json %v", err)
	}
}

func TestNewPRFailure(t *testing.T) {
	junk := []byte(`*#(*(R$#&)$#)R*(Y@#_RH`)
	_, err := newPR(junk)
	if err == nil {
		t.Error("newPR should have errored and did not")
	}
	if e, ok := err.(*json.SyntaxError); !ok {
		t.Errorf("Expected json.SyntaxError, got %T", e)
	}
}

func TestRawOptionParse(t *testing.T) {
	res, err := RawOptionParse(apibFile, 0)
	if err != nil {
		t.Fatalf("RawOptionParse failed with error: %v", err)
	}
	if res == nil {
		t.Fatal("RawOptionParse returned nil result")
	}
	fmt.Println(string(res))
}

func TestParse(t *testing.T) {
	res, err := Parse(apibFile)
	if err != nil {
		t.Fatalf("Parse failed with error: %v", err)
	}
	if res == nil {
		t.Fatal("Parse returned nil result")
	}
	// v, _ := json.MarshalIndent(res, "", "  ")
	// fmt.Println(string(v))
}

// ensure that the option parse with a 0 does the same thing as the simple parse
func TestParseEquality(t *testing.T) {
	res1, err := Parse(simpleSrc)
	if err != nil {
		t.Fatalf("Parse failed with err: %v", err)
	}
	res2, err := OptionParse(simpleSrc, 0)
	if err != nil {
		t.Fatalf("OptionParse failed with err: %v", err)
	}

	v1, err := json.Marshal(res1)
	if err != nil {
		t.Fatalf("json marshal error: %v", err)
	}
	v2, err := json.Marshal(res2)
	if err != nil {
		t.Fatalf("json marshal error: %v", err)
	}
	if string(v1) != string(v2) {
		t.Error("Results should be equal")
	}
}

func TestParseError(t *testing.T) {
	junk := "*#(*(R$#&)$#)R*(Y@#_RH"
	res, err := OptionParse(junk, -1)
	if err == nil {
		t.Errorf("OptionParse did not fail for junk input")
	}
	if res != nil {
		t.Errorf("OptionParse returned non=empty result for junk input")
	}
}

func TestFilesOptionParse(t *testing.T) {
	res, err := OptionParse(apibFile, ScRenderDescriptionsOptionKey)
	if err != nil {
		t.Errorf("OptionParse failed for key ScRenderDescriptionsOptionKey with error: %v", err)
	} else if res == nil {
		t.Errorf("OptionParse for key ScRenderDescriptionsOptionKey returned empty result")
	}

	_, err = OptionParse(namelessSrc, RequireBlueprintNameOptionKey)
	if err == nil {
		t.Errorf("strict OptionParse did not fail for key RequireBlueprintNameOptionKey")
	}

	res, err = OptionParse(apibFile, ExportSourcemapOptionKey)
	if err != nil {
		t.Errorf("OptionParse failed for ExportSourcemapOptionKey with error: %v", err)
	} else if res == nil {
		t.Errorf("OptionParse for key ExportSourcemapOptionKey returned empty result")
	}
}
<commit_msg>more tests wil print abilities<commit_after>package gosnow

import (
	"encoding/json"
	_ "fmt"
	"io/ioutil"
	"testing"
)

const simpleSrc string = `
# My API
## GET /message
 + Response 200 (text/plain)

         Hello World
`

const namelessSrc string = `
FORMAT: 1A

# Group Messages

# Message [/messages/{id}]

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

`

const warningSrc string = `
FORMAT: 1A

# Group Messages

# Message [/messages/{id}]

+ Model (text/plain)

        Hello World

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

## Retrieve Message [GET]
+ Response 200 (text/plain)

        Hello World!

`

var (
	apibFile      = "test/fixtures/sample-api.apib"
	astFile       = "test/fixtures/sample-api-ast.json"
	sourcemapFile = "test/fixtures/sample-api-sourcemap.json"
)

// replace the variables with the contents of the file they point to
func init() {
	if c, err := ioutil.ReadFile(apibFile); err != nil {
		panic("apibFile not found")
	} else {
		apibFile = string(c)
	}

	if c, err := ioutil.ReadFile(astFile); err != nil {
		panic("astFile not found")
	} else {
		astFile = string(c)
	}

	if c, err := ioutil.ReadFile(sourcemapFile); err != nil {
		panic("sourcemapFile not found")
	} else {
		sourcemapFile = string(c)
	}
}

func TestSourceAnnotationOk(t *testing.T) {
	sa := new(SourceAnnotation)
	if !sa.Ok() {
		t.Error("empty source annotation should have zero value indicating ok")
	}
	sa.Code = 2
	if sa.Ok() {
		t.Error("source annotation should have non zero value indicating not ok")
	}
}

func TestNewPR(t *testing.T) {
	_, err := newPR([]byte(`{"unrelated": "json"}`))
	if err != nil {
		t.Errorf("newPR errored for valid json %v", err)
	}
}

func TestNewPRFailure(t *testing.T) {
	junk := []byte(`*#(*(R$#&)$#)R*(Y@#_RH`)
	_, err := newPR(junk)
	if err == nil {
		t.Error("newPR should have errored and did not")
	}
	if e, ok := err.(*json.SyntaxError); !ok {
		t.Errorf("Expected json.SyntaxError, got %T", e)
	}
}

func TestRawOptionParse(t *testing.T) {
	res, err := RawOptionParse(apibFile, 0)
	if err != nil {
		t.Fatalf("RawOptionParse failed with error: %v", err)
	}
	if res == nil {
		t.Fatal("RawOptionParse returned nil result")
	}
	// fmt.Println(string(res))
}

func TestRawSourceMapOptionParse(t *testing.T) {
	res, err := RawOptionParse(apibFile, ExportSourcemapOptionKey)
	if err != nil {
		t.Errorf("RawOptionParse failed for ExportSourcemapOptionKey with error: %v", err)
	} else if res == nil {
		t.Errorf("RawOptionParse for key ExportSourcemapOptionKey returned empty result")
	}
	// fmt.Println(string(res))
}

func TestParse(t *testing.T) {
	res, err := Parse(apibFile)
	if err != nil {
		t.Fatalf("Parse failed with error: %v", err)
	}
	if res == nil {
		t.Fatal("Parse returned nil result")
	}
	// v, _ := json.MarshalIndent(res, "", "  ")
	// fmt.Println(string(v))
}

// ensure that the option parse with a 0 does the same thing as the simple parse
func TestParseEquality(t *testing.T) {
	res1, err := Parse(simpleSrc)
	if err != nil {
		t.Fatalf("Parse failed with err: %v", err)
	}
	res2, err := OptionParse(simpleSrc, 0)
	if err != nil {
		t.Fatalf("OptionParse failed with err: %v", err)
	}

	v1, err := json.Marshal(res1)
	if err != nil {
		t.Fatalf("json marshal error: %v", err)
	}
	v2, err := json.Marshal(res2)
	if err != nil {
		t.Fatalf("json marshal error: %v", err)
	}
	if string(v1) != string(v2) {
		t.Error("Results should be equal")
	}
}

func TestParseError(t *testing.T) {
	junk := "*#(*(R$#&)$#)R*(Y@#_RH"
	res, err := OptionParse(junk, -1)
	if err == nil {
		t.Errorf("OptionParse did not fail for junk input")
	}
	if res != nil {
		t.Errorf("OptionParse returned non=empty result for junk input")
	}
}

func TestFilesOptionParse(t *testing.T) {
	res, err := OptionParse(apibFile, ScRenderDescriptionsOptionKey)
	if err != nil {
		t.Errorf("OptionParse failed for key ScRenderDescriptionsOptionKey with error: %v", err)
	} else if res == nil {
		t.Errorf("OptionParse for key ScRenderDescriptionsOptionKey returned empty result")
	}
}

func TestRequiredNameOptionParse(t *testing.T) {
	_, err := OptionParse(namelessSrc, RequireBlueprintNameOptionKey)
	if err == nil {
		t.Errorf("strict OptionParse did not fail for key RequireBlueprintNameOptionKey")
	}
}

func TestSourceMapOptionParse(t *testing.T) {
	res, err := OptionParse(apibFile, ExportSourcemapOptionKey)
	if err != nil {
		t.Errorf("OptionParse failed for ExportSourcemapOptionKey with error: %v", err)
	} else if res == nil {
		t.Errorf("OptionParse for key ExportSourcemapOptionKey returned empty result")
	}
	// v, _ := json.MarshalIndent(res, "", "  ")
	// fmt.Println(string(v))
}
<|endoftext|>


<commit_before>// Copyright 2016 The Prometheus 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.

package web

import (
	"bytes"
	"net/http/httptest"
	"sort"
	"strings"
	"testing"

	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/promql"
)

var scenarios = map[string]struct {
	params string
	accept string
	code   int
	body   string
}{
	"empty": {
		params: "",
		code:   200,
		body:   ``,
	},
	"invalid params from the beginning": {
		params: "match[]=-not-a-valid-metric-name",
		code:   400,
		body: `parse error at char 1: vector selector must contain label matchers or metric name
`,
	},
	"invalid params somehwere in the middle": {
		params: "match[]=not-a-valid-metric-name",
		code:   400,
		body: `parse error at char 4: could not parse remaining input "-a-valid-metric"...
`,
	},
	"test_metric1": {
		params: "match[]=test_metric1",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
`,
	},
	"test_metric2": {
		params: "match[]=test_metric2",
		code:   200,
		body: `# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"test_metric_without_labels": {
		params: "match[]=test_metric_without_labels",
		code:   200,
		body: `# TYPE test_metric_without_labels untyped
test_metric_without_labels 1001 6000000
`,
	},
	"{foo='boo'}": {
		params: "match[]={foo='boo'}",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"two matchers": {
		params: "match[]=test_metric1&match[]=test_metric2",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"everything": {
		params: "match[]={__name__=~'.%2b'}", // '%2b' is an URL-encoded '+'.
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
# TYPE test_metric_without_labels untyped
test_metric_without_labels 1001 6000000
`,
	},
}

func TestFederation(t *testing.T) {
	suite, err := promql.NewTest(t, `
		load 1m
			test_metric1{foo="bar"}    0+100x100
			test_metric1{foo="boo"}    1+0x100
			test_metric2{foo="boo"}    1+0x100
			test_metric_without_labels 1+10x100
	`)
	if err != nil {
		t.Fatal(err)
	}
	defer suite.Close()

	if err := suite.Run(); err != nil {
		t.Fatal(err)
	}

	h := &Handler{
		storage:     suite.Storage(),
		queryEngine: suite.QueryEngine(),
		now:         func() model.Time { return 101 * 60 * 1000 }, // 101min after epoch.
	}

	for name, scenario := range scenarios {
		req := httptest.NewRequest("GET", "http://example.org/federate?"+scenario.params, nil)
		res := httptest.NewRecorder()
		h.federation(res, req)
		if got, want := res.Code, scenario.code; got != want {
			t.Errorf("Scenario %q: got code %d, want %d", name, got, want)
		}
		if got, want := normalizeBody(res.Body), scenario.body; got != want {
			t.Errorf("Scenario %q: got body %q, want %q", name, got, want)
		}
	}
}

// normalizeBody sorts the lines within a metric to make it easy to verify the body.
// (Federation is not taking care of sorting within a metric family.)
func normalizeBody(body *bytes.Buffer) string {
	var (
		lines    []string
		lastHash int
	)
	for line, err := body.ReadString('\n'); err == nil; line, err = body.ReadString('\n') {
		if line[0] == '#' && len(lines) > 0 {
			sort.Strings(lines[lastHash+1:])
			lastHash = len(lines)
		}
		lines = append(lines, line)
	}
	if len(lines) > 0 {
		sort.Strings(lines[lastHash+1:])
	}
	return strings.Join(lines, "")
}
<commit_msg>web: Inline httptest.NewRequest because it only exists in Go1.7+<commit_after>// Copyright 2016 The Prometheus 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.

package web

import (
	"bufio"
	"bytes"
	"net/http"
	"net/http/httptest"
	"sort"
	"strings"
	"testing"

	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/promql"
)

var scenarios = map[string]struct {
	params string
	accept string
	code   int
	body   string
}{
	"empty": {
		params: "",
		code:   200,
		body:   ``,
	},
	"invalid params from the beginning": {
		params: "match[]=-not-a-valid-metric-name",
		code:   400,
		body: `parse error at char 1: vector selector must contain label matchers or metric name
`,
	},
	"invalid params somehwere in the middle": {
		params: "match[]=not-a-valid-metric-name",
		code:   400,
		body: `parse error at char 4: could not parse remaining input "-a-valid-metric"...
`,
	},
	"test_metric1": {
		params: "match[]=test_metric1",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
`,
	},
	"test_metric2": {
		params: "match[]=test_metric2",
		code:   200,
		body: `# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"test_metric_without_labels": {
		params: "match[]=test_metric_without_labels",
		code:   200,
		body: `# TYPE test_metric_without_labels untyped
test_metric_without_labels 1001 6000000
`,
	},
	"{foo='boo'}": {
		params: "match[]={foo='boo'}",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"two matchers": {
		params: "match[]=test_metric1&match[]=test_metric2",
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
`,
	},
	"everything": {
		params: "match[]={__name__=~'.%2b'}", // '%2b' is an URL-encoded '+'.
		code:   200,
		body: `# TYPE test_metric1 untyped
test_metric1{foo="bar"} 10000 6000000
test_metric1{foo="boo"} 1 6000000
# TYPE test_metric2 untyped
test_metric2{foo="boo"} 1 6000000
# TYPE test_metric_without_labels untyped
test_metric_without_labels 1001 6000000
`,
	},
}

func TestFederation(t *testing.T) {
	suite, err := promql.NewTest(t, `
		load 1m
			test_metric1{foo="bar"}    0+100x100
			test_metric1{foo="boo"}    1+0x100
			test_metric2{foo="boo"}    1+0x100
			test_metric_without_labels 1+10x100
	`)
	if err != nil {
		t.Fatal(err)
	}
	defer suite.Close()

	if err := suite.Run(); err != nil {
		t.Fatal(err)
	}

	h := &Handler{
		storage:     suite.Storage(),
		queryEngine: suite.QueryEngine(),
		now:         func() model.Time { return 101 * 60 * 1000 }, // 101min after epoch.
	}

	for name, scenario := range scenarios {
		req, err := http.ReadRequest(bufio.NewReader(strings.NewReader(
			"GET http://example.org/federate?" + scenario.params + " HTTP/1.0\r\n\r\n",
		)))
		if err != nil {
			t.Fatal(err)
		}
		// HTTP/1.0 was used above to avoid needing a Host field. Change it to 1.1 here.
		req.Proto = "HTTP/1.1"
		req.ProtoMinor = 1
		req.Close = false
		// 192.0.2.0/24 is "TEST-NET" in RFC 5737 for use solely in
		// documentation and example source code and should not be
		// used publicly.
		req.RemoteAddr = "192.0.2.1:1234"
		// TODO(beorn7): Once we are completely on Go1.7, replace the lines above by the following:
		// req := httptest.NewRequest("GET", "http://example.org/federate?"+scenario.params, nil)
		res := httptest.NewRecorder()
		h.federation(res, req)
		if got, want := res.Code, scenario.code; got != want {
			t.Errorf("Scenario %q: got code %d, want %d", name, got, want)
		}
		if got, want := normalizeBody(res.Body), scenario.body; got != want {
			t.Errorf("Scenario %q: got body %q, want %q", name, got, want)
		}
	}
}

// normalizeBody sorts the lines within a metric to make it easy to verify the body.
// (Federation is not taking care of sorting within a metric family.)
func normalizeBody(body *bytes.Buffer) string {
	var (
		lines    []string
		lastHash int
	)
	for line, err := body.ReadString('\n'); err == nil; line, err = body.ReadString('\n') {
		if line[0] == '#' && len(lines) > 0 {
			sort.Strings(lines[lastHash+1:])
			lastHash = len(lines)
		}
		lines = append(lines, line)
	}
	if len(lines) > 0 {
		sort.Strings(lines[lastHash+1:])
	}
	return strings.Join(lines, "")
}
<|endoftext|>


<commit_before>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package git provides helper functions for working with git repositories.
package git

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/osutil"
)

const timeout = time.Hour // timeout for all git invocations

// Poll checkouts the specified repository/branch in dir.
// This involves fetching/resetting/cloning as necessary to recover from all possible problems.
// Returns hash of the HEAD commit in the specified branch.
func Poll(dir, repo, branch string) (string, error) {
	osutil.RunCmd(timeout, dir, "git", "reset", "--hard")
	origin, err := osutil.RunCmd(timeout, dir, "git", "remote", "get-url", "origin")
	if err != nil || strings.TrimSpace(string(origin)) != repo {
		// The repo is here, but it has wrong origin (e.g. repo in config has changed), re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	// Use origin/branch for the case the branch was force-pushed,
	// in such case branch is not the same is origin/branch and we will
	// stuck with the local version forever (git checkout won't fail).
	if _, err := osutil.RunCmd(timeout, dir, "git", "checkout", "origin/"+branch); err != nil {
		// No such branch (e.g. branch in config has changed), re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	if _, err := osutil.RunCmd(timeout, dir, "git", "fetch", "--no-tags"); err != nil {
		// Something else is wrong, re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	if _, err := osutil.RunCmd(timeout, dir, "git", "checkout", "origin/"+branch); err != nil {
		return "", err
	}
	return HeadCommit(dir)
}

func clone(dir, repo, branch string) error {
	if err := os.RemoveAll(dir); err != nil {
		return fmt.Errorf("failed to remove repo dir: %v", err)
	}
	if err := osutil.MkdirAll(dir); err != nil {
		return fmt.Errorf("failed to create repo dir: %v", err)
	}
	args := []string{
		"clone",
		repo,
		"--single-branch",
		"--branch", branch,
		dir,
	}
	_, err := osutil.RunCmd(timeout, "", "git", args...)
	return err
}

// HeadCommit returns hash of the HEAD commit of the current branch of git repository in dir.
func HeadCommit(dir string) (string, error) {
	output, err := osutil.RunCmd(timeout, dir, "git", "log", "--pretty=format:%H", "-n", "1")
	if err != nil {
		return "", err
	}
	if len(output) != 0 && output[len(output)-1] == '\n' {
		output = output[:len(output)-1]
	}
	if len(output) != 40 {
		return "", fmt.Errorf("unexpected git log output, want commit hash: %q", output)
	}
	return string(output), nil
}

// ListRecentCommits returns list of recent commit titles starting from baseCommit.
func ListRecentCommits(dir, baseCommit string) ([]string, error) {
	// On upstream kernel this produces ~11MB of output.
	// Somewhat inefficient to collect whole output in a slice
	// and then convert to string, but should be bearable.
	output, err := osutil.RunCmd(timeout, dir, "git", "log",
		"--pretty=format:%s", "--no-merges", "-n", "200000", baseCommit)
	if err != nil {
		return nil, err
	}
	return strings.Split(string(output), "\n"), nil
}

// CanonicalizeCommit returns commit title that can be used when checking
// if a particular commit is present in a git tree.
// Some trees add prefixes to commit titles during backporting,
// so we want e.g. commit "foo bar" match "BACKPORT: foo bar".
func CanonicalizeCommit(title string) string {
	for _, prefix := range commitPrefixes {
		if strings.HasPrefix(title, prefix) {
			title = title[len(prefix):]
			break
		}
	}
	return strings.TrimSpace(title)
}

var commitPrefixes = []string{
	"UPSTREAM:",
	"CHROMIUM:",
	"FROMLIST:",
	"BACKPORT:",
	"net-backports:",
}
<commit_msg>pkg/git: add another commit prefix<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package git provides helper functions for working with git repositories.
package git

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/osutil"
)

const timeout = time.Hour // timeout for all git invocations

// Poll checkouts the specified repository/branch in dir.
// This involves fetching/resetting/cloning as necessary to recover from all possible problems.
// Returns hash of the HEAD commit in the specified branch.
func Poll(dir, repo, branch string) (string, error) {
	osutil.RunCmd(timeout, dir, "git", "reset", "--hard")
	origin, err := osutil.RunCmd(timeout, dir, "git", "remote", "get-url", "origin")
	if err != nil || strings.TrimSpace(string(origin)) != repo {
		// The repo is here, but it has wrong origin (e.g. repo in config has changed), re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	// Use origin/branch for the case the branch was force-pushed,
	// in such case branch is not the same is origin/branch and we will
	// stuck with the local version forever (git checkout won't fail).
	if _, err := osutil.RunCmd(timeout, dir, "git", "checkout", "origin/"+branch); err != nil {
		// No such branch (e.g. branch in config has changed), re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	if _, err := osutil.RunCmd(timeout, dir, "git", "fetch", "--no-tags"); err != nil {
		// Something else is wrong, re-clone.
		if err := clone(dir, repo, branch); err != nil {
			return "", err
		}
	}
	if _, err := osutil.RunCmd(timeout, dir, "git", "checkout", "origin/"+branch); err != nil {
		return "", err
	}
	return HeadCommit(dir)
}

func clone(dir, repo, branch string) error {
	if err := os.RemoveAll(dir); err != nil {
		return fmt.Errorf("failed to remove repo dir: %v", err)
	}
	if err := osutil.MkdirAll(dir); err != nil {
		return fmt.Errorf("failed to create repo dir: %v", err)
	}
	args := []string{
		"clone",
		repo,
		"--single-branch",
		"--branch", branch,
		dir,
	}
	_, err := osutil.RunCmd(timeout, "", "git", args...)
	return err
}

// HeadCommit returns hash of the HEAD commit of the current branch of git repository in dir.
func HeadCommit(dir string) (string, error) {
	output, err := osutil.RunCmd(timeout, dir, "git", "log", "--pretty=format:%H", "-n", "1")
	if err != nil {
		return "", err
	}
	if len(output) != 0 && output[len(output)-1] == '\n' {
		output = output[:len(output)-1]
	}
	if len(output) != 40 {
		return "", fmt.Errorf("unexpected git log output, want commit hash: %q", output)
	}
	return string(output), nil
}

// ListRecentCommits returns list of recent commit titles starting from baseCommit.
func ListRecentCommits(dir, baseCommit string) ([]string, error) {
	// On upstream kernel this produces ~11MB of output.
	// Somewhat inefficient to collect whole output in a slice
	// and then convert to string, but should be bearable.
	output, err := osutil.RunCmd(timeout, dir, "git", "log",
		"--pretty=format:%s", "--no-merges", "-n", "200000", baseCommit)
	if err != nil {
		return nil, err
	}
	return strings.Split(string(output), "\n"), nil
}

// CanonicalizeCommit returns commit title that can be used when checking
// if a particular commit is present in a git tree.
// Some trees add prefixes to commit titles during backporting,
// so we want e.g. commit "foo bar" match "BACKPORT: foo bar".
func CanonicalizeCommit(title string) string {
	for _, prefix := range commitPrefixes {
		if strings.HasPrefix(title, prefix) {
			title = title[len(prefix):]
			break
		}
	}
	return strings.TrimSpace(title)
}

var commitPrefixes = []string{
	"UPSTREAM:",
	"CHROMIUM:",
	"FROMLIST:",
	"BACKPORT:",
	"FROMGIT:",
	"net-backports:",
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strings"
	"time"

	"github.com/go-check/check"
)

func (s *DockerTrustSuite) TestTrustedPull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-pull")

	// Try pull
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try untrusted pull to ensure we pushed the tag to the registry
	pullCmd = exec.Command(dockerBinary, "pull", "--disable-content-trust=true", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Status: Downloaded") {
		c.Fatalf("Missing expected output on trusted pull with --disable-content-trust:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestTrustedIsolatedPull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-isolatd-pull")

	// Try pull (run from isolated directory without trust information)
	pullCmd := exec.Command(dockerBinary, "--config", "/tmp/docker-isolated", "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)
}

func (s *DockerTrustSuite) TestUntrustedPull(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/trusted:latest", privateRegistryURL)
	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)
	dockerCmd(c, "push", repoName)
	dockerCmd(c, "rmi", repoName)

	// Try trusted pull on untrusted tag
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Error expected when running trusted pull with:\n%s", out)
	}

	if !strings.Contains(string(out), "no trust data available") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestPullWhenCertExpired(c *check.C) {
	c.Skip("Currently changes system time, causing instability")
	repoName := s.setupTrustedImage(c, "trusted-cert-expired")

	// Certificates have 10 years of expiration
	elevenYearsFromNow := time.Now().Add(time.Hour * 24 * 365 * 11)

	runAtDifferentDate(elevenYearsFromNow, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", repoName)
		s.trustedCmd(pullCmd)
		out, _, err := runCommandWithOutput(pullCmd)
		if err == nil {
			c.Fatalf("Error running trusted pull in the distant future: %s\n%s", err, out)
		}

		if !strings.Contains(string(out), "could not validate the path to a trusted root") {
			c.Fatalf("Missing expected output on trusted pull in the distant future:\n%s", out)
		}
	})

	runAtDifferentDate(elevenYearsFromNow, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", "--disable-content-trust", repoName)
		s.trustedCmd(pullCmd)
		out, _, err := runCommandWithOutput(pullCmd)
		if err != nil {
			c.Fatalf("Error running untrusted pull in the distant future: %s\n%s", err, out)
		}

		if !strings.Contains(string(out), "Status: Downloaded") {
			c.Fatalf("Missing expected output on untrusted pull in the distant future:\n%s", out)
		}
	})
}

func (s *DockerTrustSuite) TestTrustedPullFromBadTrustServer(c *check.C) {
	repoName := fmt.Sprintf("%v/dockerclievilpull/trusted:latest", privateRegistryURL)
	evilLocalConfigDir, err := ioutil.TempDir("", "evil-local-config-dir")
	if err != nil {
		c.Fatalf("Failed to create local temp dir")
	}

	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)

	pushCmd := exec.Command(dockerBinary, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err := runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("Error running trusted push: %s\n%s", err, out)
	}
	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try pull
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Kill the notary server, start a new "evil" one.
	s.not.Close()
	s.not, err = newTestNotary(c)
	if err != nil {
		c.Fatalf("Restarting notary server failed.")
	}

	// In order to make an evil server, lets re-init a client (with a different trust dir) and push new data.
	// tag an image and upload it to the private registry
	dockerCmd(c, "--config", evilLocalConfigDir, "tag", "busybox", repoName)

	// Push up to the new server
	pushCmd = exec.Command(dockerBinary, "--config", evilLocalConfigDir, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err = runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("Error running trusted push: %s\n%s", err, out)
	}
	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	// Now, try pulling with the original client from this new trust server. This should fail.
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Expected to fail on this pull due to different remote data: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "failed to validate data with current trusted certificates") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestTrustedPullWithExpiredSnapshot(c *check.C) {
	c.Skip("Currently changes system time, causing instability")
	repoName := fmt.Sprintf("%v/dockercliexpiredtimestamppull/trusted:latest", privateRegistryURL)
	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)

	// Push with default passphrases
	pushCmd := exec.Command(dockerBinary, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err := runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("trusted push failed: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Snapshots last for three years. This should be expired
	fourYearsLater := time.Now().Add(time.Hour * 24 * 365 * 4)

	// Should succeed because the server transparently re-signs one
	runAtDifferentDate(fourYearsLater, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", repoName)
		s.trustedCmd(pullCmd)
		out, _, err = runCommandWithOutput(pullCmd)
		if err == nil {
			c.Fatalf("Missing expected error running trusted pull with expired snapshots")
		}

		if !strings.Contains(string(out), "repository out-of-date") {
			c.Fatalf("Missing expected output on trusted pull with expired snapshot:\n%s", out)
		}
	})
}

func (s *DockerTrustSuite) TestTrustedOfflinePull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-offline-pull")

	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmdWithServer(pullCmd, "https://invalidnotaryserver")
	out, _, err := runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Expected error pulling with invalid notary server:\n%s", out)
	}

	if !strings.Contains(string(out), "error contacting notary server") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	// Do valid trusted pull to warm cache
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try pull again with invalid notary server, should use cache
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmdWithServer(pullCmd, "https://invalidnotaryserver")
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}
}
<commit_msg>Fix text not matching the actual tests<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strings"
	"time"

	"github.com/go-check/check"
)

func (s *DockerTrustSuite) TestTrustedPull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-pull")

	// Try pull
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try untrusted pull to ensure we pushed the tag to the registry
	pullCmd = exec.Command(dockerBinary, "pull", "--disable-content-trust=true", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Status: Downloaded") {
		c.Fatalf("Missing expected output on trusted pull with --disable-content-trust:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestTrustedIsolatedPull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-isolated-pull")

	// Try pull (run from isolated directory without trust information)
	pullCmd := exec.Command(dockerBinary, "--config", "/tmp/docker-isolated", "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)
}

func (s *DockerTrustSuite) TestUntrustedPull(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/trusted:latest", privateRegistryURL)
	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)
	dockerCmd(c, "push", repoName)
	dockerCmd(c, "rmi", repoName)

	// Try trusted pull on untrusted tag
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err := runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Error expected when running trusted pull with:\n%s", out)
	}

	if !strings.Contains(string(out), "no trust data available") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestPullWhenCertExpired(c *check.C) {
	c.Skip("Currently changes system time, causing instability")
	repoName := s.setupTrustedImage(c, "trusted-cert-expired")

	// Certificates have 10 years of expiration
	elevenYearsFromNow := time.Now().Add(time.Hour * 24 * 365 * 11)

	runAtDifferentDate(elevenYearsFromNow, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", repoName)
		s.trustedCmd(pullCmd)
		out, _, err := runCommandWithOutput(pullCmd)
		if err == nil {
			c.Fatalf("Error running trusted pull in the distant future: %s\n%s", err, out)
		}

		if !strings.Contains(string(out), "could not validate the path to a trusted root") {
			c.Fatalf("Missing expected output on trusted pull in the distant future:\n%s", out)
		}
	})

	runAtDifferentDate(elevenYearsFromNow, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", "--disable-content-trust", repoName)
		s.trustedCmd(pullCmd)
		out, _, err := runCommandWithOutput(pullCmd)
		if err != nil {
			c.Fatalf("Error running untrusted pull in the distant future: %s\n%s", err, out)
		}

		if !strings.Contains(string(out), "Status: Downloaded") {
			c.Fatalf("Missing expected output on untrusted pull in the distant future:\n%s", out)
		}
	})
}

func (s *DockerTrustSuite) TestTrustedPullFromBadTrustServer(c *check.C) {
	repoName := fmt.Sprintf("%v/dockerclievilpull/trusted:latest", privateRegistryURL)
	evilLocalConfigDir, err := ioutil.TempDir("", "evil-local-config-dir")
	if err != nil {
		c.Fatalf("Failed to create local temp dir")
	}

	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)

	pushCmd := exec.Command(dockerBinary, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err := runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("Error running trusted push: %s\n%s", err, out)
	}
	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try pull
	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Kill the notary server, start a new "evil" one.
	s.not.Close()
	s.not, err = newTestNotary(c)
	if err != nil {
		c.Fatalf("Restarting notary server failed.")
	}

	// In order to make an evil server, lets re-init a client (with a different trust dir) and push new data.
	// tag an image and upload it to the private registry
	dockerCmd(c, "--config", evilLocalConfigDir, "tag", "busybox", repoName)

	// Push up to the new server
	pushCmd = exec.Command(dockerBinary, "--config", evilLocalConfigDir, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err = runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("Error running trusted push: %s\n%s", err, out)
	}
	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	// Now, try pulling with the original client from this new trust server. This should fail.
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Expected to fail on this pull due to different remote data: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "failed to validate data with current trusted certificates") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}
}

func (s *DockerTrustSuite) TestTrustedPullWithExpiredSnapshot(c *check.C) {
	c.Skip("Currently changes system time, causing instability")
	repoName := fmt.Sprintf("%v/dockercliexpiredtimestamppull/trusted:latest", privateRegistryURL)
	// tag the image and upload it to the private registry
	dockerCmd(c, "tag", "busybox", repoName)

	// Push with default passphrases
	pushCmd := exec.Command(dockerBinary, "push", repoName)
	s.trustedCmd(pushCmd)
	out, _, err := runCommandWithOutput(pushCmd)
	if err != nil {
		c.Fatalf("trusted push failed: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Signing and pushing trust metadata") {
		c.Fatalf("Missing expected output on trusted push:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Snapshots last for three years. This should be expired
	fourYearsLater := time.Now().Add(time.Hour * 24 * 365 * 4)

	runAtDifferentDate(fourYearsLater, func() {
		// Try pull
		pullCmd := exec.Command(dockerBinary, "pull", repoName)
		s.trustedCmd(pullCmd)
		out, _, err = runCommandWithOutput(pullCmd)
		if err == nil {
			c.Fatalf("Missing expected error running trusted pull with expired snapshots")
		}

		if !strings.Contains(string(out), "repository out-of-date") {
			c.Fatalf("Missing expected output on trusted pull with expired snapshot:\n%s", out)
		}
	})
}

func (s *DockerTrustSuite) TestTrustedOfflinePull(c *check.C) {
	repoName := s.setupTrustedImage(c, "trusted-offline-pull")

	pullCmd := exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmdWithServer(pullCmd, "https://invalidnotaryserver")
	out, _, err := runCommandWithOutput(pullCmd)
	if err == nil {
		c.Fatalf("Expected error pulling with invalid notary server:\n%s", out)
	}

	if !strings.Contains(string(out), "error contacting notary server") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	// Do valid trusted pull to warm cache
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmd(pullCmd)
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}

	dockerCmd(c, "rmi", repoName)

	// Try pull again with invalid notary server, should use cache
	pullCmd = exec.Command(dockerBinary, "pull", repoName)
	s.trustedCmdWithServer(pullCmd, "https://invalidnotaryserver")
	out, _, err = runCommandWithOutput(pullCmd)
	if err != nil {
		c.Fatalf("Error running trusted pull: %s\n%s", err, out)
	}

	if !strings.Contains(string(out), "Tagging") {
		c.Fatalf("Missing expected output on trusted pull:\n%s", out)
	}
}
<|endoftext|>


<commit_before>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	err := c.QueryRow("SELECT * FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan()
	if err == nil {
		panic(fmt.Sprintf("SimId %x is already post-processed. Skipping.\n", c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId AND n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<commit_msg>fixed bug for recognizing when sim has already been post processed<commit_after>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	dummy := 0
	err := c.QueryRow("SELECT AgentId FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan(&dummy)
	if err == nil {
		panic(fmt.Sprintf("SimId %x is already post-processed - skipping", c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId AND n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<|endoftext|>