<commit_before>package main

import (
	"flag"
	"fmt"
	_build "neon/build"
	_ "neon/builtin"
	_ "neon/task"
	"neon/util"
	"os"
	"path/filepath"
	"time"
)

const (
	DEFAULT_BUILD_FILE = "build.yml"
)

var VERSION string

// Parse command line and return parsed options
func ParseCommandLine() (string, bool, bool, string, bool, bool, string, bool, bool, string, bool, string, bool, []string) {
	file := flag.String("file", DEFAULT_BUILD_FILE, "Build file to run")
	info := flag.Bool("info", false, "Print build information")
	version := flag.Bool("version", false, "Print neon version")
	props := flag.String("props", "", "Build properties")
	timeit := flag.Bool("time", false, "Print build duration")
	tasks := flag.Bool("tasks", false, "Print tasks list")
	task := flag.String("task", "", "Print help on given task")
	targs := flag.Bool("targets", false, "Print targets list")
	builtins := flag.Bool("builtins", false, "Print builtins list")
	builtin := flag.String("builtin", "", "Print help on given builtin")
	refs := flag.Bool("refs", false, "Print tasks and builtins reference")
	install := flag.String("install", "", "Install given plugin")
	grey := flag.Bool("grey", false, "Print on terminal without colors")
	flag.Parse()
	targets := flag.Args()
	return *file, *info, *version, *props, *timeit, *tasks, *task, *targs, *builtins,
		*builtin, *refs, *install, *grey, targets
}

// Find build file and return its path
func FindBuildFile(name string) (string, error) {
	absolute, err := filepath.Abs(name)
	if err != nil {
		return "", fmt.Errorf("getting build file path: %v", err)
	}
	file := filepath.Base(absolute)
	dir := filepath.Dir(absolute)
	for {
		path := filepath.Join(dir, file)
		if util.FileExists(path) {
			return path, nil
		} else {
			parent := filepath.Dir(dir)
			if parent == dir {
				return "", fmt.Errorf("build file not found")
			}
			dir = parent
		}
	}
}

// Program entry point
func main() {
	start := time.Now()
	file, info, version, props, timeit, tasks, task, targs, builtins, builtin, refs, install, grey, targets := ParseCommandLine()
	// options that do not require we load build file
	_build.Grey = grey
	if tasks {
		_build.PrintTasks()
		os.Exit(0)
	} else if task != "" {
		_build.PrintHelpTask(task)
		os.Exit(0)
	} else if builtins {
		_build.PrintBuiltins()
		os.Exit(0)
	} else if builtin != "" {
		_build.PrintHelpBuiltin(builtin)
		os.Exit(0)
	} else if refs {
		_build.PrintReference()
		os.Exit(0)
	} else if version {
		_build.Message(VERSION)
		os.Exit(0)
	}
	// options that do require we load build file
	path, err := FindBuildFile(file)
	PrintError(err, 1)
	build, err := _build.NewBuild(path)
	if build != nil && install != "" {
		err = build.Install(install)
		if err == nil {
			os.Exit(0)
		} else {
			_build.Message("ERROR " + err.Error())
			os.Exit(6)
		}
	}
	PrintError(err, 2)
	if props != "" {
		err = build.SetCommandLineProperties(props)
		PrintError(err, 3)
	}
	if targs {
		build.PrintTargets()
	} else if info {
		context, err := _build.NewContext(build)
		PrintError(err, 4)
		err = build.Info(context)
		PrintError(err, 4)
	} else {
		context, err := _build.NewContext(build)
		err = build.Run(context, targets)
		PrintError(err, 5)
		duration := time.Now().Sub(start)
		if timeit || duration.Seconds() > 10 {
			_build.Message("Build duration: %s", duration.String())
		}
		if err == nil {
			_build.PrintOk()
		} else {
			PrintError(err, 5)
		}
	}
}

// Print an error and exit if any
func PrintError(err error, code int) {
	if err != nil {
		_build.PrintError(err.Error())
		os.Exit(code)
	}
}
<commit_msg>main() refactoring<commit_after>package main

import (
	"flag"
	"fmt"
	_build "neon/build"
	_ "neon/builtin"
	_ "neon/task"
	"neon/util"
	"os"
	"path/filepath"
	"time"
)

const (
	DEFAULT_BUILD_FILE = "build.yml"
)

var VERSION string

// Parse command line and return parsed options
func ParseCommandLine() (string, bool, bool, string, bool, bool, string, bool, bool, string, bool, string, bool, []string) {
	file := flag.String("file", DEFAULT_BUILD_FILE, "Build file to run")
	info := flag.Bool("info", false, "Print build information")
	version := flag.Bool("version", false, "Print neon version")
	props := flag.String("props", "", "Build properties")
	timeit := flag.Bool("time", false, "Print build duration")
	tasks := flag.Bool("tasks", false, "Print tasks list")
	task := flag.String("task", "", "Print help on given task")
	targs := flag.Bool("targets", false, "Print targets list")
	builtins := flag.Bool("builtins", false, "Print builtins list")
	builtin := flag.String("builtin", "", "Print help on given builtin")
	refs := flag.Bool("refs", false, "Print tasks and builtins reference")
	install := flag.String("install", "", "Install given plugin")
	grey := flag.Bool("grey", false, "Print on terminal without colors")
	flag.Parse()
	targets := flag.Args()
	return *file, *info, *version, *props, *timeit, *tasks, *task, *targs, *builtins,
		*builtin, *refs, *install, *grey, targets
}

// Find build file and return its path
func FindBuildFile(name string) (string, error) {
	absolute, err := filepath.Abs(name)
	if err != nil {
		return "", fmt.Errorf("getting build file path: %v", err)
	}
	file := filepath.Base(absolute)
	dir := filepath.Dir(absolute)
	for {
		path := filepath.Join(dir, file)
		if util.FileExists(path) {
			return path, nil
		} else {
			parent := filepath.Dir(dir)
			if parent == dir {
				return "", fmt.Errorf("build file not found")
			}
			dir = parent
		}
	}
}

// Program entry point
func main() {
	start := time.Now()
	file, info, version, props, timeit, tasks, task, targs, builtins, builtin, refs, install, grey, targets := ParseCommandLine()
	// options that do not require we load build file
	_build.Grey = grey
	if tasks {
		_build.PrintTasks()
		return
	} else if task != "" {
		_build.PrintHelpTask(task)
		return
	} else if builtins {
		_build.PrintBuiltins()
		return
	} else if builtin != "" {
		_build.PrintHelpBuiltin(builtin)
		return
	} else if refs {
		_build.PrintReference()
		return
	} else if version {
		_build.Message(VERSION)
		return
	}
	// options that do require we load build file
	path, err := FindBuildFile(file)
	PrintError(err, 1)
	build, err := _build.NewBuild(path)
	PrintError(err, 2)
	if props != "" {
		err = build.SetCommandLineProperties(props)
		PrintError(err, 3)
	}
	if install != "" {
		err = build.Install(install)
		PrintError(err, 6)
		return
	} else if targs {
		build.PrintTargets()
		return
	} else if info {
		context, err := _build.NewContext(build)
		PrintError(err, 4)
		err = build.Info(context)
		PrintError(err, 4)
		return
	} else {
		context, err := _build.NewContext(build)
		err = build.Run(context, targets)
		PrintError(err, 5)
		duration := time.Now().Sub(start)
		if timeit || duration.Seconds() > 10 {
			_build.Message("Build duration: %s", duration.String())
		}
		PrintError(err, 5)
		_build.PrintOk()
		return
	}
}

// Print an error and exit if any
func PrintError(err error, code int) {
	if err != nil {
		_build.PrintError(err.Error())
		os.Exit(code)
	}
}
<|endoftext|>


<commit_before>// (c) Copyright IBM Corp. 2022

//go:build go1.13
// +build go1.13

package instaredis

import (
	"context"
	"strings"

	"github.com/go-redis/redis/v8"
	instana "github.com/instana/go-sensor"
	ot "github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

type commandCaptureHook struct {
	options        *redis.Options
	clusterOptions *redis.ClusterOptions
	sensor         *instana.Sensor
}

func setSpanCommands(span ot.Span, cmd redis.Cmder, cmds []redis.Cmder) {
	if cmd != nil {
		span.SetTag("redis.command", cmd.Name())

		if cmd.Err() != nil {
			span.SetTag("redis.error", cmd.Err().Error())
			span.LogFields(otlog.Object("error", cmd.Err().Error()))
		}
		return
	}

	i, end := 0, len(cmds)

	if cmds[0].Name() == "multi" && cmds[len(cmds)-1].Name() == "exec" {
		i = 1
		end = end - 1
	}

	var subCommands []string

	for ; i < end; i++ {
		subCommands = append(subCommands, cmds[i].Name())

		if cmds[i].Err() != nil {
			span.SetTag("redis.error", cmds[i].Err().Error())
			span.LogFields(otlog.Object("error", cmds[i].Err().Error()))
		}
	}

	span.SetTag("redis.command", "multi")
	span.SetTag("redis.subCommands", subCommands)
}

func (h commandCaptureHook) getConnection(ctx context.Context) string {
	if h.clusterOptions != nil {
		if len(h.clusterOptions.Addrs) > 0 {
			return h.clusterOptions.Addrs[0]
		}

		if h.clusterOptions.ClusterSlots != nil {
			cs, err := h.clusterOptions.ClusterSlots(ctx)

			if err == nil {
				return cs[0].Nodes[0].Addr
			}
		}
		return ""
	}

	connection := h.options.Addr

	if connection == "FailoverClient" {
		conn, err := h.options.Dialer(ctx, h.options.Network, "")

		if err == nil {
			return conn.RemoteAddr().String()
		}
	}
	return connection
}

func (h commandCaptureHook) handleHook(ctx context.Context, cmd redis.Cmder, cmds []redis.Cmder) {
	connection := h.getConnection(ctx)

	// if the IP provided in the Redis constructor have only ports. eg: :6379, :6380 and so on, we add localhost
	if strings.HasPrefix(connection, ":") {
		connection = "localhost" + connection
	}

	opts := []ot.StartSpanOption{
		ot.Tags{
			"redis.connection": connection,
		},
	}

	// We need an entry parent span in order to test this, so we will need a webserver or manually create an entry span
	tracer := h.sensor.Tracer()
	var span ot.Span = tracer.StartSpan("redis", opts...)

	if ps, ok := instana.SpanFromContext(ctx); ok {
		tracer = ps.Tracer()
		opts = append(opts, ot.ChildOf(ps.Context()))
		span = tracer.StartSpan("redis", opts...)
	}

	setSpanCommands(span, cmd, cmds)

	span.Finish()
}

func (h commandCaptureHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
	return ctx, nil
}

func (h commandCaptureHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error {
	h.handleHook(ctx, cmd, []redis.Cmder{})
	return nil
}

func (h commandCaptureHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) {
	return ctx, nil
}

func (h commandCaptureHook) AfterProcessPipeline(ctx context.Context, cmds []redis.Cmder) error {
	h.handleHook(ctx, nil, cmds)
	return nil
}

// WrapClient wraps the Redis client instance in order to add the instrumentation
func WrapClient(client *redis.Client, sensor *instana.Sensor) *redis.Client {
	opts := client.Options()
	client.AddHook(&commandCaptureHook{options: opts, sensor: sensor})
	return client
}

// WrapClient wraps the Redis client instance in order to add the instrumentation
func WrapClusterClient(clusterClient *redis.ClusterClient, sensor *instana.Sensor) *redis.ClusterClient {
	opts := clusterClient.Options()
	clusterClient.AddHook(&commandCaptureHook{clusterOptions: opts, sensor: sensor})
	return clusterClient
}
<commit_msg>improved how the connection is obtained<commit_after>// (c) Copyright IBM Corp. 2022

//go:build go1.13
// +build go1.13

package instaredis

import (
	"context"
	"strings"

	"github.com/go-redis/redis/v8"
	instana "github.com/instana/go-sensor"
	ot "github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

type commandCaptureHook struct {
	options        *redis.Options
	clusterOptions *redis.ClusterOptions
	sensor         *instana.Sensor
	connection     string
}

func newCommandCapture(s instana.Sensor, o *redis.Options, co *redis.ClusterOptions) *commandCaptureHook {
	var cch *commandCaptureHook

	if o != nil {
		cch = &commandCaptureHook{options: o, sensor: &s, connection: ""}
		cch.connection = cch.options.Addr
	} else {
		cch = &commandCaptureHook{clusterOptions: co, sensor: &s, connection: ""}

		if cch.clusterOptions != nil && len(cch.clusterOptions.Addrs) > 0 {
			cch.connection = cch.clusterOptions.Addrs[0]
		}
	}

	return cch
}

func setSpanCommands(span ot.Span, cmd redis.Cmder, cmds []redis.Cmder) {
	if cmd != nil {
		span.SetTag("redis.command", cmd.Name())

		if cmd.Err() != nil {
			span.SetTag("redis.error", cmd.Err().Error())
			span.LogFields(otlog.Object("error", cmd.Err().Error()))
		}
		return
	}

	i, end := 0, len(cmds)

	if cmds[0].Name() == "multi" && cmds[len(cmds)-1].Name() == "exec" {
		i = 1
		end = end - 1
	}

	var subCommands []string

	for ; i < end; i++ {
		subCommands = append(subCommands, cmds[i].Name())

		if cmds[i].Err() != nil {
			span.SetTag("redis.error", cmds[i].Err().Error())
			span.LogFields(otlog.Object("error", cmds[i].Err().Error()))
		}
	}

	span.SetTag("redis.command", "multi")
	span.SetTag("redis.subCommands", subCommands)
}

func (h commandCaptureHook) getConnection(ctx context.Context) string {
	if h.connection == "FailoverClient" {
		conn, err := h.options.Dialer(ctx, h.options.Network, "")

		if err == nil {
			h.connection = conn.RemoteAddr().String()
			return h.connection
		}
	}

	if h.connection != "" {
		return h.connection
	}

	if h.clusterOptions != nil && h.clusterOptions.ClusterSlots != nil {
		if cs, err := h.clusterOptions.ClusterSlots(ctx); err == nil {
			h.connection = cs[0].Nodes[0].Addr
			return h.connection
		}
		return ""
	}

	return ""
}

func (h commandCaptureHook) handleHook(ctx context.Context, cmd redis.Cmder, cmds []redis.Cmder) {
	connection := h.getConnection(ctx)

	// if the IP provided in the Redis constructor have only ports. eg: :6379, :6380 and so on, we add localhost
	if strings.HasPrefix(connection, ":") {
		connection = "localhost" + connection
	}

	opts := []ot.StartSpanOption{
		ot.Tags{
			"redis.connection": connection,
		},
	}

	// We need an entry parent span in order to test this, so we will need a webserver or manually create an entry span
	tracer := h.sensor.Tracer()
	var span ot.Span = tracer.StartSpan("redis", opts...)

	if ps, ok := instana.SpanFromContext(ctx); ok {
		tracer = ps.Tracer()
		opts = append(opts, ot.ChildOf(ps.Context()))
		span = tracer.StartSpan("redis", opts...)
	}

	setSpanCommands(span, cmd, cmds)

	span.Finish()
}

func (h commandCaptureHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
	return ctx, nil
}

func (h commandCaptureHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error {
	h.handleHook(ctx, cmd, []redis.Cmder{})
	return nil
}

func (h commandCaptureHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) {
	return ctx, nil
}

func (h commandCaptureHook) AfterProcessPipeline(ctx context.Context, cmds []redis.Cmder) error {
	h.handleHook(ctx, nil, cmds)
	return nil
}

// WrapClient wraps the Redis client instance in order to add the instrumentation
func WrapClient(client *redis.Client, sensor *instana.Sensor) *redis.Client {
	opts := client.Options()
	client.AddHook(newCommandCapture(*sensor, opts, nil))
	return client
}

// WrapClient wraps the Redis client instance in order to add the instrumentation
func WrapClusterClient(clusterClient *redis.ClusterClient, sensor *instana.Sensor) *redis.ClusterClient {
	opts := clusterClient.Options()
	clusterClient.AddHook(newCommandCapture(*sensor, nil, opts))
	return clusterClient
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/internal/base64util"

	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
)

const (
	// 21 Days worth of keys is the maximum per publish request (inclusive)
	maxKeysPerPublish = 21

	// only valid exposure key keyLength
	KeyLength = 16

	// Transmission risk constraints (inclusive..inclusive)
	MinTransmissionRisk = 0 // 0 indicates, no/unknown risk.
	MaxTransmissionRisk = 8

	// Intervals are defined as 10 minute periods, there are 144 of them in a day.
	// IntervalCount constraints (inclusive..inclusive)
	MinIntervalCount = 1
	MaxIntervalCount = 144

	// Self explanatory.
	// oneDay = time.Hour * 24

	// interval length
	intervalLength = 10 * time.Minute
)

// Publish represents the body of the PublishInfectedIds API call.
// Keys: Required and must have length >= 1 and <= 21 (`maxKeysPerPublish`)
// Regions: Array of regions. System defined, must match configuration.
// AppPackageName: The identifier for the mobile application.
//  - Android: The App Package AppPackageName
//  - iOS: The BundleID
// TransmissionRisk: An integer from 0-8 (inclusive) that represents
//  the transmission risk for this publish.
// Verification: The attestation payload for this request. (iOS or Android specific)
//   Base64 encoded.
// VerificationAuthorityName: a string that should be verified against the code provider.
//  Note: This project doesn't directly include a diagnosis code verification System
//        but does provide the ability to configure one in `serverevn.ServerEnv`
//
// The following fields are deprecated, but accepted for backwards-compatability:
// DeviceVerificationPayload: (attestation)
// Platform: "ios" or "android"
type Publish struct {
	Keys                []ExposureKey `json:"temporaryExposureKeys"`
	Regions             []string      `json:"regions"`
	AppPackageName      string        `json:"appPackageName"`
	VerificationPayload string        `json:"verificationPayload"`
	HMACKey             string        `json:"hmackey"`
	Padding             string        `json:"padding"`

	Platform                  string `json:"platform"`                  // DEPRECATED
	DeviceVerificationPayload string `json:"deviceVerificationPayload"` // DEPRECATED
}

// ApplyTransmissionRiskOverrides modifies the transmission risk values in the publish request
// based on the provided TransmissionRiskVector.
// In the live system, the TransmissionRiskVector values come from a trusted public health authority
// and are embedded in the verification certificate (JWT) transmitted on the publish request.
func (p *Publish) ApplyTransmissionRiskOverrides(overrides verifyapi.TransmissionRiskVector) {
	if len(overrides) == 0 {
		return
	}
	// The default sort order for TransmissionRiskVector is descending by SinceRollingPeriod.
	sort.Sort(overrides)
	// Sort the keys with the largest start interval first (descending), same as overrides.
	sort.Slice(p.Keys, func(i int, j int) bool {
		return p.Keys[i].IntervalNumber > p.Keys[j].IntervalNumber
	})

	overrideIdx := 0
	for i, eKey := range p.Keys {
		// Advance the overrideIdx until the current key is covered or we exhaust the
		// override index.
		for overrideIdx < len(overrides) &&
			eKey.IntervalNumber+eKey.IntervalCount <= overrides[overrideIdx].SinceRollingPeriod {
			overrideIdx++
		}

		// If we've run out of overrides to apply, then we have to break free.
		if overrideIdx >= len(overrides) {
			break
		}

		// Check to see if this key is in the current override.
		// If the key was EVERY valid during the SinceRollingPeriod then the override applies.
		if eKey.IntervalNumber+eKey.IntervalCount >= overrides[overrideIdx].SinceRollingPeriod {
			p.Keys[i].TransmissionRisk = overrides[overrideIdx].TranismissionRisk
			// don't advance overrideIdx, there might be additional keys in this override.
		}
	}
}

// ExposureKey is the 16 byte key, the start time of the key and the
// duration of the key. A duration of 0 means 24 hours.
// - ALL fields are REQUIRED and must meet the constraints below.
// Key must be the base64 (RFC 4648) encoded 16 byte exposure key from the device.
// - Base64 encoding should include padding, as per RFC 4648
// - if the key is not exactly 16 bytes in length, the request will be failed
// - that is, the whole batch will fail.
// IntervalNumber must be "reasonable" as in the system won't accept keys that
//   are scheduled to start in the future or that are too far in the past, which
//   is configurable per installation.
// IntervalCount must >= `minIntervalCount` and <= `maxIntervalCount`
//   1 - 144 inclusive.
// transmissionRisk must be >= 0 and <= 8.
type ExposureKey struct {
	Key              string `json:"key"`
	IntervalNumber   int32  `json:"rollingStartNumber"`
	IntervalCount    int32  `json:"rollingPeriod"`
	TransmissionRisk int    `json:"transmissionRisk"`
}

// ExposureKeys represents a set of ExposureKey objects as input to
// export file generation utility.
// Keys: Required and must have length >= 1.
type ExposureKeys struct {
	Keys []ExposureKey `json:"temporaryExposureKeys"`
}

// Exposure represents the record as stored in the database
// TODO(mikehelmick) - refactor this so that there is a public
// Exposure struct that doesn't have public fields and an
// internal struct that does. Separate out the database model
// from direct access.
// Mark records as writable/nowritable - is exposure key encrypted.
type Exposure struct {
	ExposureKey      []byte    `db:"exposure_key"`
	TransmissionRisk int       `db:"transmission_risk"`
	AppPackageName   string    `db:"app_package_name"`
	Regions          []string  `db:"regions"`
	IntervalNumber   int32     `db:"interval_number"`
	IntervalCount    int32     `db:"interval_count"`
	CreatedAt        time.Time `db:"created_at"`
	LocalProvenance  bool      `db:"local_provenance"`
	FederationSyncID int64     `db:"sync_id"`
}

// IntervalNumber calculates the exposure notification system interval
// number based on the input time.
func IntervalNumber(t time.Time) int32 {
	return int32(t.UTC().Unix()) / int32(intervalLength.Seconds())
}

// TruncateWindow truncates a time based on the size of the creation window.
func TruncateWindow(t time.Time, d time.Duration) time.Time {
	return t.Truncate(d)
}

// Transformer represents a configured Publish -> Exposure[] transformer.
type Transformer struct {
	maxExposureKeys     int
	maxIntervalStartAge time.Duration // How many intervals old does this server accept?
	truncateWindow      time.Duration
	debugAllowRestOfDay bool // raises end time of keys to the end of day, but doesn't embargo. For e2e testing only.
}

// NewTransformer creates a transformer for turning publish API requests into
// records for insertion into the database. On the call to TransformPublish
// all data is validated according to the transformer that is used.
func NewTransformer(maxExposureKeys int, maxIntervalStartAge time.Duration, truncateWindow time.Duration, allowRestOfDay bool) (*Transformer, error) {
	if maxExposureKeys < 0 || maxExposureKeys > maxKeysPerPublish {
		return nil, fmt.Errorf("maxExposureKeys must be > 0 and <= %v, got %v", maxKeysPerPublish, maxExposureKeys)
	}
	return &Transformer{
		maxExposureKeys:     maxExposureKeys,
		maxIntervalStartAge: maxIntervalStartAge,
		truncateWindow:      truncateWindow,
		debugAllowRestOfDay: allowRestOfDay,
	}, nil
}

// TransformExposureKey converts individual key data to an exposure entity.
// Validations during the transform include:
//
// * exposure keys are exactly 16 bytes in length after base64 decoding
// * minInterval <= interval number <= maxInterval
// * MinIntervalCount <= interval count <= MaxIntervalCount
//
func TransformExposureKey(exposureKey ExposureKey, appPackageName string, upcaseRegions []string, createdAt time.Time, minIntervalNumber, maxIntervalNumber int32) (*Exposure, error) {
	binKey, err := base64util.DecodeString(exposureKey.Key)
	if err != nil {
		return nil, err
	}

	// Validate individual pieces of the exposure key
	if len(binKey) != KeyLength {
		return nil, fmt.Errorf("invalid key length, %v, must be %v", len(binKey), KeyLength)
	}
	if ic := exposureKey.IntervalCount; ic < MinIntervalCount || ic > MaxIntervalCount {
		return nil, fmt.Errorf("invalid interval count, %v, must be >= %v && <= %v", ic, MinIntervalCount, MaxIntervalCount)
	}

	// Validate the IntervalNumber.
	if exposureKey.IntervalNumber < minIntervalNumber {
		return nil, fmt.Errorf("interval number %v is too old, must be >= %v", exposureKey.IntervalNumber, minIntervalNumber)
	}
	if exposureKey.IntervalNumber >= maxIntervalNumber {
		return nil, fmt.Errorf("interval number %v is in the future, must be < %v", exposureKey.IntervalNumber, maxIntervalNumber)
	}

	// Validate that the key is no longer effective.
	if exposureKey.IntervalNumber+exposureKey.IntervalCount > maxIntervalNumber {
		return nil, fmt.Errorf("interval number %v + interval count %v represents a key that is still valid, must end <= %v",
			exposureKey.IntervalNumber, exposureKey.IntervalCount, maxIntervalNumber)
	}

	if tr := exposureKey.TransmissionRisk; tr < MinTransmissionRisk || tr > MaxTransmissionRisk {
		return nil, fmt.Errorf("invalid transmission risk: %v, must be >= %v && <= %v", tr, MinTransmissionRisk, MaxTransmissionRisk)
	}

	return &Exposure{
		ExposureKey:      binKey,
		TransmissionRisk: exposureKey.TransmissionRisk,
		AppPackageName:   appPackageName,
		Regions:          upcaseRegions,
		IntervalNumber:   exposureKey.IntervalNumber,
		IntervalCount:    exposureKey.IntervalCount,
		CreatedAt:        createdAt,
		LocalProvenance:  true,
	}, nil
}

// TransformPublish converts incoming key data to a list of exposure entities.
// The data in the request is validated during the transform, including:
//
// * 0 exposure Keys in the requests
// * > Transformer.maxExposureKeys in the request
//
func (t *Transformer) TransformPublish(inData *Publish, batchTime time.Time) ([]*Exposure, error) {
	// Validate the number of keys that want to be published.
	if len(inData.Keys) == 0 {
		return nil, fmt.Errorf("no exposure keys in publish request")
	}
	if len(inData.Keys) > t.maxExposureKeys {
		return nil, fmt.Errorf("too many exposure keys in publish: %v, max of %v is allowed", len(inData.Keys), t.maxExposureKeys)
	}

	createdAt := TruncateWindow(batchTime, t.truncateWindow)
	entities := make([]*Exposure, 0, len(inData.Keys))

	// An exposure key must have an interval >= minInterval (max configured age)
	minIntervalNumber := IntervalNumber(batchTime.Add(-1 * t.maxIntervalStartAge))
	// And have an interval <= maxInterval (configured allowed clock skew)
	maxIntervalNumber := IntervalNumber(batchTime)

	// If, for testing, we are accepting keys that are valid the rest of the day:
	// adjust the maxIntervalNumber to be the end of the UTC day.
	if t.debugAllowRestOfDay {
		maxIntervalNumber = IntervalNumber(batchTime.Add(24 * time.Hour).Truncate(24 * time.Hour))
	}

	// Regions are a multi-value property, uppercase them for storage.
	// There is no set of "valid" regions overall, but it is defined
	// elsewhere by what regions an authorized application may write to.
	// See `authorizedapp.Config`
	upcaseRegions := make([]string, len(inData.Regions))
	for i, r := range inData.Regions {
		upcaseRegions[i] = strings.ToUpper(r)
	}

	for _, exposureKey := range inData.Keys {
		exposure, err := TransformExposureKey(exposureKey, inData.AppPackageName, upcaseRegions, createdAt, minIntervalNumber, maxIntervalNumber)
		if err != nil {
			return nil, fmt.Errorf("invalid publish data: %v", err)
		}
		entities = append(entities, exposure)
	}

	// Ensure that the uploaded keys are for a consecutive time period. No
	// overlaps and no gaps.
	// 1) Sort by interval number.
	sort.Slice(entities, func(i int, j int) bool {
		return entities[i].IntervalNumber < entities[j].IntervalNumber
	})
	// 2) Walk the slice and verify no gaps/overlaps.
	// We know the slice isn't empty, seed w/ the first interval.
	nextInterval := entities[0].IntervalNumber
	for _, ex := range entities {
		if ex.IntervalNumber < nextInterval {
			return nil, fmt.Errorf("exposure keys have overlapping intervals")
		}
		nextInterval = ex.IntervalNumber + ex.IntervalCount
	}

	return entities, nil
}
<commit_msg>[master] Fix spelling errors (#539)<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/internal/base64util"

	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
)

const (
	// 21 Days worth of keys is the maximum per publish request (inclusive)
	maxKeysPerPublish = 21

	// only valid exposure key keyLength
	KeyLength = 16

	// Transmission risk constraints (inclusive..inclusive)
	MinTransmissionRisk = 0 // 0 indicates, no/unknown risk.
	MaxTransmissionRisk = 8

	// Intervals are defined as 10 minute periods, there are 144 of them in a day.
	// IntervalCount constraints (inclusive..inclusive)
	MinIntervalCount = 1
	MaxIntervalCount = 144

	// Self explanatory.
	// oneDay = time.Hour * 24

	// interval length
	intervalLength = 10 * time.Minute
)

// Publish represents the body of the PublishInfectedIds API call.
// Keys: Required and must have length >= 1 and <= 21 (`maxKeysPerPublish`)
// Regions: Array of regions. System defined, must match configuration.
// AppPackageName: The identifier for the mobile application.
//  - Android: The App Package AppPackageName
//  - iOS: The BundleID
// TransmissionRisk: An integer from 0-8 (inclusive) that represents
//  the transmission risk for this publish.
// Verification: The attestation payload for this request. (iOS or Android specific)
//   Base64 encoded.
// VerificationAuthorityName: a string that should be verified against the code provider.
//  Note: This project doesn't directly include a diagnosis code verification System
//        but does provide the ability to configure one in `serverevn.ServerEnv`
//
// The following fields are deprecated, but accepted for backwards-compatibility:
// DeviceVerificationPayload: (attestation)
// Platform: "ios" or "android"
type Publish struct {
	Keys                []ExposureKey `json:"temporaryExposureKeys"`
	Regions             []string      `json:"regions"`
	AppPackageName      string        `json:"appPackageName"`
	VerificationPayload string        `json:"verificationPayload"`
	HMACKey             string        `json:"hmackey"`
	Padding             string        `json:"padding"`

	Platform                  string `json:"platform"`                  // DEPRECATED
	DeviceVerificationPayload string `json:"deviceVerificationPayload"` // DEPRECATED
}

// ApplyTransmissionRiskOverrides modifies the transmission risk values in the publish request
// based on the provided TransmissionRiskVector.
// In the live system, the TransmissionRiskVector values come from a trusted public health authority
// and are embedded in the verification certificate (JWT) transmitted on the publish request.
func (p *Publish) ApplyTransmissionRiskOverrides(overrides verifyapi.TransmissionRiskVector) {
	if len(overrides) == 0 {
		return
	}
	// The default sort order for TransmissionRiskVector is descending by SinceRollingPeriod.
	sort.Sort(overrides)
	// Sort the keys with the largest start interval first (descending), same as overrides.
	sort.Slice(p.Keys, func(i int, j int) bool {
		return p.Keys[i].IntervalNumber > p.Keys[j].IntervalNumber
	})

	overrideIdx := 0
	for i, eKey := range p.Keys {
		// Advance the overrideIdx until the current key is covered or we exhaust the
		// override index.
		for overrideIdx < len(overrides) &&
			eKey.IntervalNumber+eKey.IntervalCount <= overrides[overrideIdx].SinceRollingPeriod {
			overrideIdx++
		}

		// If we've run out of overrides to apply, then we have to break free.
		if overrideIdx >= len(overrides) {
			break
		}

		// Check to see if this key is in the current override.
		// If the key was EVERY valid during the SinceRollingPeriod then the override applies.
		if eKey.IntervalNumber+eKey.IntervalCount >= overrides[overrideIdx].SinceRollingPeriod {
			p.Keys[i].TransmissionRisk = overrides[overrideIdx].TranismissionRisk
			// don't advance overrideIdx, there might be additional keys in this override.
		}
	}
}

// ExposureKey is the 16 byte key, the start time of the key and the
// duration of the key. A duration of 0 means 24 hours.
// - ALL fields are REQUIRED and must meet the constraints below.
// Key must be the base64 (RFC 4648) encoded 16 byte exposure key from the device.
// - Base64 encoding should include padding, as per RFC 4648
// - if the key is not exactly 16 bytes in length, the request will be failed
// - that is, the whole batch will fail.
// IntervalNumber must be "reasonable" as in the system won't accept keys that
//   are scheduled to start in the future or that are too far in the past, which
//   is configurable per installation.
// IntervalCount must >= `minIntervalCount` and <= `maxIntervalCount`
//   1 - 144 inclusive.
// transmissionRisk must be >= 0 and <= 8.
type ExposureKey struct {
	Key              string `json:"key"`
	IntervalNumber   int32  `json:"rollingStartNumber"`
	IntervalCount    int32  `json:"rollingPeriod"`
	TransmissionRisk int    `json:"transmissionRisk"`
}

// ExposureKeys represents a set of ExposureKey objects as input to
// export file generation utility.
// Keys: Required and must have length >= 1.
type ExposureKeys struct {
	Keys []ExposureKey `json:"temporaryExposureKeys"`
}

// Exposure represents the record as stored in the database
// TODO(mikehelmick) - refactor this so that there is a public
// Exposure struct that doesn't have public fields and an
// internal struct that does. Separate out the database model
// from direct access.
// Mark records as writable/nowritable - is exposure key encrypted.
type Exposure struct {
	ExposureKey      []byte    `db:"exposure_key"`
	TransmissionRisk int       `db:"transmission_risk"`
	AppPackageName   string    `db:"app_package_name"`
	Regions          []string  `db:"regions"`
	IntervalNumber   int32     `db:"interval_number"`
	IntervalCount    int32     `db:"interval_count"`
	CreatedAt        time.Time `db:"created_at"`
	LocalProvenance  bool      `db:"local_provenance"`
	FederationSyncID int64     `db:"sync_id"`
}

// IntervalNumber calculates the exposure notification system interval
// number based on the input time.
func IntervalNumber(t time.Time) int32 {
	return int32(t.UTC().Unix()) / int32(intervalLength.Seconds())
}

// TruncateWindow truncates a time based on the size of the creation window.
func TruncateWindow(t time.Time, d time.Duration) time.Time {
	return t.Truncate(d)
}

// Transformer represents a configured Publish -> Exposure[] transformer.
type Transformer struct {
	maxExposureKeys     int
	maxIntervalStartAge time.Duration // How many intervals old does this server accept?
	truncateWindow      time.Duration
	debugAllowRestOfDay bool // raises end time of keys to the end of day, but doesn't embargo. For e2e testing only.
}

// NewTransformer creates a transformer for turning publish API requests into
// records for insertion into the database. On the call to TransformPublish
// all data is validated according to the transformer that is used.
func NewTransformer(maxExposureKeys int, maxIntervalStartAge time.Duration, truncateWindow time.Duration, allowRestOfDay bool) (*Transformer, error) {
	if maxExposureKeys < 0 || maxExposureKeys > maxKeysPerPublish {
		return nil, fmt.Errorf("maxExposureKeys must be > 0 and <= %v, got %v", maxKeysPerPublish, maxExposureKeys)
	}
	return &Transformer{
		maxExposureKeys:     maxExposureKeys,
		maxIntervalStartAge: maxIntervalStartAge,
		truncateWindow:      truncateWindow,
		debugAllowRestOfDay: allowRestOfDay,
	}, nil
}

// TransformExposureKey converts individual key data to an exposure entity.
// Validations during the transform include:
//
// * exposure keys are exactly 16 bytes in length after base64 decoding
// * minInterval <= interval number <= maxInterval
// * MinIntervalCount <= interval count <= MaxIntervalCount
//
func TransformExposureKey(exposureKey ExposureKey, appPackageName string, upcaseRegions []string, createdAt time.Time, minIntervalNumber, maxIntervalNumber int32) (*Exposure, error) {
	binKey, err := base64util.DecodeString(exposureKey.Key)
	if err != nil {
		return nil, err
	}

	// Validate individual pieces of the exposure key
	if len(binKey) != KeyLength {
		return nil, fmt.Errorf("invalid key length, %v, must be %v", len(binKey), KeyLength)
	}
	if ic := exposureKey.IntervalCount; ic < MinIntervalCount || ic > MaxIntervalCount {
		return nil, fmt.Errorf("invalid interval count, %v, must be >= %v && <= %v", ic, MinIntervalCount, MaxIntervalCount)
	}

	// Validate the IntervalNumber.
	if exposureKey.IntervalNumber < minIntervalNumber {
		return nil, fmt.Errorf("interval number %v is too old, must be >= %v", exposureKey.IntervalNumber, minIntervalNumber)
	}
	if exposureKey.IntervalNumber >= maxIntervalNumber {
		return nil, fmt.Errorf("interval number %v is in the future, must be < %v", exposureKey.IntervalNumber, maxIntervalNumber)
	}

	// Validate that the key is no longer effective.
	if exposureKey.IntervalNumber+exposureKey.IntervalCount > maxIntervalNumber {
		return nil, fmt.Errorf("interval number %v + interval count %v represents a key that is still valid, must end <= %v",
			exposureKey.IntervalNumber, exposureKey.IntervalCount, maxIntervalNumber)
	}

	if tr := exposureKey.TransmissionRisk; tr < MinTransmissionRisk || tr > MaxTransmissionRisk {
		return nil, fmt.Errorf("invalid transmission risk: %v, must be >= %v && <= %v", tr, MinTransmissionRisk, MaxTransmissionRisk)
	}

	return &Exposure{
		ExposureKey:      binKey,
		TransmissionRisk: exposureKey.TransmissionRisk,
		AppPackageName:   appPackageName,
		Regions:          upcaseRegions,
		IntervalNumber:   exposureKey.IntervalNumber,
		IntervalCount:    exposureKey.IntervalCount,
		CreatedAt:        createdAt,
		LocalProvenance:  true,
	}, nil
}

// TransformPublish converts incoming key data to a list of exposure entities.
// The data in the request is validated during the transform, including:
//
// * 0 exposure Keys in the requests
// * > Transformer.maxExposureKeys in the request
//
func (t *Transformer) TransformPublish(inData *Publish, batchTime time.Time) ([]*Exposure, error) {
	// Validate the number of keys that want to be published.
	if len(inData.Keys) == 0 {
		return nil, fmt.Errorf("no exposure keys in publish request")
	}
	if len(inData.Keys) > t.maxExposureKeys {
		return nil, fmt.Errorf("too many exposure keys in publish: %v, max of %v is allowed", len(inData.Keys), t.maxExposureKeys)
	}

	createdAt := TruncateWindow(batchTime, t.truncateWindow)
	entities := make([]*Exposure, 0, len(inData.Keys))

	// An exposure key must have an interval >= minInterval (max configured age)
	minIntervalNumber := IntervalNumber(batchTime.Add(-1 * t.maxIntervalStartAge))
	// And have an interval <= maxInterval (configured allowed clock skew)
	maxIntervalNumber := IntervalNumber(batchTime)

	// If, for testing, we are accepting keys that are valid the rest of the day:
	// adjust the maxIntervalNumber to be the end of the UTC day.
	if t.debugAllowRestOfDay {
		maxIntervalNumber = IntervalNumber(batchTime.Add(24 * time.Hour).Truncate(24 * time.Hour))
	}

	// Regions are a multi-value property, uppercase them for storage.
	// There is no set of "valid" regions overall, but it is defined
	// elsewhere by what regions an authorized application may write to.
	// See `authorizedapp.Config`
	upcaseRegions := make([]string, len(inData.Regions))
	for i, r := range inData.Regions {
		upcaseRegions[i] = strings.ToUpper(r)
	}

	for _, exposureKey := range inData.Keys {
		exposure, err := TransformExposureKey(exposureKey, inData.AppPackageName, upcaseRegions, createdAt, minIntervalNumber, maxIntervalNumber)
		if err != nil {
			return nil, fmt.Errorf("invalid publish data: %v", err)
		}
		entities = append(entities, exposure)
	}

	// Ensure that the uploaded keys are for a consecutive time period. No
	// overlaps and no gaps.
	// 1) Sort by interval number.
	sort.Slice(entities, func(i int, j int) bool {
		return entities[i].IntervalNumber < entities[j].IntervalNumber
	})
	// 2) Walk the slice and verify no gaps/overlaps.
	// We know the slice isn't empty, seed w/ the first interval.
	nextInterval := entities[0].IntervalNumber
	for _, ex := range entities {
		if ex.IntervalNumber < nextInterval {
			return nil, fmt.Errorf("exposure keys have overlapping intervals")
		}
		nextInterval = ex.IntervalNumber + ex.IntervalCount
	}

	return entities, nil
}
<|endoftext|>


<commit_before>package gorm

import (
	"errors"
	"fmt"
	"go/ast"
	"strings"
	"time"

	"reflect"
	"regexp"
)

type Scope struct {
	Value         interface{}
	indirectValue *reflect.Value
	Search        *search
	Sql           string
	SqlVars       []interface{}
	db            *DB
	skipLeft      bool
	primaryKey    string
	instanceId    string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

// NewScope create scope for callbacks, including DB's search information
func (db *DB) NewScope(value interface{}) *Scope {
	db.Value = value
	return &Scope{db: db, Search: db.search, Value: value}
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.db.parent, Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	return scope.db.new()
}

// DB get *sql.DB
func (scope *Scope) DB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	return scope.Dialect().Quote(str)
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.err(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if scope.primaryKey != "" {
		return scope.primaryKey
	}

	scope.primaryKey = ToSnake(GetPrimaryKey(scope.Value))
	return scope.primaryKey
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	return isBlank(reflect.ValueOf(scope.PrimaryKeyValue()))
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if scope.IndirectValue().Kind() == reflect.Struct {
		if field := scope.IndirectValue().FieldByName(SnakeToUpperCamel(scope.PrimaryKey())); field.IsValid() {
			return field.Interface()
		}
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(name string) bool {
	_, result := scope.FieldValueByName(name)
	return result
}

// FieldValueByName to get column's value and existence
func (scope *Scope) FieldValueByName(name string) (interface{}, bool) {
	return FieldValueByName(name, scope.Value)
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column string, value interface{}) bool {
	if scope.Value == nil {
		return false
	}

	return setFieldValue(scope.IndirectValue().FieldByName(SnakeToUpperCamel(column)), value)
}

// CallMethod invoke method with necessary argument
func (scope *Scope) CallMethod(name string) {
	if scope.Value == nil {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			fi := fm.Interface()
			if f, ok := fi.(func()); ok {
				f()
			} else if f, ok := fi.(func(s *Scope)); ok {
				f(scope)
			} else if f, ok := fi.(func(s *DB)); ok {
				f(scope.db.new())
			} else if f, ok := fi.(func() error); ok {
				scope.Err(f())
			} else if f, ok := fi.(func(s *Scope) error); ok {
				scope.Err(f(scope))
			} else if f, ok := fi.(func(s *DB) error); ok {
				scope.Err(f(scope.db.new()))
			} else {
				scope.Err(errors.New(fmt.Sprintf("unsupported function %v", name)))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			call(values.Index(i).Addr().Interface())
		}
	} else {
		call(scope.Value)
	}
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	scope.SqlVars = append(scope.SqlVars, value)
	return scope.Dialect().BinVar(len(scope.SqlVars))
}

// TableName get table name
var pluralMapKeys = []*regexp.Regexp{regexp.MustCompile("ch$"), regexp.MustCompile("ss$"), regexp.MustCompile("sh$"), regexp.MustCompile("day$"), regexp.MustCompile("y$"), regexp.MustCompile("x$"), regexp.MustCompile("([^s])s?$")}
var pluralMapValues = []string{"ches", "sses", "shes", "days", "ies", "xes", "${1}s"}

func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		if scope.Value == nil {
			scope.Err(errors.New("can't get table name"))
			return ""
		}

		data := scope.IndirectValue()
		if data.Kind() == reflect.Slice {
			elem := data.Type().Elem()
			if elem.Kind() == reflect.Ptr {
				elem = elem.Elem()
			}
			data = reflect.New(elem).Elem()
		}

		if fm := data.MethodByName("TableName"); fm.IsValid() {
			if v := fm.Call([]reflect.Value{}); len(v) > 0 {
				if result, ok := v[0].Interface().(string); ok {
					return result
				}
			}
		}

		str := ToSnake(data.Type().Name())

		if scope.db == nil || !scope.db.parent.singularTable {
			for index, reg := range pluralMapKeys {
				if reg.MatchString(str) {
					return reg.ReplaceAllString(str, pluralMapValues[index])
				}
			}
		}

		return str
	}
}

func (scope *Scope) QuotedTableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		keys := strings.Split(scope.TableName(), ".")
		for i, v := range keys {
			keys[i] = scope.Quote(v)
		}
		return strings.Join(keys, ".")
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) FieldByName(name string) (field *Field, ok bool) {
	var f reflect.StructField
	if scope.Value != nil {
		if scope.IndirectValue().Kind() == reflect.Struct {
			if f, ok = scope.IndirectValue().Type().FieldByName(SnakeToUpperCamel(name)); ok {
				field = scope.fieldFromStruct(f)
			}
		}
	}
	return
}

func (scope *Scope) fieldFromStruct(fieldStruct reflect.StructField) *Field {
	var field Field
	field.Name = fieldStruct.Name
	field.DBName = ToSnake(fieldStruct.Name)

	value := scope.IndirectValue().FieldByName(fieldStruct.Name)
	indirectValue := reflect.Indirect(value)
	field.Value = value.Interface()
	field.IsBlank = isBlank(value)

	// Search for primary key tag identifier
	settings := parseTagSetting(fieldStruct.Tag.Get("gorm"))
	if scope.PrimaryKey() == field.DBName {
		field.IsPrimaryKey = true
	}

	field.Tag = fieldStruct.Tag
	field.SqlTag = scope.sqlTagForField(&field)

	if !field.IsIgnored {
		// parse association
		if !indirectValue.IsValid() {
			indirectValue = reflect.New(value.Type())
		}
		typ := indirectValue.Type()
		foreignKey := SnakeToUpperCamel(settings["FOREIGNKEY"])
		associationForeignKey := SnakeToUpperCamel(settings["ASSOCIATIONFOREIGNKEY"])
		many2many := settings["MANY2MANY"]
		scopeTyp := scope.IndirectValue().Type()

		switch indirectValue.Kind() {
		case reflect.Slice:
			typ = typ.Elem()

			if typ.Kind() == reflect.Struct {
				if foreignKey == "" {
					foreignKey = scopeTyp.Name() + "Id"
				}
				if associationForeignKey == "" {
					associationForeignKey = typ.Name() + "Id"
				}

				// if not many to many, foreign key could be null
				if many2many == "" {
					if !reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						foreignKey = ""
					}
				}

				field.Relationship = &relationship{
					JoinTable:             many2many,
					ForeignKey:            foreignKey,
					AssociationForeignKey: associationForeignKey,
					Kind: "has_many",
				}

				if many2many != "" {
					field.Relationship.Kind = "many_to_many"
				}
			}
		case reflect.Struct:
			if !field.IsTime() && !field.IsScanner() {
				if foreignKey == "" && scope.HasColumn(field.Name+"Id") {
					field.Relationship = &relationship{ForeignKey: field.Name + "Id", Kind: "belongs_to"}
				} else if scope.HasColumn(foreignKey) {
					field.Relationship = &relationship{ForeignKey: foreignKey, Kind: "belongs_to"}
				} else {
					if foreignKey == "" {
						foreignKey = scopeTyp.Name() + "Id"
					}
					if reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						field.Relationship = &relationship{ForeignKey: foreignKey, Kind: "has_one"}
					}
				}
			}
		}
	}
	return &field
}

// Fields get value's fields
func (scope *Scope) Fields() (fields []*Field) {
	if scope.IndirectValue().IsValid() {
		scopeTyp := scope.IndirectValue().Type()
		for i := 0; i < scopeTyp.NumField(); i++ {
			fieldStruct := scopeTyp.Field(i)
			if !ast.IsExported(fieldStruct.Name) {
				continue
			}
			fields = append(fields, scope.fieldFromStruct(fieldStruct))
		}
	}
	return
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(NowFunc())

	if !scope.HasError() {
		result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...)
		if scope.Err(err) == nil {
			if count, err := result.RowsAffected(); err == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope.db.InstantSet(name, value)
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (interface{}, bool) {
	return scope.db.Get(name)
}

// InstanceId get InstanceId for scope
func (scope *Scope) InstanceId() string {
	if scope.instanceId == "" {
		scope.instanceId = fmt.Sprintf("%v", &scope)
	}
	return scope.instanceId
}

func (scope *Scope) InstanceSet(name string, value interface{}) *Scope {
	return scope.Set(name+scope.InstanceId(), value)
}

func (scope *Scope) InstanceGet(name string) (interface{}, bool) {
	return scope.Get(name + scope.InstanceId())
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.DB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.InstanceSet("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.InstanceGet("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				db.Commit()
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}
<commit_msg>Use the original DB when create a new Scope<commit_after>package gorm

import (
	"errors"
	"fmt"
	"go/ast"
	"strings"
	"time"

	"reflect"
	"regexp"
)

type Scope struct {
	Value         interface{}
	indirectValue *reflect.Value
	Search        *search
	Sql           string
	SqlVars       []interface{}
	db            *DB
	skipLeft      bool
	primaryKey    string
	instanceId    string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

// NewScope create scope for callbacks, including DB's search information
func (db *DB) NewScope(value interface{}) *Scope {
	db.Value = value
	return &Scope{db: db, Search: db.search, Value: value}
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.db, Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	return scope.db.new()
}

// DB get *sql.DB
func (scope *Scope) DB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	return scope.Dialect().Quote(str)
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.err(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if scope.primaryKey != "" {
		return scope.primaryKey
	}

	scope.primaryKey = ToSnake(GetPrimaryKey(scope.Value))
	return scope.primaryKey
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	return isBlank(reflect.ValueOf(scope.PrimaryKeyValue()))
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if scope.IndirectValue().Kind() == reflect.Struct {
		if field := scope.IndirectValue().FieldByName(SnakeToUpperCamel(scope.PrimaryKey())); field.IsValid() {
			return field.Interface()
		}
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(name string) bool {
	_, result := scope.FieldValueByName(name)
	return result
}

// FieldValueByName to get column's value and existence
func (scope *Scope) FieldValueByName(name string) (interface{}, bool) {
	return FieldValueByName(name, scope.Value)
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column string, value interface{}) bool {
	if scope.Value == nil {
		return false
	}

	return setFieldValue(scope.IndirectValue().FieldByName(SnakeToUpperCamel(column)), value)
}

// CallMethod invoke method with necessary argument
func (scope *Scope) CallMethod(name string) {
	if scope.Value == nil {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			fi := fm.Interface()
			if f, ok := fi.(func()); ok {
				f()
			} else if f, ok := fi.(func(s *Scope)); ok {
				f(scope)
			} else if f, ok := fi.(func(s *DB)); ok {
				f(scope.db.new())
			} else if f, ok := fi.(func() error); ok {
				scope.Err(f())
			} else if f, ok := fi.(func(s *Scope) error); ok {
				scope.Err(f(scope))
			} else if f, ok := fi.(func(s *DB) error); ok {
				scope.Err(f(scope.db.new()))
			} else {
				scope.Err(errors.New(fmt.Sprintf("unsupported function %v", name)))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			call(values.Index(i).Addr().Interface())
		}
	} else {
		call(scope.Value)
	}
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	scope.SqlVars = append(scope.SqlVars, value)
	return scope.Dialect().BinVar(len(scope.SqlVars))
}

// TableName get table name
var pluralMapKeys = []*regexp.Regexp{regexp.MustCompile("ch$"), regexp.MustCompile("ss$"), regexp.MustCompile("sh$"), regexp.MustCompile("day$"), regexp.MustCompile("y$"), regexp.MustCompile("x$"), regexp.MustCompile("([^s])s?$")}
var pluralMapValues = []string{"ches", "sses", "shes", "days", "ies", "xes", "${1}s"}

func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		if scope.Value == nil {
			scope.Err(errors.New("can't get table name"))
			return ""
		}

		data := scope.IndirectValue()
		if data.Kind() == reflect.Slice {
			elem := data.Type().Elem()
			if elem.Kind() == reflect.Ptr {
				elem = elem.Elem()
			}
			data = reflect.New(elem).Elem()
		}

		if fm := data.MethodByName("TableName"); fm.IsValid() {
			if v := fm.Call([]reflect.Value{}); len(v) > 0 {
				if result, ok := v[0].Interface().(string); ok {
					return result
				}
			}
		}

		str := ToSnake(data.Type().Name())

		if scope.db == nil || !scope.db.parent.singularTable {
			for index, reg := range pluralMapKeys {
				if reg.MatchString(str) {
					return reg.ReplaceAllString(str, pluralMapValues[index])
				}
			}
		}

		return str
	}
}

func (scope *Scope) QuotedTableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		keys := strings.Split(scope.TableName(), ".")
		for i, v := range keys {
			keys[i] = scope.Quote(v)
		}
		return strings.Join(keys, ".")
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) FieldByName(name string) (field *Field, ok bool) {
	var f reflect.StructField
	if scope.Value != nil {
		if scope.IndirectValue().Kind() == reflect.Struct {
			if f, ok = scope.IndirectValue().Type().FieldByName(SnakeToUpperCamel(name)); ok {
				field = scope.fieldFromStruct(f)
			}
		}
	}
	return
}

func (scope *Scope) fieldFromStruct(fieldStruct reflect.StructField) *Field {
	var field Field
	field.Name = fieldStruct.Name
	field.DBName = ToSnake(fieldStruct.Name)

	value := scope.IndirectValue().FieldByName(fieldStruct.Name)
	indirectValue := reflect.Indirect(value)
	field.Value = value.Interface()
	field.IsBlank = isBlank(value)

	// Search for primary key tag identifier
	settings := parseTagSetting(fieldStruct.Tag.Get("gorm"))
	if scope.PrimaryKey() == field.DBName {
		field.IsPrimaryKey = true
	}

	field.Tag = fieldStruct.Tag
	field.SqlTag = scope.sqlTagForField(&field)

	if !field.IsIgnored {
		// parse association
		if !indirectValue.IsValid() {
			indirectValue = reflect.New(value.Type())
		}
		typ := indirectValue.Type()
		foreignKey := SnakeToUpperCamel(settings["FOREIGNKEY"])
		associationForeignKey := SnakeToUpperCamel(settings["ASSOCIATIONFOREIGNKEY"])
		many2many := settings["MANY2MANY"]
		scopeTyp := scope.IndirectValue().Type()

		switch indirectValue.Kind() {
		case reflect.Slice:
			typ = typ.Elem()

			if typ.Kind() == reflect.Struct {
				if foreignKey == "" {
					foreignKey = scopeTyp.Name() + "Id"
				}
				if associationForeignKey == "" {
					associationForeignKey = typ.Name() + "Id"
				}

				// if not many to many, foreign key could be null
				if many2many == "" {
					if !reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						foreignKey = ""
					}
				}

				field.Relationship = &relationship{
					JoinTable:             many2many,
					ForeignKey:            foreignKey,
					AssociationForeignKey: associationForeignKey,
					Kind: "has_many",
				}

				if many2many != "" {
					field.Relationship.Kind = "many_to_many"
				}
			}
		case reflect.Struct:
			if !field.IsTime() && !field.IsScanner() {
				if foreignKey == "" && scope.HasColumn(field.Name+"Id") {
					field.Relationship = &relationship{ForeignKey: field.Name + "Id", Kind: "belongs_to"}
				} else if scope.HasColumn(foreignKey) {
					field.Relationship = &relationship{ForeignKey: foreignKey, Kind: "belongs_to"}
				} else {
					if foreignKey == "" {
						foreignKey = scopeTyp.Name() + "Id"
					}
					if reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						field.Relationship = &relationship{ForeignKey: foreignKey, Kind: "has_one"}
					}
				}
			}
		}
	}
	return &field
}

// Fields get value's fields
func (scope *Scope) Fields() (fields []*Field) {
	if scope.IndirectValue().IsValid() {
		scopeTyp := scope.IndirectValue().Type()
		for i := 0; i < scopeTyp.NumField(); i++ {
			fieldStruct := scopeTyp.Field(i)
			if !ast.IsExported(fieldStruct.Name) {
				continue
			}
			fields = append(fields, scope.fieldFromStruct(fieldStruct))
		}
	}
	return
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(NowFunc())

	if !scope.HasError() {
		result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...)
		if scope.Err(err) == nil {
			if count, err := result.RowsAffected(); err == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope.db.InstantSet(name, value)
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (interface{}, bool) {
	return scope.db.Get(name)
}

// InstanceId get InstanceId for scope
func (scope *Scope) InstanceId() string {
	if scope.instanceId == "" {
		scope.instanceId = fmt.Sprintf("%v", &scope)
	}
	return scope.instanceId
}

func (scope *Scope) InstanceSet(name string, value interface{}) *Scope {
	return scope.Set(name+scope.InstanceId(), value)
}

func (scope *Scope) InstanceGet(name string) (interface{}, bool) {
	return scope.Get(name + scope.InstanceId())
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.DB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.InstanceSet("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.InstanceGet("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				db.Commit()
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"github.com/abh/geodns/countries"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

func getQuestionName(z *Zone, req *dns.Msg) string {
	lx := dns.SplitLabels(req.Question[0].Name)
	ql := lx[0 : len(lx)-z.LenLabels]
	return strings.ToLower(strings.Join(ql, "."))
}

var geoIP = setupGeoIP()

func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {

	qtype := req.Question[0].Qtype

	logPrintf("[zone %s] incoming %s %s %d from %s\n", z.Origin, req.Question[0].Name,
		dns.TypeToString[qtype], req.MsgHdr.Id, w.RemoteAddr())

	qCounter.Add(1)

	logPrintln("Got request", req)

	label := getQuestionName(z, req)

	var ip string
	var edns *dns.EDNS0_SUBNET
	var opt_rr *dns.OPT

	for _, extra := range req.Extra {
		log.Println("Extra", extra)
		for _, o := range extra.(*dns.OPT).Option {
			opt_rr = extra.(*dns.OPT)
			switch e := o.(type) {
			case *dns.EDNS0_NSID:
				// do stuff with e.Nsid
			case *dns.EDNS0_SUBNET:
				log.Println("Got edns", e.Address, e.Family, e.SourceNetmask, e.SourceScope)
				if e.Address != nil {
					log.Println("Setting edns to", e)
					edns = e
					ip = e.Address.String()
				}
			}
		}
	}

	var country string
	if geoIP != nil {
		if len(ip) == 0 { // no edns subnet
			ip, _, _ = net.SplitHostPort(w.RemoteAddr().String())
		}
		country = strings.ToLower(geoIP.GetCountry(ip))
		logPrintln("Country:", ip, country)
	}

	m := new(dns.Msg)
	m.SetReply(req)
	if e := m.IsEdns0(); e != nil {
		m.SetEdns0(4096, e.Do())
	}
	m.Authoritative = true

	// TODO: set scope to 0 if there are no alternate responses
	if edns != nil {
		log.Println("family", edns.Family)
		if edns.Family != 0 {
			log.Println("edns response!")
			edns.SourceScope = 16
			m.Extra = append(m.Extra, opt_rr)
		}
	}

	labels, labelQtype := z.findLabels(label, country, qTypes{dns.TypeMF, dns.TypeCNAME, qtype})
	if labelQtype == 0 {
		labelQtype = qtype
	}

	if labels == nil {

		if label == "_status" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			m.Answer = statusRR(z)
			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		if label == "_country" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
			h.Name = "_country." + z.Origin + "."

			m.Answer = []dns.RR{&dns.TXT{Hdr: h,
				Txt: []string{
					w.RemoteAddr().String(),
					ip,
					string(country),
					string(countries.CountryContinent[country]),
				},
			}}

			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		// return NXDOMAIN
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		m.Ns = []dns.RR{z.SoaRR()}

		w.WriteMsg(m)
		return
	}

	if servers := labels.Picker(labelQtype, labels.MaxHosts); servers != nil {
		var rrs []dns.RR
		for _, record := range servers {
			rr := record.RR
			rr.Header().Name = req.Question[0].Name
			rrs = append(rrs, rr)
		}
		m.Answer = rrs
	}

	if len(m.Answer) == 0 {
		m.Ns = append(m.Ns, z.SoaRR())
	}

	logPrintln(m)

	err := w.WriteMsg(m)
	if err != nil {
		// if Pack'ing fails the Write fails. Return SERVFAIL.
		log.Println("Error writing packet", m)
		dns.HandleFailed(w, req)
	}
	return
}

func statusRR(z *Zone) []dns.RR {
	h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
	h.Name = "_status." + z.Origin + "."

	status := map[string]string{"v": VERSION, "id": serverId}

	hostname, err := os.Hostname()
	if err == nil {
		status["h"] = hostname
	}
	status["up"] = strconv.Itoa(int(time.Since(timeStarted).Seconds()))
	status["qs"] = qCounter.String()

	js, err := json.Marshal(status)

	return []dns.RR{&dns.TXT{Hdr: h, Txt: []string{string(js)}}}
}

func setupServerFunc(Zone *Zone) func(dns.ResponseWriter, *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		serve(w, r, Zone)
	}
}

func listenAndServe(ip string, Zones *Zones) {

	prots := []string{"udp", "tcp"}

	for _, prot := range prots {
		go func(p string) {
			server := &dns.Server{Addr: ip, Net: p}

			log.Printf("Opening on %s %s", ip, p)
			if err := server.ListenAndServe(); err != nil {
				log.Fatalf("geodns: failed to setup %s %s: %s", ip, p, err)
			}
			log.Fatalf("geodns: ListenAndServe unexpectedly returned")
		}(prot)
	}

}
<commit_msg>Disable EDNS logging<commit_after>package main

import (
	"encoding/json"
	"github.com/abh/geodns/countries"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

func getQuestionName(z *Zone, req *dns.Msg) string {
	lx := dns.SplitLabels(req.Question[0].Name)
	ql := lx[0 : len(lx)-z.LenLabels]
	return strings.ToLower(strings.Join(ql, "."))
}

var geoIP = setupGeoIP()

func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {

	qtype := req.Question[0].Qtype

	logPrintf("[zone %s] incoming %s %s %d from %s\n", z.Origin, req.Question[0].Name,
		dns.TypeToString[qtype], req.MsgHdr.Id, w.RemoteAddr())

	qCounter.Add(1)

	logPrintln("Got request", req)

	label := getQuestionName(z, req)

	var ip string
	var edns *dns.EDNS0_SUBNET
	var opt_rr *dns.OPT

	for _, extra := range req.Extra {
		for _, o := range extra.(*dns.OPT).Option {
			opt_rr = extra.(*dns.OPT)
			switch e := o.(type) {
			case *dns.EDNS0_NSID:
				// do stuff with e.Nsid
			case *dns.EDNS0_SUBNET:
				logPrintln("Got edns", e.Address, e.Family, e.SourceNetmask, e.SourceScope)
				if e.Address != nil {
					edns = e
					ip = e.Address.String()
				}
			}
		}
	}

	var country string
	if geoIP != nil {
		if len(ip) == 0 { // no edns subnet
			ip, _, _ = net.SplitHostPort(w.RemoteAddr().String())
		}
		country = strings.ToLower(geoIP.GetCountry(ip))
		logPrintln("Country:", ip, country)
	}

	m := new(dns.Msg)
	m.SetReply(req)
	if e := m.IsEdns0(); e != nil {
		m.SetEdns0(4096, e.Do())
	}
	m.Authoritative = true

	// TODO: set scope to 0 if there are no alternate responses
	if edns != nil {
		if edns.Family != 0 {
			edns.SourceScope = 16
			m.Extra = append(m.Extra, opt_rr)
		}
	}

	labels, labelQtype := z.findLabels(label, country, qTypes{dns.TypeMF, dns.TypeCNAME, qtype})
	if labelQtype == 0 {
		labelQtype = qtype
	}

	if labels == nil {

		if label == "_status" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			m.Answer = statusRR(z)
			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		if label == "_country" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
			h.Name = "_country." + z.Origin + "."

			m.Answer = []dns.RR{&dns.TXT{Hdr: h,
				Txt: []string{
					w.RemoteAddr().String(),
					ip,
					string(country),
					string(countries.CountryContinent[country]),
				},
			}}

			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		// return NXDOMAIN
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		m.Ns = []dns.RR{z.SoaRR()}

		w.WriteMsg(m)
		return
	}

	if servers := labels.Picker(labelQtype, labels.MaxHosts); servers != nil {
		var rrs []dns.RR
		for _, record := range servers {
			rr := record.RR
			rr.Header().Name = req.Question[0].Name
			rrs = append(rrs, rr)
		}
		m.Answer = rrs
	}

	if len(m.Answer) == 0 {
		m.Ns = append(m.Ns, z.SoaRR())
	}

	logPrintln(m)

	err := w.WriteMsg(m)
	if err != nil {
		// if Pack'ing fails the Write fails. Return SERVFAIL.
		log.Println("Error writing packet", m)
		dns.HandleFailed(w, req)
	}
	return
}

func statusRR(z *Zone) []dns.RR {
	h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
	h.Name = "_status." + z.Origin + "."

	status := map[string]string{"v": VERSION, "id": serverId}

	hostname, err := os.Hostname()
	if err == nil {
		status["h"] = hostname
	}
	status["up"] = strconv.Itoa(int(time.Since(timeStarted).Seconds()))
	status["qs"] = qCounter.String()

	js, err := json.Marshal(status)

	return []dns.RR{&dns.TXT{Hdr: h, Txt: []string{string(js)}}}
}

func setupServerFunc(Zone *Zone) func(dns.ResponseWriter, *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		serve(w, r, Zone)
	}
}

func listenAndServe(ip string, Zones *Zones) {

	prots := []string{"udp", "tcp"}

	for _, prot := range prots {
		go func(p string) {
			server := &dns.Server{Addr: ip, Net: p}

			log.Printf("Opening on %s %s", ip, p)
			if err := server.ListenAndServe(); err != nil {
				log.Fatalf("geodns: failed to setup %s %s: %s", ip, p, err)
			}
			log.Fatalf("geodns: ListenAndServe unexpectedly returned")
		}(prot)
	}

}
<|endoftext|>


<commit_before>package engine

import (
	"bufio"
	"container/ring"
	"fmt"
	"io"
	"sync"
)

type Output struct {
	sync.Mutex
	dests []io.Writer
	tasks sync.WaitGroup
}

// NewOutput returns a new Output object with no destinations attached.
// Writing to an empty Output will cause the written data to be discarded.
func NewOutput() *Output {
	return &Output{}
}

// Add attaches a new destination to the Output. Any data subsequently written
// to the output will be written to the new destination in addition to all the others.
// This method is thread-safe.
// FIXME: Add cannot fail
func (o *Output) Add(dst io.Writer) error {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	o.dests = append(o.dests, dst)
	return nil
}

// AddPipe creates an in-memory pipe with io.Pipe(), adds its writing end as a destination,
// and returns its reading end for consumption by the caller.
// This is a rough equivalent similar to Cmd.StdoutPipe() in the standard os/exec package.
// This method is thread-safe.
func (o *Output) AddPipe() (io.Reader, error) {
	r, w := io.Pipe()
	o.Add(w)
	return r, nil
}

// AddTail starts a new goroutine which will read all subsequent data written to the output,
// line by line, and append the last `n` lines to `dst`.
func (o *Output) AddTail(dst *[]string, n int) error {
	src, err := o.AddPipe()
	if err != nil {
		return err
	}
	o.tasks.Add(1)
	go func() {
		defer o.tasks.Done()
		Tail(src, n, dst)
	}()
	return nil
}

// AddString starts a new goroutine which will read all subsequent data written to the output,
// line by line, and store the last line into `dst`.
func (o *Output) AddString(dst *string) error {
	src, err := o.AddPipe()
	if err != nil {
		return err
	}
	o.tasks.Add(1)
	go func() {
		defer o.tasks.Done()
		lines := make([]string, 0, 1)
		Tail(src, 1, &lines)
		if len(lines) == 0 {
			*dst = ""
		} else {
			*dst = lines[0]
		}
	}()
	return nil
}

// Write writes the same data to all registered destinations.
// This method is thread-safe.
func (o *Output) Write(p []byte) (n int, err error) {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	var firstErr error
	for _, dst := range o.dests {
		_, err := dst.Write(p)
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return len(p), err
}

// Close unregisters all destinations and waits for all background
// AddTail and AddString tasks to complete.
// The Close method of each destination is called if it exists.
func (o *Output) Close() error {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	var firstErr error
	for _, dst := range o.dests {
		if closer, ok := dst.(io.WriteCloser); ok {
			err := closer.Close()
			if err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	o.tasks.Wait()
	return firstErr
}

type Input struct {
	src io.Reader
	sync.Mutex
}

// NewInput returns a new Input object with no source attached.
// Reading to an empty Input will return io.EOF.
func NewInput() *Input {
	return &Input{}
}

// Read reads from the input in a thread-safe way.
func (i *Input) Read(p []byte) (n int, err error) {
	i.Mutex.Lock()
	defer i.Mutex.Unlock()
	if i.src == nil {
		return 0, io.EOF
	}
	return i.src.Read(p)
}

// Add attaches a new source to the input.
// Add can only be called once per input. Subsequent calls will
// return an error.
func (i *Input) Add(src io.Reader) error {
	i.Mutex.Lock()
	defer i.Mutex.Unlock()
	if i.src != nil {
		return fmt.Errorf("Maximum number of sources reached: 1")
	}
	i.src = src
	return nil
}

// Tail reads from `src` line per line, and returns the last `n` lines as an array.
// A ring buffer is used to only store `n` lines at any time.
func Tail(src io.Reader, n int, dst *[]string) {
	scanner := bufio.NewScanner(src)
	r := ring.New(n)
	for scanner.Scan() {
		if n == 0 {
			continue
		}
		r.Value = scanner.Text()
		r = r.Next()
	}
	r.Do(func(v interface{}) {
		if v == nil {
			return
		}
		*dst = append(*dst, v.(string))
	})
}
<commit_msg>revert 'firstErr'<commit_after>package engine

import (
	"bufio"
	"container/ring"
	"fmt"
	"io"
	"sync"
)

type Output struct {
	sync.Mutex
	dests []io.Writer
	tasks sync.WaitGroup
}

// NewOutput returns a new Output object with no destinations attached.
// Writing to an empty Output will cause the written data to be discarded.
func NewOutput() *Output {
	return &Output{}
}

// Add attaches a new destination to the Output. Any data subsequently written
// to the output will be written to the new destination in addition to all the others.
// This method is thread-safe.
// FIXME: Add cannot fail
func (o *Output) Add(dst io.Writer) error {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	o.dests = append(o.dests, dst)
	return nil
}

// AddPipe creates an in-memory pipe with io.Pipe(), adds its writing end as a destination,
// and returns its reading end for consumption by the caller.
// This is a rough equivalent similar to Cmd.StdoutPipe() in the standard os/exec package.
// This method is thread-safe.
func (o *Output) AddPipe() (io.Reader, error) {
	r, w := io.Pipe()
	o.Add(w)
	return r, nil
}

// AddTail starts a new goroutine which will read all subsequent data written to the output,
// line by line, and append the last `n` lines to `dst`.
func (o *Output) AddTail(dst *[]string, n int) error {
	src, err := o.AddPipe()
	if err != nil {
		return err
	}
	o.tasks.Add(1)
	go func() {
		defer o.tasks.Done()
		Tail(src, n, dst)
	}()
	return nil
}

// AddString starts a new goroutine which will read all subsequent data written to the output,
// line by line, and store the last line into `dst`.
func (o *Output) AddString(dst *string) error {
	src, err := o.AddPipe()
	if err != nil {
		return err
	}
	o.tasks.Add(1)
	go func() {
		defer o.tasks.Done()
		lines := make([]string, 0, 1)
		Tail(src, 1, &lines)
		if len(lines) == 0 {
			*dst = ""
		} else {
			*dst = lines[0]
		}
	}()
	return nil
}

// Write writes the same data to all registered destinations.
// This method is thread-safe.
func (o *Output) Write(p []byte) (n int, err error) {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	var firstErr error
	for _, dst := range o.dests {
		_, err := dst.Write(p)
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return len(p), firstErr
}

// Close unregisters all destinations and waits for all background
// AddTail and AddString tasks to complete.
// The Close method of each destination is called if it exists.
func (o *Output) Close() error {
	o.Mutex.Lock()
	defer o.Mutex.Unlock()
	var firstErr error
	for _, dst := range o.dests {
		if closer, ok := dst.(io.WriteCloser); ok {
			err := closer.Close()
			if err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	o.tasks.Wait()
	return firstErr
}

type Input struct {
	src io.Reader
	sync.Mutex
}

// NewInput returns a new Input object with no source attached.
// Reading to an empty Input will return io.EOF.
func NewInput() *Input {
	return &Input{}
}

// Read reads from the input in a thread-safe way.
func (i *Input) Read(p []byte) (n int, err error) {
	i.Mutex.Lock()
	defer i.Mutex.Unlock()
	if i.src == nil {
		return 0, io.EOF
	}
	return i.src.Read(p)
}

// Add attaches a new source to the input.
// Add can only be called once per input. Subsequent calls will
// return an error.
func (i *Input) Add(src io.Reader) error {
	i.Mutex.Lock()
	defer i.Mutex.Unlock()
	if i.src != nil {
		return fmt.Errorf("Maximum number of sources reached: 1")
	}
	i.src = src
	return nil
}

// Tail reads from `src` line per line, and returns the last `n` lines as an array.
// A ring buffer is used to only store `n` lines at any time.
func Tail(src io.Reader, n int, dst *[]string) {
	scanner := bufio.NewScanner(src)
	r := ring.New(n)
	for scanner.Scan() {
		if n == 0 {
			continue
		}
		r.Value = scanner.Text()
		r = r.Next()
	}
	r.Do(func(v interface{}) {
		if v == nil {
			return
		}
		*dst = append(*dst, v.(string))
	})
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "net/http"
  "os"
  "io/ioutil"
)

func main(){
  var url string
  fmt.Println(">> url: (Make sure to use http://)")
  fmt.Scanf("%s", &url)
  data, err := http.Get(url)
  if err != nil{
    os.Exit(1)
  }else{
    defer data.Body.Close()
    contents, err := ioutil.ReadAll(data.Body)
    if err != nil{
      fmt.Println("%s\n", err)
    }
    fmt.Println("%s\n", string(contents))
  }
}
<commit_msg>Added error message if user type in no url<commit_after>package main

import (
  "fmt"
  "net/http"
  "os"
  "io/ioutil"
)

func main(){
  var url string
  fmt.Println(">> url: (Make sure to use http://)")
  fmt.Scanf("%s", &url)
  if url == "" {
    fmt.Println("Please type in an url")
  }
  data, err := http.Get(url)
  if err != nil{
    os.Exit(1)
  }else{
    defer data.Body.Close()
    contents, err := ioutil.ReadAll(data.Body)
    if err != nil{
      fmt.Println("%s\n", err)
    }
    fmt.Println("%s\n", string(contents))
  }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tcontainer

import (
	"github.com/trivago/tgo/ttesting"
	//"runtime"
	"testing"
	//"time"
)

func TestBytePool(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	tinyMin := pool.Get(1)
	expect.Equal(tiny, cap(tinyMin))
	expect.Equal(1, len(tinyMin))

	tinyMax := pool.Get(tiny)
	expect.Equal(tiny, len(tinyMax))

	smallMin := pool.Get(tiny + 1)
	expect.Equal(small, cap(smallMin))
	expect.Equal(tiny+1, len(smallMin))

	smallMax := pool.Get(small)
	expect.Equal(small, len(smallMax))

	mediumMin := pool.Get(small + 1)
	expect.Equal(medium, cap(mediumMin))
	expect.Equal(small+1, len(mediumMin))

	mediumMax := pool.Get(medium)
	expect.Equal(medium, len(mediumMax))

	largeMin := pool.Get(medium + 1)
	expect.Equal(large, cap(largeMin))
	expect.Equal(medium+1, len(largeMin))

	largeMax := pool.Get(large)
	expect.Equal(large, len(largeMax))

	hugeMin := pool.Get(large + 1)
	expect.Equal(huge, cap(hugeMin))
	expect.Equal(large+1, len(hugeMin))

	hugeMax := pool.Get(huge)
	expect.Equal(huge, len(hugeMax))
}

/*func TestBytePool(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	tinyMin := pool.Get(1)
	expect.Equal(pool.tiny.unit, cap(tinyMin))
	expect.Equal(1, len(tinyMin))

	tinyMax := pool.Get(pool.tiny.max)
	expect.Equal(pool.tiny.max, cap(tinyMax))
	expect.Equal(pool.tiny.max, len(tinyMax))

	smallMin := pool.Get(pool.tiny.max + 1)
	expect.Equal(pool.small.unit, cap(smallMin))
	expect.Equal(pool.tiny.max+1, len(smallMin))

	smallMax := pool.Get(pool.small.max)
	expect.Equal(pool.small.max, cap(smallMax))
	expect.Equal(pool.small.max, len(smallMax))

	mediumMin := pool.Get(pool.small.max + 1)
	expect.Equal(pool.medium.unit, cap(mediumMin))
	expect.Equal(pool.small.max+1, len(mediumMin))

	mediumMax := pool.Get(pool.medium.max)
	expect.Equal(pool.medium.max, cap(mediumMax))
	expect.Equal(pool.medium.max, len(mediumMax))

	largeMin := pool.Get(pool.medium.max + 1)
	expect.Equal(pool.large.unit, cap(largeMin))
	expect.Equal(pool.medium.max+1, len(largeMin))

	largeMax := pool.Get(pool.large.max)
	expect.Equal(pool.large.max, cap(largeMax))
	expect.Equal(pool.large.max, len(largeMax))

	huge := pool.Get(pool.large.max + 1)
	expect.Equal(pool.large.max+1, cap(huge))
	expect.Equal(pool.large.max+1, len(huge))
}

func allocateWaste(pool *BytePool, expect ttesting.Expect) {
	data := pool.Get(32)
	for i := 0; i < 32; i++ {
		data[i] = byte(i)
	}

	expect.Equal(int32(-1), *pool.tiny.slabs[0].top)
}

func TestBytePoolRecycle(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	expect.Nil(pool.tiny.slabs[0])
	allocateWaste(&pool, expect)

	expect.NonBlocking(time.Second, func() {
		for *pool.tiny.slabs[0].top < 0 {
			runtime.Gosched()
			runtime.GC()
		}
	})

	expect.Equal(int32(0), *pool.tiny.slabs[0].top)

	data := pool.Get(32)
	for i := 0; i < 32; i++ {
		expect.Equal(byte(i), data[i])
	}
}*/
<commit_msg>Added bytepool tests<commit_after>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tcontainer

import (
	"bytes"
	"encoding/binary"
	"github.com/trivago/tgo/ttesting"
	"math/rand"
	"sync"
	"testing"
	"time"
)

func TestBytePool(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	tinyMin := pool.Get(1)
	expect.Equal(tiny, cap(tinyMin))
	expect.Equal(1, len(tinyMin))

	tinyMax := pool.Get(tiny)
	expect.Equal(tiny, len(tinyMax))

	smallMin := pool.Get(tiny + 1)
	expect.Equal(small, cap(smallMin))
	expect.Equal(tiny+1, len(smallMin))

	smallMax := pool.Get(small)
	expect.Equal(small, len(smallMax))

	mediumMin := pool.Get(small + 1)
	expect.Equal(medium, cap(mediumMin))
	expect.Equal(small+1, len(mediumMin))

	mediumMax := pool.Get(medium)
	expect.Equal(medium, len(mediumMax))

	largeMin := pool.Get(medium + 1)
	expect.Equal(large, cap(largeMin))
	expect.Equal(medium+1, len(largeMin))

	largeMax := pool.Get(large)
	expect.Equal(large, len(largeMax))

	hugeMin := pool.Get(large + 1)
	expect.Equal(huge, cap(hugeMin))
	expect.Equal(large+1, len(hugeMin))

	hugeMax := pool.Get(huge)
	expect.Equal(huge, len(hugeMax))
}

func TestBytePoolParallel(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	start := new(sync.WaitGroup)
	done := make(chan int)

	allocate := func() {
		start.Wait()
		for {
			select {
			case <-done:
				return
			default:
				pool.Get(rand.Intn(huge))
			}
		}
	}

	start.Add(1)
	for i := 0; i < 100; i++ {
		go expect.NoPanic(allocate)
	}
	start.Done()

	time.Sleep(2 * time.Second)
	close(done)
}

func TestBytePoolUnique(t *testing.T) {
	expect := ttesting.NewExpect(t)
	pool := NewBytePool()

	numTests := tinyCount * 10
	chunks := make([][]byte, numTests)

	for i := 0; i < numTests; i++ {
		data := pool.Get(8)
		binary.PutVarint(data, int64(i))
		chunks[i] = data
	}

	for i := 0; i < numTests; i++ {
		num, err := binary.ReadVarint(bytes.NewReader(chunks[i]))
		expect.NoError(err)
		expect.Equal(int64(i), num)
	}
}
<|endoftext|>


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

package gosnmpquerier

import (
	"time"

	"github.com/soniah/gosnmp"
)

type SyncQuerier interface {
	ExecuteQuery(query Query) Query
	Get(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
	GetNext(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
	Walk(destination, community, oid string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
}

type syncQuerier struct {
	Input        chan QueryWithOutputChannel
	asyncQuerier *AsyncQuerier
}

func NewSyncQuerier(contention, numErrors int, resetTime time.Duration) *syncQuerier {
	querier := syncQuerier{
		Input:        make(chan QueryWithOutputChannel),
		asyncQuerier: NewAsyncQuerier(contention, numErrors, resetTime),
	}
	go querier.processAndDispatchQueries()
	return &querier
}

func (querier *syncQuerier) ExecuteQuery(query Query) Query {
	output := make(chan Query)
	querier.Input <- QueryWithOutputChannel{query, output}
	return <-output
}

func (querier *syncQuerier) Get(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(GET, destination, community, oids, timeout, retries)
}

func (querier *syncQuerier) GetNext(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(GETNEXT, destination, community, oids, timeout, retries)
}

func (querier *syncQuerier) Walk(destination, community, oid string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(WALK, destination, community, []string{oid}, timeout, retries)
}

func (querier *syncQuerier) executeCommand(command OpSnmp, destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	query := querier.makeQuery(command, destination, community, oids, timeout, retries)
	processedQuery := querier.ExecuteQuery(query)
	return processedQuery.Response, processedQuery.Error
}

func (querier *syncQuerier) makeQuery(command OpSnmp, destination, community string, oids []string, timeout time.Duration, retries int) Query {
	return Query{
		Cmd:         command,
		Community:   community,
		Oids:        oids,
		Timeout:     timeout,
		Retries:     retries,
		Destination: destination,
	}
}

func (querier *syncQuerier) processAndDispatchQueries() {
	m := make(map[int]chan Query)
	i := 0
	for {
		select {
		case queryWithOutputChannel := <-querier.Input:
			queryWithOutputChannel.query.Id = i
			i += 1
			m[queryWithOutputChannel.query.Id] = queryWithOutputChannel.responseChannel
			querier.asyncQuerier.Input <- queryWithOutputChannel.query
		case processedQuery := <-querier.asyncQuerier.Output:
			m[processedQuery.Id] <- processedQuery
			delete(m, processedQuery.Id)
		}
	}
}
<commit_msg>Drop queries when the channel is full<commit_after>// Copyright 2014 The GoSNMPQuerier Authors. All rights reserved.  Use of this
// source code is governed by a MIT-style license that can be found in the
// LICENSE file.

package gosnmpquerier

import (
	"errors"
	"time"

	"github.com/soniah/gosnmp"
)

type SyncQuerier interface {
	ExecuteQuery(query Query) Query
	Get(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
	GetNext(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
	Walk(destination, community, oid string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error)
}

type syncQuerier struct {
	Input        chan QueryWithOutputChannel
	asyncQuerier *AsyncQuerier
}

func NewSyncQuerier(contention, numErrors int, resetTime time.Duration) *syncQuerier {
	querier := syncQuerier{
		Input:        make(chan QueryWithOutputChannel),
		asyncQuerier: NewAsyncQuerier(contention, numErrors, resetTime),
	}
	go querier.processAndDispatchQueries()
	return &querier
}

func (querier *syncQuerier) ExecuteQuery(query Query) Query {
	output := make(chan Query)

	select {
	case querier.Input <- QueryWithOutputChannel{query, output}:
	case <-time.After(1 * time.Millisecond):
		query.Error = errors.New("ExecuteQuery error to much queued queries")
		return query
	}
	return <-output
}

func (querier *syncQuerier) Get(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(GET, destination, community, oids, timeout, retries)
}

func (querier *syncQuerier) GetNext(destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(GETNEXT, destination, community, oids, timeout, retries)
}

func (querier *syncQuerier) Walk(destination, community, oid string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	return querier.executeCommand(WALK, destination, community, []string{oid}, timeout, retries)
}

func (querier *syncQuerier) executeCommand(command OpSnmp, destination, community string, oids []string, timeout time.Duration, retries int) ([]gosnmp.SnmpPDU, error) {
	query := querier.makeQuery(command, destination, community, oids, timeout, retries)
	processedQuery := querier.ExecuteQuery(query)
	return processedQuery.Response, processedQuery.Error
}

func (querier *syncQuerier) makeQuery(command OpSnmp, destination, community string, oids []string, timeout time.Duration, retries int) Query {
	return Query{
		Cmd:         command,
		Community:   community,
		Oids:        oids,
		Timeout:     timeout,
		Retries:     retries,
		Destination: destination,
	}
}

func (querier *syncQuerier) processAndDispatchQueries() {
	m := make(map[int]chan Query)
	i := 0
	for {
		select {
		case queryWithOutputChannel := <-querier.Input:
			queryWithOutputChannel.query.Id = i
			i += 1
			m[queryWithOutputChannel.query.Id] = queryWithOutputChannel.responseChannel
			querier.asyncQuerier.Input <- queryWithOutputChannel.query
		case processedQuery := <-querier.asyncQuerier.Output:
			m[processedQuery.Id] <- processedQuery
			delete(m, processedQuery.Id)
		}
	}
}
<|endoftext|>


<commit_before>package table

import (
	"bytes"
	"fmt"
)

var (
	ST_STANDARD = NewStyle("Standard")
)

type Style struct {
	handle int
	Name string // 2
	FontName string // 3
	BigFontName string // 4
}

func NewStyle(name string) *Style {
	st := new(Style)
	st.Name = name
	st.FontName = "arial.ttf"
	return st
}

func (st *Style) IsSymbolTable() bool {
	return true
}

func (st *Style) String() string {
	var otp bytes.Buffer
	otp.WriteString("0\nSTYLE\n")
	otp.WriteString(fmt.Sprintf("5\n%x\n", st.handle))
	otp.WriteString("100\nAcDbSymbostableRecord\n100\nAcDbTextStyleTableRecord\n")
	otp.WriteString(fmt.Sprintf("2\n%s\n", st.Name))
	otp.WriteString("70\n0\n")
	otp.WriteString("40\n0.0\n")
	otp.WriteString("41\n1.0\n")
	otp.WriteString("50\n0.0\n")
	otp.WriteString("71\n0\n")
	otp.WriteString("42\n0.2\n")
	otp.WriteString(fmt.Sprintf("3\n%s\n", st.FontName))
	otp.WriteString(fmt.Sprintf("4\n%s\n", st.BigFontName))
	otp.WriteString("1001\nACAD\n")
	otp.WriteString("1000\nArial\n")
	otp.WriteString("1071\n34\n")
	return otp.String()
}

func (st *Style) Handle() int {
	return st.handle
}
func (st *Style) SetHandle(v *int) {
	st.handle = *v
	(*v)++
}

<commit_msg>change Style: delete 1001, 1000, 1071 flags<commit_after>package table

import (
	"bytes"
	"fmt"
)

var (
	ST_STANDARD = NewStyle("Standard")
)

type Style struct {
	handle int
	Name string // 2
	FontName string // 3
	BigFontName string // 4
}

func NewStyle(name string) *Style {
	st := new(Style)
	st.Name = name
	st.FontName = "arial.ttf"
	return st
}

func (st *Style) IsSymbolTable() bool {
	return true
}

func (st *Style) String() string {
	var otp bytes.Buffer
	otp.WriteString("0\nSTYLE\n")
	otp.WriteString(fmt.Sprintf("5\n%x\n", st.handle))
	otp.WriteString("100\nAcDbSymbostableRecord\n100\nAcDbTextStyleTableRecord\n")
	otp.WriteString(fmt.Sprintf("2\n%s\n", st.Name))
	otp.WriteString("70\n0\n")
	otp.WriteString("40\n0.0\n")
	otp.WriteString("41\n1.0\n")
	otp.WriteString("50\n0.0\n")
	otp.WriteString("71\n0\n")
	otp.WriteString("42\n0.2\n")
	otp.WriteString(fmt.Sprintf("3\n%s\n", st.FontName))
	otp.WriteString(fmt.Sprintf("4\n%s\n", st.BigFontName))
	return otp.String()
}

func (st *Style) Handle() int {
	return st.handle
}
func (st *Style) SetHandle(v *int) {
	st.handle = *v
	(*v)++
}

<|endoftext|>