<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"sort"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/utils/set"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/process"
)

var logger = loggo.GetLogger("juju.process.context")

// TODO(ericsnow) Normalize method names across all the arch. layers
// (e.g. List->AllProcesses, Get->Processes, Set->AddProcess).

// APIClient represents the API needs of a Context.
type APIClient interface {
	// List requests the list of registered process IDs from state.
	List() ([]string, error)
	// Get requests the process info for the given ID.
	Get(ids ...string) ([]*process.Info, error)
	// Set sends a request to update state with the provided processes.
	Set(procs ...*process.Info) error
	// AllDefinitions returns the process definitions found in the
	// unit's metadata.
	AllDefinitions() ([]charm.Process, error)
}

// TODO(ericsnow) Rename Get and Set to more specifically describe what
// they are for.

// omponent provides the hook context data specific to workload processes.
type Component interface {
	// Get returns the process info corresponding to the given ID.
	Get(procName string) (*process.Info, error)
	// Set records the process info in the hook context.
	Set(procName string, info *process.Info) error
	// List returns the list of registered process IDs.
	List() ([]string, error)
	// ListDefinitions returns the charm-defined processes.
	ListDefinitions() ([]charm.Process, error)
	// Flush pushes the hook context data out to state.
	Flush() error
}

// Context is the workload process portion of the hook context.
type Context struct {
	api       APIClient
	processes map[string]*process.Info
	updates   map[string]*process.Info
	ids       set.Strings
}

// NewContext returns a new jujuc.ContextComponent for workload processes.
func NewContext(api APIClient, procs ...*process.Info) *Context {
	ids := set.NewStrings()
	processes := make(map[string]*process.Info)
	for _, proc := range procs {
		processes[proc.Name] = proc
		ids.Add(proc.ID())
	}
	return &Context{
		processes: processes,
		api:       api,
		ids:       ids,
	}
}

// NewContextAPI returns a new jujuc.ContextComponent for workload processes.
func NewContextAPI(api APIClient) (*Context, error) {
	ids, err := api.List()
	if err != nil {
		return nil, errors.Trace(err)
	}

	ctx := NewContext(api)
	for _, id := range ids {
		ctx.processes[id] = nil
		ctx.ids.Add(id)
	}
	return ctx, nil
}

// HookContext is the portion of jujuc.Context used in this package.
type HookContext interface {
	// Component implements jujuc.Context.
	Component(string) (Component, error)
}

// ContextComponent returns the hook context for the workload
// process component.
func ContextComponent(ctx HookContext) (Component, error) {
	compCtx, err := ctx.Component(process.ComponentName)
	if errors.IsNotFound(err) {
		return nil, errors.Errorf("component %q not registered", process.ComponentName)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	if compCtx == nil {
		return nil, errors.Errorf("component %q disabled", process.ComponentName)
	}
	return compCtx, nil
}

func (c *Context) addProc(id string, original *process.Info) error {
	var proc *process.Info
	if original != nil {
		info := *original
		info.Name = id
		proc = &info
	}
	if _, ok := c.processes[id]; !ok {
		c.processes[id] = proc
	} else {
		if proc == nil {
			return errors.Errorf("update can't be nil")
		}
		c.set(id, proc)
	}
	return nil
}

// Processes returns the processes known to the context.
func (c *Context) Processes() ([]*process.Info, error) {
	var procs []*process.Info
	for id, info := range mergeProcMaps(c.processes, c.updates) {
		if info == nil {
			fetched, err := c.api.Get(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			info = fetched[0]
			c.processes[id] = info
		}
		procs = append(procs, info)
	}
	return procs, nil
}

func mergeProcMaps(procs, updates map[string]*process.Info) map[string]*process.Info {
	// At this point procs and updates have already been checked for
	// nil values so we won't see any here.
	result := make(map[string]*process.Info)
	for k, v := range procs {
		result[k] = v
	}
	for k, v := range updates {
		result[k] = v
	}
	return result
}

// TODO(ericsnow) Should be build in refreshes?

// Get returns the process info corresponding to the given ID.
func (c *Context) Get(procName string) (*process.Info, error) {
	actual, ok := c.updates[procName]
	if !ok {
		actual, ok = c.processes[procName]
		if !ok {
			return nil, errors.NotFoundf("%s", procName)
		}
	}
	if actual == nil {
		fetched, err := c.api.Get(procName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		actual = fetched[0]
		c.processes[procName] = actual
	}
	return actual, nil
}

// List returns the names of all registered processes.
func (c *Context) List() ([]string, error) {
	ids := make([]string, len(c.ids))
	copy(ids, c.ids.Values())
	sort.Strings(ids)
	return ids, nil
}

// Set records the process info in the hook context.
func (c *Context) Set(procName string, info *process.Info) error {
	if procName != info.Name {
		return errors.Errorf("mismatch on name: %s != %s", procName, info.Name)
	}
	// TODO(ericsnow) We are likely missing mechanisim for local persistence.

	c.set(procName, info)
	return nil
}

func (c *Context) set(id string, pInfo *process.Info) {
	if c.updates == nil {
		c.updates = make(map[string]*process.Info)
	}
	var info process.Info
	info = *pInfo
	c.updates[id] = &info
	c.ids.Add(id)
}

// ListDefinitions returns the unit's charm-defined processes.
func (c *Context) ListDefinitions() ([]charm.Process, error) {
	definitions, err := c.api.AllDefinitions()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return definitions, nil
}

// TODO(ericsnow) The context machinery is not actually using this yet.

// Flush implements jujuc.ContextComponent. In this case that means all
// added and updated process.Info in the hook context are pushed to
// Juju state via the API.
func (c *Context) Flush() error {
	if len(c.updates) == 0 {
		return nil
	}

	var updates []*process.Info
	for _, info := range c.updates {
		updates = append(updates, info)
	}
	if err := c.api.Set(updates...); err != nil {
		return errors.Trace(err)
	}

	for k, v := range c.updates {
		c.processes[k] = v
	}
	c.updates = nil
	return nil
}
<commit_msg>Get rid of Context.ids.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"sort"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/process"
)

var logger = loggo.GetLogger("juju.process.context")

// TODO(ericsnow) Normalize method names across all the arch. layers
// (e.g. List->AllProcesses, Get->Processes, Set->AddProcess).

// APIClient represents the API needs of a Context.
type APIClient interface {
	// List requests the list of registered process IDs from state.
	List() ([]string, error)
	// Get requests the process info for the given ID.
	Get(ids ...string) ([]*process.Info, error)
	// Set sends a request to update state with the provided processes.
	Set(procs ...*process.Info) error
	// AllDefinitions returns the process definitions found in the
	// unit's metadata.
	AllDefinitions() ([]charm.Process, error)
}

// TODO(ericsnow) Rename Get and Set to more specifically describe what
// they are for.

// omponent provides the hook context data specific to workload processes.
type Component interface {
	// Get returns the process info corresponding to the given ID.
	Get(procName string) (*process.Info, error)
	// Set records the process info in the hook context.
	Set(procName string, info *process.Info) error
	// List returns the list of registered process IDs.
	List() ([]string, error)
	// ListDefinitions returns the charm-defined processes.
	ListDefinitions() ([]charm.Process, error)
	// Flush pushes the hook context data out to state.
	Flush() error
}

// Context is the workload process portion of the hook context.
type Context struct {
	api       APIClient
	processes map[string]*process.Info
	updates   map[string]*process.Info
}

// NewContext returns a new jujuc.ContextComponent for workload processes.
func NewContext(api APIClient, procs ...*process.Info) *Context {
	processes := make(map[string]*process.Info)
	for _, proc := range procs {
		processes[proc.Name] = proc
	}
	return &Context{
		processes: processes,
		api:       api,
	}
}

// NewContextAPI returns a new jujuc.ContextComponent for workload processes.
func NewContextAPI(api APIClient) (*Context, error) {
	ids, err := api.List()
	if err != nil {
		return nil, errors.Trace(err)
	}

	ctx := NewContext(api)
	for _, id := range ids {
		ctx.processes[id] = nil
	}
	return ctx, nil
}

// HookContext is the portion of jujuc.Context used in this package.
type HookContext interface {
	// Component implements jujuc.Context.
	Component(string) (Component, error)
}

// ContextComponent returns the hook context for the workload
// process component.
func ContextComponent(ctx HookContext) (Component, error) {
	compCtx, err := ctx.Component(process.ComponentName)
	if errors.IsNotFound(err) {
		return nil, errors.Errorf("component %q not registered", process.ComponentName)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	if compCtx == nil {
		return nil, errors.Errorf("component %q disabled", process.ComponentName)
	}
	return compCtx, nil
}

func (c *Context) addProc(id string, original *process.Info) error {
	var proc *process.Info
	if original != nil {
		info := *original
		info.Name = id
		proc = &info
	}
	if _, ok := c.processes[id]; !ok {
		c.processes[id] = proc
	} else {
		if proc == nil {
			return errors.Errorf("update can't be nil")
		}
		c.set(id, proc)
	}
	return nil
}

// Processes returns the processes known to the context.
func (c *Context) Processes() ([]*process.Info, error) {
	var procs []*process.Info
	for id, info := range mergeProcMaps(c.processes, c.updates) {
		if info == nil {
			fetched, err := c.api.Get(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			info = fetched[0]
			c.processes[id] = info
		}
		procs = append(procs, info)
	}
	return procs, nil
}

func mergeProcMaps(procs, updates map[string]*process.Info) map[string]*process.Info {
	// At this point procs and updates have already been checked for
	// nil values so we won't see any here.
	result := make(map[string]*process.Info)
	for k, v := range procs {
		result[k] = v
	}
	for k, v := range updates {
		result[k] = v
	}
	return result
}

// TODO(ericsnow) Should be build in refreshes?

// Get returns the process info corresponding to the given ID.
func (c *Context) Get(procName string) (*process.Info, error) {
	actual, ok := c.updates[procName]
	if !ok {
		actual, ok = c.processes[procName]
		if !ok {
			return nil, errors.NotFoundf("%s", procName)
		}
	}
	if actual == nil {
		fetched, err := c.api.Get(procName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		actual = fetched[0]
		c.processes[procName] = actual
	}
	return actual, nil
}

// List returns the names of all registered processes.
func (c *Context) List() ([]string, error) {
	procs, err := c.Processes()
	if err != nil {
		return nil, errors.Trace(err)
	}
	if len(procs) == 0 {
		return nil, nil
	}
	var ids []string
	for _, proc := range procs {
		ids = append(ids, proc.ID())
	}
	sort.Strings(ids)
	return ids, nil
}

// Set records the process info in the hook context.
func (c *Context) Set(procName string, info *process.Info) error {
	if procName != info.Name {
		return errors.Errorf("mismatch on name: %s != %s", procName, info.Name)
	}
	// TODO(ericsnow) We are likely missing mechanisim for local persistence.

	c.set(procName, info)
	return nil
}

func (c *Context) set(id string, pInfo *process.Info) {
	if c.updates == nil {
		c.updates = make(map[string]*process.Info)
	}
	var info process.Info
	info = *pInfo
	c.updates[id] = &info
}

// ListDefinitions returns the unit's charm-defined processes.
func (c *Context) ListDefinitions() ([]charm.Process, error) {
	definitions, err := c.api.AllDefinitions()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return definitions, nil
}

// TODO(ericsnow) The context machinery is not actually using this yet.

// Flush implements jujuc.ContextComponent. In this case that means all
// added and updated process.Info in the hook context are pushed to
// Juju state via the API.
func (c *Context) Flush() error {
	if len(c.updates) == 0 {
		return nil
	}

	var updates []*process.Info
	for _, info := range c.updates {
		updates = append(updates, info)
	}
	if err := c.api.Set(updates...); err != nil {
		return errors.Trace(err)
	}

	for k, v := range c.updates {
		c.processes[k] = v
	}
	c.updates = nil
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"github.com/juju/errors"

	"github.com/juju/juju/process"
)

// APIClient represents the API needs of a Context.
type APIClient interface {
	// List requests the list of registered process IDs from state.
	List() ([]string, error)
	// Get requests the process info for the given ID.
	Get(ids ...string) ([]*process.Info, error)
	// Set sends a request to update state with the provided processes.
	Set(procs ...*process.Info) error
}

// omponent provides the hook context data specific to workload processes.
type Component interface {
	// Get returns the process info corresponding to the given ID.
	Get(procName string) (*process.Info, error)
	// Set records the process info in the hook context.
	Set(procName string, info *process.Info) error
}

// Context is the workload process portion of the hook context.
type Context struct {
	api       APIClient
	processes map[string]*process.Info
	updates   map[string]*process.Info
}

// NewContext returns a new jujuc.ContextComponent for workload processes.
func NewContext(api APIClient, procs ...*process.Info) *Context {
	processes := make(map[string]*process.Info)
	for _, proc := range procs {
		processes[proc.Name] = proc
	}
	return &Context{
		processes: processes,
		api:       api,
	}
}

// NewContextAPI returns a new jujuc.ContextComponent for workload processes.
func NewContextAPI(api APIClient) (*Context, error) {
	ids, err := api.List()
	if err != nil {
		return nil, errors.Trace(err)
	}

	ctx := NewContext(api)
	for _, id := range ids {
		ctx.processes[id] = nil
	}
	return ctx, nil
}

// HookContext is the portion of jujuc.Context used in this package.
type HookContext interface {
	// Component implements jujuc.Context.
	Component(string) (Component, error)
}

// ContextComponent returns the hook context for the workload
// process component.
func ContextComponent(ctx HookContext) (Component, error) {
	compCtx, err := ctx.Component(process.ComponentName)
	if errors.IsNotFound(err) {
		return nil, errors.Errorf("component %q not registered", process.ComponentName)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	if compCtx == nil {
		return nil, errors.Errorf("component %q disabled", process.ComponentName)
	}
	return compCtx, nil
}

func (c *Context) addProc(id string, original *process.Info) error {
	var proc *process.Info
	if original != nil {
		info := *original
		info.Name = id
		proc = &info
	}
	if _, ok := c.processes[id]; !ok {
		c.processes[id] = proc
	} else {
		if proc == nil {
			return errors.Errorf("update can't be nil")
		}
		c.set(id, proc)
	}
	return nil
}

// Processes returns the processes known to the context.
func (c *Context) Processes() ([]*process.Info, error) {
	var procs []*process.Info
	for id, info := range mergeProcMaps(c.processes, c.updates) {
		if info == nil {
			fetched, err := c.api.Get(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			info = fetched[0]
			c.processes[id] = info
		}
		procs = append(procs, info)
	}
	return procs, nil
}

func mergeProcMaps(procs, updates map[string]*process.Info) map[string]*process.Info {
	// At this point procs and updates have already been checked for
	// nil values so we won't see any here.
	result := make(map[string]*process.Info)
	for k, v := range procs {
		result[k] = v
	}
	for k, v := range updates {
		result[k] = v
	}
	return result
}

// Get returns the process info corresponding to the given ID.
func (c *Context) Get(procName string) (*process.Info, error) {
	actual, ok := c.updates[procName]
	if !ok {
		actual, ok = c.processes[procName]
		if !ok {
			return nil, errors.NotFoundf("%s", procName)
		}
	}
	if actual == nil {
		fetched, err := c.api.Get(procName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		actual = fetched[0]
		c.processes[procName] = actual
	}
	return actual, nil
}

// Set records the process info in the hook context.
func (c *Context) Set(procName string, info *process.Info) error {
	if procName != info.Name {
		return errors.Errorf("mismatch on name: %s != %s", procName, info.Name)
	}
	// TODO(ericsnow) We are likely missing mechanisim for local persistence.

	c.set(procName, info)
	return nil
}

func (c *Context) set(id string, pInfo *process.Info) {
	if c.updates == nil {
		c.updates = make(map[string]*process.Info)
	}
	var info process.Info
	info = *pInfo
	c.updates[id] = &info
}

// Flush implements jujuc.ContextComponent. In this case that means all
// added and updated process.Info in the hook context are pushed to
// Juju state via the API.
func (c *Context) Flush() error {
	if len(c.updates) == 0 {
		return nil
	}

	var updates []*process.Info
	for _, info := range c.updates {
		updates = append(updates, info)
	}
	if err := c.api.Set(updates...); err != nil {
		return errors.Trace(err)
	}

	for k, v := range c.updates {
		c.processes[k] = v
	}
	c.updates = nil
	return nil
}
<commit_msg>Add a TODO.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"github.com/juju/errors"

	"github.com/juju/juju/process"
)

// APIClient represents the API needs of a Context.
type APIClient interface {
	// List requests the list of registered process IDs from state.
	List() ([]string, error)
	// Get requests the process info for the given ID.
	Get(ids ...string) ([]*process.Info, error)
	// Set sends a request to update state with the provided processes.
	Set(procs ...*process.Info) error
}

// omponent provides the hook context data specific to workload processes.
type Component interface {
	// Get returns the process info corresponding to the given ID.
	Get(procName string) (*process.Info, error)
	// Set records the process info in the hook context.
	Set(procName string, info *process.Info) error
}

// Context is the workload process portion of the hook context.
type Context struct {
	api       APIClient
	processes map[string]*process.Info
	updates   map[string]*process.Info
}

// NewContext returns a new jujuc.ContextComponent for workload processes.
func NewContext(api APIClient, procs ...*process.Info) *Context {
	processes := make(map[string]*process.Info)
	for _, proc := range procs {
		processes[proc.Name] = proc
	}
	return &Context{
		processes: processes,
		api:       api,
	}
}

// NewContextAPI returns a new jujuc.ContextComponent for workload processes.
func NewContextAPI(api APIClient) (*Context, error) {
	ids, err := api.List()
	if err != nil {
		return nil, errors.Trace(err)
	}

	ctx := NewContext(api)
	for _, id := range ids {
		ctx.processes[id] = nil
	}
	return ctx, nil
}

// HookContext is the portion of jujuc.Context used in this package.
type HookContext interface {
	// Component implements jujuc.Context.
	Component(string) (Component, error)
}

// ContextComponent returns the hook context for the workload
// process component.
func ContextComponent(ctx HookContext) (Component, error) {
	compCtx, err := ctx.Component(process.ComponentName)
	if errors.IsNotFound(err) {
		return nil, errors.Errorf("component %q not registered", process.ComponentName)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	if compCtx == nil {
		return nil, errors.Errorf("component %q disabled", process.ComponentName)
	}
	return compCtx, nil
}

func (c *Context) addProc(id string, original *process.Info) error {
	var proc *process.Info
	if original != nil {
		info := *original
		info.Name = id
		proc = &info
	}
	if _, ok := c.processes[id]; !ok {
		c.processes[id] = proc
	} else {
		if proc == nil {
			return errors.Errorf("update can't be nil")
		}
		c.set(id, proc)
	}
	return nil
}

// Processes returns the processes known to the context.
func (c *Context) Processes() ([]*process.Info, error) {
	var procs []*process.Info
	for id, info := range mergeProcMaps(c.processes, c.updates) {
		if info == nil {
			fetched, err := c.api.Get(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			info = fetched[0]
			c.processes[id] = info
		}
		procs = append(procs, info)
	}
	return procs, nil
}

func mergeProcMaps(procs, updates map[string]*process.Info) map[string]*process.Info {
	// At this point procs and updates have already been checked for
	// nil values so we won't see any here.
	result := make(map[string]*process.Info)
	for k, v := range procs {
		result[k] = v
	}
	for k, v := range updates {
		result[k] = v
	}
	return result
}

// Get returns the process info corresponding to the given ID.
func (c *Context) Get(procName string) (*process.Info, error) {
	actual, ok := c.updates[procName]
	if !ok {
		actual, ok = c.processes[procName]
		if !ok {
			return nil, errors.NotFoundf("%s", procName)
		}
	}
	if actual == nil {
		fetched, err := c.api.Get(procName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		actual = fetched[0]
		c.processes[procName] = actual
	}
	return actual, nil
}

// Set records the process info in the hook context.
func (c *Context) Set(procName string, info *process.Info) error {
	if procName != info.Name {
		return errors.Errorf("mismatch on name: %s != %s", procName, info.Name)
	}
	// TODO(ericsnow) We are likely missing mechanisim for local persistence.

	c.set(procName, info)
	return nil
}

func (c *Context) set(id string, pInfo *process.Info) {
	if c.updates == nil {
		c.updates = make(map[string]*process.Info)
	}
	var info process.Info
	info = *pInfo
	c.updates[id] = &info
}

// TODO(ericsnow) The context machinery is not actually using this yet.

// Flush implements jujuc.ContextComponent. In this case that means all
// added and updated process.Info in the hook context are pushed to
// Juju state via the API.
func (c *Context) Flush() error {
	if len(c.updates) == 0 {
		return nil
	}

	var updates []*process.Info
	for _, info := range c.updates {
		updates = append(updates, info)
	}
	if err := c.api.Set(updates...); err != nil {
		return errors.Trace(err)
	}

	for k, v := range c.updates {
		c.processes[k] = v
	}
	c.updates = nil
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"time"

	humanize "github.com/dustin/go-humanize"

	"github.com/golang/protobuf/proto"
	"github.com/luci/luci-go/client/isolate"
	"github.com/luci/luci-go/common/eventlog"
	logpb "github.com/luci/luci-go/common/eventlog/proto"
	"github.com/luci/luci-go/common/isolated"
	"github.com/luci/luci-go/common/isolatedclient"
	"github.com/maruel/subcommands"
)

const (
	// archiveThreshold is the size (in bytes) used to determine whether to add
	// files to a tar archive before uploading. Files smaller than this size will
	// be combined into archives before being uploaded to the server.
	archiveThreshold = 100e3 // 100kB

	// archiveSizeTrigger is the desired size of the created archives. Once
	// archives reach this size, they will be closed and prepared for upload.
	archiveSizeTrigger = 10e6
)

var cmdExpArchive = &subcommands.Command{
	UsageLine: "exparchive <options>",
	ShortDesc: "EXPERIMENTAL parses a .isolate file to create a .isolated file, and uploads it and all referenced files to an isolate server",
	LongDesc:  "All the files listed in the .isolated file are put in the isolate server cache. Small files are combined together in a tar archive before uploading.",
	CommandRun: func() subcommands.CommandRun {
		c := &expArchiveRun{}
		c.commonServerFlags.Init()
		c.isolateFlags.Init(&c.Flags)
		return c
	},
}

// expArchiveRun contains the logic for the experimental archive subcommand.
// It implements subcommand.CommandRun
type expArchiveRun struct {
	commonServerFlags // Provides the GetFlags method.
	isolateFlags      isolateFlags
}

// Item represents a file or symlink referenced by an isolate file.
type Item struct {
	Path    string
	RelPath string
	Size    int64
	Mode    os.FileMode

	Digest isolated.HexDigest
}

// main contains the core logic for experimental archive.
func (c *expArchiveRun) main() error {
	start := time.Now()
	archiveOpts := &c.isolateFlags.ArchiveOptions
	// Parse the incoming isolate file.
	deps, rootDir, isol, err := isolate.ProcessIsolate(archiveOpts)
	if err != nil {
		return fmt.Errorf("failed to process isolate: %v", err)
	}
	log.Printf("Isolate referenced %d deps", len(deps))

	// Create the isolated client which connects to the isolate server.
	authCl, err := c.createAuthClient()
	if err != nil {
		return err
	}
	client := isolatedclient.New(nil, authCl, c.isolatedFlags.ServerURL, c.isolatedFlags.Namespace, nil)

	checker := NewChecker(client)

	// Walk each of the deps, partioning the results into symlinks and files categorised by size.
	var links, archiveFiles, indivFiles []*Item
	var archiveSize, indivSize int64 // Cumulative size of archived/individual files.
	for _, dep := range deps {
		// Try to walk dep. If dep is a file (or symlink), the inner function is called exactly once.
		err := filepath.Walk(filepath.Clean(dep), func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				return nil
			}

			relPath, err := filepath.Rel(rootDir, path)
			if err != nil {
				return err
			}

			item := &Item{
				Path:    path,
				RelPath: relPath,
				Mode:    info.Mode(),
				Size:    info.Size(),
			}

			switch {
			case item.Mode&os.ModeSymlink == os.ModeSymlink:
				links = append(links, item)
			case item.Size < archiveThreshold:
				archiveFiles = append(archiveFiles, item)
				archiveSize += item.Size
			default:
				indivFiles = append(indivFiles, item)
				indivSize += item.Size
			}
			return nil
		})
		if err != nil {
			return err
		}
	}

	// Construct a map of the files that constitute the isolate.
	files := make(map[string]isolated.File)

	log.Printf("Isolate expanded to %d files (total size %s) and %d symlinks", len(archiveFiles)+len(indivFiles), humanize.Bytes(uint64(archiveSize+indivSize)), len(links))
	log.Printf("\t%d files (%s) to be isolated individually", len(indivFiles), humanize.Bytes(uint64(indivSize)))
	log.Printf("\t%d files (%s) to be isolated in archives", len(archiveFiles), humanize.Bytes(uint64(archiveSize)))

	// Handle the symlinks.
	for _, item := range links {
		l, err := os.Readlink(item.Path)
		if err != nil {
			return fmt.Errorf("unable to resolve symlink for %q: %v", item.Path, err)
		}
		files[item.RelPath] = isolated.SymLink(l)
	}

	// Handle the small to-be-archived files.
	archiveCallback := func(tar *Tar) {
		log.Printf("Created tar archive %q (%s)", tar.Digest, humanize.Bytes(uint64(len(tar.Content))))
		log.Printf("\tcontains %d files (total %s)", tar.FileCount, humanize.Bytes(uint64(tar.FileSize)))
		// Mint an item for this tar.
		item := &Item{
			Path:    fmt.Sprintf(".%s.tar", tar.Digest),
			RelPath: fmt.Sprintf(".%s.tar", tar.Digest),
			Size:    int64(len(tar.Content)),
			Mode:    0644, // Read
			Digest:  tar.Digest,
		}
		files[item.RelPath] = isolated.TarFile(item.Digest, int(item.Mode), item.Size)
		checker.AddItem(item, false, func(item *Item, ps *isolatedclient.PushState) {
			if ps == nil {
				tar.Release()
				return
			}
			// XXX(djd): Actually upload the tar file.
			log.Printf("XXX: %q needs to be uploaded", item.RelPath)
			tar.Release()
		})
	}

	archiver := NewTarAchiver(archiveCallback)
	for _, item := range archiveFiles {
		if err := archiver.AddItem(item); err != nil {
			return err
		}
	}
	archiver.Flush()

	// Handle the large individually-uploaded files.
	for _, item := range indivFiles {
		d, err := hashFile(item.Path)
		if err != nil {
			return err
		}
		item.Digest = d
		files[item.RelPath] = isolated.BasicFile(item.Digest, int(item.Mode), item.Size)
		checker.AddItem(item, false, func(item *Item, ps *isolatedclient.PushState) {
			if ps == nil {
				return
			}
			// XXX(djd): Actually upload the file.
			log.Printf("XXX: %q needs to be uploaded", item.RelPath)
		})
	}

	// Marshal the isolated file into JSON.
	isol.Files = files
	isolJSON, err := json.Marshal(isol)
	if err != nil {
		return err
	}
	// TODO(djd): actually check/upload the isolated.

	// Make sure that all pending items have been checked.
	if err := checker.Close(); err != nil {
		return err
	}

	// Write the isolated file, and emit its digest to stdout.
	if err := ioutil.WriteFile(archiveOpts.Isolated, isolJSON, 0644); err != nil {
		return err
	}
	isolatedHash := string(isolated.HashBytes(isolJSON))
	fmt.Printf("%s\t%s\n", isolatedHash, filepath.Base(archiveOpts.Isolated))

	end := time.Now()

	if endpoint := eventlogEndpoint(c.isolateFlags.EventlogEndpoint); endpoint != "" {
		ctx := context.Background()
		logger := eventlog.NewClient(ctx, endpoint)

		// TODO(mcgreevy): fill out more stats in archiveDetails.
		archiveDetails := &logpb.IsolateClientEvent_ArchiveDetails{
			IsolateHash: []string{isolatedHash},
		}
		if err := logStats(ctx, logger, start, end, archiveDetails); err != nil {
			log.Printf("Failed to log to eventlog: %v", err)
		}
	}

	return errors.New("experimental archive is not implemented")
}

func (c *expArchiveRun) parseFlags(args []string) error {
	if len(args) != 0 {
		return errors.New("position arguments not expected")
	}
	if err := c.commonServerFlags.Parse(); err != nil {
		return err
	}
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err := c.isolateFlags.Parse(cwd, RequireIsolateFile&RequireIsolatedFile); err != nil {
		return err
	}
	return nil
}

func (c *expArchiveRun) Run(a subcommands.Application, args []string, _ subcommands.Env) int {
	fmt.Fprintln(a.GetErr(), "WARNING: this command is experimental")
	if err := c.parseFlags(args); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	if len(c.isolateFlags.ArchiveOptions.Blacklist) != 0 {
		fmt.Fprintf(a.GetErr(), "%s: blacklist is not supported\n", a.GetName())
		return 1
	}
	if err := c.main(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

func hashFile(path string) (isolated.HexDigest, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return isolated.Hash(f)
}

func logStats(ctx context.Context, logger *eventlog.Client, start, end time.Time, archiveDetails *logpb.IsolateClientEvent_ArchiveDetails) error {
	event := logger.NewLogEvent(ctx, eventlog.Point())
	event.InfraEvent.IsolateClientEvent = &logpb.IsolateClientEvent{
		Binary: &logpb.Binary{
			Name:          proto.String("isolate"),
			VersionNumber: proto.String(version),
		},
		Operation:      logpb.IsolateClientEvent_ARCHIVE.Enum(),
		ArchiveDetails: archiveDetails,
		// TODO(mcgreevy): fill out Master, Builder, BuildId, Slave.
		StartTsUsec: proto.Int64(int64(start.UnixNano() / 1e3)),
		EndTsUsec:   proto.Int64(int64(end.UnixNano() / 1e3)),
	}
	return logger.LogSync(ctx, event)
}

func eventlogEndpoint(endpointFlag string) string {
	switch endpointFlag {
	case "test":
		return eventlog.TestEndpoint
	case "prod":
		return eventlog.ProdEndpoint
	default:
		return endpointFlag
	}
}
<commit_msg>client/isolate: wire up object uploading in exparchive<commit_after>// Copyright 2015 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"time"

	humanize "github.com/dustin/go-humanize"

	"github.com/golang/protobuf/proto"
	"github.com/luci/luci-go/client/isolate"
	"github.com/luci/luci-go/common/eventlog"
	logpb "github.com/luci/luci-go/common/eventlog/proto"
	"github.com/luci/luci-go/common/isolated"
	"github.com/luci/luci-go/common/isolatedclient"
	"github.com/maruel/subcommands"
)

const (
	// archiveThreshold is the size (in bytes) used to determine whether to add
	// files to a tar archive before uploading. Files smaller than this size will
	// be combined into archives before being uploaded to the server.
	archiveThreshold = 100e3 // 100kB

	// archiveSizeTrigger is the desired size of the created archives. Once
	// archives reach this size, they will be closed and prepared for upload.
	archiveSizeTrigger = 10e6
)

var cmdExpArchive = &subcommands.Command{
	UsageLine: "exparchive <options>",
	ShortDesc: "EXPERIMENTAL parses a .isolate file to create a .isolated file, and uploads it and all referenced files to an isolate server",
	LongDesc:  "All the files listed in the .isolated file are put in the isolate server cache. Small files are combined together in a tar archive before uploading.",
	CommandRun: func() subcommands.CommandRun {
		c := &expArchiveRun{}
		c.commonServerFlags.Init()
		c.isolateFlags.Init(&c.Flags)
		return c
	},
}

// expArchiveRun contains the logic for the experimental archive subcommand.
// It implements subcommand.CommandRun
type expArchiveRun struct {
	commonServerFlags // Provides the GetFlags method.
	isolateFlags      isolateFlags
}

// Item represents a file or symlink referenced by an isolate file.
type Item struct {
	Path    string
	RelPath string
	Size    int64
	Mode    os.FileMode

	Digest isolated.HexDigest
}

// main contains the core logic for experimental archive.
func (c *expArchiveRun) main() error {
	// TODO(djd): This func is long and has a lot of internal complexity (like,
	// such as, archiveCallback). Refactor.

	start := time.Now()
	archiveOpts := &c.isolateFlags.ArchiveOptions
	// Parse the incoming isolate file.
	deps, rootDir, isol, err := isolate.ProcessIsolate(archiveOpts)
	if err != nil {
		return fmt.Errorf("failed to process isolate: %v", err)
	}
	log.Printf("Isolate referenced %d deps", len(deps))

	// Set up a background context which is cancelled when this function returns.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create the isolated client which connects to the isolate server.
	authCl, err := c.createAuthClient()
	if err != nil {
		return err
	}
	client := isolatedclient.New(nil, authCl, c.isolatedFlags.ServerURL, c.isolatedFlags.Namespace, nil)

	// Set up a checker and pair of uploaders. One uploader is for in-memory
	// files, and one is for on-disk files (we want to limit the latter to only
	// one upload at a time).
	// TODO(djd): Make NewChecker take a context arg.
	checker := NewChecker(client)
	memUploader, fileUploader := NewUploader(ctx, client, 10), NewUploader(ctx, client, 1)

	// Walk each of the deps, partioning the results into symlinks and files categorised by size.
	var links, archiveFiles, indivFiles []*Item
	var archiveSize, indivSize int64 // Cumulative size of archived/individual files.
	for _, dep := range deps {
		// Try to walk dep. If dep is a file (or symlink), the inner function is called exactly once.
		err := filepath.Walk(filepath.Clean(dep), func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				return nil
			}

			relPath, err := filepath.Rel(rootDir, path)
			if err != nil {
				return err
			}

			item := &Item{
				Path:    path,
				RelPath: relPath,
				Mode:    info.Mode(),
				Size:    info.Size(),
			}

			switch {
			case item.Mode&os.ModeSymlink == os.ModeSymlink:
				links = append(links, item)
			case item.Size < archiveThreshold:
				archiveFiles = append(archiveFiles, item)
				archiveSize += item.Size
			default:
				indivFiles = append(indivFiles, item)
				indivSize += item.Size
			}
			return nil
		})
		if err != nil {
			return err
		}
	}

	// Construct a map of the files that constitute the isolate.
	files := make(map[string]isolated.File)

	log.Printf("Isolate expanded to %d files (total size %s) and %d symlinks", len(archiveFiles)+len(indivFiles), humanize.Bytes(uint64(archiveSize+indivSize)), len(links))
	log.Printf("\t%d files (%s) to be isolated individually", len(indivFiles), humanize.Bytes(uint64(indivSize)))
	log.Printf("\t%d files (%s) to be isolated in archives", len(archiveFiles), humanize.Bytes(uint64(archiveSize)))

	// Handle the symlinks.
	for _, item := range links {
		l, err := os.Readlink(item.Path)
		if err != nil {
			return fmt.Errorf("unable to resolve symlink for %q: %v", item.Path, err)
		}
		files[item.RelPath] = isolated.SymLink(l)
	}

	// Handle the small to-be-archived files.
	archiveCallback := func(tar *Tar) {
		log.Printf("Created tar archive %q (%s)", tar.Digest, humanize.Bytes(uint64(len(tar.Content))))
		log.Printf("\tcontains %d files (total %s)", tar.FileCount, humanize.Bytes(uint64(tar.FileSize)))
		// Mint an item for this tar.
		item := &Item{
			Path:    fmt.Sprintf(".%s.tar", tar.Digest),
			RelPath: fmt.Sprintf(".%s.tar", tar.Digest),
			Size:    int64(len(tar.Content)),
			Mode:    0644, // Read
			Digest:  tar.Digest,
		}
		files[item.RelPath] = isolated.TarFile(item.Digest, int(item.Mode), item.Size)
		checker.AddItem(item, false, func(item *Item, ps *isolatedclient.PushState) {
			if ps == nil {
				tar.Release()
				return
			}
			log.Printf("QUEUED %q for upload", item.RelPath)
			memUploader.UploadBytes(item.RelPath, tar.Content, ps, func() {
				log.Printf("UPLOADED %q", item.RelPath)
				tar.Release()
			})
		})
	}

	archiver := NewTarAchiver(archiveCallback)
	for _, item := range archiveFiles {
		if err := archiver.AddItem(item); err != nil {
			return err
		}
	}
	archiver.Flush()

	// Handle the large individually-uploaded files.
	for _, item := range indivFiles {
		d, err := hashFile(item.Path)
		if err != nil {
			return err
		}
		item.Digest = d
		files[item.RelPath] = isolated.BasicFile(item.Digest, int(item.Mode), item.Size)
		checker.AddItem(item, false, func(item *Item, ps *isolatedclient.PushState) {
			if ps == nil {
				return
			}
			log.Printf("QUEUED %q for upload", item.RelPath)
			fileUploader.UploadFile(item, ps, func() {
				log.Printf("UPLOADED %q", item.RelPath)
			})
		})
	}

	// Marshal the isolated file into JSON, and create an Item to describe it.
	isol.Files = files
	isolJSON, err := json.Marshal(isol)
	if err != nil {
		return err
	}
	isolItem := &Item{
		Path:    archiveOpts.Isolated,
		RelPath: filepath.Base(archiveOpts.Isolated),
		Digest:  isolated.HashBytes(isolJSON),
		Size:    int64(len(isolJSON)),
	}

	// Check and upload isolate JSON.
	checker.AddItem(isolItem, true, func(item *Item, ps *isolatedclient.PushState) {
		if ps == nil {
			return
		}
		log.Printf("QUEUED %q for upload", item.RelPath)
		memUploader.UploadBytes(item.RelPath, isolJSON, ps, func() {
			log.Printf("UPLOADED %q", item.RelPath)
		})
	})

	// Make sure that all pending items have been checked.
	if err := checker.Close(); err != nil {
		return err
	}

	// Make sure that all the uploads have completed successfully.
	if err := fileUploader.Close(); err != nil {
		return err
	}
	if err := memUploader.Close(); err != nil {
		return err
	}

	// Write the isolated file, and emit its digest to stdout.
	if err := ioutil.WriteFile(archiveOpts.Isolated, isolJSON, 0644); err != nil {
		return err
	}
	fmt.Printf("%s\t%s\n", isolItem.Digest, filepath.Base(archiveOpts.Isolated))

	end := time.Now()

	if endpoint := eventlogEndpoint(c.isolateFlags.EventlogEndpoint); endpoint != "" {
		logger := eventlog.NewClient(ctx, endpoint)

		// TODO(mcgreevy): fill out more stats in archiveDetails.
		archiveDetails := &logpb.IsolateClientEvent_ArchiveDetails{
			IsolateHash: []string{string(isolItem.Digest)},
		}
		if err := logStats(ctx, logger, start, end, archiveDetails); err != nil {
			log.Printf("Failed to log to eventlog: %v", err)
		}
	}

	return nil
}

func (c *expArchiveRun) parseFlags(args []string) error {
	if len(args) != 0 {
		return errors.New("position arguments not expected")
	}
	if err := c.commonServerFlags.Parse(); err != nil {
		return err
	}
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err := c.isolateFlags.Parse(cwd, RequireIsolateFile&RequireIsolatedFile); err != nil {
		return err
	}
	return nil
}

func (c *expArchiveRun) Run(a subcommands.Application, args []string, _ subcommands.Env) int {
	fmt.Fprintln(a.GetErr(), "WARNING: this command is experimental")
	if err := c.parseFlags(args); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	if len(c.isolateFlags.ArchiveOptions.Blacklist) != 0 {
		fmt.Fprintf(a.GetErr(), "%s: blacklist is not supported\n", a.GetName())
		return 1
	}
	if err := c.main(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

func hashFile(path string) (isolated.HexDigest, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return isolated.Hash(f)
}

func logStats(ctx context.Context, logger *eventlog.Client, start, end time.Time, archiveDetails *logpb.IsolateClientEvent_ArchiveDetails) error {
	event := logger.NewLogEvent(ctx, eventlog.Point())
	event.InfraEvent.IsolateClientEvent = &logpb.IsolateClientEvent{
		Binary: &logpb.Binary{
			Name:          proto.String("isolate"),
			VersionNumber: proto.String(version),
		},
		Operation:      logpb.IsolateClientEvent_ARCHIVE.Enum(),
		ArchiveDetails: archiveDetails,
		// TODO(mcgreevy): fill out Master, Builder, BuildId, Slave.
		StartTsUsec: proto.Int64(int64(start.UnixNano() / 1e3)),
		EndTsUsec:   proto.Int64(int64(end.UnixNano() / 1e3)),
	}
	return logger.LogSync(ctx, event)
}

func eventlogEndpoint(endpointFlag string) string {
	switch endpointFlag {
	case "test":
		return eventlog.TestEndpoint
	case "prod":
		return eventlog.ProdEndpoint
	default:
		return endpointFlag
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Authors of Cilium
//
// 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 main

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/pkg/api"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/fqdn"
	"github.com/cilium/cilium/pkg/fqdn/dnsproxy"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	policyApi "github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/proxy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/proxy/logger"
	"github.com/cilium/cilium/pkg/u8proto"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/strfmt"

	"github.com/miekg/dns"
)

// bootstrapFQDN initializes the toFQDNs related subsystems: DNSPoller,
// d.dnsRuleGen, and the DNS proxy.
// dnsRuleGen and DNSPoller will use the default resolver and, implicitly, the
// default DNS cache. The proxy binds to all interfaces, and uses the
// configured DNS proxy port (this may be 0 and so OS-assigned).
func (d *Daemon) bootstrapFQDN(restoredEndpoints *endpointRestoreState) (err error) {
	cfg := fqdn.Config{
		MinTTL:         option.Config.ToFQDNsMinTTL,
		Cache:          fqdn.DefaultDNSCache,
		LookupDNSNames: fqdn.DNSLookupDefaultResolver,
		AddGeneratedRules: func(generatedRules []*policyApi.Rule) error {
			// Insert the new rules into the policy repository. We need them to
			// replace the previous set. This requires the labels to match (including
			// the ToFQDN-UUID one).
			_, err := d.PolicyAdd(generatedRules, &AddOptions{Replace: true, Generated: true})
			return err
		},
		PollerResponseNotify: func(lookupTime time.Time, qname string, response *fqdn.DNSIPRecords) {
			// Do nothing if this option is off
			if !option.Config.ToFQDNsEnablePollerEvents {
				return
			}

			// FIXME: Not always true but we don't have the protocol information here
			protocol := accesslog.TransportProtocol(u8proto.ProtoIDs["udp"])

			record := logger.LogRecord{
				LogRecord: accesslog.LogRecord{
					Type:              accesslog.TypeResponse,
					ObservationPoint:  accesslog.Ingress,
					IPVersion:         accesslog.VersionIPv4,
					TransportProtocol: protocol,
					Timestamp:         time.Now().UTC().Format(time.RFC3339Nano),
					NodeAddressInfo: accesslog.NodeAddressInfo{
						IPv4: node.GetExternalIPv4().String(),
						IPv6: node.GetIPv6().String(),
					},
				},
			}

			logger.LogTags.Verdict(accesslog.VerdictForwarded, "DNSPoller")(&record)
			logger.LogTags.DNS(&accesslog.LogRecordDNS{
				Query:             qname,
				IPs:               response.IPs,
				TTL:               uint32(response.TTL),
				CNAMEs:            nil,
				ObservationSource: accesslog.DNSSourceAgentPoller,
			})(&record)
			record.Log()
		}}

	d.dnsRuleGen = fqdn.NewRuleGen(cfg)
	d.dnsPoller = fqdn.NewDNSPoller(cfg, d.dnsRuleGen)
	if option.Config.ToFQDNsEnablePoller {
		fqdn.StartDNSPoller(d.dnsPoller)
	}

	// Prefill the cache with DNS lookups from restored endpoints. This is needed
	// to maintain continuity of which IPs are allowed.
	// Note: This is TTL aware, and expired data will not be used (e.g. when
	// restoring after a long delay).
	for _, restoredEP := range restoredEndpoints.restored {
		// Upgrades from old ciliums have this nil
		if restoredEP.DNSHistory != nil {
			fqdn.DefaultDNSCache.UpdateFromCache(restoredEP.DNSHistory)
		}
	}

	// Once we stop returning errors from StartDNSProxy this should live in
	// StartProxySupport
	proxy.DefaultDNSProxy, err = dnsproxy.StartDNSProxy("", uint16(option.Config.ToFQDNsProxyPort),
		// LookupEPByIP
		func(endpointIP net.IP) (endpointID string, err error) {
			e := endpointmanager.LookupIPv4(endpointIP.String())
			if e == nil {
				return "", fmt.Errorf("Cannot find endpoint with IP %s", endpointIP.String())
			}

			return e.StringID(), nil
		},
		// NotifyOnDNSMsg handles DNS data in the daemon by emitting monitor
		// events, proxy metrics and storing DNS data in the DNS cache. This may
		// result in rule generation.
		// It will:
		// - Report a monitor error event and proxy metrics when the proxy sees an
		//   error, and when it can't process something in this function
		// - Report the verdict in a monitor event and emit proxy metrics
		// - Insert the DNS data into the cache when msg is a DNS response and we
		//   can lookup the endpoint related to it
		// srcAddr and dstAddr should match the packet reported on (i.e. the
		// endpoint is srcAddr for requests, and dstAddr for responses).
		func(lookupTime time.Time, srcAddr, dstAddr string, msg *dns.Msg, protocol string, allowed bool, proxyErr error) error {
			var protoID = u8proto.ProtoIDs[strings.ToLower(protocol)]

			var verdict accesslog.FlowVerdict
			var reason string
			switch {
			case proxyErr != nil:
				verdict = accesslog.VerdictError
				reason = "Error: " + proxyErr.Error()
			case allowed:
				verdict = accesslog.VerdictForwarded
				reason = "Allowed by policy"
			case !allowed:
				verdict = accesslog.VerdictDenied
				reason = "Denied by policy"
			}

			var epAddr string     // the address of the endpoint that originated the request
			var serverAddr string // the address of the DNS target
			var ingress = msg.Response
			var flowType accesslog.FlowType
			if ingress {
				flowType = accesslog.TypeResponse
				epAddr = dstAddr
				serverAddr = srcAddr
			} else {
				flowType = accesslog.TypeRequest
				epAddr = srcAddr
				serverAddr = dstAddr
			}

			var serverPort int
			_, serverPortStr, err := net.SplitHostPort(serverAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
			} else {
				if serverPort, err = strconv.Atoi(serverPortStr); err != nil {
					log.WithError(err).WithField(logfields.Port, serverPortStr).Error("cannot parse destination port")
				}
			}

			var ep *endpoint.Endpoint
			epIP, _, err := net.SplitHostPort(epAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
				ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, accesslog.VerdictError)
				return err
			}
			ep = endpointmanager.LookupIPv4(epIP)
			if ep == nil {
				// This is a hard fail. We cannot proceed because record.Log requires a
				// non-nil ep, and we also don't want to insert this data into the
				// cache if we don't know that an endpoint asked for it (this is
				// asserted via ep != nil here and msg.Response && msg.Rcode ==
				// dns.RcodeSuccess below).
				err := fmt.Errorf("Cannot find matching endpoint for IPs %s or %s", srcAddr, dstAddr)
				log.WithError(err).Error("cannot find matching endpoint")
				ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, accesslog.VerdictError)
				return err
			}

			qname, responseIPs, TTL, CNAMEs, err := dnsproxy.ExtractMsgDetails(msg)
			if err != nil {
				// This error is ok because all these values are used for reporting, or filling in the cache.
				log.WithError(err).Error("cannot extract DNS message details")
			}

			ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, verdict)
			record := logger.NewLogRecord(proxy.DefaultEndpointInfoRegistry, ep, flowType, ingress,
				func(lr *logger.LogRecord) { lr.LogRecord.TransportProtocol = accesslog.TransportProtocol(protoID) },
				logger.LogTags.Verdict(verdict, reason),
				logger.LogTags.Addressing(logger.AddressingInfo{
					SrcIPPort:   srcAddr,
					DstIPPort:   dstAddr,
					SrcIdentity: 0, // 0 more correctly finds src and dst EP data
				}),
				logger.LogTags.DNS(&accesslog.LogRecordDNS{
					Query:             qname,
					IPs:               responseIPs,
					TTL:               TTL,
					CNAMEs:            CNAMEs,
					ObservationSource: accesslog.DNSSourceProxy,
				}),
			)
			record.Log()

			if msg.Response && msg.Rcode == dns.RcodeSuccess {
				// This must happen before the ruleGen update below, to ensure that
				// this data is included in the serialized Endpoint object.
				// Note: We need to fixup minTTL to be consistent with how we insert it
				// elsewhere i.e. we don't want to lose the lower bound for DNS data
				// TTL if we reboot twice.
				log.WithField(logfields.EndpointID, ep.ID).Debug("Recording DNS lookup in endpoint specific cache")
				effectiveTTL := int(TTL)
				if effectiveTTL < option.Config.ToFQDNsMinTTL {
					effectiveTTL = option.Config.ToFQDNsMinTTL
				}
				ep.DNSHistory.Update(lookupTime, qname, responseIPs, effectiveTTL)
				log.Debug("Updating DNS name in cache from response to to query")
				err = d.dnsRuleGen.UpdateGenerateDNS(lookupTime, map[string]*fqdn.DNSIPRecords{
					qname: {
						IPs: responseIPs,
						TTL: int(effectiveTTL),
					}})
				if err != nil {
					log.WithError(err).Error("error updating internal DNS cache for rule generation")
				}
			}

			return nil
		})
	proxy.DefaultDNSProxy.SetRejectReply(option.Config.FQDNRejectResponse)
	return err // filled by StartDNSProxy
}

type getFqdnCache struct {
	daemon *Daemon
}

func NewGetFqdnCacheHandler(d *Daemon) GetFqdnCacheHandler {
	return &getFqdnCache{daemon: d}
}

func (h *getFqdnCache) Handle(params GetFqdnCacheParams) middleware.Responder {
	// endpoints we want data from
	endpoints := endpointmanager.GetEndpoints()

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

type getFqdnCacheID struct {
	daemon *Daemon
}

func NewGetFqdnCacheIDHandler(d *Daemon) GetFqdnCacheIDHandler {
	return &getFqdnCacheID{daemon: d}
}

func (h *getFqdnCacheID) Handle(params GetFqdnCacheIDParams) middleware.Responder {
	var endpoints []*endpoint.Endpoint
	if params.ID != "" {
		ep, err := endpointmanager.Lookup(params.ID)
		switch {
		case err != nil:
			return api.Error(GetFqdnCacheIDBadRequestCode, err)
		case ep == nil:
			return api.Error(GetFqdnCacheIDNotFoundCode, fmt.Errorf("Cannot find endpoint %s", params.ID))
		default:
			endpoints = []*endpoint.Endpoint{ep}
		}
	}

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

// extractDNSLookups returns API models.DNSLookup copies of DNS data in each
// endpoint's DNSHistory. These are filtered by CIDRStr and matchPatternStr if
// they are non-empty.
func extractDNSLookups(endpoints []*endpoint.Endpoint, CIDRStr, matchPatternStr string) (lookups []*models.DNSLookup, err error) {
	cidrMatcher := func(ip net.IP) bool { return true }
	if CIDRStr != "" {
		_, cidr, err := net.ParseCIDR(CIDRStr)
		if err != nil {
			return nil, err
		}
		cidrMatcher = func(ip net.IP) bool { return cidr.Contains(ip) }
	}

	nameMatcher := func(name string) bool { return true }
	if matchPatternStr != "" {
		matcher, err := matchpattern.Validate(matchPatternStr)
		if err != nil {
			return nil, err
		}
		nameMatcher = func(name string) bool { return matcher.MatchString(name) }
	}

	for _, ep := range endpoints {
		for _, lookup := range ep.DNSHistory.Dump() {
			if !nameMatcher(lookup.Name) {
				continue
			}

			// The API model needs strings
			IPStrings := make([]string, 0, len(lookup.IPs))

			// only proceed if any IP matches the cidr selector
			anIPMatches := false
			for _, ip := range lookup.IPs {
				anIPMatches = anIPMatches || cidrMatcher(ip)
				IPStrings = append(IPStrings, ip.String())
			}
			if !anIPMatches {
				continue
			}

			lookups = append(lookups, &models.DNSLookup{
				Fqdn:           lookup.Name,
				Ips:            IPStrings,
				LookupTime:     strfmt.DateTime(lookup.LookupTime),
				TTL:            int64(lookup.TTL),
				ExpirationTime: strfmt.DateTime(lookup.ExpirationTime),
				EndpointID:     int64(ep.ID),
			})
		}
	}

	return lookups, nil
}
<commit_msg>fqdn: Do not report stats on nil Endpoint<commit_after>// Copyright 2019 Authors of Cilium
//
// 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 main

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/pkg/api"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/fqdn"
	"github.com/cilium/cilium/pkg/fqdn/dnsproxy"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	policyApi "github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/proxy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/proxy/logger"
	"github.com/cilium/cilium/pkg/u8proto"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/strfmt"

	"github.com/miekg/dns"
)

// bootstrapFQDN initializes the toFQDNs related subsystems: DNSPoller,
// d.dnsRuleGen, and the DNS proxy.
// dnsRuleGen and DNSPoller will use the default resolver and, implicitly, the
// default DNS cache. The proxy binds to all interfaces, and uses the
// configured DNS proxy port (this may be 0 and so OS-assigned).
func (d *Daemon) bootstrapFQDN(restoredEndpoints *endpointRestoreState) (err error) {
	cfg := fqdn.Config{
		MinTTL:         option.Config.ToFQDNsMinTTL,
		Cache:          fqdn.DefaultDNSCache,
		LookupDNSNames: fqdn.DNSLookupDefaultResolver,
		AddGeneratedRules: func(generatedRules []*policyApi.Rule) error {
			// Insert the new rules into the policy repository. We need them to
			// replace the previous set. This requires the labels to match (including
			// the ToFQDN-UUID one).
			_, err := d.PolicyAdd(generatedRules, &AddOptions{Replace: true, Generated: true})
			return err
		},
		PollerResponseNotify: func(lookupTime time.Time, qname string, response *fqdn.DNSIPRecords) {
			// Do nothing if this option is off
			if !option.Config.ToFQDNsEnablePollerEvents {
				return
			}

			// FIXME: Not always true but we don't have the protocol information here
			protocol := accesslog.TransportProtocol(u8proto.ProtoIDs["udp"])

			record := logger.LogRecord{
				LogRecord: accesslog.LogRecord{
					Type:              accesslog.TypeResponse,
					ObservationPoint:  accesslog.Ingress,
					IPVersion:         accesslog.VersionIPv4,
					TransportProtocol: protocol,
					Timestamp:         time.Now().UTC().Format(time.RFC3339Nano),
					NodeAddressInfo: accesslog.NodeAddressInfo{
						IPv4: node.GetExternalIPv4().String(),
						IPv6: node.GetIPv6().String(),
					},
				},
			}

			logger.LogTags.Verdict(accesslog.VerdictForwarded, "DNSPoller")(&record)
			logger.LogTags.DNS(&accesslog.LogRecordDNS{
				Query:             qname,
				IPs:               response.IPs,
				TTL:               uint32(response.TTL),
				CNAMEs:            nil,
				ObservationSource: accesslog.DNSSourceAgentPoller,
			})(&record)
			record.Log()
		}}

	d.dnsRuleGen = fqdn.NewRuleGen(cfg)
	d.dnsPoller = fqdn.NewDNSPoller(cfg, d.dnsRuleGen)
	if option.Config.ToFQDNsEnablePoller {
		fqdn.StartDNSPoller(d.dnsPoller)
	}

	// Prefill the cache with DNS lookups from restored endpoints. This is needed
	// to maintain continuity of which IPs are allowed.
	// Note: This is TTL aware, and expired data will not be used (e.g. when
	// restoring after a long delay).
	for _, restoredEP := range restoredEndpoints.restored {
		// Upgrades from old ciliums have this nil
		if restoredEP.DNSHistory != nil {
			fqdn.DefaultDNSCache.UpdateFromCache(restoredEP.DNSHistory)
		}
	}

	// Once we stop returning errors from StartDNSProxy this should live in
	// StartProxySupport
	proxy.DefaultDNSProxy, err = dnsproxy.StartDNSProxy("", uint16(option.Config.ToFQDNsProxyPort),
		// LookupEPByIP
		func(endpointIP net.IP) (endpointID string, err error) {
			e := endpointmanager.LookupIPv4(endpointIP.String())
			if e == nil {
				return "", fmt.Errorf("Cannot find endpoint with IP %s", endpointIP.String())
			}

			return e.StringID(), nil
		},
		// NotifyOnDNSMsg handles DNS data in the daemon by emitting monitor
		// events, proxy metrics and storing DNS data in the DNS cache. This may
		// result in rule generation.
		// It will:
		// - Report a monitor error event and proxy metrics when the proxy sees an
		//   error, and when it can't process something in this function
		// - Report the verdict in a monitor event and emit proxy metrics
		// - Insert the DNS data into the cache when msg is a DNS response and we
		//   can lookup the endpoint related to it
		// srcAddr and dstAddr should match the packet reported on (i.e. the
		// endpoint is srcAddr for requests, and dstAddr for responses).
		func(lookupTime time.Time, srcAddr, dstAddr string, msg *dns.Msg, protocol string, allowed bool, proxyErr error) error {
			var protoID = u8proto.ProtoIDs[strings.ToLower(protocol)]

			var verdict accesslog.FlowVerdict
			var reason string
			switch {
			case proxyErr != nil:
				verdict = accesslog.VerdictError
				reason = "Error: " + proxyErr.Error()
			case allowed:
				verdict = accesslog.VerdictForwarded
				reason = "Allowed by policy"
			case !allowed:
				verdict = accesslog.VerdictDenied
				reason = "Denied by policy"
			}

			var epAddr string     // the address of the endpoint that originated the request
			var serverAddr string // the address of the DNS target
			var ingress = msg.Response
			var flowType accesslog.FlowType
			if ingress {
				flowType = accesslog.TypeResponse
				epAddr = dstAddr
				serverAddr = srcAddr
			} else {
				flowType = accesslog.TypeRequest
				epAddr = srcAddr
				serverAddr = dstAddr
			}

			var serverPort int
			_, serverPortStr, err := net.SplitHostPort(serverAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
			} else {
				if serverPort, err = strconv.Atoi(serverPortStr); err != nil {
					log.WithError(err).WithField(logfields.Port, serverPortStr).Error("cannot parse destination port")
				}
			}

			var ep *endpoint.Endpoint
			epIP, _, err := net.SplitHostPort(epAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
				ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, accesslog.VerdictError)
				return err
			}
			ep = endpointmanager.LookupIPv4(epIP)
			if ep == nil {
				// This is a hard fail. We cannot proceed because record.Log requires a
				// non-nil ep, and we also don't want to insert this data into the
				// cache if we don't know that an endpoint asked for it (this is
				// asserted via ep != nil here and msg.Response && msg.Rcode ==
				// dns.RcodeSuccess below).
				err := fmt.Errorf("Cannot find matching endpoint for IPs %s or %s", srcAddr, dstAddr)
				log.WithError(err).Error("cannot find matching endpoint")
				return err
			}

			qname, responseIPs, TTL, CNAMEs, err := dnsproxy.ExtractMsgDetails(msg)
			if err != nil {
				// This error is ok because all these values are used for reporting, or filling in the cache.
				log.WithError(err).Error("cannot extract DNS message details")
			}

			ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, verdict)
			record := logger.NewLogRecord(proxy.DefaultEndpointInfoRegistry, ep, flowType, ingress,
				func(lr *logger.LogRecord) { lr.LogRecord.TransportProtocol = accesslog.TransportProtocol(protoID) },
				logger.LogTags.Verdict(verdict, reason),
				logger.LogTags.Addressing(logger.AddressingInfo{
					SrcIPPort:   srcAddr,
					DstIPPort:   dstAddr,
					SrcIdentity: 0, // 0 more correctly finds src and dst EP data
				}),
				logger.LogTags.DNS(&accesslog.LogRecordDNS{
					Query:             qname,
					IPs:               responseIPs,
					TTL:               TTL,
					CNAMEs:            CNAMEs,
					ObservationSource: accesslog.DNSSourceProxy,
				}),
			)
			record.Log()

			if msg.Response && msg.Rcode == dns.RcodeSuccess {
				// This must happen before the ruleGen update below, to ensure that
				// this data is included in the serialized Endpoint object.
				// Note: We need to fixup minTTL to be consistent with how we insert it
				// elsewhere i.e. we don't want to lose the lower bound for DNS data
				// TTL if we reboot twice.
				log.WithField(logfields.EndpointID, ep.ID).Debug("Recording DNS lookup in endpoint specific cache")
				effectiveTTL := int(TTL)
				if effectiveTTL < option.Config.ToFQDNsMinTTL {
					effectiveTTL = option.Config.ToFQDNsMinTTL
				}
				ep.DNSHistory.Update(lookupTime, qname, responseIPs, effectiveTTL)
				log.Debug("Updating DNS name in cache from response to to query")
				err = d.dnsRuleGen.UpdateGenerateDNS(lookupTime, map[string]*fqdn.DNSIPRecords{
					qname: {
						IPs: responseIPs,
						TTL: int(effectiveTTL),
					}})
				if err != nil {
					log.WithError(err).Error("error updating internal DNS cache for rule generation")
				}
			}

			return nil
		})
	proxy.DefaultDNSProxy.SetRejectReply(option.Config.FQDNRejectResponse)
	return err // filled by StartDNSProxy
}

type getFqdnCache struct {
	daemon *Daemon
}

func NewGetFqdnCacheHandler(d *Daemon) GetFqdnCacheHandler {
	return &getFqdnCache{daemon: d}
}

func (h *getFqdnCache) Handle(params GetFqdnCacheParams) middleware.Responder {
	// endpoints we want data from
	endpoints := endpointmanager.GetEndpoints()

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

type getFqdnCacheID struct {
	daemon *Daemon
}

func NewGetFqdnCacheIDHandler(d *Daemon) GetFqdnCacheIDHandler {
	return &getFqdnCacheID{daemon: d}
}

func (h *getFqdnCacheID) Handle(params GetFqdnCacheIDParams) middleware.Responder {
	var endpoints []*endpoint.Endpoint
	if params.ID != "" {
		ep, err := endpointmanager.Lookup(params.ID)
		switch {
		case err != nil:
			return api.Error(GetFqdnCacheIDBadRequestCode, err)
		case ep == nil:
			return api.Error(GetFqdnCacheIDNotFoundCode, fmt.Errorf("Cannot find endpoint %s", params.ID))
		default:
			endpoints = []*endpoint.Endpoint{ep}
		}
	}

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

// extractDNSLookups returns API models.DNSLookup copies of DNS data in each
// endpoint's DNSHistory. These are filtered by CIDRStr and matchPatternStr if
// they are non-empty.
func extractDNSLookups(endpoints []*endpoint.Endpoint, CIDRStr, matchPatternStr string) (lookups []*models.DNSLookup, err error) {
	cidrMatcher := func(ip net.IP) bool { return true }
	if CIDRStr != "" {
		_, cidr, err := net.ParseCIDR(CIDRStr)
		if err != nil {
			return nil, err
		}
		cidrMatcher = func(ip net.IP) bool { return cidr.Contains(ip) }
	}

	nameMatcher := func(name string) bool { return true }
	if matchPatternStr != "" {
		matcher, err := matchpattern.Validate(matchPatternStr)
		if err != nil {
			return nil, err
		}
		nameMatcher = func(name string) bool { return matcher.MatchString(name) }
	}

	for _, ep := range endpoints {
		for _, lookup := range ep.DNSHistory.Dump() {
			if !nameMatcher(lookup.Name) {
				continue
			}

			// The API model needs strings
			IPStrings := make([]string, 0, len(lookup.IPs))

			// only proceed if any IP matches the cidr selector
			anIPMatches := false
			for _, ip := range lookup.IPs {
				anIPMatches = anIPMatches || cidrMatcher(ip)
				IPStrings = append(IPStrings, ip.String())
			}
			if !anIPMatches {
				continue
			}

			lookups = append(lookups, &models.DNSLookup{
				Fqdn:           lookup.Name,
				Ips:            IPStrings,
				LookupTime:     strfmt.DateTime(lookup.LookupTime),
				TTL:            int64(lookup.TTL),
				ExpirationTime: strfmt.DateTime(lookup.ExpirationTime),
				EndpointID:     int64(ep.ID),
			})
		}
	}

	return lookups, nil
}
<|endoftext|>


<commit_before>package dalga

// TODO backoff

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/cenkalti/dalga/vendor/github.com/go-sql-driver/mysql"
)

var debugging = flag.Bool("debug", false, "turn on debug messages")

func debug(args ...interface{}) {
	if *debugging {
		log.Println(args...)
	}
}

type Dalga struct {
	config   Config
	db       *sql.DB
	table    *table
	listener net.Listener
	client   http.Client
	// to wake up publisher when a new job is scheduled or cancelled
	notify chan struct{}
	// will be closed when dalga is ready to accept requests
	ready chan struct{}
	// will be closed by Shutdown method
	shutdown chan struct{}
	// to stop publisher goroutine
	stopPublisher chan struct{}
	// will be closed when publisher goroutine is stopped
	publisherStopped chan struct{}
}

func New(config Config) *Dalga {
	d := &Dalga{
		config:           config,
		notify:           make(chan struct{}, 1),
		ready:            make(chan struct{}),
		shutdown:         make(chan struct{}),
		stopPublisher:    make(chan struct{}),
		publisherStopped: make(chan struct{}),
	}
	d.client.Timeout = time.Duration(config.Endpoint.Timeout) * time.Second
	return d
}

// Run Dalga. This function is blocking. Returns nil if Shutdown is called.
func (d *Dalga) Run() error {
	if err := d.connectDB(); err != nil {
		return err
	}
	defer d.db.Close()

	var err error
	d.listener, err = net.Listen("tcp", d.config.Listen.Addr())
	if err != nil {
		return err
	}
	log.Println("Listening", d.listener.Addr())

	close(d.ready)

	go d.publisher()
	defer func() {
		close(d.stopPublisher)
		<-d.publisherStopped
	}()

	if err = d.serveHTTP(); err != nil {
		select {
		case _, ok := <-d.shutdown:
			if !ok {
				// shutdown in progress, do not return error
				return nil
			}
		default:
		}
	}
	return err
}

// Shutdown running Dalga.
func (d *Dalga) Shutdown() error {
	close(d.shutdown)
	return d.listener.Close()
}

// NotifyReady returns a channel that will be closed when Dalga is running.
func (d *Dalga) NotifyReady() <-chan struct{} {
	return d.ready
}

func (d *Dalga) connectDB() error {
	var err error
	d.db, err = sql.Open("mysql", d.config.MySQL.DSN())
	if err != nil {
		return err
	}
	if err = d.db.Ping(); err != nil {
		return err
	}
	log.Print("Connected to MySQL")
	d.table = &table{d.db, d.config.MySQL.Table}
	return nil
}

// CreateTable creates the table for storing jobs.
func (d *Dalga) CreateTable() error {
	db, err := sql.Open("mysql", d.config.MySQL.DSN())
	if err != nil {
		return err
	}
	defer db.Close()
	t := &table{db, d.config.MySQL.Table}
	return t.Create()
}

// GetJob returns the job with description routing key.
func (d *Dalga) GetJob(description, routingKey string) (*Job, error) {
	return d.table.Get(description, routingKey)
}

// ScheduleJob inserts a new job to the table or replaces existing one.
// Returns the created or replaced job.
func (d *Dalga) ScheduleJob(description, routingKey string, interval uint32, oneOff bool) (*Job, error) {
	job := newJob(description, routingKey, interval, oneOff)
	err := d.table.Insert(job)
	if err != nil {
		return nil, err
	}
	d.notifyPublisher("new job")
	debug("Job is scheduled:", job)
	return job, nil
}

// TriggerJob publishes the job to RabbitMQ immediately and resets the next run time of the job.
func (d *Dalga) TriggerJob(description, routingKey string) (*Job, error) {
	job, err := d.GetJob(description, routingKey)
	if err != nil {
		return nil, err
	}
	job.NextRun = time.Now().UTC()
	if err := d.table.Insert(job); err != nil {
		return nil, err
	}
	d.notifyPublisher("job is triggered")
	debug("Job is triggered:", job)
	return job, nil
}

// CancelJob deletes the job with description and routing key.
func (d *Dalga) CancelJob(description, routingKey string) error {
	err := d.table.Delete(description, routingKey)
	if err != nil {
		return err
	}
	d.notifyPublisher("job cancelled")
	debug("Job is cancelled")
	return nil
}

func (d *Dalga) notifyPublisher(debugMessage string) {
	select {
	case d.notify <- struct{}{}:
		debug("notifying publisher:", debugMessage)
	default:
	}
}

// publisher runs a loop that reads the next Job from the queue and publishes it.
func (d *Dalga) publisher() {
	defer close(d.publisherStopped)

	for {
		debug("---")

		var after <-chan time.Time

		job, err := d.table.Front()
		if err != nil {
			if err == sql.ErrNoRows {
				debug("No scheduled jobs in the table")
			} else if myErr, ok := err.(*mysql.MySQLError); ok && myErr.Number == 1146 {
				// Table doesn't exist
				log.Fatal(myErr)
			} else {
				log.Print(err)
				time.Sleep(time.Second)
				continue
			}
		} else {
			remaining := job.Remaining()
			after = time.After(remaining)
			debug("Next job:", job, "Remaining:", remaining)
		}

		// Sleep until the next job's run time or the webserver's wakes us up.
		select {
		case <-after:
			debug("Job sleep time finished")
			if err = d.publish(job); err != nil {
				log.Print(err)
				time.Sleep(time.Second)
			}
		case <-d.notify:
			debug("Woken up from sleep by notification")
			continue
		case <-d.stopPublisher:
			debug("Came quit message")
			return
		}
	}
}

// publish makes a POST request to the endpoint and updates the Job's next run time.
func (d *Dalga) publish(j *Job) error {
	debug("publish", *j)

	var add time.Duration
	if j.Interval == 0 {
		add = time.Duration(d.config.Endpoint.Timeout) * time.Second
	} else {
		add = j.Interval
	}

	j.NextRun = time.Now().UTC().Add(add)

	if err := d.table.UpdateNextRun(j); err != nil {
		return err
	}

	go func() {
		if err := d.postJob(j); err != nil {
			log.Print(err)
		}
	}()

	return nil
}

func (d *Dalga) postJob(j *Job) error {
	resp, err := d.client.Post(d.config.Endpoint.BaseURL+j.Path, "text/plain", strings.NewReader(j.Body))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("endpoint error: %d", resp.StatusCode)
	}

	if j.Interval == 0 {
		debug("deleting one-off job")
		return d.table.Delete(j.Path, j.Body)
	}

	return nil
}
<commit_msg>Do not do multiple POSTs for the same job at the same time<commit_after>package dalga

// TODO backoff

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/cenkalti/dalga/vendor/github.com/go-sql-driver/mysql"
)

var debugging = flag.Bool("debug", false, "turn on debug messages")

func debug(args ...interface{}) {
	if *debugging {
		log.Println(args...)
	}
}

type Dalga struct {
	config     Config
	db         *sql.DB
	table      *table
	listener   net.Listener
	client     http.Client
	activeJobs map[string]struct{}
	m          sync.Mutex
	// to wake up publisher when a new job is scheduled or cancelled
	notify chan struct{}
	// will be closed when dalga is ready to accept requests
	ready chan struct{}
	// will be closed by Shutdown method
	shutdown chan struct{}
	// to stop publisher goroutine
	stopPublisher chan struct{}
	// will be closed when publisher goroutine is stopped
	publisherStopped chan struct{}
}

func New(config Config) *Dalga {
	d := &Dalga{
		config:           config,
		activeJobs:       make(map[string]struct{}),
		notify:           make(chan struct{}, 1),
		ready:            make(chan struct{}),
		shutdown:         make(chan struct{}),
		stopPublisher:    make(chan struct{}),
		publisherStopped: make(chan struct{}),
	}
	d.client.Timeout = time.Duration(config.Endpoint.Timeout) * time.Second
	return d
}

// Run Dalga. This function is blocking. Returns nil if Shutdown is called.
func (d *Dalga) Run() error {
	if err := d.connectDB(); err != nil {
		return err
	}
	defer d.db.Close()

	var err error
	d.listener, err = net.Listen("tcp", d.config.Listen.Addr())
	if err != nil {
		return err
	}
	log.Println("Listening", d.listener.Addr())

	close(d.ready)

	go d.publisher()
	defer func() {
		close(d.stopPublisher)
		<-d.publisherStopped
	}()

	if err = d.serveHTTP(); err != nil {
		select {
		case _, ok := <-d.shutdown:
			if !ok {
				// shutdown in progress, do not return error
				return nil
			}
		default:
		}
	}
	return err
}

// Shutdown running Dalga.
func (d *Dalga) Shutdown() error {
	close(d.shutdown)
	return d.listener.Close()
}

// NotifyReady returns a channel that will be closed when Dalga is running.
func (d *Dalga) NotifyReady() <-chan struct{} {
	return d.ready
}

func (d *Dalga) connectDB() error {
	var err error
	d.db, err = sql.Open("mysql", d.config.MySQL.DSN())
	if err != nil {
		return err
	}
	if err = d.db.Ping(); err != nil {
		return err
	}
	log.Print("Connected to MySQL")
	d.table = &table{d.db, d.config.MySQL.Table}
	return nil
}

// CreateTable creates the table for storing jobs.
func (d *Dalga) CreateTable() error {
	db, err := sql.Open("mysql", d.config.MySQL.DSN())
	if err != nil {
		return err
	}
	defer db.Close()
	t := &table{db, d.config.MySQL.Table}
	return t.Create()
}

// GetJob returns the job with description routing key.
func (d *Dalga) GetJob(description, routingKey string) (*Job, error) {
	return d.table.Get(description, routingKey)
}

// ScheduleJob inserts a new job to the table or replaces existing one.
// Returns the created or replaced job.
func (d *Dalga) ScheduleJob(description, routingKey string, interval uint32, oneOff bool) (*Job, error) {
	job := newJob(description, routingKey, interval, oneOff)
	err := d.table.Insert(job)
	if err != nil {
		return nil, err
	}
	d.notifyPublisher("new job")
	debug("Job is scheduled:", job)
	return job, nil
}

// TriggerJob publishes the job to RabbitMQ immediately and resets the next run time of the job.
func (d *Dalga) TriggerJob(description, routingKey string) (*Job, error) {
	job, err := d.GetJob(description, routingKey)
	if err != nil {
		return nil, err
	}
	job.NextRun = time.Now().UTC()
	if err := d.table.Insert(job); err != nil {
		return nil, err
	}
	d.notifyPublisher("job is triggered")
	debug("Job is triggered:", job)
	return job, nil
}

// CancelJob deletes the job with description and routing key.
func (d *Dalga) CancelJob(description, routingKey string) error {
	err := d.table.Delete(description, routingKey)
	if err != nil {
		return err
	}
	d.notifyPublisher("job cancelled")
	debug("Job is cancelled")
	return nil
}

func (d *Dalga) notifyPublisher(debugMessage string) {
	select {
	case d.notify <- struct{}{}:
		debug("notifying publisher:", debugMessage)
	default:
	}
}

// publisher runs a loop that reads the next Job from the queue and publishes it.
func (d *Dalga) publisher() {
	defer close(d.publisherStopped)

	for {
		debug("---")

		var after <-chan time.Time

		job, err := d.table.Front()
		if err != nil {
			if err == sql.ErrNoRows {
				debug("No scheduled jobs in the table")
			} else if myErr, ok := err.(*mysql.MySQLError); ok && myErr.Number == 1146 {
				// Table doesn't exist
				log.Fatal(myErr)
			} else {
				log.Print(err)
				time.Sleep(time.Second)
				continue
			}
		} else {
			remaining := job.Remaining()
			after = time.After(remaining)
			debug("Next job:", job, "Remaining:", remaining)
		}

		// Sleep until the next job's run time or the webserver's wakes us up.
		select {
		case <-after:
			debug("Job sleep time finished")
			if err = d.publish(job); err != nil {
				log.Print(err)
				time.Sleep(time.Second)
			}
		case <-d.notify:
			debug("Woken up from sleep by notification")
			continue
		case <-d.stopPublisher:
			debug("Came quit message")
			return
		}
	}
}

// publish makes a POST request to the endpoint and updates the Job's next run time.
func (d *Dalga) publish(j *Job) error {
	debug("publish", *j)

	var add time.Duration
	if j.Interval == 0 {
		add = d.client.Timeout
	} else {
		add = j.Interval
	}

	j.NextRun = time.Now().UTC().Add(add)

	if err := d.table.UpdateNextRun(j); err != nil {
		return err
	}

	go func() {
		// Do not do multiple POSTs for the same job at the same time.
		key := j.Path + j.Body
		d.m.Lock()
		if _, ok := d.activeJobs[key]; ok {
			d.m.Unlock()
			return
		}
		d.activeJobs[key] = struct{}{}
		d.m.Unlock()

		if err := d.postJob(j); err != nil {
			log.Print(err)
		}

		d.m.Lock()
		delete(d.activeJobs, key)
		d.m.Unlock()
	}()

	return nil
}

func (d *Dalga) postJob(j *Job) error {
	resp, err := d.client.Post(d.config.Endpoint.BaseURL+j.Path, "text/plain", strings.NewReader(j.Body))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("endpoint error: %d", resp.StatusCode)
	}

	if j.Interval == 0 {
		debug("deleting one-off job")
		return d.table.Delete(j.Path, j.Body)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"os"
)

type UDPRecv struct {
	Data       []byte // UDP payload
	RemoteAddr *net.UDPAddr
}

// listenUDP goroutine has its own internal read buf that it gives a slice into
// when it receives a packet. It passes the slice and the source address into
// the channel.
// TODO: make this take an Iface so it can log statistics
func listenUDP(conn UDPReadWrite, c chan UDPRecv) error {
	colors := []string{"magenta", "yellow", "cyan", "white:blue", "black:white"}
	ansi_colors := make([]string, len(colors))
	for i, color := range colors {
		ansi_colors[i] = ansi.ColorCode(color)
	}
	ansi_reset := ansi.ColorCode("reset")

	read_buf := make([]byte, BUF_SIZE)
	for {
		log.Printf("Listening on conn %p", conn)
		count, remote_addr, err := conn.ReadFromUDP(read_buf)
		if err != nil {
			log.Print(err)
			return err
		}
		c <- UDPRecv{
			Data:       read_buf[:count],
			RemoteAddr: remote_addr,
		}

		if DEBUG_LEVEL >= 1 {
			fmt.Print(ansi_colors[0], "R", ansi_reset)
		}

	}
}

func listenTun(tundev *os.File, read_buf []byte, c chan int) error {
	for {
		count, err := tundev.Read(read_buf)
		if err != nil {
			log.Print(err)
			return err
		}
		c <- count
	}
}
<commit_msg>Fixed ansi color imports<commit_after>package main

import (
	"fmt"
	"github.com/mgutz/ansi"
	"log"
	"net"
	"os"
)

type UDPRecv struct {
	Data       []byte // UDP payload
	RemoteAddr *net.UDPAddr
}

// listenUDP goroutine has its own internal read buf that it gives a slice into
// when it receives a packet. It passes the slice and the source address into
// the channel.
// TODO: make this take an Iface so it can log statistics
func listenUDP(conn UDPReadWrite, c chan UDPRecv) error {
	colors := []string{"magenta", "yellow", "cyan", "white:blue", "black:white"}
	ansi_colors := make([]string, len(colors))
	for i, color := range colors {
		ansi_colors[i] = ansi.ColorCode(color)
	}
	ansi_reset := ansi.ColorCode("reset")

	read_buf := make([]byte, BUF_SIZE)
	for {
		log.Printf("Listening on conn %p", conn)
		count, remote_addr, err := conn.ReadFromUDP(read_buf)
		if err != nil {
			log.Print(err)
			return err
		}
		c <- UDPRecv{
			Data:       read_buf[:count],
			RemoteAddr: remote_addr,
		}

		if DEBUG_LEVEL >= 1 {
			fmt.Print(ansi_colors[0], "R", ansi_reset)
		}

	}
}

func listenTun(tundev *os.File, read_buf []byte, c chan int) error {
	for {
		count, err := tundev.Read(read_buf)
		if err != nil {
			log.Print(err)
			return err
		}
		c <- count
	}
}
<|endoftext|>


<commit_before>package deployments

import (
	"testing"

	"github.com/asobti/kube-monkey/config"
	"k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	IDENTIFIER = "kube-monkey-id"
	NAME       = "deployment_name"
	NAMESPACE  = metav1.NamespaceDefault
)

func newDeployment(name string, labels map[string]string) v1beta1.Deployment {

	return v1beta1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: NAMESPACE,
			Labels:    labels,
		},
	}
}

func TestNew(t *testing.T) {

	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "1",
		},
	)
	depl, err := New(&v1depl)

	if err != nil {
		t.Error(err)
	}

	if ns := depl.Namespace(); ns != NAMESPACE {
		t.Errorf("Unexpected deployment Namespace, got %s", ns)
	}

	if k := depl.Kind(); k != "v1beta1.Deployment" {
		t.Errorf("Unexpected deployment Kindepl, got %s", k)
	}

	if n := depl.Name(); n != NAME {
		t.Errorf("Unexpected deployment Name, got %s", n)
	}

	if i := depl.Identifier(); i != IDENTIFIER {
		t.Errorf("Unexpected deployment Identifier, got %s", i)
	}

	if m := depl.Mtbf(); m != 1 {
		t.Errorf("Unexpected deployment Mtbf, got %d", m)
	}

}

func TestInvalidIdentifier(t *testing.T) {
	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.MtbfLabelKey: "1",
		},
	)
	_, err := New(&v1depl)

	if err == nil {
		t.Error("Expected an error if config.IdentLabelKey label doesn't exist")
	}
}

func TestInvalidMtbf(t *testing.T) {
	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
		},
	)
	_, err := New(&v1depl)

	if err == nil {
		t.Error("Expected an error if config.MtbfLabelKey label doesn't exist")
	}

	v1depl = newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "string",
		},
	)
	_, err = New(&v1depl)

	if err == nil {
		t.Error("Expected an error if config.MtbfLabelKey label can't be converted a Int type")
	}
	v1depl = newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "0",
		},
	)
	_, err = New(&v1depl)

	if err == nil {
		t.Error("Expected an error if config.MtbfLabelKey label is lower than 1")
	}
}
<commit_msg>Port UT to testify/assert<commit_after>package deployments

import (
	"testing"

	"github.com/asobti/kube-monkey/config"
	"github.com/stretchr/testify/assert"
	"k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	IDENTIFIER = "kube-monkey-id"
	NAME       = "deployment_name"
	NAMESPACE  = metav1.NamespaceDefault
)

func newDeployment(name string, labels map[string]string) v1beta1.Deployment {

	return v1beta1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: NAMESPACE,
			Labels:    labels,
		},
	}
}

func TestNew(t *testing.T) {

	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "1",
		},
	)
	depl, err := New(&v1depl)

	assert.NoError(t, err)
	assert.Equal(t, "v1beta1.Deployment", depl.Kind())
	assert.Equal(t, NAME, depl.Name())
	assert.Equal(t, NAMESPACE, depl.Namespace())
	assert.Equal(t, IDENTIFIER, depl.Identifier())
	assert.Equal(t, 1, depl.Mtbf())
}

func TestInvalidIdentifier(t *testing.T) {
	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.MtbfLabelKey: "1",
		},
	)
	_, err := New(&v1depl)

	assert.Errorf(t, err, "Expected an error if "+config.IdentLabelKey+" label doesn't exist")
}

func TestInvalidMtbf(t *testing.T) {
	v1depl := newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
		},
	)
	_, err := New(&v1depl)

	assert.Errorf(t, err, "Expected an error if "+config.MtbfLabelKey+" label doesn't exist")

	v1depl = newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "string",
		},
	)
	_, err = New(&v1depl)

	assert.Errorf(t, err, "Expected an error if "+config.MtbfLabelKey+" label can't be converted a Int type")

	v1depl = newDeployment(
		NAME,
		map[string]string{
			config.IdentLabelKey: IDENTIFIER,
			config.MtbfLabelKey:  "0",
		},
	)
	_, err = New(&v1depl)

	assert.Errorf(t, err, "Expected an error if "+config.MtbfLabelKey+" label is lower than 1")
}
<|endoftext|>


<commit_before>// Package api has type definitions for webdav
package api

import (
	"encoding/xml"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	// Wed, 27 Sep 2017 14:28:34 GMT
	timeFormat = time.RFC1123
	// The same as time.RFC1123 with optional leading zeros on the date
	// see https://github.com/ncw/rclone/issues/2574
	noZerosRFC1123 = "Mon, _2 Jan 2006 15:04:05 MST"
)

// Multistatus contains responses returned from an HTTP 207 return code
type Multistatus struct {
	Responses []Response `xml:"response"`
}

// Response contains an Href the response it about and its properties
type Response struct {
	Href  string `xml:"href"`
	Props Prop   `xml:"propstat"`
}

// Prop is the properties of a response
//
// This is a lazy way of decoding the multiple <s:propstat> in the
// response.
//
// The response might look like this
//
// <d:response>
//   <d:href>/remote.php/webdav/Nextcloud%20Manual.pdf</d:href>
//   <d:propstat>
//     <d:prop>
//       <d:getlastmodified>Tue, 19 Dec 2017 22:02:36 GMT</d:getlastmodified>
//       <d:getcontentlength>4143665</d:getcontentlength>
//       <d:resourcetype/>
//       <d:getetag>"048d7be4437ff7deeae94db50ff3e209"</d:getetag>
//       <d:getcontenttype>application/pdf</d:getcontenttype>
//     </d:prop>
//     <d:status>HTTP/1.1 200 OK</d:status>
//   </d:propstat>
//   <d:propstat>
//     <d:prop>
//       <d:quota-used-bytes/>
//       <d:quota-available-bytes/>
//     </d:prop>
//     <d:status>HTTP/1.1 404 Not Found</d:status>
//   </d:propstat>
// </d:response>
//
// So we elide the array of <d:propstat> and within that the array of
// <d:prop> into one struct.
//
// Note that status collects all the status values for which we just
// check the first is OK.
type Prop struct {
	Status   []string  `xml:"DAV: status"`
	Name     string    `xml:"DAV: prop>displayname,omitempty"`
	Type     *xml.Name `xml:"DAV: prop>resourcetype>collection,omitempty"`
	Size     int64     `xml:"DAV: prop>getcontentlength,omitempty"`
	Modified Time      `xml:"DAV: prop>getlastmodified,omitempty"`
}

// Parse a status of the form "HTTP/1.1 200 OK" or "HTTP/1.1 200"
var parseStatus = regexp.MustCompile(`^HTTP/[0-9.]+\s+(\d+)`)

// StatusOK examines the Status and returns an OK flag
func (p *Prop) StatusOK() bool {
	// Assume OK if no statuses received
	if len(p.Status) == 0 {
		return true
	}
	match := parseStatus.FindStringSubmatch(p.Status[0])
	if len(match) < 2 {
		return false
	}
	code, err := strconv.Atoi(match[1])
	if err != nil {
		return false
	}
	if code >= 200 && code < 300 {
		return true
	}
	return false
}

// PropValue is a tagged name and value
type PropValue struct {
	XMLName xml.Name `xml:""`
	Value   string   `xml:",chardata"`
}

// Error is used to desribe webdav errors
//
// <d:error xmlns:d="DAV:" xmlns:s="http://sabredav.org/ns">
//   <s:exception>Sabre\DAV\Exception\NotFound</s:exception>
//   <s:message>File with name Photo could not be located</s:message>
// </d:error>
type Error struct {
	Exception  string `xml:"exception,omitempty"`
	Message    string `xml:"message,omitempty"`
	Status     string
	StatusCode int
}

// Error returns a string for the error and statistifes the error interface
func (e *Error) Error() string {
	var out []string
	if e.Message != "" {
		out = append(out, e.Message)
	}
	if e.Exception != "" {
		out = append(out, e.Exception)
	}
	if e.Status != "" {
		out = append(out, e.Status)
	}
	if len(out) == 0 {
		return "Webdav Error"
	}
	return strings.Join(out, ": ")
}

// Time represents represents date and time information for the
// webdav API marshalling to and from timeFormat
type Time time.Time

// MarshalXML turns a Time into XML
func (t *Time) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	timeString := (*time.Time)(t).Format(timeFormat)
	return e.EncodeElement(timeString, start)
}

// Possible time formats to parse the time with
var timeFormats = []string{
	timeFormat,     // Wed, 27 Sep 2017 14:28:34 GMT (as per RFC)
	time.RFC1123Z,  // Fri, 05 Jan 2018 14:14:38 +0000 (as used by mydrive.ch)
	time.UnixDate,  // Wed May 17 15:31:58 UTC 2017 (as used in an internal server)
	noZerosRFC1123, // Fri, 7 Sep 2018 08:49:58 GMT (as used by server in #2574)
}

// UnmarshalXML turns XML into a Time
func (t *Time) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var v string
	err := d.DecodeElement(&v, &start)
	if err != nil {
		return err
	}

	// If time is missing then return the epoch
	if v == "" {
		*t = Time(time.Unix(0, 0))
		return nil
	}

	// Parse the time format in multiple possible ways
	var newT time.Time
	for _, timeFormat := range timeFormats {
		newT, err = time.Parse(timeFormat, v)
		if err == nil {
			*t = Time(newT)
			break
		}
	}
	return err
}
<commit_msg>WebDav - Add RFC3339 date format - fixes #2712<commit_after>// Package api has type definitions for webdav
package api

import (
	"encoding/xml"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	// Wed, 27 Sep 2017 14:28:34 GMT
	timeFormat = time.RFC1123
	// The same as time.RFC1123 with optional leading zeros on the date
	// see https://github.com/ncw/rclone/issues/2574
	noZerosRFC1123 = "Mon, _2 Jan 2006 15:04:05 MST"
)

// Multistatus contains responses returned from an HTTP 207 return code
type Multistatus struct {
	Responses []Response `xml:"response"`
}

// Response contains an Href the response it about and its properties
type Response struct {
	Href  string `xml:"href"`
	Props Prop   `xml:"propstat"`
}

// Prop is the properties of a response
//
// This is a lazy way of decoding the multiple <s:propstat> in the
// response.
//
// The response might look like this
//
// <d:response>
//   <d:href>/remote.php/webdav/Nextcloud%20Manual.pdf</d:href>
//   <d:propstat>
//     <d:prop>
//       <d:getlastmodified>Tue, 19 Dec 2017 22:02:36 GMT</d:getlastmodified>
//       <d:getcontentlength>4143665</d:getcontentlength>
//       <d:resourcetype/>
//       <d:getetag>"048d7be4437ff7deeae94db50ff3e209"</d:getetag>
//       <d:getcontenttype>application/pdf</d:getcontenttype>
//     </d:prop>
//     <d:status>HTTP/1.1 200 OK</d:status>
//   </d:propstat>
//   <d:propstat>
//     <d:prop>
//       <d:quota-used-bytes/>
//       <d:quota-available-bytes/>
//     </d:prop>
//     <d:status>HTTP/1.1 404 Not Found</d:status>
//   </d:propstat>
// </d:response>
//
// So we elide the array of <d:propstat> and within that the array of
// <d:prop> into one struct.
//
// Note that status collects all the status values for which we just
// check the first is OK.
type Prop struct {
	Status   []string  `xml:"DAV: status"`
	Name     string    `xml:"DAV: prop>displayname,omitempty"`
	Type     *xml.Name `xml:"DAV: prop>resourcetype>collection,omitempty"`
	Size     int64     `xml:"DAV: prop>getcontentlength,omitempty"`
	Modified Time      `xml:"DAV: prop>getlastmodified,omitempty"`
}

// Parse a status of the form "HTTP/1.1 200 OK" or "HTTP/1.1 200"
var parseStatus = regexp.MustCompile(`^HTTP/[0-9.]+\s+(\d+)`)

// StatusOK examines the Status and returns an OK flag
func (p *Prop) StatusOK() bool {
	// Assume OK if no statuses received
	if len(p.Status) == 0 {
		return true
	}
	match := parseStatus.FindStringSubmatch(p.Status[0])
	if len(match) < 2 {
		return false
	}
	code, err := strconv.Atoi(match[1])
	if err != nil {
		return false
	}
	if code >= 200 && code < 300 {
		return true
	}
	return false
}

// PropValue is a tagged name and value
type PropValue struct {
	XMLName xml.Name `xml:""`
	Value   string   `xml:",chardata"`
}

// Error is used to desribe webdav errors
//
// <d:error xmlns:d="DAV:" xmlns:s="http://sabredav.org/ns">
//   <s:exception>Sabre\DAV\Exception\NotFound</s:exception>
//   <s:message>File with name Photo could not be located</s:message>
// </d:error>
type Error struct {
	Exception  string `xml:"exception,omitempty"`
	Message    string `xml:"message,omitempty"`
	Status     string
	StatusCode int
}

// Error returns a string for the error and statistifes the error interface
func (e *Error) Error() string {
	var out []string
	if e.Message != "" {
		out = append(out, e.Message)
	}
	if e.Exception != "" {
		out = append(out, e.Exception)
	}
	if e.Status != "" {
		out = append(out, e.Status)
	}
	if len(out) == 0 {
		return "Webdav Error"
	}
	return strings.Join(out, ": ")
}

// Time represents represents date and time information for the
// webdav API marshalling to and from timeFormat
type Time time.Time

// MarshalXML turns a Time into XML
func (t *Time) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	timeString := (*time.Time)(t).Format(timeFormat)
	return e.EncodeElement(timeString, start)
}

// Possible time formats to parse the time with
var timeFormats = []string{
	timeFormat,     // Wed, 27 Sep 2017 14:28:34 GMT (as per RFC)
	time.RFC1123Z,  // Fri, 05 Jan 2018 14:14:38 +0000 (as used by mydrive.ch)
	time.UnixDate,  // Wed May 17 15:31:58 UTC 2017 (as used in an internal server)
	noZerosRFC1123, // Fri, 7 Sep 2018 08:49:58 GMT (as used by server in #2574)
        time.RFC3339,   // Wed, 31 Oct 2018 13:57:11 CET (as used by komfortcloud.de)
}

// UnmarshalXML turns XML into a Time
func (t *Time) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var v string
	err := d.DecodeElement(&v, &start)
	if err != nil {
		return err
	}

	// If time is missing then return the epoch
	if v == "" {
		*t = Time(time.Unix(0, 0))
		return nil
	}

	// Parse the time format in multiple possible ways
	var newT time.Time
	for _, timeFormat := range timeFormats {
		newT, err = time.Parse(timeFormat, v)
		if err == nil {
			*t = Time(newT)
			break
		}
	}
	return err
}
<|endoftext|>