<commit_before>// Copyright 2015 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.

// execprog executes a single program or a set of programs
// and optionally prints information about execution.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/pkg/host"
	"github.com/google/syzkaller/pkg/ipc"
	"github.com/google/syzkaller/pkg/ipc/ipcconfig"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/tool"
	"github.com/google/syzkaller/prog"
	_ "github.com/google/syzkaller/sys"
)

var (
	flagOS        = flag.String("os", runtime.GOOS, "target os")
	flagArch      = flag.String("arch", runtime.GOARCH, "target arch")
	flagCoverFile = flag.String("coverfile", "", "write coverage to the file")
	flagRepeat    = flag.Int("repeat", 1, "repeat execution that many times (0 for infinite loop)")
	flagProcs     = flag.Int("procs", 1, "number of parallel processes to execute programs")
	flagOutput    = flag.Bool("output", false, "write programs and results to stdout")
	flagHints     = flag.Bool("hints", false, "do a hints-generation run")
	flagFaultCall = flag.Int("fault_call", -1, "inject fault into this call (0-based)")
	flagFaultNth  = flag.Int("fault_nth", 0, "inject fault on n-th operation (0-based)")
	flagEnable    = flag.String("enable", "none", "enable only listed additional features")
	flagDisable   = flag.String("disable", "none", "enable all additional features except listed")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: execprog [flags] file-with-programs+\n")
		flag.PrintDefaults()
		csource.PrintAvailableFeaturesFlags()
	}
	defer tool.Init()()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	featuresFlags, err := csource.ParseFeaturesFlags(*flagEnable, *flagDisable, true)
	if err != nil {
		log.Fatalf("%v", err)
	}

	target, err := prog.GetTarget(*flagOS, *flagArch)
	if err != nil {
		log.Fatalf("%v", err)
	}
	entries := loadPrograms(target, flag.Args())
	if len(entries) == 0 {
		return
	}
	features, err := host.Check(target)
	if err != nil {
		log.Fatalf("%v", err)
	}
	if *flagOutput {
		for _, feat := range features.Supported() {
			log.Logf(0, "%-24v: %v", feat.Name, feat.Reason)
		}
	}
	config, execOpts := createConfig(target, features, featuresFlags)
	if err = host.Setup(target, features, featuresFlags, config.Executor); err != nil {
		log.Fatal(err)
	}
	var gateCallback func()
	if features[host.FeatureLeak].Enabled {
		gateCallback = func() {
			output, err := osutil.RunCmd(10*time.Minute, "", config.Executor, "leak")
			if err != nil {
				os.Stdout.Write(output)
				os.Exit(1)
			}
		}
	}
	ctx := &Context{
		entries:  entries,
		config:   config,
		execOpts: execOpts,
		gate:     ipc.NewGate(2**flagProcs, gateCallback),
		shutdown: make(chan struct{}),
		repeat:   *flagRepeat,
	}
	var wg sync.WaitGroup
	wg.Add(*flagProcs)
	for p := 0; p < *flagProcs; p++ {
		pid := p
		go func() {
			defer wg.Done()
			ctx.run(pid)
		}()
	}
	osutil.HandleInterrupts(ctx.shutdown)
	wg.Wait()
}

type Context struct {
	entries   []*prog.LogEntry
	config    *ipc.Config
	execOpts  *ipc.ExecOpts
	gate      *ipc.Gate
	shutdown  chan struct{}
	logMu     sync.Mutex
	posMu     sync.Mutex
	repeat    int
	pos       int
	lastPrint time.Time
}

func (ctx *Context) run(pid int) {
	env, err := ipc.MakeEnv(ctx.config, pid)
	if err != nil {
		log.Fatalf("failed to create ipc env: %v", err)
	}
	defer env.Close()
	for {
		select {
		case <-ctx.shutdown:
			return
		default:
		}
		idx := ctx.getProgramIndex()
		if ctx.repeat > 0 && idx >= len(ctx.entries)*ctx.repeat {
			return
		}
		entry := ctx.entries[idx%len(ctx.entries)]
		ctx.execute(pid, env, entry)
	}
}

func (ctx *Context) execute(pid int, env *ipc.Env, entry *prog.LogEntry) {
	// Limit concurrency window.
	ticket := ctx.gate.Enter()
	defer ctx.gate.Leave(ticket)

	callOpts := ctx.execOpts
	if *flagFaultCall == -1 && entry.Fault {
		newOpts := *ctx.execOpts
		newOpts.Flags |= ipc.FlagInjectFault
		newOpts.FaultCall = entry.FaultCall
		newOpts.FaultNth = entry.FaultNth
		callOpts = &newOpts
	}
	if *flagOutput {
		ctx.logProgram(pid, entry.P, callOpts)
	}
	output, info, hanged, err := env.Exec(callOpts, entry.P)
	if ctx.config.Flags&ipc.FlagDebug != 0 || err != nil {
		log.Logf(0, "result: hanged=%v err=%v\n\n%s", hanged, err, output)
	}
	if info != nil {
		ctx.printCallResults(info)
		if *flagHints {
			ctx.printHints(entry.P, info)
		}
		if *flagCoverFile != "" {
			ctx.dumpCoverage(*flagCoverFile, info)
		}
	} else {
		log.Logf(1, "RESULT: no calls executed")
	}
}

func (ctx *Context) logProgram(pid int, p *prog.Prog, callOpts *ipc.ExecOpts) {
	strOpts := ""
	if callOpts.Flags&ipc.FlagInjectFault != 0 {
		strOpts = fmt.Sprintf(" (fault-call:%v fault-nth:%v)",
			callOpts.FaultCall, callOpts.FaultNth)
	}
	data := p.Serialize()
	ctx.logMu.Lock()
	log.Logf(0, "executing program %v%v:\n%s", pid, strOpts, data)
	ctx.logMu.Unlock()
}

func (ctx *Context) printCallResults(info *ipc.ProgInfo) {
	for i, inf := range info.Calls {
		if inf.Flags&ipc.CallExecuted == 0 {
			continue
		}
		flags := ""
		if inf.Flags&ipc.CallFinished == 0 {
			flags += " unfinished"
		}
		if inf.Flags&ipc.CallBlocked != 0 {
			flags += " blocked"
		}
		if inf.Flags&ipc.CallFaultInjected != 0 {
			flags += " faulted"
		}
		log.Logf(1, "CALL %v: signal %v, coverage %v errno %v%v",
			i, len(inf.Signal), len(inf.Cover), inf.Errno, flags)
	}
}

func (ctx *Context) printHints(p *prog.Prog, info *ipc.ProgInfo) {
	ncomps, ncandidates := 0, 0
	for i := range p.Calls {
		if *flagOutput {
			fmt.Printf("call %v:\n", i)
		}
		comps := info.Calls[i].Comps
		for v, args := range comps {
			ncomps += len(args)
			if *flagOutput {
				fmt.Printf("comp 0x%x:", v)
				for arg := range args {
					fmt.Printf(" 0x%x", arg)
				}
				fmt.Printf("\n")
			}
		}
		p.MutateWithHints(i, comps, func(p *prog.Prog) {
			ncandidates++
			if *flagOutput {
				log.Logf(1, "PROGRAM:\n%s", p.Serialize())
			}
		})
	}
	log.Logf(0, "ncomps=%v ncandidates=%v", ncomps, ncandidates)
}

func (ctx *Context) dumpCallCoverage(coverFile string, info *ipc.CallInfo) {
	if len(info.Cover) == 0 {
		return
	}
	buf := new(bytes.Buffer)
	for _, pc := range info.Cover {
		fmt.Fprintf(buf, "0x%x\n", cover.RestorePC(pc, 0xffffffff))
	}
	err := osutil.WriteFile(coverFile, buf.Bytes())
	if err != nil {
		log.Fatalf("failed to write coverage file: %v", err)
	}
}

func (ctx *Context) dumpCoverage(coverFile string, info *ipc.ProgInfo) {
	for i, inf := range info.Calls {
		log.Logf(0, "call #%v: signal %v, coverage %v", i, len(inf.Signal), len(inf.Cover))
		ctx.dumpCallCoverage(fmt.Sprintf("%v.%v", coverFile, i), &inf)
	}
	log.Logf(0, "extra: signal %v, coverage %v", len(info.Extra.Signal), len(info.Extra.Cover))
	ctx.dumpCallCoverage(fmt.Sprintf("%v.extra", coverFile), &info.Extra)
}

func (ctx *Context) getProgramIndex() int {
	ctx.posMu.Lock()
	idx := ctx.pos
	ctx.pos++
	if idx%len(ctx.entries) == 0 && time.Since(ctx.lastPrint) > 5*time.Second {
		log.Logf(0, "executed programs: %v", idx)
		ctx.lastPrint = time.Now()
	}
	ctx.posMu.Unlock()
	return idx
}

func loadPrograms(target *prog.Target, files []string) []*prog.LogEntry {
	var entries []*prog.LogEntry
	for _, fn := range files {
		data, err := ioutil.ReadFile(fn)
		if err != nil {
			log.Fatalf("failed to read log file: %v", err)
		}
		entries = append(entries, target.ParseLog(data)...)
	}
	log.Logf(0, "parsed %v programs", len(entries))
	return entries
}

func createConfig(target *prog.Target, features *host.Features, featuresFlags csource.Features) (
	*ipc.Config, *ipc.ExecOpts) {
	config, execOpts, err := ipcconfig.Default(target)
	if err != nil {
		log.Fatalf("%v", err)
	}
	if config.Flags&ipc.FlagSignal != 0 {
		execOpts.Flags |= ipc.FlagCollectCover
	}
	if *flagCoverFile != "" {
		config.Flags |= ipc.FlagSignal
		execOpts.Flags |= ipc.FlagCollectCover
		execOpts.Flags &^= ipc.FlagDedupCover
	}
	if *flagHints {
		if execOpts.Flags&ipc.FlagCollectCover != 0 {
			execOpts.Flags ^= ipc.FlagCollectCover
		}
		execOpts.Flags |= ipc.FlagCollectComps
	}
	if features[host.FeatureExtraCoverage].Enabled {
		config.Flags |= ipc.FlagExtraCover
	}
	if *flagFaultCall >= 0 {
		execOpts.Flags |= ipc.FlagInjectFault
		execOpts.FaultCall = *flagFaultCall
		execOpts.FaultNth = *flagFaultNth
	}
	if featuresFlags["tun"].Enabled && features[host.FeatureNetInjection].Enabled {
		config.Flags |= ipc.FlagEnableTun
	}
	if featuresFlags["net_dev"].Enabled && features[host.FeatureNetDevices].Enabled {
		config.Flags |= ipc.FlagEnableNetDev
	}
	if featuresFlags["net_reset"].Enabled {
		config.Flags |= ipc.FlagEnableNetReset
	}
	if featuresFlags["cgroups"].Enabled {
		config.Flags |= ipc.FlagEnableCgroups
	}
	if featuresFlags["close_fds"].Enabled {
		config.Flags |= ipc.FlagEnableCloseFds
	}
	if featuresFlags["devlink_pci"].Enabled && features[host.FeatureDevlinkPCI].Enabled {
		config.Flags |= ipc.FlagEnableDevlinkPCI
	}
	if featuresFlags["vhci"].Enabled && features[host.FeatureVhciInjection].Enabled {
		config.Flags |= ipc.FlagEnableVhciInjection
	}
	if featuresFlags["wifi"].Enabled && features[host.FeatureWifiEmulation].Enabled {
		config.Flags |= ipc.FlagEnableWifi
	}
	return config, execOpts
}
<commit_msg>tools/syz-execprog: mimic syz-fuzzer logic for executor failures<commit_after>// Copyright 2015 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.

// execprog executes a single program or a set of programs
// and optionally prints information about execution.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/pkg/host"
	"github.com/google/syzkaller/pkg/ipc"
	"github.com/google/syzkaller/pkg/ipc/ipcconfig"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/tool"
	"github.com/google/syzkaller/prog"
	_ "github.com/google/syzkaller/sys"
)

var (
	flagOS        = flag.String("os", runtime.GOOS, "target os")
	flagArch      = flag.String("arch", runtime.GOARCH, "target arch")
	flagCoverFile = flag.String("coverfile", "", "write coverage to the file")
	flagRepeat    = flag.Int("repeat", 1, "repeat execution that many times (0 for infinite loop)")
	flagProcs     = flag.Int("procs", 1, "number of parallel processes to execute programs")
	flagOutput    = flag.Bool("output", false, "write programs and results to stdout")
	flagHints     = flag.Bool("hints", false, "do a hints-generation run")
	flagFaultCall = flag.Int("fault_call", -1, "inject fault into this call (0-based)")
	flagFaultNth  = flag.Int("fault_nth", 0, "inject fault on n-th operation (0-based)")
	flagEnable    = flag.String("enable", "none", "enable only listed additional features")
	flagDisable   = flag.String("disable", "none", "enable all additional features except listed")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: execprog [flags] file-with-programs+\n")
		flag.PrintDefaults()
		csource.PrintAvailableFeaturesFlags()
	}
	defer tool.Init()()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	featuresFlags, err := csource.ParseFeaturesFlags(*flagEnable, *flagDisable, true)
	if err != nil {
		log.Fatalf("%v", err)
	}

	target, err := prog.GetTarget(*flagOS, *flagArch)
	if err != nil {
		log.Fatalf("%v", err)
	}
	entries := loadPrograms(target, flag.Args())
	if len(entries) == 0 {
		return
	}
	features, err := host.Check(target)
	if err != nil {
		log.Fatalf("%v", err)
	}
	if *flagOutput {
		for _, feat := range features.Supported() {
			log.Logf(0, "%-24v: %v", feat.Name, feat.Reason)
		}
	}
	config, execOpts := createConfig(target, features, featuresFlags)
	if err = host.Setup(target, features, featuresFlags, config.Executor); err != nil {
		log.Fatal(err)
	}
	var gateCallback func()
	if features[host.FeatureLeak].Enabled {
		gateCallback = func() {
			output, err := osutil.RunCmd(10*time.Minute, "", config.Executor, "leak")
			if err != nil {
				os.Stdout.Write(output)
				os.Exit(1)
			}
		}
	}
	ctx := &Context{
		entries:  entries,
		config:   config,
		execOpts: execOpts,
		gate:     ipc.NewGate(2**flagProcs, gateCallback),
		shutdown: make(chan struct{}),
		repeat:   *flagRepeat,
	}
	var wg sync.WaitGroup
	wg.Add(*flagProcs)
	for p := 0; p < *flagProcs; p++ {
		pid := p
		go func() {
			defer wg.Done()
			ctx.run(pid)
		}()
	}
	osutil.HandleInterrupts(ctx.shutdown)
	wg.Wait()
}

type Context struct {
	entries   []*prog.LogEntry
	config    *ipc.Config
	execOpts  *ipc.ExecOpts
	gate      *ipc.Gate
	shutdown  chan struct{}
	logMu     sync.Mutex
	posMu     sync.Mutex
	repeat    int
	pos       int
	lastPrint time.Time
}

func (ctx *Context) run(pid int) {
	env, err := ipc.MakeEnv(ctx.config, pid)
	if err != nil {
		log.Fatalf("failed to create ipc env: %v", err)
	}
	defer env.Close()
	for {
		select {
		case <-ctx.shutdown:
			return
		default:
		}
		idx := ctx.getProgramIndex()
		if ctx.repeat > 0 && idx >= len(ctx.entries)*ctx.repeat {
			return
		}
		entry := ctx.entries[idx%len(ctx.entries)]
		ctx.execute(pid, env, entry)
	}
}

func (ctx *Context) execute(pid int, env *ipc.Env, entry *prog.LogEntry) {
	// Limit concurrency window.
	ticket := ctx.gate.Enter()
	defer ctx.gate.Leave(ticket)

	callOpts := ctx.execOpts
	if *flagFaultCall == -1 && entry.Fault {
		newOpts := *ctx.execOpts
		newOpts.Flags |= ipc.FlagInjectFault
		newOpts.FaultCall = entry.FaultCall
		newOpts.FaultNth = entry.FaultNth
		callOpts = &newOpts
	}
	if *flagOutput {
		ctx.logProgram(pid, entry.P, callOpts)
	}
	// This mimics the syz-fuzzer logic. This is important for reproduction.
	for try := 0; ; try++ {
		output, info, hanged, err := env.Exec(callOpts, entry.P)
		if err != nil && err != prog.ErrExecBufferTooSmall {
			if try > 10 {
				log.Fatalf("executor failed %v times: %v\n%s", try, err, output)
			}
			// Don't print err/output in this case as it may contain "SYZFAIL" and we want to fail yet.
			log.Logf(1, "executor failed, retrying")
			time.Sleep(time.Second)
			continue
		}
		if ctx.config.Flags&ipc.FlagDebug != 0 || err != nil {
			log.Logf(0, "result: hanged=%v err=%v\n\n%s", hanged, err, output)
		}
		if info != nil {
			ctx.printCallResults(info)
			if *flagHints {
				ctx.printHints(entry.P, info)
			}
			if *flagCoverFile != "" {
				ctx.dumpCoverage(*flagCoverFile, info)
			}
		} else {
			log.Logf(1, "RESULT: no calls executed")
		}
		break
	}
}

func (ctx *Context) logProgram(pid int, p *prog.Prog, callOpts *ipc.ExecOpts) {
	strOpts := ""
	if callOpts.Flags&ipc.FlagInjectFault != 0 {
		strOpts = fmt.Sprintf(" (fault-call:%v fault-nth:%v)",
			callOpts.FaultCall, callOpts.FaultNth)
	}
	data := p.Serialize()
	ctx.logMu.Lock()
	log.Logf(0, "executing program %v%v:\n%s", pid, strOpts, data)
	ctx.logMu.Unlock()
}

func (ctx *Context) printCallResults(info *ipc.ProgInfo) {
	for i, inf := range info.Calls {
		if inf.Flags&ipc.CallExecuted == 0 {
			continue
		}
		flags := ""
		if inf.Flags&ipc.CallFinished == 0 {
			flags += " unfinished"
		}
		if inf.Flags&ipc.CallBlocked != 0 {
			flags += " blocked"
		}
		if inf.Flags&ipc.CallFaultInjected != 0 {
			flags += " faulted"
		}
		log.Logf(1, "CALL %v: signal %v, coverage %v errno %v%v",
			i, len(inf.Signal), len(inf.Cover), inf.Errno, flags)
	}
}

func (ctx *Context) printHints(p *prog.Prog, info *ipc.ProgInfo) {
	ncomps, ncandidates := 0, 0
	for i := range p.Calls {
		if *flagOutput {
			fmt.Printf("call %v:\n", i)
		}
		comps := info.Calls[i].Comps
		for v, args := range comps {
			ncomps += len(args)
			if *flagOutput {
				fmt.Printf("comp 0x%x:", v)
				for arg := range args {
					fmt.Printf(" 0x%x", arg)
				}
				fmt.Printf("\n")
			}
		}
		p.MutateWithHints(i, comps, func(p *prog.Prog) {
			ncandidates++
			if *flagOutput {
				log.Logf(1, "PROGRAM:\n%s", p.Serialize())
			}
		})
	}
	log.Logf(0, "ncomps=%v ncandidates=%v", ncomps, ncandidates)
}

func (ctx *Context) dumpCallCoverage(coverFile string, info *ipc.CallInfo) {
	if len(info.Cover) == 0 {
		return
	}
	buf := new(bytes.Buffer)
	for _, pc := range info.Cover {
		fmt.Fprintf(buf, "0x%x\n", cover.RestorePC(pc, 0xffffffff))
	}
	err := osutil.WriteFile(coverFile, buf.Bytes())
	if err != nil {
		log.Fatalf("failed to write coverage file: %v", err)
	}
}

func (ctx *Context) dumpCoverage(coverFile string, info *ipc.ProgInfo) {
	for i, inf := range info.Calls {
		log.Logf(0, "call #%v: signal %v, coverage %v", i, len(inf.Signal), len(inf.Cover))
		ctx.dumpCallCoverage(fmt.Sprintf("%v.%v", coverFile, i), &inf)
	}
	log.Logf(0, "extra: signal %v, coverage %v", len(info.Extra.Signal), len(info.Extra.Cover))
	ctx.dumpCallCoverage(fmt.Sprintf("%v.extra", coverFile), &info.Extra)
}

func (ctx *Context) getProgramIndex() int {
	ctx.posMu.Lock()
	idx := ctx.pos
	ctx.pos++
	if idx%len(ctx.entries) == 0 && time.Since(ctx.lastPrint) > 5*time.Second {
		log.Logf(0, "executed programs: %v", idx)
		ctx.lastPrint = time.Now()
	}
	ctx.posMu.Unlock()
	return idx
}

func loadPrograms(target *prog.Target, files []string) []*prog.LogEntry {
	var entries []*prog.LogEntry
	for _, fn := range files {
		data, err := ioutil.ReadFile(fn)
		if err != nil {
			log.Fatalf("failed to read log file: %v", err)
		}
		entries = append(entries, target.ParseLog(data)...)
	}
	log.Logf(0, "parsed %v programs", len(entries))
	return entries
}

func createConfig(target *prog.Target, features *host.Features, featuresFlags csource.Features) (
	*ipc.Config, *ipc.ExecOpts) {
	config, execOpts, err := ipcconfig.Default(target)
	if err != nil {
		log.Fatalf("%v", err)
	}
	if config.Flags&ipc.FlagSignal != 0 {
		execOpts.Flags |= ipc.FlagCollectCover
	}
	if *flagCoverFile != "" {
		config.Flags |= ipc.FlagSignal
		execOpts.Flags |= ipc.FlagCollectCover
		execOpts.Flags &^= ipc.FlagDedupCover
	}
	if *flagHints {
		if execOpts.Flags&ipc.FlagCollectCover != 0 {
			execOpts.Flags ^= ipc.FlagCollectCover
		}
		execOpts.Flags |= ipc.FlagCollectComps
	}
	if features[host.FeatureExtraCoverage].Enabled {
		config.Flags |= ipc.FlagExtraCover
	}
	if *flagFaultCall >= 0 {
		execOpts.Flags |= ipc.FlagInjectFault
		execOpts.FaultCall = *flagFaultCall
		execOpts.FaultNth = *flagFaultNth
	}
	if featuresFlags["tun"].Enabled && features[host.FeatureNetInjection].Enabled {
		config.Flags |= ipc.FlagEnableTun
	}
	if featuresFlags["net_dev"].Enabled && features[host.FeatureNetDevices].Enabled {
		config.Flags |= ipc.FlagEnableNetDev
	}
	if featuresFlags["net_reset"].Enabled {
		config.Flags |= ipc.FlagEnableNetReset
	}
	if featuresFlags["cgroups"].Enabled {
		config.Flags |= ipc.FlagEnableCgroups
	}
	if featuresFlags["close_fds"].Enabled {
		config.Flags |= ipc.FlagEnableCloseFds
	}
	if featuresFlags["devlink_pci"].Enabled && features[host.FeatureDevlinkPCI].Enabled {
		config.Flags |= ipc.FlagEnableDevlinkPCI
	}
	if featuresFlags["vhci"].Enabled && features[host.FeatureVhciInjection].Enabled {
		config.Flags |= ipc.FlagEnableVhciInjection
	}
	if featuresFlags["wifi"].Enabled && features[host.FeatureWifiEmulation].Enabled {
		config.Flags |= ipc.FlagEnableWifi
	}
	return config, execOpts
}
<|endoftext|>


<commit_before>// +build linux

/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 netns

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"gopkg.in/fsnotify/fsnotify.v1"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
	"github.com/skydive-project/skydive/topology/probes/netlink"
	"github.com/vishvananda/netns"
)

// Probe describes a netlink probe in a network namespace
type Probe struct {
	common.RWMutex
	Graph       *graph.Graph
	Root        *graph.Node
	Probe       *netlink.Probe
	pathToNetNS map[string]*NetNs
	netNsProbes map[string]*netNsProbe
	rootNs      *NetNs
	watcher     *fsnotify.Watcher
	pending     chan string
	exclude     []string
}

// NetNs describes a network namespace path associated with a device / inode
type NetNs struct {
	path string
	dev  uint64
	ino  uint64
}

// extends the original struct to add use count number
type netNsProbe struct {
	*netlink.NetNsProbe
	useCount int
}

func getNetNSName(path string) string {
	s := strings.Split(path, "/")
	return s[len(s)-1]
}

func (ns *NetNs) String() string {
	return fmt.Sprintf("%d,%d", ns.dev, ns.ino)
}

// Equal compares two NetNs objects
func (ns *NetNs) Equal(o *NetNs) bool {
	return (ns.dev == o.dev && ns.ino == o.ino)
}

func (u *Probe) checkNamespace(path string) error {
	// When a new network namespace has been seen by inotify, the path to
	// the namespace may still be a regular file, not a bind mount to the
	// file in /proc/<pid>/tasks/<tid>/ns/net yet, so we wait a bit for the
	// bind mount to be set up

	return common.Retry(func() error {
		var stats, parentStats syscall.Stat_t
		fd, err := syscall.Open(path, syscall.O_RDONLY, 0)
		if err != nil {
			return err
		}

		err = syscall.Fstat(fd, &stats)
		syscall.Close(fd)
		if err != nil {
			return err
		}

		if (stats.Mode & syscall.S_IFLNK) > 0 {
			return nil
		}

		if parent := filepath.Dir(path); parent != "" {
			if err := syscall.Stat(parent, &parentStats); err == nil {
				if stats.Dev == parentStats.Dev {
					return fmt.Errorf("%s does not seem to be a valid namespace", path)
				}
			}
		}

		return nil
	}, 30, time.Millisecond*100)
}

// Register a new network namespace path
func (u *Probe) Register(path string, name string) (*graph.Node, error) {
	logging.GetLogger().Debugf("Register network namespace: %s", path)

	if err := u.checkNamespace(path); err != nil {
		return nil, err
	}

	var stats syscall.Stat_t
	if err := syscall.Stat(path, &stats); err != nil {
		return nil, fmt.Errorf("Failed to stat namespace %s: %s", path, err)
	}

	newns := &NetNs{path: path, dev: stats.Dev, ino: stats.Ino}

	// avoid hard link to root ns
	if u.rootNs.Equal(newns) {
		return u.Root, nil
	}

	u.Lock()
	defer u.Unlock()

	_, ok := u.pathToNetNS[path]
	if !ok {
		u.pathToNetNS[path] = newns
	}

	nsString := newns.String()
	if probe, ok := u.netNsProbes[nsString]; ok {
		probe.useCount++
		logging.GetLogger().Debugf("Increasing counter for namespace %s to %d", nsString, probe.useCount)
		return probe.Root, nil
	}

	u.Graph.Lock()

	logging.GetLogger().Debugf("Network namespace added: %s", nsString)
	metadata := graph.Metadata{
		"Name":   name,
		"Type":   "netns",
		"Path":   path,
		"Inode":  int64(newns.ino),
		"Device": int64(newns.dev),
	}

	n := u.Graph.NewNode(graph.GenID(), metadata)
	topology.AddOwnershipLink(u.Graph, u.Root, n, nil)

	u.Graph.Unlock()

	logging.GetLogger().Debugf("Registering namespace: %s", nsString)

	var probe *netlink.NetNsProbe
	err := common.Retry(func() error {
		var err error
		probe, err = u.Probe.Register(path, n)
		if err != nil {
			return fmt.Errorf("Could not register netlink probe within namespace: %s", err)
		}
		return nil
	}, 100, 10*time.Millisecond)
	if err != nil {
		return nil, err
	}

	u.netNsProbes[nsString] = &netNsProbe{NetNsProbe: probe, useCount: 1}

	return n, nil
}

// Unregister a network namespace path
func (u *Probe) Unregister(path string) {
	logging.GetLogger().Debugf("Unregister network Namespace: %s", path)

	u.Lock()
	defer u.Unlock()

	ns, ok := u.pathToNetNS[path]
	if !ok {
		return
	}

	delete(u.pathToNetNS, path)
	nsString := ns.String()
	probe, ok := u.netNsProbes[nsString]
	if !ok {
		logging.GetLogger().Debugf("No existing network namespace found: %s (%s)", nsString)
		return
	}

	if probe.useCount > 1 {
		probe.useCount--
		logging.GetLogger().Debugf("Decremented counter for namespace %s to %d", nsString, probe.useCount)
		return
	}

	u.Probe.Unregister(path)
	logging.GetLogger().Debugf("Network namespace deleted: %s", nsString)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	for _, child := range u.Graph.LookupChildren(probe.Root, nil, nil) {
		u.Graph.DelNode(child)
	}
	u.Graph.DelNode(probe.Root)

	delete(u.netNsProbes, nsString)
}

func (u *Probe) initializeRunPath(path string) {
	for {
		if _, err := os.Stat(path); err == nil {
			break
		}
		time.Sleep(time.Second)
	}

	if err := u.watcher.Add(path); err != nil {
		logging.GetLogger().Errorf("Unable to Watch %s: %s", path, err)
	}

	files, _ := ioutil.ReadDir(path)
LOOP:
	for _, f := range files {
		fullpath, name := path+"/"+f.Name(), f.Name()

		if u.isPathExcluded(fullpath) {
			continue LOOP
		}

		if _, err := u.Register(fullpath, name); err != nil {
			logging.GetLogger().Errorf("Failed to register namespace %s: %s", fullpath, err)
			continue
		}
	}
	logging.GetLogger().Debugf("Probe initialized %s", path)
}

func (u *Probe) start() {
	logging.GetLogger().Debugf("Probe initialized")
LOOP:
	for {
		select {
		case path := <-u.pending:
			go u.initializeRunPath(path)
		case ev := <-u.watcher.Events:
			if u.isPathExcluded(ev.Name) {
				continue LOOP
			}
			if ev.Op&fsnotify.Create == fsnotify.Create {
				if _, err := u.Register(ev.Name, getNetNSName(ev.Name)); err != nil {
					logging.GetLogger().Errorf("Failed to register namespace %s: %s", ev.Name, err)
					continue
				}
			}
			if ev.Op&fsnotify.Remove == fsnotify.Remove {
				u.Unregister(ev.Name)
			}

		case err := <-u.watcher.Errors:
			logging.GetLogger().Errorf("Error while watching network namespace: %s", err)
		}
	}
}

// Watch add a path to the inotify watcher
func (u *Probe) Watch(path string) {
	u.pending <- path
}

// Start the probe
func (u *Probe) Start() {
	go u.start()
}

// Stop the probe
func (u *Probe) Stop() {
	u.Probe.Stop()
}

func (u *Probe) isPathExcluded(path string) bool {
	u.RLock()
	defer u.RUnlock()

	for _, e := range u.exclude {
		if e == path {
			return true
		}
	}
	return false
}

// Exclude specify path to not process
func (u *Probe) Exclude(paths ...string) {
	u.Lock()
	u.exclude = append(u.exclude, paths...)
	u.Unlock()
}

// NewProbe creates a new network namespace probe
func NewProbe(g *graph.Graph, n *graph.Node, nlProbe *netlink.Probe) (*Probe, error) {
	if uid := os.Geteuid(); uid != 0 {
		return nil, errors.New("NetNS probe has to be run as root")
	}

	ns, err := netns.Get()
	if err != nil {
		return nil, errors.New("Failed to get root namespace")
	}
	defer ns.Close()

	var stats syscall.Stat_t
	if err = syscall.Fstat(int(ns), &stats); err != nil {
		return nil, errors.New("Failed to stat root namespace")
	}
	rootNs := &NetNs{dev: stats.Dev, ino: stats.Ino}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("Unable to create a new Watcher: %s", err)
	}

	nsProbe := &Probe{
		Graph:       g,
		Root:        n,
		Probe:       nlProbe,
		pathToNetNS: make(map[string]*NetNs),
		netNsProbes: make(map[string]*netNsProbe),
		rootNs:      rootNs,
		watcher:     watcher,
		pending:     make(chan string, 10),
	}

	if path := config.GetString("netns.run_path"); path != "" {
		nsProbe.Watch(path)
	}

	return nsProbe, nil
}
<commit_msg>netns: allow to start not being root<commit_after>// +build linux

/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 netns

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"gopkg.in/fsnotify/fsnotify.v1"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/graph"
	"github.com/skydive-project/skydive/topology/probes/netlink"
	"github.com/vishvananda/netns"
)

// Probe describes a netlink probe in a network namespace
type Probe struct {
	common.RWMutex
	Graph       *graph.Graph
	Root        *graph.Node
	Probe       *netlink.Probe
	pathToNetNS map[string]*NetNs
	netNsProbes map[string]*netNsProbe
	rootNs      *NetNs
	watcher     *fsnotify.Watcher
	pending     chan string
	exclude     []string
}

// NetNs describes a network namespace path associated with a device / inode
type NetNs struct {
	path string
	dev  uint64
	ino  uint64
}

// extends the original struct to add use count number
type netNsProbe struct {
	*netlink.NetNsProbe
	useCount int
}

func getNetNSName(path string) string {
	s := strings.Split(path, "/")
	return s[len(s)-1]
}

func (ns *NetNs) String() string {
	return fmt.Sprintf("%d,%d", ns.dev, ns.ino)
}

// Equal compares two NetNs objects
func (ns *NetNs) Equal(o *NetNs) bool {
	return (ns.dev == o.dev && ns.ino == o.ino)
}

func (u *Probe) checkNamespace(path string) error {
	// When a new network namespace has been seen by inotify, the path to
	// the namespace may still be a regular file, not a bind mount to the
	// file in /proc/<pid>/tasks/<tid>/ns/net yet, so we wait a bit for the
	// bind mount to be set up

	return common.Retry(func() error {
		var stats, parentStats syscall.Stat_t
		fd, err := syscall.Open(path, syscall.O_RDONLY, 0)
		if err != nil {
			return err
		}

		err = syscall.Fstat(fd, &stats)
		syscall.Close(fd)
		if err != nil {
			return err
		}

		if (stats.Mode & syscall.S_IFLNK) > 0 {
			return nil
		}

		if parent := filepath.Dir(path); parent != "" {
			if err := syscall.Stat(parent, &parentStats); err == nil {
				if stats.Dev == parentStats.Dev {
					return fmt.Errorf("%s does not seem to be a valid namespace", path)
				}
			}
		}

		return nil
	}, 30, time.Millisecond*100)
}

// Register a new network namespace path
func (u *Probe) Register(path string, name string) (*graph.Node, error) {
	logging.GetLogger().Debugf("Register network namespace: %s", path)

	if err := u.checkNamespace(path); err != nil {
		return nil, err
	}

	var stats syscall.Stat_t
	if err := syscall.Stat(path, &stats); err != nil {
		return nil, fmt.Errorf("Failed to stat namespace %s: %s", path, err)
	}

	newns := &NetNs{path: path, dev: stats.Dev, ino: stats.Ino}

	// avoid hard link to root ns
	if u.rootNs.Equal(newns) {
		return u.Root, nil
	}

	u.Lock()
	defer u.Unlock()

	_, ok := u.pathToNetNS[path]
	if !ok {
		u.pathToNetNS[path] = newns
	}

	nsString := newns.String()
	if probe, ok := u.netNsProbes[nsString]; ok {
		probe.useCount++
		logging.GetLogger().Debugf("Increasing counter for namespace %s to %d", nsString, probe.useCount)
		return probe.Root, nil
	}

	u.Graph.Lock()

	logging.GetLogger().Debugf("Network namespace added: %s", nsString)
	metadata := graph.Metadata{
		"Name":   name,
		"Type":   "netns",
		"Path":   path,
		"Inode":  int64(newns.ino),
		"Device": int64(newns.dev),
	}

	n := u.Graph.NewNode(graph.GenID(), metadata)
	topology.AddOwnershipLink(u.Graph, u.Root, n, nil)

	u.Graph.Unlock()

	logging.GetLogger().Debugf("Registering namespace: %s", nsString)

	var probe *netlink.NetNsProbe
	err := common.Retry(func() error {
		var err error
		probe, err = u.Probe.Register(path, n)
		if err != nil {
			return fmt.Errorf("Could not register netlink probe within namespace: %s", err)
		}
		return nil
	}, 100, 10*time.Millisecond)
	if err != nil {
		return nil, err
	}

	u.netNsProbes[nsString] = &netNsProbe{NetNsProbe: probe, useCount: 1}

	return n, nil
}

// Unregister a network namespace path
func (u *Probe) Unregister(path string) {
	logging.GetLogger().Debugf("Unregister network Namespace: %s", path)

	u.Lock()
	defer u.Unlock()

	ns, ok := u.pathToNetNS[path]
	if !ok {
		return
	}

	delete(u.pathToNetNS, path)
	nsString := ns.String()
	probe, ok := u.netNsProbes[nsString]
	if !ok {
		logging.GetLogger().Debugf("No existing network namespace found: %s (%s)", nsString)
		return
	}

	if probe.useCount > 1 {
		probe.useCount--
		logging.GetLogger().Debugf("Decremented counter for namespace %s to %d", nsString, probe.useCount)
		return
	}

	u.Probe.Unregister(path)
	logging.GetLogger().Debugf("Network namespace deleted: %s", nsString)

	u.Graph.Lock()
	defer u.Graph.Unlock()

	for _, child := range u.Graph.LookupChildren(probe.Root, nil, nil) {
		u.Graph.DelNode(child)
	}
	u.Graph.DelNode(probe.Root)

	delete(u.netNsProbes, nsString)
}

func (u *Probe) initializeRunPath(path string) {
	for {
		if _, err := os.Stat(path); err == nil {
			break
		}
		time.Sleep(time.Second)
	}

	if err := u.watcher.Add(path); err != nil {
		logging.GetLogger().Errorf("Unable to Watch %s: %s", path, err)
	}

	files, _ := ioutil.ReadDir(path)
LOOP:
	for _, f := range files {
		fullpath, name := path+"/"+f.Name(), f.Name()

		if u.isPathExcluded(fullpath) {
			continue LOOP
		}

		if _, err := u.Register(fullpath, name); err != nil {
			logging.GetLogger().Errorf("Failed to register namespace %s: %s", fullpath, err)
			continue
		}
	}
	logging.GetLogger().Debugf("Probe initialized %s", path)
}

func (u *Probe) start() {
	logging.GetLogger().Debugf("Probe initialized")
LOOP:
	for {
		select {
		case path := <-u.pending:
			go u.initializeRunPath(path)
		case ev := <-u.watcher.Events:
			if u.isPathExcluded(ev.Name) {
				continue LOOP
			}
			if ev.Op&fsnotify.Create == fsnotify.Create {
				if _, err := u.Register(ev.Name, getNetNSName(ev.Name)); err != nil {
					logging.GetLogger().Errorf("Failed to register namespace %s: %s", ev.Name, err)
					continue
				}
			}
			if ev.Op&fsnotify.Remove == fsnotify.Remove {
				u.Unregister(ev.Name)
			}

		case err := <-u.watcher.Errors:
			logging.GetLogger().Errorf("Error while watching network namespace: %s", err)
		}
	}
}

// Watch add a path to the inotify watcher
func (u *Probe) Watch(path string) {
	u.pending <- path
}

// Start the probe
func (u *Probe) Start() {
	go u.start()
}

// Stop the probe
func (u *Probe) Stop() {
	u.Probe.Stop()
}

func (u *Probe) isPathExcluded(path string) bool {
	u.RLock()
	defer u.RUnlock()

	for _, e := range u.exclude {
		if e == path {
			return true
		}
	}
	return false
}

// Exclude specify path to not process
func (u *Probe) Exclude(paths ...string) {
	u.Lock()
	u.exclude = append(u.exclude, paths...)
	u.Unlock()
}

// NewProbe creates a new network namespace probe
func NewProbe(g *graph.Graph, n *graph.Node, nlProbe *netlink.Probe) (*Probe, error) {
	if uid := os.Geteuid(); uid != 0 {
		logging.GetLogger().Warning("NetNS probe needs root permissions, network namespace might not be repported")
	}

	ns, err := netns.Get()
	if err != nil {
		return nil, errors.New("Failed to get root namespace")
	}
	defer ns.Close()

	var stats syscall.Stat_t
	if err = syscall.Fstat(int(ns), &stats); err != nil {
		return nil, errors.New("Failed to stat root namespace")
	}
	rootNs := &NetNs{dev: stats.Dev, ino: stats.Ino}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("Unable to create a new Watcher: %s", err)
	}

	nsProbe := &Probe{
		Graph:       g,
		Root:        n,
		Probe:       nlProbe,
		pathToNetNS: make(map[string]*NetNs),
		netNsProbes: make(map[string]*netNsProbe),
		rootNs:      rootNs,
		watcher:     watcher,
		pending:     make(chan string, 10),
	}

	if path := config.GetString("netns.run_path"); path != "" {
		nsProbe.Watch(path)
	}

	return nsProbe, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"encoding/json"
	"fmt"
	"os"
	"testing"

	"github.com/buildkite/agent/env"
	"github.com/stretchr/testify/assert"
)

func TestPipelineParserParsesYaml(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello \"friend\""}]}`, string(j))
}

func TestPipelineParserSupportsYamlMergesAndAnchors(t *testing.T) {
	complexYAML := `---
base_step: &base_step
  type: script
  agent_query_rules:
    - queue=default

steps:
  - <<: *base_step
    name: ':docker: building image'
    command: docker build .
    agents:
      queue: default`

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte(complexYAML)}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"base_step":{"agent_query_rules":["queue=default"],"type":"script"},"steps":[{"agent_query_rules":["queue=default"],"agents":{"queue":"default"},"command":"docker build .","name":":docker: building image","type":"script"}]}`, string(j))
}

func TestPipelineParserReturnsYamlParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.yml", Pipeline: []byte("steps: %blah%")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.yml: found character that cannot start any token`, fmt.Sprintf("%s", err))
}

func TestPipelineParserReturnsJsonParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.json", Pipeline: []byte("{")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.json: line 1: did not find expected node content`, fmt.Sprintf("%s", err))
}

func TestPipelineParserParsesJson(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "thing.json",
		Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonObjects(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonArrays(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  [ { \"foo\": \"bye ${ENV_VAR_FRIEND}\" } ]\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `[{"foo":"bye \"friend\""}]`, string(j))
}

func TestPipelineParserPreservesBools(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    async: true")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"async":true,"trigger":"hello"}]}`, string(j))
}

func TestPipelineParserPreservesInts(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - label: hello\n    parallelism: 10")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello","parallelism":10}]}`, string(j))
}

func TestPipelineParserPreservesNull(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - wait: ~")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"wait":null}]}`, string(j))
}

func TestPipelineParserPreservesFloats(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 3.142")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"llamas":3.142,"trigger":"hello"}]}`, string(j))
}

func TestPipelineParserHandlesDates(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 2002-08-15T17:18:23.18-06:00")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"llamas":"2002-08-15T17:18:23.18-06:00","trigger":"hello"}]}`, string(j))
}

func TestPipelineParserInterpolatesKeysAsWellAsValues(t *testing.T) {
	var pipeline = `{
		"env": {
			"${FROM_ENV}TEST1": "MyTest",
			"TEST2": "${FROM_ENV}"
		}
	}`

	var decoded struct {
		Env map[string]string `json:"env"`
	}

	environ := env.FromSlice([]string{`FROM_ENV=llamas`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `MyTest`, decoded.Env["llamasTEST1"])
	assert.Equal(t, `llamas`, decoded.Env["TEST2"])
}

func TestPipelineParserLoadsGlobalEnvBlockFirst(t *testing.T) {
	var pipeline = `{
		"env": {
			"TEAM1": "England",
			"TEAM2": "Australia",
			"HEADLINE": "${TEAM1} smashes ${TEAM2} to win the ashes in ${YEAR_FROM_SHELL}!!"
		},
		"steps": [{
			"command": "echo ${HEADLINE}"
		}]
	}`

	var decoded struct {
		Env   map[string]string `json:"env"`
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	environ := env.FromSlice([]string{`YEAR_FROM_SHELL=1912`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `England`, decoded.Env["TEAM1"])
	assert.Equal(t, `England smashes Australia to win the ashes in 1912!!`, decoded.Env["HEADLINE"])
	assert.Equal(t, `echo England smashes Australia to win the ashes in 1912!!`, decoded.Steps[0].Command)
}

func decodeIntoStruct(into interface{}, from interface{}) error {
	b, err := json.Marshal(from)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, into)
}

func TestPipelineParserLoadsSystemEnvironment(t *testing.T) {
	var pipeline = `{
		"steps": [{
			"command": "echo ${LLAMAS_ROCK?}"
		}]
	}`

	var decoded struct {
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	_, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err == nil {
		t.Fatalf("Expected $LLAMAS_ROCK: not set")
	}

	os.Setenv("LLAMAS_ROCK", "absolutely")
	defer os.Unsetenv("LLAMAS_ROCK")

	result2, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result2)
	if err != nil {
		t.Fatal(err)
	}

	if decoded.Steps[0].Command != "echo absolutely" {
		t.Fatalf("Unexpected: %q", decoded.Steps[0].Command)
	}
}
<commit_msg>Add a test for no interpolation<commit_after>package agent

import (
	"encoding/json"
	"fmt"
	"os"
	"testing"

	"github.com/buildkite/agent/env"
	"github.com/stretchr/testify/assert"
)

func TestPipelineParserParsesYaml(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello \"friend\""}]}`, string(j))
}

func TestPipelineParserParsesYamlWithNoInterpolation(t *testing.T) {
	result, err := PipelineParser{
		Filename:        "awesome.yml",
		Pipeline:        []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		NoInterpolation: true,
	}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello ${ENV_VAR_FRIEND}"}]}`, string(j))
}

func TestPipelineParserSupportsYamlMergesAndAnchors(t *testing.T) {
	complexYAML := `---
base_step: &base_step
  type: script
  agent_query_rules:
    - queue=default

steps:
  - <<: *base_step
    name: ':docker: building image'
    command: docker build .
    agents:
      queue: default`

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte(complexYAML)}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"base_step":{"agent_query_rules":["queue=default"],"type":"script"},"steps":[{"agent_query_rules":["queue=default"],"agents":{"queue":"default"},"command":"docker build .","name":":docker: building image","type":"script"}]}`, string(j))
}

func TestPipelineParserReturnsYamlParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.yml", Pipeline: []byte("steps: %blah%")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.yml: found character that cannot start any token`, fmt.Sprintf("%s", err))
}

func TestPipelineParserReturnsJsonParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.json", Pipeline: []byte("{")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.json: line 1: did not find expected node content`, fmt.Sprintf("%s", err))
}

func TestPipelineParserParsesJson(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "thing.json",
		Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonObjects(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonArrays(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  [ { \"foo\": \"bye ${ENV_VAR_FRIEND}\" } ]\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `[{"foo":"bye \"friend\""}]`, string(j))
}

func TestPipelineParserPreservesBools(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    async: true")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"async":true,"trigger":"hello"}]}`, string(j))
}

func TestPipelineParserPreservesInts(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - label: hello\n    parallelism: 10")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello","parallelism":10}]}`, string(j))
}

func TestPipelineParserPreservesNull(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - wait: ~")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"wait":null}]}`, string(j))
}

func TestPipelineParserPreservesFloats(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 3.142")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"llamas":3.142,"trigger":"hello"}]}`, string(j))
}

func TestPipelineParserHandlesDates(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 2002-08-15T17:18:23.18-06:00")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"llamas":"2002-08-15T17:18:23.18-06:00","trigger":"hello"}]}`, string(j))
}

func TestPipelineParserInterpolatesKeysAsWellAsValues(t *testing.T) {
	var pipeline = `{
		"env": {
			"${FROM_ENV}TEST1": "MyTest",
			"TEST2": "${FROM_ENV}"
		}
	}`

	var decoded struct {
		Env map[string]string `json:"env"`
	}

	environ := env.FromSlice([]string{`FROM_ENV=llamas`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `MyTest`, decoded.Env["llamasTEST1"])
	assert.Equal(t, `llamas`, decoded.Env["TEST2"])
}

func TestPipelineParserLoadsGlobalEnvBlockFirst(t *testing.T) {
	var pipeline = `{
		"env": {
			"TEAM1": "England",
			"TEAM2": "Australia",
			"HEADLINE": "${TEAM1} smashes ${TEAM2} to win the ashes in ${YEAR_FROM_SHELL}!!"
		},
		"steps": [{
			"command": "echo ${HEADLINE}"
		}]
	}`

	var decoded struct {
		Env   map[string]string `json:"env"`
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	environ := env.FromSlice([]string{`YEAR_FROM_SHELL=1912`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `England`, decoded.Env["TEAM1"])
	assert.Equal(t, `England smashes Australia to win the ashes in 1912!!`, decoded.Env["HEADLINE"])
	assert.Equal(t, `echo England smashes Australia to win the ashes in 1912!!`, decoded.Steps[0].Command)
}

func decodeIntoStruct(into interface{}, from interface{}) error {
	b, err := json.Marshal(from)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, into)
}

func TestPipelineParserLoadsSystemEnvironment(t *testing.T) {
	var pipeline = `{
		"steps": [{
			"command": "echo ${LLAMAS_ROCK?}"
		}]
	}`

	var decoded struct {
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	_, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err == nil {
		t.Fatalf("Expected $LLAMAS_ROCK: not set")
	}

	os.Setenv("LLAMAS_ROCK", "absolutely")
	defer os.Unsetenv("LLAMAS_ROCK")

	result2, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result2)
	if err != nil {
		t.Fatal(err)
	}

	if decoded.Steps[0].Command != "echo absolutely" {
		t.Fatalf("Unexpected: %q", decoded.Steps[0].Command)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"github.com/miekg/dns"
	"io/ioutil"
	"log"
	"net"
	"path"
	"sort"
	"strconv"
	"strings"
)

type Options struct {
	Serial int
	Ttl    int
}

type Record struct {
	RR     dns.RR
	Weight int
}

type Records []Record

func (s Records) Len() int      { return len(s) }
func (s Records) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type RecordsByWeight struct{ Records }

func (s RecordsByWeight) Less(i, j int) bool { return s.Records[i].Weight > s.Records[j].Weight }

type Label struct {
	Label    string
	MaxHosts int
	Ttl      int
	Records  map[uint16]Records
	Weight   map[uint16]int
}

type labels map[string]*Label

type Zones map[string]*Zone

type Zone struct {
	Origin    string
	Labels    labels
	LenLabels int
	Options   Options
}

var (
	listen  = flag.String("listen", ":8053", "set the listener address")
	flaglog = flag.Bool("log", false, "be more verbose")
	flagrun = flag.Bool("run", false, "run server")
)

func (z *Zone) findLabels(s, cc string, qtype uint16) *Label {

	selectors := []string{}

	if len(cc) > 0 {
		if len(s) > 0 {
			cc = s + "." + cc
		}
		selectors = append(selectors, cc)
	}
	// TODO(ask) Add continent, see https://github.com/abh/geodns/issues/1
	selectors = append(selectors, s)

	for _, name := range selectors {
		if label, ok := z.Labels[name]; ok {
			// return the label if it has the right records
			if label.Records[qtype] != nil {
				return label
			}
		}
	}
	return z.Labels[s]
}

func main() {

	log.SetPrefix("geodns ")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	dirName := "dns"

	dir, err := ioutil.ReadDir(dirName)
	if err != nil {
		panic(err)
	}

	Zones := make(Zones)

	for i, file := range dir {
		fileName := file.Name()
		if !strings.HasSuffix(strings.ToLower(fileName), ".json") {
			continue
		}
		zoneName := fileName[0:strings.LastIndex(fileName, ".")]
		log.Println("FILE:", i, file, zoneName)
		config := readZoneFile(zoneName, path.Join(dirName, fileName))
		Zones[zoneName] = config
	}

	log.Println("ZONES", Zones)

	runServe(&Zones)
}

func readZoneFile(zoneName, fileName string) *Zone {

	b, err := ioutil.ReadFile(fileName)
	if err != nil {
		panic(err)
	}

	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	Zone.Origin = zoneName

	if err == nil {
		var objmap map[string]interface{}
		err := json.Unmarshal(b, &objmap)
		if err != nil {
			panic(err)
		}
		//log.Println(objmap)

		var data map[string]interface{}

		for k, v := range objmap {
			log.Printf("k: %s v: %#v, T: %T\n", k, v, v)

			switch k {
			case "ttl", "serial":
				switch option := k; option {
				case "ttl":
					Zone.Options.Ttl = int(v.(float64))
				case "serial":
					Zone.Options.Serial = int(v.(float64))
				}
				continue

			case "data":
				data = v.(map[string]interface{})
			}
		}

		setupZoneData(data, Zone)

	}

	//log.Printf("ZO T: %T %s\n", Zones["0.us"], Zones["0.us"])

	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))

	return Zone
}

func setupZoneData(data map[string]interface{}, Zone *Zone) {

	var recordTypes = map[string]uint16{
		"a":    dns.TypeA,
		"aaaa": dns.TypeAAAA,
		"ns":   dns.TypeNS,
	}

	for dk, dv := range data {

		log.Printf("K %s V %s TYPE-V %T\n", dk, dv, dv)

		Zone.Labels[dk] = new(Label)
		label := Zone.Labels[dk]

		// BUG(ask) Read 'ttl' value in label data

		for rType, dnsType := range recordTypes {

			var rdata = dv.(map[string]interface{})[rType]

			if rdata == nil {
				//log.Printf("No %s records for label %s\n", rType, dk)
				continue
			}

			log.Printf("rdata %s TYPE-R %T\n", rdata, rdata)

			records := make(map[string][]interface{})

			switch rdata.(type) {
			case map[string]interface{}:
				// Handle map[ns2.example.net:<nil> ns1.example.net:<nil>]
				tmp := make([]interface{}, 0)
				for rdata_k, rdata_v := range rdata.(map[string]interface{}) {
					if rdata_v == nil {
						rdata_v = ""
					}
					tmp = append(tmp, []string{rdata_k, rdata_v.(string)})
				}
				records[rType] = tmp
			default:
				records[rType] = rdata.([]interface{})
			}

			//log.Printf("RECORDS %s TYPE-REC %T\n", Records, Records)

			if label.Records == nil {
				label.Records = make(map[uint16]Records)
				label.Weight = make(map[uint16]int)
			}

			label.Records[dnsType] = make(Records, len(records[rType]))

			for i := 0; i < len(records[rType]); i++ {

				log.Printf("RT %T %#v\n", records[rType][i], records[rType][i])

				record := new(Record)

				var h dns.RR_Header
				// log.Println("TTL OPTIONS", Zone.Options.Ttl)
				h.Ttl = uint32(Zone.Options.Ttl)
				h.Class = dns.ClassINET
				h.Rrtype = dnsType

				switch dnsType {
				case dns.TypeA, dns.TypeAAAA:
					rec := records[rType][i].([]interface{})
					ip := rec[0].(string)
					var err error
					switch rec[1].(type) {
					case string:
						record.Weight, err = strconv.Atoi(rec[1].(string))
						if err != nil {
							panic("Error converting weight to integer")
						}
						label.Weight[dnsType] += record.Weight
					case float64:
						record.Weight = int(rec[1].(float64))
					}
					switch dnsType {
					case dns.TypeA:
						rr := &dns.RR_A{Hdr: h}
						rr.A = net.ParseIP(ip)
						if rr.A == nil {
							panic("Bad A record")
						}
						record.RR = rr
					case dns.TypeAAAA:
						rr := &dns.RR_AAAA{Hdr: h}
						rr.AAAA = net.ParseIP(ip)
						if rr.AAAA == nil {
							panic("Bad AAAA record")
						}
						record.RR = rr
					}
				case dns.TypeNS:
					rec := records[rType][i]
					rr := &dns.RR_NS{Hdr: h}

					switch rec.(type) {
					case string:
						rr.Ns = rec.(string)
					case []string:
						recl := rec.([]string)
						log.Println("RECL:", recl)
						rr.Ns = recl[0]
						if len(recl[1]) > 0 {
							log.Println("NS records with names syntax not supported")
						}
					default:
						log.Printf("Data: %T %#v\n", rec, rec)
						panic("Unrecognized NS format/syntax")
					}

					if h.Ttl < 43000 {
						h.Ttl = 43200
					}
					record.RR = rr

				default:
					log.Println("type:", rType)
					panic("Don't know how to handle this type")
				}

				if record.RR == nil {
					panic("record.RR is nil")
				}

				label.Records[dnsType][i] = *record
			}
			if label.Weight[dnsType] > 0 {
				sort.Sort(RecordsByWeight{label.Records[dnsType]})
			}
		}
	}

	setupSOA(Zone)

	//log.Println(Zones[k])
}

func setupSOA(Zone *Zone) {
	label := Zone.Labels[""]

	primaryNs := "ns"

	if record, ok := label.Records[dns.TypeNS]; ok {
		primaryNs = record[0].RR.(*dns.RR_NS).Ns
	}

	s := Zone.Origin + ". 3600 IN SOA " +
		primaryNs + " support.bitnames.com. " +
		strconv.Itoa(Zone.Options.Serial) +
		" 5400 5400 2419200 " +
		strconv.Itoa(Zone.Options.Ttl)

	log.Println("SOA: ", s)

	rr, err := dns.NewRR(s)

	if err != nil {
		log.Println("SOA Error", err)
		panic("Could not setup SOA")
	}

	record := Record{RR: rr}

	label.Records[dns.TypeSOA] = make([]Record, 1)
	label.Records[dns.TypeSOA][0] = record

}
<commit_msg>Add filename and microseconds to output while developing/testing<commit_after>package main

import (
	"encoding/json"
	"flag"
	"github.com/miekg/dns"
	"io/ioutil"
	"log"
	"net"
	"path"
	"sort"
	"strconv"
	"strings"
)

type Options struct {
	Serial int
	Ttl    int
}

type Record struct {
	RR     dns.RR
	Weight int
}

type Records []Record

func (s Records) Len() int      { return len(s) }
func (s Records) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type RecordsByWeight struct{ Records }

func (s RecordsByWeight) Less(i, j int) bool { return s.Records[i].Weight > s.Records[j].Weight }

type Label struct {
	Label    string
	MaxHosts int
	Ttl      int
	Records  map[uint16]Records
	Weight   map[uint16]int
}

type labels map[string]*Label

type Zones map[string]*Zone

type Zone struct {
	Origin    string
	Labels    labels
	LenLabels int
	Options   Options
}

var (
	listen  = flag.String("listen", ":8053", "set the listener address")
	flaglog = flag.Bool("log", false, "be more verbose")
	flagrun = flag.Bool("run", false, "run server")
)

func (z *Zone) findLabels(s, cc string, qtype uint16) *Label {

	selectors := []string{}

	if len(cc) > 0 {
		if len(s) > 0 {
			cc = s + "." + cc
		}
		selectors = append(selectors, cc)
	}
	// TODO(ask) Add continent, see https://github.com/abh/geodns/issues/1
	selectors = append(selectors, s)

	for _, name := range selectors {
		if label, ok := z.Labels[name]; ok {
			// return the label if it has the right records
			if label.Records[qtype] != nil {
				return label
			}
		}
	}
	return z.Labels[s]
}

func main() {

	log.SetPrefix("geodns ")
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	dirName := "dns"

	dir, err := ioutil.ReadDir(dirName)
	if err != nil {
		panic(err)
	}

	Zones := make(Zones)

	for i, file := range dir {
		fileName := file.Name()
		if !strings.HasSuffix(strings.ToLower(fileName), ".json") {
			continue
		}
		zoneName := fileName[0:strings.LastIndex(fileName, ".")]
		log.Println("FILE:", i, file, zoneName)
		config := readZoneFile(zoneName, path.Join(dirName, fileName))
		Zones[zoneName] = config
	}

	log.Println("ZONES", Zones)

	runServe(&Zones)
}

func readZoneFile(zoneName, fileName string) *Zone {

	b, err := ioutil.ReadFile(fileName)
	if err != nil {
		panic(err)
	}

	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	Zone.Origin = zoneName

	if err == nil {
		var objmap map[string]interface{}
		err := json.Unmarshal(b, &objmap)
		if err != nil {
			panic(err)
		}
		//log.Println(objmap)

		var data map[string]interface{}

		for k, v := range objmap {
			log.Printf("k: %s v: %#v, T: %T\n", k, v, v)

			switch k {
			case "ttl", "serial":
				switch option := k; option {
				case "ttl":
					Zone.Options.Ttl = int(v.(float64))
				case "serial":
					Zone.Options.Serial = int(v.(float64))
				}
				continue

			case "data":
				data = v.(map[string]interface{})
			}
		}

		setupZoneData(data, Zone)

	}

	//log.Printf("ZO T: %T %s\n", Zones["0.us"], Zones["0.us"])

	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))

	return Zone
}

func setupZoneData(data map[string]interface{}, Zone *Zone) {

	var recordTypes = map[string]uint16{
		"a":    dns.TypeA,
		"aaaa": dns.TypeAAAA,
		"ns":   dns.TypeNS,
	}

	for dk, dv := range data {

		log.Printf("K %s V %s TYPE-V %T\n", dk, dv, dv)

		Zone.Labels[dk] = new(Label)
		label := Zone.Labels[dk]

		// BUG(ask) Read 'ttl' value in label data

		for rType, dnsType := range recordTypes {

			var rdata = dv.(map[string]interface{})[rType]

			if rdata == nil {
				//log.Printf("No %s records for label %s\n", rType, dk)
				continue
			}

			log.Printf("rdata %s TYPE-R %T\n", rdata, rdata)

			records := make(map[string][]interface{})

			switch rdata.(type) {
			case map[string]interface{}:
				// Handle map[ns2.example.net:<nil> ns1.example.net:<nil>]
				tmp := make([]interface{}, 0)
				for rdata_k, rdata_v := range rdata.(map[string]interface{}) {
					if rdata_v == nil {
						rdata_v = ""
					}
					tmp = append(tmp, []string{rdata_k, rdata_v.(string)})
				}
				records[rType] = tmp
			default:
				records[rType] = rdata.([]interface{})
			}

			//log.Printf("RECORDS %s TYPE-REC %T\n", Records, Records)

			if label.Records == nil {
				label.Records = make(map[uint16]Records)
				label.Weight = make(map[uint16]int)
			}

			label.Records[dnsType] = make(Records, len(records[rType]))

			for i := 0; i < len(records[rType]); i++ {

				log.Printf("RT %T %#v\n", records[rType][i], records[rType][i])

				record := new(Record)

				var h dns.RR_Header
				// log.Println("TTL OPTIONS", Zone.Options.Ttl)
				h.Ttl = uint32(Zone.Options.Ttl)
				h.Class = dns.ClassINET
				h.Rrtype = dnsType

				switch dnsType {
				case dns.TypeA, dns.TypeAAAA:
					rec := records[rType][i].([]interface{})
					ip := rec[0].(string)
					var err error
					switch rec[1].(type) {
					case string:
						record.Weight, err = strconv.Atoi(rec[1].(string))
						if err != nil {
							panic("Error converting weight to integer")
						}
						label.Weight[dnsType] += record.Weight
					case float64:
						record.Weight = int(rec[1].(float64))
					}
					switch dnsType {
					case dns.TypeA:
						rr := &dns.RR_A{Hdr: h}
						rr.A = net.ParseIP(ip)
						if rr.A == nil {
							panic("Bad A record")
						}
						record.RR = rr
					case dns.TypeAAAA:
						rr := &dns.RR_AAAA{Hdr: h}
						rr.AAAA = net.ParseIP(ip)
						if rr.AAAA == nil {
							panic("Bad AAAA record")
						}
						record.RR = rr
					}
				case dns.TypeNS:
					rec := records[rType][i]
					rr := &dns.RR_NS{Hdr: h}

					switch rec.(type) {
					case string:
						rr.Ns = rec.(string)
					case []string:
						recl := rec.([]string)
						log.Println("RECL:", recl)
						rr.Ns = recl[0]
						if len(recl[1]) > 0 {
							log.Println("NS records with names syntax not supported")
						}
					default:
						log.Printf("Data: %T %#v\n", rec, rec)
						panic("Unrecognized NS format/syntax")
					}

					if h.Ttl < 43000 {
						h.Ttl = 43200
					}
					record.RR = rr

				default:
					log.Println("type:", rType)
					panic("Don't know how to handle this type")
				}

				if record.RR == nil {
					panic("record.RR is nil")
				}

				label.Records[dnsType][i] = *record
			}
			if label.Weight[dnsType] > 0 {
				sort.Sort(RecordsByWeight{label.Records[dnsType]})
			}
		}
	}

	setupSOA(Zone)

	//log.Println(Zones[k])
}

func setupSOA(Zone *Zone) {
	label := Zone.Labels[""]

	primaryNs := "ns"

	if record, ok := label.Records[dns.TypeNS]; ok {
		primaryNs = record[0].RR.(*dns.RR_NS).Ns
	}

	s := Zone.Origin + ". 3600 IN SOA " +
		primaryNs + " support.bitnames.com. " +
		strconv.Itoa(Zone.Options.Serial) +
		" 5400 5400 2419200 " +
		strconv.Itoa(Zone.Options.Ttl)

	log.Println("SOA: ", s)

	rr, err := dns.NewRR(s)

	if err != nil {
		log.Println("SOA Error", err)
		panic("Could not setup SOA")
	}

	record := Record{RR: rr}

	label.Records[dns.TypeSOA] = make([]Record, 1)
	label.Records[dns.TypeSOA][0] = record

}
<|endoftext|>


<commit_before>package s3gof3r

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"sync"
	"syscall"
)

type getter struct {
	url    url.URL
	client *http.Client
	b      *Bucket
	bufsz  int64
	err    error
	wg     sync.WaitGroup

	cur_chunk_id   int
	cur_chunk      *chunk
	content_length int64
	chunk_total    int
	read_ch        chan *chunk
	get_ch         chan *chunk

	//	get   chan *bytes.Buffer
	//	give  chan *bytes.Buffer
	bp *bp

	q_wait map[int]*chunk

	concurrency int
	nTry        int
	closed      bool
	c           *Config

	md5 hash.Hash
}

type chunk struct {
	id     int
	header http.Header
	start  int64
	size   int64
	b      *bytes.Buffer
	len    int64
}

func newGetter(p_url url.URL, c *Config, b *Bucket) (io.ReadCloser, http.Header, error) {
	// initialize getter
	g := new(getter)
	g.url = p_url
	g.bufsz = c.PartSize
	g.get_ch = make(chan *chunk)
	g.read_ch = make(chan *chunk)
	g.nTry = c.NTry
	g.q_wait = make(map[int]*chunk)
	g.b = b
	g.c = c
	g.client = c.Client
	g.md5 = md5.New()

	// get content length
	resp, err := g.retryRequest("HEAD", p_url.String(), nil)
	if err != nil {
		return nil, nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, nil, newRespError(resp)
	}
	g.content_length = resp.ContentLength
	g.chunk_total = int((g.content_length + g.bufsz - 1) / g.bufsz) // round up, integer division
	g.concurrency = min(c.Concurrency, g.chunk_total)
	log.Println("chunk total: ", g.chunk_total)
	log.Println("content length : ", g.content_length)
	log.Println("concurrency: ", g.concurrency)

	g.bp = NewBufferPool(g.bufsz)

	for i := 0; i < g.concurrency; i++ {
		go g.worker()
	}
	go g.init_chunks()
	return g, resp.Header, nil
}

func (g *getter) retryRequest(method, urlStr string, body io.ReadSeeker) (resp *http.Response, err error) {
	for i := 0; i < g.nTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		g.b.Sign(req)
		resp, err = g.client.Do(req)
		if err == nil {
			return
		}
		log.Println(err)
		if body != nil {
			body.Seek(0, 0)
		}
	}
	return
}

func (g *getter) init_chunks() {
	id := 0
	for i := int64(0); i < g.content_length; {
		size := min64(g.bufsz, g.content_length-i)
		c := &chunk{
			id: id,
			header: http.Header{
				"Range": {fmt.Sprintf("bytes=%d-%d",
					i, i+size-1)},
			},
			start: i,
			size:  size,
			b:     nil,
			len:   0}
		i += size
		id++
		g.wg.Add(1)
		g.get_ch <- c
	}
	close(g.get_ch)
}

func (g *getter) worker() {
	for c := range g.get_ch {
		g.retryGetChunk(c)
	}

}

func (g *getter) retryGetChunk(c *chunk) {
	defer g.wg.Done()
	var err error
	c.b = <-g.bp.get
	for i := 0; i < g.nTry; i++ {
		err = g.getChunk(c)
		if err == nil {
			return
		}
		log.Printf("Error on attempt %d: retrying chunk: %v, Error: %s", i, c, err)
	}
	g.err = err
}

func (g *getter) getChunk(c *chunk) error {
	// ensure buffer is empty
	c.b.Reset()

	r, err := http.NewRequest("GET", g.url.String(), nil)
	if err != nil {
		return err
	}
	//r.Close = true
	r.Header = c.header
	g.b.Sign(r)
	resp, err := g.client.Do(r)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 206 {
		return newRespError(resp)
	}
	n, err := c.b.ReadFrom(resp.Body)
	if err != nil {
		return err
	}
	if n != c.size {
		return fmt.Errorf("Chunk %d: Expected %d bytes, received %d",
			c.id, c.size, n)
	}
	g.read_ch <- c
	return nil
}

func (g *getter) Read(p []byte) (int, error) {
	var err error
	if g.closed {
		return 0, syscall.EINVAL
	}
	if g.err != nil {
		return 0, g.err
	}
	if g.cur_chunk == nil {
		g.cur_chunk = g.find_next_chunk()
	}
	// write to md5 hash in parallel with output
	tr := io.TeeReader(g.cur_chunk.b, g.md5)
	//n, err := g.cur_chunk.b.Read(p)
	n, err := tr.Read(p)

	// Empty buffer, move on to next
	if err == io.EOF {
		// Do not send EOF for each chunk.
		if g.cur_chunk.id == g.chunk_total-1 && g.cur_chunk.b.Len() == 0 {
			return n, err // end of stream, send eof
		}
		g.bp.give <- g.cur_chunk.b // recycle buffer
		g.cur_chunk = nil
		g.cur_chunk_id++
		return n - 1, nil // subtract EOF
	}
	return n, err
}

func (g *getter) find_next_chunk() (cur_chunk *chunk) {
	for {
		// first check q_wait
		if cur_chunk, ok := g.q_wait[g.cur_chunk_id]; ok {
			delete(g.q_wait, g.cur_chunk_id)
			return cur_chunk
		}
		// if next chunk not in q_wait, read from channel
		c := <-g.read_ch
		g.q_wait[c.id] = c
	}
}

func (g *getter) Close() error {
	if g.closed {
		return syscall.EINVAL
	}
	if g.err != nil {
		return g.err
	}
	g.wg.Wait()
	close(g.read_ch)
	g.closed = true
	log.Println("makes:", g.bp.makes)
	if g.c.Md5Check {
		if err := g.checkMd5(); err != nil {
			return err
		}
	}
	return nil
}

func (g *getter) checkMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", g.md5.Sum(nil))
	md5Path := fmt.Sprint(".md5", g.url.Path, ".md5")
	md5Url := g.b.Url(md5Path, g.c)
	log.Println("md5: ", calcMd5)
	log.Println("md5Path: ", md5Path)
	resp, err := g.retryRequest("GET", md5Url.String(), nil)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("MD5 check failed: %s not found: %s", md5Url.String(), newRespError(resp))
	}
	givenMd5, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	if calcMd5 != string(givenMd5) {
		return fmt.Errorf("MD5 mismatch. given:%s calculated:%s", givenMd5, calcMd5)
	}
	return
}
<commit_msg>Don't blow up memory on download.<commit_after>package s3gof3r

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"sync"
	"syscall"
	"time"
)

type getter struct {
	url    url.URL
	client *http.Client
	b      *Bucket
	bufsz  int64
	err    error
	wg     sync.WaitGroup

	cur_chunk_id   int
	cur_chunk      *chunk
	content_length int64
	chunk_total    int
	read_ch        chan *chunk
	get_ch         chan *chunk

	//	get   chan *bytes.Buffer
	//	give  chan *bytes.Buffer
	bp *bp

	q_wait map[int]*chunk

	concurrency int
	nTry        int
	closed      bool
	c           *Config

	md5 hash.Hash
}

type chunk struct {
	id     int
	header http.Header
	start  int64
	size   int64
	b      *bytes.Buffer
	len    int64
}

func newGetter(p_url url.URL, c *Config, b *Bucket) (io.ReadCloser, http.Header, error) {
	// initialize getter
	g := new(getter)
	g.url = p_url
	g.bufsz = c.PartSize
	g.get_ch = make(chan *chunk)
	g.read_ch = make(chan *chunk)
	g.nTry = c.NTry
	g.q_wait = make(map[int]*chunk)
	g.b = b
	g.c = c
	g.client = c.Client
	g.md5 = md5.New()

	// get content length
	resp, err := g.retryRequest("HEAD", p_url.String(), nil)
	if err != nil {
		return nil, nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, nil, newRespError(resp)
	}
	g.content_length = resp.ContentLength
	g.chunk_total = int((g.content_length + g.bufsz - 1) / g.bufsz) // round up, integer division
	g.concurrency = min(c.Concurrency, g.chunk_total)
	log.Println("chunk total: ", g.chunk_total)
	log.Println("content length : ", g.content_length)
	log.Println("concurrency: ", g.concurrency)

	g.bp = NewBufferPool(g.bufsz)

	for i := 0; i < g.concurrency; i++ {
		go g.worker()
	}
	go g.init_chunks()
	return g, resp.Header, nil
}

func (g *getter) retryRequest(method, urlStr string, body io.ReadSeeker) (resp *http.Response, err error) {
	for i := 0; i < g.nTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		g.b.Sign(req)
		resp, err = g.client.Do(req)
		if err == nil {
			return
		}
		log.Println(err)
		if body != nil {
			body.Seek(0, 0)
		}
	}
	return
}

func (g *getter) init_chunks() {
	id := 0
	for i := int64(0); i < g.content_length; {
		for len(g.q_wait) > g.concurrency {
			time.Sleep(500 * time.Millisecond)
		}
		size := min64(g.bufsz, g.content_length-i)
		c := &chunk{
			id: id,
			header: http.Header{
				"Range": {fmt.Sprintf("bytes=%d-%d",
					i, i+size-1)},
			},
			start: i,
			size:  size,
			b:     nil,
			len:   0}
		i += size
		id++
		g.wg.Add(1)
		g.get_ch <- c
	}
	close(g.get_ch)
}

func (g *getter) worker() {
	for c := range g.get_ch {
		g.retryGetChunk(c)
	}

}

func (g *getter) retryGetChunk(c *chunk) {
	defer g.wg.Done()
	var err error
	c.b = <-g.bp.get
	for i := 0; i < g.nTry; i++ {
		err = g.getChunk(c)
		if err == nil {
			return
		}
		log.Printf("Error on attempt %d: retrying chunk: %v, Error: %s", i, c, err)
	}
	g.err = err
}

func (g *getter) getChunk(c *chunk) error {
	// ensure buffer is empty
	c.b.Reset()

	r, err := http.NewRequest("GET", g.url.String(), nil)
	if err != nil {
		return err
	}
	//r.Close = true
	r.Header = c.header
	g.b.Sign(r)
	resp, err := g.client.Do(r)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 206 {
		return newRespError(resp)
	}
	n, err := c.b.ReadFrom(resp.Body)
	if err != nil {
		return err
	}
	if n != c.size {
		return fmt.Errorf("Chunk %d: Expected %d bytes, received %d",
			c.id, c.size, n)
	}
	g.read_ch <- c
	return nil
}

func (g *getter) Read(p []byte) (int, error) {
	var err error
	if g.closed {
		return 0, syscall.EINVAL
	}
	if g.err != nil {
		return 0, g.err
	}
	if g.cur_chunk == nil {
		g.cur_chunk = g.find_next_chunk()
	}
	// write to md5 hash in parallel with output
	tr := io.TeeReader(g.cur_chunk.b, g.md5)
	//n, err := g.cur_chunk.b.Read(p)
	n, err := tr.Read(p)

	// Empty buffer, move on to next
	if err == io.EOF {
		// Do not send EOF for each chunk.
		if g.cur_chunk.id == g.chunk_total-1 && g.cur_chunk.b.Len() == 0 {
			return n, err // end of stream, send eof
		}
		g.bp.give <- g.cur_chunk.b // recycle buffer
		g.cur_chunk = nil
		g.cur_chunk_id++
		return n - 1, nil // subtract EOF
	}
	return n, err
}

func (g *getter) find_next_chunk() (cur_chunk *chunk) {
	for {
		// first check q_wait
		if cur_chunk, ok := g.q_wait[g.cur_chunk_id]; ok {
			delete(g.q_wait, g.cur_chunk_id)
			return cur_chunk
		}
		// if next chunk not in q_wait, read from channel
		c := <-g.read_ch
		g.q_wait[c.id] = c
	}
}

func (g *getter) Close() error {
	if g.closed {
		return syscall.EINVAL
	}
	if g.err != nil {
		return g.err
	}
	g.wg.Wait()
	close(g.read_ch)
	g.closed = true
	log.Println("makes:", g.bp.makes)
	if g.c.Md5Check {
		if err := g.checkMd5(); err != nil {
			return err
		}
	}
	return nil
}

func (g *getter) checkMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", g.md5.Sum(nil))
	md5Path := fmt.Sprint(".md5", g.url.Path, ".md5")
	md5Url := g.b.Url(md5Path, g.c)
	log.Println("md5: ", calcMd5)
	log.Println("md5Path: ", md5Path)
	resp, err := g.retryRequest("GET", md5Url.String(), nil)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("MD5 check failed: %s not found: %s", md5Url.String(), newRespError(resp))
	}
	givenMd5, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	if calcMd5 != string(givenMd5) {
		return fmt.Errorf("MD5 mismatch. given:%s calculated:%s", givenMd5, calcMd5)
	}
	return
}
<|endoftext|>


<commit_before><commit_msg>Add GetPanic() function<commit_after><|endoftext|>


<commit_before>package misc

import (
    //"unsafe"
    "encoding/binary"
    "jvmgo/jvm/rtda"
)

func init() {
    _unsafe(allocateMemory, "allocateMemory",   "(J)J")
    _unsafe(putLong,        "putLong",          "(JJ)V")
    _unsafe(getByte,        "getByte",          "(J)B")
}

// public native long allocateMemory(long l);
// (J)J
func allocateMemory(frame *rtda.Frame) {
    stack := frame.OperandStack()
    size := stack.PopLong()
    stack.PopRef() // this

    address := allocate(size)
    stack.PushLong(address)
}

// public native void putLong(long l, long l1);
// (JJ)V
func putLong(frame *rtda.Frame) {
    stack := frame.OperandStack()
    value := stack.PopLong()
    address := stack.PopLong()
    stack.PopRef() // this

    mem := memoryAt(address)
    binary.BigEndian.PutUint64(mem, uint64(value))
}

// public native byte getByte(long l);
// (J)B
func getByte(frame *rtda.Frame) {
    stack := frame.OperandStack()
    address := stack.PopLong()
    stack.PopRef() // this

    mem := memoryAt(address)
    b := mem[address]
    stack.PushInt(int32(b))
}
<commit_msg>native: Unsafe.freeMemory()<commit_after>package misc

import (
    //"unsafe"
    "encoding/binary"
    "jvmgo/jvm/rtda"
)

func init() {
    _unsafe(allocateMemory, "allocateMemory",   "(J)J")
    _unsafe(freeMemory,     "freeMemory",       "(J)V")
    _unsafe(putLong,        "putLong",          "(JJ)V")
    _unsafe(getByte,        "getByte",          "(J)B")
}

// public native long allocateMemory(long l);
// (J)J
func allocateMemory(frame *rtda.Frame) {
    stack := frame.OperandStack()
    size := stack.PopLong()
    stack.PopRef() // this

    address := allocate(size)
    stack.PushLong(address)
}

// public native void freeMemory(long l);
// (J)V
func freeMemory(frame *rtda.Frame) {
    stack := frame.OperandStack()
    address := stack.PopLong()
    stack.PopRef() // this
    free(address)
}

// public native void putLong(long l, long l1);
// (JJ)V
func putLong(frame *rtda.Frame) {
    stack := frame.OperandStack()
    value := stack.PopLong()
    address := stack.PopLong()
    stack.PopRef() // this

    mem := memoryAt(address)
    binary.BigEndian.PutUint64(mem, uint64(value))
}

// public native byte getByte(long l);
// (J)B
func getByte(frame *rtda.Frame) {
    stack := frame.OperandStack()
    address := stack.PopLong()
    stack.PopRef() // this

    mem := memoryAt(address)
    b := mem[address]
    stack.PushInt(int32(b))
}
<|endoftext|>


<commit_before>package portping

import (
	"testing"
	"fmt"
)

// TODO use localhost
const host = "192.168.1.10"
const localhost = "localhost"

func ping(host string, port int) bool {
	err := Ping(host, port)
	fmt.Printf("%s:%d %v", host, port, err)
	return err == nil
}

func Test_ping_open_port(t*testing.T) {
	// TODO listen on a port to ensure it's open
	open_port := 22

	if !ping(host, open_port) {
		t.Errorf("should be open")
	}
}

func Test_ping_unopen_port(t*testing.T) {
	// TODO choose some random port guaranteed closed
	unopen_port := 123

	if ping(host, unopen_port) {
		t.Errorf("should be closed")
	}
}

func Test_ping_nonexistent_host(t*testing.T) {
	if ping("nonexistent.janosgyerik.com", 80) {
		t.Errorf("should be closed")
	}
}

func Test_ping_negative_port(t*testing.T) {
	if ping(localhost, -1) {
		t.Errorf("should be closed")
	}
}

func Test_ping_too_high_port(t*testing.T) {
	if ping(localhost, 123456) {
		t.Errorf("should be closed")
	}
}
<commit_msg>open local port for testing, extract vars and constants<commit_after>package portping

import (
	"testing"
	"fmt"
	"net"
	"log"
)

const testHost = "localhost"

// TODO hopefully unused. Better ideas?
const testPort = 1234

const knownNonexistentHost = "nonexistent.janosgyerik.com"

func acceptN(host string, port int, count int) {
	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		log.Fatal(err)
	}
	defer ln.Close()
	for i := 0; i < count; i++ {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatal(err)
		}
		conn.Close()
	}
}

func ping(host string, port int) bool {
	err := Ping(host, port)
	fmt.Printf("%s:%d %v", host, port, err)
	return err == nil
}

func Test_ping_open_port(t*testing.T) {
	go acceptN(testHost, testPort, 1)

	if !ping(testHost, testPort) {
		t.Errorf("should be open")
	}

	// for sanity: acceptN should have shut down already
	if ping(testHost, testPort) {
		t.Errorf("should be closed")
	}
}

func Test_ping_unopen_port(t*testing.T) {
	if ping(testHost, testPort) {
		t.Errorf("should be closed")
	}
}

func Test_ping_nonexistent_host(t*testing.T) {
	if ping(knownNonexistentHost, testPort) {
		t.Errorf("should be closed")
	}
}

func Test_ping_negative_port(t*testing.T) {
	if ping(testHost, -1) {
		t.Errorf("should be closed")
	}
}

func Test_ping_too_high_port(t*testing.T) {
	if ping(testHost, 123456) {
		t.Errorf("should be closed")
	}
}
<|endoftext|>


<commit_before>package postgres

import (
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/aodin/sol"
	"github.com/aodin/sol/types"
)

const travisCI = "host=localhost port=5432 dbname=sol_test user=postgres sslmode=disable"

var testconn *sol.DB
var once sync.Once

// getConn returns a postgres connection pool
func getConn(t *testing.T) *sol.DB {
	// Check if an ENV VAR has been set, otherwise, use travis
	credentials := os.Getenv("SOL_TEST_POSTGRES")
	if credentials == "" {
		credentials = travisCI
	}

	once.Do(func() {
		var err error
		if testconn, err = sol.Open("postgres", credentials); err != nil {
			t.Fatalf("Failed to open connection: %s", err)
		}
		testconn.SetMaxOpenConns(20)
	})
	return testconn
}

var things = sol.Table("things",
	sol.Column("name", types.Varchar()),
	sol.Column("created_at", Timestamp().Default(Now)),
)

type thing struct {
	Name      string
	CreatedAt time.Time `db:",omitempty"`
}

var itemsA = Table("items_a",
	sol.Column("id", Serial()),
	sol.Column("name", types.Varchar()),
)

var itemsB = Table("items_b",
	sol.Column("id", Serial()),
	sol.Column("name", types.Varchar()),
	sol.PrimaryKey("id"),
)

var itemsFK = Table("items_fk",
	sol.ForeignKey("id", itemsB, types.Integer().NotNull()),
	sol.Column("name", types.Varchar()),
)

type item struct {
	ID   uint64 `db:",omitempty"`
	Name string
}

func (i item) Exists() bool {
	return i.ID != 0
}

var meetings = Table("meetings",
	sol.Column("uuid", UUID().NotNull().Unique().Default(GenerateV4)),
	sol.Column("time", TimestampRange()),
)

// TestPostGres performs the standard integration test
func TestPostGres(t *testing.T) {
	conn := getConn(t) // TODO close
	sol.IntegrationTest(t, conn, false)
}

func TestPostGres_Create(t *testing.T) {
	expect := sol.NewTester(t, &PostGres{})

	expect.SQL(
		itemsFK.Create(),
		`CREATE TABLE items_fk (
  id INTEGER NOT NULL REFERENCES items_b(id),
  name VARCHAR
);`,
	)
}

// TestPostGres_Select tests a variety of SelectStmt features against the
// postgres database
func TestPostGres_Select(t *testing.T) {
	conn := getConn(t) // TODO close

	tx, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")
	defer tx.Rollback()

	// TODO temp tables
	require.Nil(t,
		tx.Query(itemsA.Create().Temporary().IfNotExists()),
		`Create table "%s" should not error`, itemsA.Name(),
	)

	require.Nil(t,
		tx.Query(itemsB.Create().Temporary().IfNotExists()),
		`Create table "%s" should not error`, itemsB.Name(),
	)

	a := item{Name: "A"}
	require.Nil(t,
		tx.Query(itemsA.Insert().Values(a).Returning(), &a),
		`Insert into table "%s" within a transaction should not error`,
		itemsA.Name(),
	)
	if a.ID == 0 {
		t.Fatal("Failed to set primary key of item during INSERT")
	}

	require.Nil(t,
		tx.Query(itemsB.Insert().Values([]item{{Name: "A"}, {Name: "B"}})),
		`Insert into table "%s" within a transaction should not error`,
		itemsB.Name(),
	)

	stmt := itemsB.Select().InnerJoin(
		itemsA,
		itemsB.C("name").Equals(itemsA.C("name")),
	).Where(
		itemsA.C("id").Equals(a.ID),
	)

	var selected item
	tx.Query(stmt, &selected)
	if selected.ID == 0 {
		t.Fatal("Failed to SELECT item through joined table")
	}
}

// TestPostGres_Transaction tests the transactional operations of PostGres,
// including Commit, Rollback, and Close
func TestPostGres_Transaction(t *testing.T) {
	conn := getConn(t) // TODO close

	require.Nil(t, conn.Query(things.Drop().IfExists()))
	require.Nil(t,
		conn.Query(things.Create()),
		`Create table "things" should not error`,
	)

	// Start a transaction with the intent to commit
	tx1, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")

	require.Nil(t,
		tx1.Query(things.Insert().Values(thing{Name: "A"})),
		`Insert into table "things" within a transaction should not error`,
	)

	require.Nil(t, tx1.Commit(), "Committing a transaction should not error")

	var first []thing
	require.Nil(t, conn.Query(things.Select(), &first))
	assert.Equal(t, 1, len(first), "Thing A should have been committed")

	// Start a transaction with the intent to rollback
	tx2, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")

	require.Nil(t, tx2.Query(things.Insert().Values(thing{Name: "B"})))
	require.Nil(t, tx2.Rollback(), "Transaction Rollback should not error")

	var second []thing
	require.Nil(t, conn.Query(things.Select(), &second))
	assert.Equal(t, 1, len(second), "Thing B should not have been committed")

	// Create functions that fail, panic, and are successful
	func() {
		tx3, err := conn.Begin()
		require.Nil(t, err)
		defer tx3.Close()

		require.Nil(t, tx3.Query(things.Insert().Values(thing{Name: "C"})))
		// Without a call to IsSuccessful, the transaction should rollback
	}()

	var third []thing
	require.Nil(t, conn.Query(things.Select(), &third))
	assert.Equal(t, 1, len(third), "Thing C should not have been committed")

	func() {
		tx4, err := conn.Begin()
		require.Nil(t, err)
		defer tx4.Close()

		require.Nil(t, tx4.Query(things.Insert().Values(thing{Name: "D"})))
		tx4.IsSuccessful()
	}()

	var fourth []thing
	require.Nil(t, conn.Query(things.Select(), &fourth))
	assert.Equal(t, 2, len(fourth), "Thing D should have been committed")

	assert.Panics(t,
		func() {
			tx5, err := conn.Begin()
			require.Nil(t, err)
			defer tx5.Close()

			require.Nil(t, tx5.Query(things.Insert().Values(thing{Name: "E"})))
			panic("I'm panicking")
		},
	)

	var fifth []thing
	require.Nil(t, conn.Query(things.Select(), &fifth))
	assert.Equal(t, 2, len(fifth), "Thing E should have been committed")
}
<commit_msg>Test pq's NullTime implementation<commit_after>package postgres

import (
	"os"
	"sync"
	"testing"
	"time"

	"github.com/lib/pq"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/aodin/sol"
	"github.com/aodin/sol/types"
)

const travisCI = "host=localhost port=5432 dbname=sol_test user=postgres sslmode=disable"

var testconn *sol.DB
var once sync.Once

// getConn returns a postgres connection pool
func getConn(t *testing.T) *sol.DB {
	// Check if an ENV VAR has been set, otherwise, use travis
	credentials := os.Getenv("SOL_TEST_POSTGRES")
	if credentials == "" {
		credentials = travisCI
	}

	once.Do(func() {
		var err error
		if testconn, err = sol.Open("postgres", credentials); err != nil {
			t.Fatalf("Failed to open connection: %s", err)
		}
		testconn.SetMaxOpenConns(20)
	})
	return testconn
}

var things = sol.Table("things",
	sol.Column("name", types.Varchar()),
	sol.Column("created_at", Timestamp().Default(Now)),
)

var thingsWithoutNow = sol.Table("things",
	sol.Column("name", types.Varchar()),
	sol.Column("created_at", Timestamp()),
)

type thing struct {
	Name      string
	CreatedAt time.Time `db:",omitempty"`
}

var itemsA = Table("items_a",
	sol.Column("id", Serial()),
	sol.Column("name", types.Varchar()),
)

var itemsB = Table("items_b",
	sol.Column("id", Serial()),
	sol.Column("name", types.Varchar()),
	sol.PrimaryKey("id"),
)

var itemsFK = Table("items_fk",
	sol.ForeignKey("id", itemsB, types.Integer().NotNull()),
	sol.Column("name", types.Varchar()),
)

type item struct {
	ID   uint64 `db:",omitempty"`
	Name string
}

func (i item) Exists() bool {
	return i.ID != 0
}

var meetings = Table("meetings",
	sol.Column("uuid", UUID().NotNull().Unique().Default(GenerateV4)),
	sol.Column("time", TimestampRange()),
)

// TestPostGres performs the standard integration test
func TestPostGres(t *testing.T) {
	conn := getConn(t) // TODO close
	sol.IntegrationTest(t, conn, false)
}

func TestPostGres_NullTime(t *testing.T) {
	conn := getConn(t) // TODO close

	tx, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")
	defer tx.Rollback()

	// TODO temp tables
	require.Nil(t,
		tx.Query(thingsWithoutNow.Create().Temporary().IfNotExists()),
		`Create table "%s" should not error`, thingsWithoutNow.Name(),
	)

	type nullthing struct {
		Name      string
		CreatedAt pq.NullTime `db:"created_at"`
	}

	nonzero := nullthing{
		Name:      "a",
		CreatedAt: pq.NullTime{Valid: true, Time: time.Now()},
	}
	tx.Query(thingsWithoutNow.Insert().Values(nonzero))

	var things []nullthing
	tx.Query(thingsWithoutNow.Select(), &things)
	require.Equal(t, 1, len(things))
	assert.True(t, things[0].CreatedAt.Valid)
	assert.False(t, things[0].CreatedAt.Time.IsZero())
}

func TestPostGres_Create(t *testing.T) {
	expect := sol.NewTester(t, &PostGres{})

	expect.SQL(
		itemsFK.Create(),
		`CREATE TABLE items_fk (
  id INTEGER NOT NULL REFERENCES items_b(id),
  name VARCHAR
);`,
	)
}

// TestPostGres_Select tests a variety of SelectStmt features against the
// postgres database
func TestPostGres_Select(t *testing.T) {
	conn := getConn(t) // TODO close

	tx, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")
	defer tx.Rollback()

	// TODO temp tables
	require.Nil(t,
		tx.Query(itemsA.Create().Temporary().IfNotExists()),
		`Create table "%s" should not error`, itemsA.Name(),
	)

	require.Nil(t,
		tx.Query(itemsB.Create().Temporary().IfNotExists()),
		`Create table "%s" should not error`, itemsB.Name(),
	)

	a := item{Name: "A"}
	require.Nil(t,
		tx.Query(itemsA.Insert().Values(a).Returning(), &a),
		`Insert into table "%s" within a transaction should not error`,
		itemsA.Name(),
	)
	if a.ID == 0 {
		t.Fatal("Failed to set primary key of item during INSERT")
	}

	require.Nil(t,
		tx.Query(itemsB.Insert().Values([]item{{Name: "A"}, {Name: "B"}})),
		`Insert into table "%s" within a transaction should not error`,
		itemsB.Name(),
	)

	stmt := itemsB.Select().InnerJoin(
		itemsA,
		itemsB.C("name").Equals(itemsA.C("name")),
	).Where(
		itemsA.C("id").Equals(a.ID),
	)

	var selected item
	tx.Query(stmt, &selected)
	if selected.ID == 0 {
		t.Fatal("Failed to SELECT item through joined table")
	}
}

// TestPostGres_Transaction tests the transactional operations of PostGres,
// including Commit, Rollback, and Close
func TestPostGres_Transaction(t *testing.T) {
	conn := getConn(t) // TODO close

	require.Nil(t, conn.Query(things.Drop().IfExists()))
	require.Nil(t,
		conn.Query(things.Create()),
		`Create table "things" should not error`,
	)

	// Start a transaction with the intent to commit
	tx1, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")

	require.Nil(t,
		tx1.Query(things.Insert().Values(thing{Name: "A"})),
		`Insert into table "things" within a transaction should not error`,
	)

	require.Nil(t, tx1.Commit(), "Committing a transaction should not error")

	var first []thing
	require.Nil(t, conn.Query(things.Select(), &first))
	assert.Equal(t, 1, len(first), "Thing A should have been committed")

	// Start a transaction with the intent to rollback
	tx2, err := conn.Begin()
	require.Nil(t, err, "Creating a new transaction should not error")

	require.Nil(t, tx2.Query(things.Insert().Values(thing{Name: "B"})))
	require.Nil(t, tx2.Rollback(), "Transaction Rollback should not error")

	var second []thing
	require.Nil(t, conn.Query(things.Select(), &second))
	assert.Equal(t, 1, len(second), "Thing B should not have been committed")

	// Create functions that fail, panic, and are successful
	func() {
		tx3, err := conn.Begin()
		require.Nil(t, err)
		defer tx3.Close()

		require.Nil(t, tx3.Query(things.Insert().Values(thing{Name: "C"})))
		// Without a call to IsSuccessful, the transaction should rollback
	}()

	var third []thing
	require.Nil(t, conn.Query(things.Select(), &third))
	assert.Equal(t, 1, len(third), "Thing C should not have been committed")

	func() {
		tx4, err := conn.Begin()
		require.Nil(t, err)
		defer tx4.Close()

		require.Nil(t, tx4.Query(things.Insert().Values(thing{Name: "D"})))
		tx4.IsSuccessful()
	}()

	var fourth []thing
	require.Nil(t, conn.Query(things.Select(), &fourth))
	assert.Equal(t, 2, len(fourth), "Thing D should have been committed")

	assert.Panics(t,
		func() {
			tx5, err := conn.Begin()
			require.Nil(t, err)
			defer tx5.Close()

			require.Nil(t, tx5.Query(things.Insert().Values(thing{Name: "E"})))
			panic("I'm panicking")
		},
	)

	var fifth []thing
	require.Nil(t, conn.Query(things.Select(), &fifth))
	assert.Equal(t, 2, len(fifth), "Thing E should have been committed")
}
<|endoftext|>