<commit_before>package lib

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

//Line represents line contents and settings.
type Line struct {
	*Configuration
	Line    string   //String of the entire line
	Columns []string //Parsed columns
	ID      int64    //Calculated uniqueobject ID.  Used for path calculation.
	LastID  int64    //Remeber last object ID processed.  Prevents duplicates.
	Dir     string   //Directory of object file
	Path    string   //Full path of object calc. from objectID
	File    *os.File //File to write the line out to

	lineCount  int64 //Increment for each line
	successful int64 //Sucessfully copied files
	failed     int64 //Failed to copy count.
	duplicates int64 //duplicates found in the sort file
	skipped    int64 //offset rows (future other skips)
}

//ProcessLine copies file to output.
//Returns false in the event of error or empty
func (l *Line) ProcessLine() (b bool) {
	var err error
	var id int

	//Is the line empty?  Skip line.
	if len(l.Columns) == 1 && len(l.Columns[0]) == 0 {
		return false
	}

	//Get object ID from column
	id, err = strconv.Atoi(l.Columns[l.ColObjectID])
	l.ID = int64(id)
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Check to see if object was already processed.
	//We do this by comparing objectID, which represents a unique object
	if l.LastID == l.ID {
		l.duplicates++
		log.Println("Skipping duplicate.", l.ID)
		writeLine(l.Line, l.outDups)
		return false
	}

	//Remember this line to compare with the next line to check for duplicates
	l.LastID = l.ID

	//Get full path for file in
	l.GetInPath()
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Create new Line for line out.  Copy values from line in.
	lo := *l
	//What object out are we on?  Should be successful plus offset
	current := l.successful + l.OutFileRenameIntOffset
	lo.ID = int64(current)

	//init out file
	lo.OutLineFile()
	defer lo.File.Close()

	//Get full path for file out.
	lo.Path, err = lo.GetOutPath()
	if err != nil {
		//Write line in to error line.
		return l.errorLine(&l.Line, err)
	}

	//fmt.Println("inFullPath:", l.Path, "outFullPath:", lo.Path)
	//Create path for out.
	MkdirAll(lo.Dir)

	//Copy file
	err = l.copy(l.Path, lo.Path)
	//Write line to out file if successful
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Columns
	err = lo.GenLineFromColumns()
	if err != nil {
		return l.errorLine(&l.Line, err)
	}
	writeLine(lo.Line+lo.Delimiter+lo.Path, lo.File)

	return true
}

//GetInPath Get full path for file in
func (l *Line) GetInPath() (fullPath string, err error) {
	//Get extension for in file.
	var inExtension string
	if l.InFileExt == "" {
		inExtension = l.Columns[l.ColFileExtIn]
	} else {
		inExtension = l.InFileExt
	}

	//Use a static path OR path from flat file dump row
	var parentPath string
	if l.InDir != "" {
		parentPath = l.InDir
	} else {
		parentPath = l.Columns[l.ColFileName]
	}

	var subpath string
	subpath, err = l.GetPathFromID()
	if err != nil {
		return "", err
	}

	//Full path for file in.
	//99,999,999 plus hex codes.
	//File names change to a hex convention after 99,999,999.
	//Convention: "A" + the object ID converted into HEX
	var fileName string
	if l.ID > 99999999 {
		fileName = "A" + strconv.FormatInt(l.ID, 16)
		fileName = strings.ToUpper(fileName)
	} else {
		fileName = strconv.FormatInt(l.ID, 10)
	}

	fullPath = filepath.Join(parentPath, subpath, fileName) + inExtension
	l.Path = fullPath
	return fullPath, nil
}

//GetOutPath Get full path for out file
func (l *Line) GetOutPath() (out string, err error) {
	//full file path out.
	var filename string

	if l.OutFileRenameInt == true {
		//Use ID as file name
		filename = strconv.FormatInt(l.ID, 10)
	} else {
		//Get the file name from the in file
		filename = l.Columns[l.ColFileName]
	}

	//Extension for out file.
	//Use static extension.  If blank, assume this value is provided in row's column
	outExtension := l.OutFileExt
	if outExtension == "" {
		outExtension = l.Columns[l.ColFileExtOut]
	}

	//Are we batching?  If so, add batch name to path.
	if l.OutAutoBatch == true {
		l.Dir = l.GetBatch()
	}

	//Create full path for out file
	//Create parent path first.
	if l.OutXtenderStructure == true {
		var subpath string
		subpath, err = l.GetPathFromID()
		if err != nil {
			return "", err
		}
		//Add the subpath to the current dir.
		l.Dir = filepath.Join(l.Dir, subpath)
	}

	//Add the parent Out directory.  This give the final directory.
	l.Dir = filepath.Join(l.OutDir, l.Dir)

	//Full path including file name.
	l.Path = filepath.Join(l.Dir, filename) + outExtension

	return l.Path, nil
}

//OutLineFile Create the "OutLines" file.
//Should call defer os.Close
func (l *Line) OutLineFile() {
	outFile := l.OutLinesName
	b := l.GetBatch()

	if b != "" {
		outFile = b + "_" + l.OutLinesName
	}
	outPath := filepath.Join(l.OutDir, outFile)
	var err error
	l.File, err = os.OpenFile(outPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		log.Fatal(err)
	}
}

//GetBatch Return the batch name, including incrementer
func (l *Line) GetBatch() string {
	//If batch is false, there is no batch.  Return blank string
	if l.OutAutoBatch == false {
		return ""
	}

	var batchCount int64
	batchCount = l.ID / l.OutAutoBatchCount

	//Zero pad our batch
	pad := strconv.Itoa(l.OutAutoBatchZeroPad)
	bcn := fmt.Sprintf("%0"+pad+"d", batchCount)

	return l.OutAutoBatchName + bcn
}

//GenLineFromColumns Instead of copying the line from flat in, grab only some columns and write those to out.
func (l *Line) GenLineFromColumns() (err error) {

	var line string
	cols := strings.Split(l.OutLinesColumns, ",")
	for _, v := range cols {
		i, err := strconv.Atoi(v)
		if err != nil {
			return err
		}

		//Fencepost
		if line == "" {
			line += l.Columns[i]
		} else {
			line += l.Delimiter + l.Columns[i]
		}
	}

	l.Line = line

	return nil
}

//GetPathFromID Calculate the ApplicationXtender from a given object id, s
//See README for examples and further explaination.
func (l *Line) GetPathFromID() (p string, e error) {
	//For each step of the path, we will calculate that portion of the path
	for i := l.DirDepth; i > 0; i-- {
		//Get maximum of how many objects there could be at this level.
		powered := math.Pow(float64(l.FolderSize), float64(i))
		//divide ID by max objects, and then Mod the id by how many objects per folder.
		subpath := int(math.Mod((float64(l.ID) / powered), float64(l.FolderSize)))
		//Add this portion of the path to the directory string.  Iterate next level until done.
		p = filepath.Join(p, strconv.Itoa(subpath))
		//fmt.Println("Path:", p, "Id value:", id, "Subpath:", subpath, "I:", i, "Powered: ", powered)
	}
	return p, e
}

//writeLine writes given string to given file with a newline appended at the end.
func writeLine(s string, f *os.File) {
	f.WriteString(s + "\n")
}

//errorLine
//helper function for when encountering error lines.  Logs the line number with
//error and and writes the line to the error line file.
func (l *Line) errorLine(line *string, err error) (b bool) {
	writeLine(*line, l.errorLines)
	log.Println("Line:", l.lineCount, err)
	return false
}

//copy copies file from in to out.
func (l *Line) copy(in string, out string) (err error) {
	i, err := os.Open(in)
	if err != nil {
		l.failed++
		message := fmt.Sprint("File does not exist. ", in)
		fmt.Println(message)
		return errors.New(message)
	}
	defer i.Close()

	o, err := os.Create(out)
	if err != nil {
		log.Println("Cannot create file out.  Stopping execution", out)
		panic(err)
	}
	defer o.Close()

	_, err = io.Copy(o, i)
	if err != nil {
		log.Println("Copying failed.  Stopping execution", out)
		panic(err)
	} else {
		l.successful++
	}

	return nil
}
<commit_msg>Update line.go<commit_after>package lib

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

//Line represents line contents and settings.
type Line struct {
	*Configuration
	Line    string   //String of the entire line
	Columns []string //Parsed columns
	ID      int64    //Calculated uniqueobject ID.  Used for path calculation.
	LastID  int64    //Remeber last object ID processed.  Prevents duplicates.
	Dir     string   //Directory of object file
	Path    string   //Full path of object calc. from objectID
	File    *os.File //File to write the line out to

	lineCount  int64 //Increment for each line
	successful int64 //Sucessfully copied files
	failed     int64 //Failed to copy count.
	duplicates int64 //duplicates found in the sort file
	skipped    int64 //offset rows (future other skips)
}

//ProcessLine copies file to output.
//Returns false in the event of error or empty
func (l *Line) ProcessLine() (b bool) {
	var err error
	var id int

	//Is the line empty?  Skip line.
	if len(l.Columns) == 1 && len(l.Columns[0]) == 0 {
		return false
	}

	//Get object ID from column
	id, err = strconv.Atoi(l.Columns[l.ColObjectID])
	l.ID = int64(id)
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Check to see if object was already processed.
	//We do this by comparing objectID, which represents a unique object
	if l.LastID == l.ID {
		l.duplicates++
		log.Println("Skipping duplicate.", l.ID)
		writeLine(l.Line, l.outDups)
		return false
	}

	//Remember this line to compare with the next line to check for duplicates
	l.LastID = l.ID

	//Get full path for file in
	l.GetInPath()
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Create new Line for line out.  Copy values from line in.
	lo := *l
	//What object out are we on?  Should be successful plus offset
	current := l.successful + l.OutFileRenameIntOffset

	//*** Added 2/16/17 PCS
	if l.OutFileRenameInt != true {
		//Use the object ID in instead since we aren't changing the names
		current = l.ID
	}

	lo.ID = int64(current)

	//init out file
	lo.OutLineFile()
	defer lo.File.Close()

	//Get full path for file out.
	lo.Path, err = lo.GetOutPath()
	if err != nil {
		//Write line in to error line.
		return l.errorLine(&l.Line, err)
	}

	//fmt.Println("inFullPath:", l.Path, "outFullPath:", lo.Path)
	//Create path for out.
	MkdirAll(lo.Dir)

	//Copy file
	err = l.copy(l.Path, lo.Path)
	//Write line to out file if successful
	if err != nil {
		return l.errorLine(&l.Line, err)
	}

	//Columns
	err = lo.GenLineFromColumns()
	if err != nil {
		return l.errorLine(&l.Line, err)
	}
	writeLine(lo.Line+lo.Delimiter+lo.Path, lo.File)

	return true
}

//GetInPath Get full path for file in
func (l *Line) GetInPath() (fullPath string, err error) {
	//Get extension for in file.
	var inExtension string
	if l.InFileExt == "" {
		inExtension = l.Columns[l.ColFileExtIn]
	} else {
		inExtension = l.InFileExt
	}

	//Use a static path OR path from flat file dump row
	var parentPath string
	if l.InDir != "" {
		parentPath = l.InDir
	} else {
		parentPath = l.Columns[l.ColFileName]
	}

	var subpath string
	subpath, err = l.GetPathFromID()
	if err != nil {
		return "", err
	}

	//Full path for file in.
	//99,999,999 plus hex codes.
	//File names change to a hex convention after 99,999,999.
	//Convention: "A" + the object ID converted into HEX
	var fileName string
	if l.ID > 99999999 {
		fileName = "A" + strconv.FormatInt(l.ID, 16)
		fileName = strings.ToUpper(fileName)
	} else {
		fileName = strconv.FormatInt(l.ID, 10)
	}

	fullPath = filepath.Join(parentPath, subpath, fileName) + inExtension
	l.Path = fullPath
	return fullPath, nil
}

//GetOutPath Get full path for out file
func (l *Line) GetOutPath() (out string, err error) {
	//full file path out.
	var filename string

	if l.OutFileRenameInt == true {
		//Use ID as file name
		filename = strconv.FormatInt(l.ID, 10)
	} else {
		//Get the file name from the in file
		filename = l.Columns[l.ColFileName]
	}

	//Extension for out file.
	//Use static extension.  If blank, assume this value is provided in row's column
	outExtension := l.OutFileExt
	if outExtension == "" {
		outExtension = l.Columns[l.ColFileExtOut]
	}

	//Are we batching?  If so, add batch name to path.
	if l.OutAutoBatch == true {
		l.Dir = l.GetBatch()
	}

	//Create full path for out file
	//Create parent path first.
	if l.OutXtenderStructure == true {
		var subpath string
		subpath, err = l.GetPathFromID()
		if err != nil {
			return "", err
		}
		//Add the subpath to the current dir.
		l.Dir = filepath.Join(l.Dir, subpath)
	}

	//Add the parent Out directory.  This give the final directory.
	l.Dir = filepath.Join(l.OutDir, l.Dir)

	//Full path including file name.
	l.Path = filepath.Join(l.Dir, filename) + outExtension

	return l.Path, nil
}

//OutLineFile Create the "OutLines" file.
//Should call defer os.Close
func (l *Line) OutLineFile() {
	outFile := l.OutLinesName
	b := l.GetBatch()

	if b != "" {
		outFile = b + "_" + l.OutLinesName
	}
	outPath := filepath.Join(l.OutDir, outFile)
	var err error
	l.File, err = os.OpenFile(outPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		log.Fatal(err)
	}
}

//GetBatch Return the batch name, including incrementer
func (l *Line) GetBatch() string {
	//If batch is false, there is no batch.  Return blank string
	if l.OutAutoBatch == false {
		return ""
	}

	var batchCount int64
	batchCount = l.ID / l.OutAutoBatchCount

	//Zero pad our batch
	pad := strconv.Itoa(l.OutAutoBatchZeroPad)
	bcn := fmt.Sprintf("%0"+pad+"d", batchCount)

	return l.OutAutoBatchName + bcn
}

//GenLineFromColumns Instead of copying the line from flat in, grab only some columns and write those to out.
func (l *Line) GenLineFromColumns() (err error) {

	var line string
	cols := strings.Split(l.OutLinesColumns, ",")
	for _, v := range cols {
		i, err := strconv.Atoi(v)
		if err != nil {
			return err
		}

		//Fencepost
		if line == "" {
			line += l.Columns[i]
		} else {
			line += l.Delimiter + l.Columns[i]
		}
	}

	l.Line = line

	return nil
}

//GetPathFromID Calculate the ApplicationXtender from a given object id, s
//See README for examples and further explaination.
func (l *Line) GetPathFromID() (p string, e error) {
	//For each step of the path, we will calculate that portion of the path
	for i := l.DirDepth; i > 0; i-- {
		//Get maximum of how many objects there could be at this level.
		powered := math.Pow(float64(l.FolderSize), float64(i))
		//divide ID by max objects, and then Mod the id by how many objects per folder.
		subpath := int(math.Mod((float64(l.ID) / powered), float64(l.FolderSize)))
		//Add this portion of the path to the directory string.  Iterate next level until done.
		p = filepath.Join(p, strconv.Itoa(subpath))
		//fmt.Println("i:", i)
		//fmt.Println("Path:", p, "Id value:", l.ID, "Subpath:", subpath, "I:", i, "Powered: ", powered)
	}
	//fmt.Println("******** End GetPathFromID function *******")
	return p, e
}

//writeLine writes given string to given file with a newline appended at the end.
func writeLine(s string, f *os.File) {
	f.WriteString(s + "\n")
}

//errorLine
//helper function for when encountering error lines.  Logs the line number with
//error and and writes the line to the error line file.
func (l *Line) errorLine(line *string, err error) (b bool) {
	writeLine(*line, l.errorLines)
	log.Println("Line:", l.lineCount, err)
	return false
}

//copy copies file from in to out.
func (l *Line) copy(in string, out string) (err error) {
	i, err := os.Open(in)
	if err != nil {
		l.failed++
		message := fmt.Sprint("File does not exist. ", in)
		fmt.Println(message)
		return errors.New(message)
	}
	defer i.Close()

	o, err := os.Create(out)
	if err != nil {
		log.Println("Cannot create file out.  Stopping execution", out)
		panic(err)
	}
	defer o.Close()

	_, err = io.Copy(o, i)
	if err != nil {
		log.Println("Copying failed.  Stopping execution", out)
		panic(err)
	} else {
		l.successful++
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package priorities

import (
	"fmt"

	"github.com/golang/glog"
	"k8s.io/api/core/v1"
	priorityutil "k8s.io/kubernetes/pkg/scheduler/algorithm/priorities/util"
	schedulerapi "k8s.io/kubernetes/pkg/scheduler/api"
	schedulercache "k8s.io/kubernetes/pkg/scheduler/cache"
)

// ResourceAllocationPriority contains information to calculate resource allocation priority.
type ResourceAllocationPriority struct {
	Name   string
	scorer func(requested, allocable *schedulercache.Resource, includeVolumes bool, requestedVolumes int, allocatableVolumes int) int64
}

// PriorityMap priorities nodes according to the resource allocations on the node.
// It will use `scorer` function to calculate the score.
func (r *ResourceAllocationPriority) PriorityMap(
	pod *v1.Pod,
	meta interface{},
	nodeInfo *schedulercache.NodeInfo) (schedulerapi.HostPriority, error) {
	node := nodeInfo.Node()
	if node == nil {
		return schedulerapi.HostPriority{}, fmt.Errorf("node not found")
	}
	allocatable := nodeInfo.AllocatableResource()

	var requested schedulercache.Resource
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		requested = *priorityMeta.nonZeroRequest
	} else {
		// We couldn't parse metadata - fallback to computing it.
		requested = *getNonZeroRequests(pod)
	}

	requested.MilliCPU += nodeInfo.NonZeroRequest().MilliCPU
	requested.Memory += nodeInfo.NonZeroRequest().Memory
	var score int64
	// Check if the pod has volumes and this could be added to scorer function for balanced resource allocation.
	if len(pod.Spec.Volumes) >= 0 && nodeInfo.TransientInfo != nil {
		score = r.scorer(&requested, &allocatable, true, nodeInfo.TransientInfo.TransNodeInfo.RequestedVolumes, nodeInfo.TransientInfo.TransNodeInfo.AllocatableVolumesCount)
	} else {
		score = r.scorer(&requested, &allocatable, false, 0, 0)
	}

	if glog.V(10) {
		glog.Infof(
			"%v -> %v: %v, capacity %d millicores %d memory bytes, total request %d millicores %d memory bytes, score %d",
			pod.Name, node.Name, r.Name,
			allocatable.MilliCPU, allocatable.Memory,
			requested.MilliCPU, requested.Memory,
			score,
		)
	}

	return schedulerapi.HostPriority{
		Host:  node.Name,
		Score: int(score),
	}, nil
}

func getNonZeroRequests(pod *v1.Pod) *schedulercache.Resource {
	result := &schedulercache.Resource{}
	for i := range pod.Spec.Containers {
		container := &pod.Spec.Containers[i]
		cpu, memory := priorityutil.GetNonzeroRequests(&container.Resources.Requests)
		result.MilliCPU += cpu
		result.Memory += memory
	}
	return result
}
<commit_msg>Ouput volumes (total capacity and requests) too along with cpu and memory when the feature BalanceAttachedNodeVolumes is used.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package priorities

import (
	"fmt"

	"github.com/golang/glog"
	"k8s.io/api/core/v1"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/kubernetes/pkg/features"
	priorityutil "k8s.io/kubernetes/pkg/scheduler/algorithm/priorities/util"
	schedulerapi "k8s.io/kubernetes/pkg/scheduler/api"
	schedulercache "k8s.io/kubernetes/pkg/scheduler/cache"
)

// ResourceAllocationPriority contains information to calculate resource allocation priority.
type ResourceAllocationPriority struct {
	Name   string
	scorer func(requested, allocable *schedulercache.Resource, includeVolumes bool, requestedVolumes int, allocatableVolumes int) int64
}

// PriorityMap priorities nodes according to the resource allocations on the node.
// It will use `scorer` function to calculate the score.
func (r *ResourceAllocationPriority) PriorityMap(
	pod *v1.Pod,
	meta interface{},
	nodeInfo *schedulercache.NodeInfo) (schedulerapi.HostPriority, error) {
	node := nodeInfo.Node()
	if node == nil {
		return schedulerapi.HostPriority{}, fmt.Errorf("node not found")
	}
	allocatable := nodeInfo.AllocatableResource()

	var requested schedulercache.Resource
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		requested = *priorityMeta.nonZeroRequest
	} else {
		// We couldn't parse metadata - fallback to computing it.
		requested = *getNonZeroRequests(pod)
	}

	requested.MilliCPU += nodeInfo.NonZeroRequest().MilliCPU
	requested.Memory += nodeInfo.NonZeroRequest().Memory
	var score int64
	// Check if the pod has volumes and this could be added to scorer function for balanced resource allocation.
	if len(pod.Spec.Volumes) >= 0 && utilfeature.DefaultFeatureGate.Enabled(features.BalanceAttachedNodeVolumes) && nodeInfo.TransientInfo != nil {
		score = r.scorer(&requested, &allocatable, true, nodeInfo.TransientInfo.TransNodeInfo.RequestedVolumes, nodeInfo.TransientInfo.TransNodeInfo.AllocatableVolumesCount)
	} else {
		score = r.scorer(&requested, &allocatable, false, 0, 0)
	}

	if glog.V(10) {
		if len(pod.Spec.Volumes) >= 0 && utilfeature.DefaultFeatureGate.Enabled(features.BalanceAttachedNodeVolumes) && nodeInfo.TransientInfo != nil {
			glog.Infof(
				"%v -> %v: %v, capacity %d millicores %d memory bytes, %d volumes, total request %d millicores %d memory bytes %d volumes, score %d",
				pod.Name, node.Name, r.Name,
				allocatable.MilliCPU, allocatable.Memory, nodeInfo.TransientInfo.TransNodeInfo.AllocatableVolumesCount,
				requested.MilliCPU, requested.Memory,
				nodeInfo.TransientInfo.TransNodeInfo.RequestedVolumes,
				score,
			)
		} else {
			glog.Infof(
				"%v -> %v: %v, capacity %d millicores %d memory bytes, total request %d millicores %d memory bytes, score %d",
				pod.Name, node.Name, r.Name,
				allocatable.MilliCPU, allocatable.Memory,
				requested.MilliCPU, requested.Memory,
				score,
			)
		}
	}

	return schedulerapi.HostPriority{
		Host:  node.Name,
		Score: int(score),
	}, nil
}

func getNonZeroRequests(pod *v1.Pod) *schedulercache.Resource {
	result := &schedulercache.Resource{}
	for i := range pod.Spec.Containers {
		container := &pod.Spec.Containers[i]
		cpu, memory := priorityutil.GetNonzeroRequests(&container.Resources.Requests)
		result.MilliCPU += cpu
		result.Memory += memory
	}
	return result
}
<|endoftext|>


<commit_before>package deploymentconfigs

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/golang/glog"

	"k8s.io/client-go/tools/cache"
	kapi "k8s.io/kubernetes/pkg/api"

	"github.com/openshift/origin/pkg/client"
	deployapi "github.com/openshift/origin/pkg/deploy/api"
	triggerapi "github.com/openshift/origin/pkg/image/api/v1/trigger"
	"github.com/openshift/origin/pkg/image/trigger"
)

func indicesForContainerNames(containers []kapi.Container, names []string) []int {
	var index []int
	for _, name := range names {
		for i, container := range containers {
			if name == container.Name {
				index = append(index, i)
			}
		}
	}
	return index
}

func namesInclude(names []string, name string) bool {
	for _, n := range names {
		if name == n {
			return true
		}
	}
	return false
}

// calculateDeploymentConfigTrigger resolves a particular trigger against the deployment config and extracts a list of triggers.
// It will silently ignore triggers that do not point to valid container. It returns empty if no triggers can be found.
func calculateDeploymentConfigTrigger(t deployapi.DeploymentTriggerPolicy, dc *deployapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	if t.ImageChangeParams == nil {
		return nil
	}
	from := t.ImageChangeParams.From
	if from.Kind != "ImageStreamTag" || len(from.Name) == 0 {
		return nil
	}

	var triggers []triggerapi.ObjectFieldTrigger

	// add one trigger per init container and container
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.Containers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.containers[@name='%s'].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.InitContainers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.initContainers[@name='%s'].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	return triggers
}

// calculateDeploymentConfigTriggers transforms a deployment config into a set of image change triggers. If retrieveChanges
// is true an array of the latest state of the triggers will be returned.
func calculateDeploymentConfigTriggers(dc *deployapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	var triggers []triggerapi.ObjectFieldTrigger
	for _, t := range dc.Spec.Triggers {
		addedTriggers := calculateDeploymentConfigTrigger(t, dc)
		triggers = append(triggers, addedTriggers...)
	}
	return triggers
}

// deploymentConfigTriggerIndexer converts deployment config events into entries for the trigger cache, and
// also calculates the latest state of the changes on the object.
type deploymentConfigTriggerIndexer struct {
	prefix string
}

func NewDeploymentConfigTriggerIndexer(prefix string) trigger.Indexer {
	return deploymentConfigTriggerIndexer{prefix: prefix}
}

func (i deploymentConfigTriggerIndexer) Index(obj, old interface{}) (string, *trigger.CacheEntry, cache.DeltaType, error) {
	var (
		triggers []triggerapi.ObjectFieldTrigger
		dc       *deployapi.DeploymentConfig
		change   cache.DeltaType
	)
	switch {
	case obj != nil && old == nil:
		// added
		dc = obj.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Added
	case old != nil && obj == nil:
		// deleted
		dc = old.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Deleted
	default:
		// updated
		dc = obj.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		oldTriggers := calculateDeploymentConfigTriggers(old.(*deployapi.DeploymentConfig))
		switch {
		case len(oldTriggers) == 0:
			change = cache.Added
		case !reflect.DeepEqual(oldTriggers, triggers):
			change = cache.Updated
		}
	}

	if len(triggers) > 0 {
		key := i.prefix + dc.Namespace + "/" + dc.Name
		return key, &trigger.CacheEntry{
			Key:       key,
			Namespace: dc.Namespace,
			Triggers:  triggers,
		}, change, nil
	}
	return "", nil, change, nil
}

// DeploymentConfigReactor converts image trigger changes into updates on deployments.
type DeploymentConfigReactor struct {
	Client client.DeploymentConfigsNamespacer
}

// UpdateDeploymentConfigImages sets the latest image value from all triggers onto each container, returning false if
// one or more triggers could not be resolved yet or an error. The returned dc will be copied if mutated.
func UpdateDeploymentConfigImages(dc *deployapi.DeploymentConfig, tagRetriever trigger.TagRetriever) (*deployapi.DeploymentConfig, bool, error) {
	var updated *deployapi.DeploymentConfig

	// copy the object and reset dc to the copy
	copyObject := func() {
		if updated != nil {
			return
		}
		copied, err := kapi.Scheme.Copy(dc)
		if err != nil {
			return
		}
		dc = copied.(*deployapi.DeploymentConfig)
		updated = dc
	}

	for i, t := range dc.Spec.Triggers {
		p := t.ImageChangeParams
		if p == nil || p.From.Kind != "ImageStreamTag" || len(p.LastTriggeredImage) > 0 {
			continue
		}
		if !p.Automatic {
			continue
		}

		namespace := p.From.Namespace
		if len(namespace) == 0 {
			namespace = dc.Namespace
		}

		ref, _, ok := tagRetriever.ImageStreamTag(namespace, p.From.Name)
		if !ok && len(p.LastTriggeredImage) == 0 {
			glog.V(4).Infof("trigger %#v in deployment %s is not resolveable", p, dc.Name)
			return nil, false, nil
		}
		if len(ref) == 0 {
			ref = p.LastTriggeredImage
		}

		if p.LastTriggeredImage != ref {
			copyObject()
			p = dc.Spec.Triggers[i].ImageChangeParams
			p.LastTriggeredImage = ref
		}

		for i, c := range dc.Spec.Template.Spec.InitContainers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.InitContainers[i]
			container.Image = ref
		}

		for i, c := range dc.Spec.Template.Spec.Containers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.Containers[i]
			container.Image = ref
		}
	}
	return updated, true, nil
}

// ImageChanged is passed a deployment config and a set of changes.
func (r *DeploymentConfigReactor) ImageChanged(obj interface{}, tagRetriever trigger.TagRetriever) error {
	dc := obj.(*deployapi.DeploymentConfig)
	copied, err := kapi.Scheme.DeepCopy(dc)
	if err != nil {
		return err
	}
	newDC := copied.(*deployapi.DeploymentConfig)

	updated, resolvable, err := UpdateDeploymentConfigImages(newDC, tagRetriever)
	if err != nil {
		return err
	}
	if !resolvable {
		if glog.V(4) {
			glog.Infof("Ignoring changes to deployment config %s, has unresolved images: %s", dc.Name, printDeploymentTriggers(newDC.Spec.Triggers))
		}
		return nil
	}
	if updated == nil {
		glog.V(4).Infof("Deployment config %s has not changed", dc.Name)
		return nil
	}
	glog.V(4).Infof("Deployment config %s has changed", dc.Name)
	_, err = r.Client.DeploymentConfigs(dc.Namespace).Update(updated)
	return err
}

func printDeploymentTriggers(triggers []deployapi.DeploymentTriggerPolicy) string {
	var values []string
	for _, t := range triggers {
		if t.ImageChangeParams == nil {
			continue
		}
		values = append(values, fmt.Sprintf("[from=%s last=%s]", t.ImageChangeParams.From.Name, t.ImageChangeParams.LastTriggeredImage))
	}
	return strings.Join(values, ", ")
}
<commit_msg>Fix chained DC IST trigger<commit_after>package deploymentconfigs

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/golang/glog"

	"k8s.io/client-go/tools/cache"
	kapi "k8s.io/kubernetes/pkg/api"

	"github.com/openshift/origin/pkg/client"
	deployapi "github.com/openshift/origin/pkg/deploy/api"
	triggerapi "github.com/openshift/origin/pkg/image/api/v1/trigger"
	"github.com/openshift/origin/pkg/image/trigger"
)

func indicesForContainerNames(containers []kapi.Container, names []string) []int {
	var index []int
	for _, name := range names {
		for i, container := range containers {
			if name == container.Name {
				index = append(index, i)
			}
		}
	}
	return index
}

func namesInclude(names []string, name string) bool {
	for _, n := range names {
		if name == n {
			return true
		}
	}
	return false
}

// calculateDeploymentConfigTrigger resolves a particular trigger against the deployment config and extracts a list of triggers.
// It will silently ignore triggers that do not point to valid container. It returns empty if no triggers can be found.
func calculateDeploymentConfigTrigger(t deployapi.DeploymentTriggerPolicy, dc *deployapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	if t.ImageChangeParams == nil {
		return nil
	}
	from := t.ImageChangeParams.From
	if from.Kind != "ImageStreamTag" || len(from.Name) == 0 {
		return nil
	}

	var triggers []triggerapi.ObjectFieldTrigger

	// add one trigger per init container and container
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.Containers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.containers[@name='%s'].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.InitContainers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.initContainers[@name='%s'].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	return triggers
}

// calculateDeploymentConfigTriggers transforms a deployment config into a set of image change triggers. If retrieveChanges
// is true an array of the latest state of the triggers will be returned.
func calculateDeploymentConfigTriggers(dc *deployapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	var triggers []triggerapi.ObjectFieldTrigger
	for _, t := range dc.Spec.Triggers {
		addedTriggers := calculateDeploymentConfigTrigger(t, dc)
		triggers = append(triggers, addedTriggers...)
	}
	return triggers
}

// deploymentConfigTriggerIndexer converts deployment config events into entries for the trigger cache, and
// also calculates the latest state of the changes on the object.
type deploymentConfigTriggerIndexer struct {
	prefix string
}

func NewDeploymentConfigTriggerIndexer(prefix string) trigger.Indexer {
	return deploymentConfigTriggerIndexer{prefix: prefix}
}

func (i deploymentConfigTriggerIndexer) Index(obj, old interface{}) (string, *trigger.CacheEntry, cache.DeltaType, error) {
	var (
		triggers []triggerapi.ObjectFieldTrigger
		dc       *deployapi.DeploymentConfig
		change   cache.DeltaType
	)
	switch {
	case obj != nil && old == nil:
		// added
		dc = obj.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Added
	case old != nil && obj == nil:
		// deleted
		dc = old.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Deleted
	default:
		// updated
		dc = obj.(*deployapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		oldTriggers := calculateDeploymentConfigTriggers(old.(*deployapi.DeploymentConfig))
		switch {
		case len(oldTriggers) == 0:
			change = cache.Added
		case !reflect.DeepEqual(oldTriggers, triggers):
			change = cache.Updated
		}
	}

	if len(triggers) > 0 {
		key := i.prefix + dc.Namespace + "/" + dc.Name
		return key, &trigger.CacheEntry{
			Key:       key,
			Namespace: dc.Namespace,
			Triggers:  triggers,
		}, change, nil
	}
	return "", nil, change, nil
}

// DeploymentConfigReactor converts image trigger changes into updates on deployments.
type DeploymentConfigReactor struct {
	Client client.DeploymentConfigsNamespacer
}

// UpdateDeploymentConfigImages sets the latest image value from all triggers onto each container, returning false if
// one or more triggers could not be resolved yet or an error. The returned dc will be copied if mutated.
func UpdateDeploymentConfigImages(dc *deployapi.DeploymentConfig, tagRetriever trigger.TagRetriever) (*deployapi.DeploymentConfig, bool, error) {
	var updated *deployapi.DeploymentConfig

	// copy the object and reset dc to the copy
	copyObject := func() {
		if updated != nil {
			return
		}
		copied, err := kapi.Scheme.Copy(dc)
		if err != nil {
			return
		}
		dc = copied.(*deployapi.DeploymentConfig)
		updated = dc
	}

	for i, t := range dc.Spec.Triggers {
		p := t.ImageChangeParams
		if p == nil || p.From.Kind != "ImageStreamTag" {
			continue
		}
		if !p.Automatic {
			continue
		}

		namespace := p.From.Namespace
		if len(namespace) == 0 {
			namespace = dc.Namespace
		}

		ref, _, ok := tagRetriever.ImageStreamTag(namespace, p.From.Name)
		if !ok && len(p.LastTriggeredImage) == 0 {
			glog.V(4).Infof("trigger %#v in deployment %s is not resolveable", p, dc.Name)
			return nil, false, nil
		}
		if ref == p.LastTriggeredImage {
			continue
		}

		if len(ref) == 0 {
			ref = p.LastTriggeredImage
		}

		if p.LastTriggeredImage != ref {
			copyObject()
			p = dc.Spec.Triggers[i].ImageChangeParams
			p.LastTriggeredImage = ref
		}

		for i, c := range dc.Spec.Template.Spec.InitContainers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.InitContainers[i]
			container.Image = ref
		}

		for i, c := range dc.Spec.Template.Spec.Containers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.Containers[i]
			container.Image = ref
		}
	}
	return updated, true, nil
}

// ImageChanged is passed a deployment config and a set of changes.
func (r *DeploymentConfigReactor) ImageChanged(obj interface{}, tagRetriever trigger.TagRetriever) error {
	dc := obj.(*deployapi.DeploymentConfig)
	copied, err := kapi.Scheme.DeepCopy(dc)
	if err != nil {
		return err
	}
	newDC := copied.(*deployapi.DeploymentConfig)

	updated, resolvable, err := UpdateDeploymentConfigImages(newDC, tagRetriever)
	if err != nil {
		return err
	}
	if !resolvable {
		if glog.V(4) {
			glog.Infof("Ignoring changes to deployment config %s, has unresolved images: %s", dc.Name, printDeploymentTriggers(newDC.Spec.Triggers))
		}
		return nil
	}
	if updated == nil {
		glog.V(4).Infof("Deployment config %s has not changed", dc.Name)
		return nil
	}
	glog.V(4).Infof("Deployment config %s has changed", dc.Name)
	_, err = r.Client.DeploymentConfigs(dc.Namespace).Update(updated)
	return err
}

func printDeploymentTriggers(triggers []deployapi.DeploymentTriggerPolicy) string {
	var values []string
	for _, t := range triggers {
		if t.ImageChangeParams == nil {
			continue
		}
		values = append(values, fmt.Sprintf("[from=%s last=%s]", t.ImageChangeParams.From.Name, t.ImageChangeParams.LastTriggeredImage))
	}
	return strings.Join(values, ", ")
}
<|endoftext|>


<commit_before>package logger

import (
    "io"
    "log"
    "sync"
    "sync/atomic"
)

type Level int32

const (
    Debug Level = iota
    Info
    Error
    Always
)

type loggerObject struct {
    Flags   int
    Level   int32
    Logger *log.Logger
    Writer  io.Writer
}

type logger struct {
    Jam  []loggerObject
    Once   sync.Once
}

var instance logger

func GetLevel(index int) Level {
    var ret Level = Always
    logger := getLogger()

    if index < len(logger.Jam) {
        ret = Level(atomic.LoadInt32(&logger.Jam[index].Level))
    }

    return ret
}

func getLogger() *logger {
    instance.Once.Do(func() {
        for i := range instance.Jam {
            instance.Jam[i].Logger = log.New(instance.Jam[i].Writer,
                                             "",
                                             instance.Jam[i].Flags)
        }
    })
    return &instance
}

func Init(flags int, level Level, writer io.Writer) {
    instance.Jam = append(instance.Jam, loggerObject{ Flags: flags, Level: int32(level), Writer: writer })
}

func SetLevel(index int, level Level) {
    logger := getLogger()

    if index < len(logger.Jam) {
        atomic.StoreInt32(&logger.Jam[index].Level, int32(level))
    }
}

func println(level Level, v ...interface{}) {
    logger := getLogger()

    for i := range logger.Jam {
        if level >= GetLevel(i) {
            prefix := ""

            switch level {
            case Debug:
                prefix = "[DBG]"
            case Info:
                prefix = "[INF]"
            case Error:
                prefix = "[ERR]"
            case Always:
                prefix = "[ALW]"
            }

            v = append([]interface{}{prefix}, v...)
            logger.Jam[i].Logger.Println(v...)
        }
    }
}

func PrintlnAlways(v ...interface{}) {
    println(Always, v...)
}

func PrintlnDebug(v ...interface{}) {
    println(Debug, v...)
}

func PrintlnError(v ...interface{}) {
    println(Error, v...)
}

func PrintlnInfo(v ...interface{}) {
    println(Info, v...)
}
<commit_msg>Fixed prepending of multiple log level tags to log message<commit_after>package logger

import (
    "io"
    "log"
    "sync"
    "sync/atomic"
)

type Level int32

const (
    Debug Level = iota
    Info
    Error
    Always
)

type loggerObject struct {
    Flags   int
    Level   int32
    Logger *log.Logger
    Writer  io.Writer
}

type logger struct {
    Jam  []loggerObject
    Once   sync.Once
}

var instance logger

func GetLevel(index int) Level {
    var ret Level = Always
    logger := getLogger()

    if index < len(logger.Jam) {
        ret = Level(atomic.LoadInt32(&logger.Jam[index].Level))
    }

    return ret
}

func getLogger() *logger {
    instance.Once.Do(func() {
        for i := range instance.Jam {
            instance.Jam[i].Logger = log.New(instance.Jam[i].Writer,
                                             "",
                                             instance.Jam[i].Flags)
        }
    })
    return &instance
}

func Init(flags int, level Level, writer io.Writer) {
    instance.Jam = append(instance.Jam, loggerObject{ Flags: flags, Level: int32(level), Writer: writer })
}

func SetLevel(index int, level Level) {
    logger := getLogger()

    if index < len(logger.Jam) {
        atomic.StoreInt32(&logger.Jam[index].Level, int32(level))
    }
}

func println(level Level, v ...interface{}) {
    logger := getLogger()

    for i := range logger.Jam {
        if level >= GetLevel(i) {
            prefix := ""

            switch level {
            case Debug:
                prefix = "[DBG]"
            case Info:
                prefix = "[INF]"
            case Error:
                prefix = "[ERR]"
            case Always:
                prefix = "[ALW]"
            }

            m := append([]interface{}{prefix}, v...)
            logger.Jam[i].Logger.Println(m...)
        }
    }
}

func PrintlnAlways(v ...interface{}) {
    println(Always, v...)
}

func PrintlnDebug(v ...interface{}) {
    println(Debug, v...)
}

func PrintlnError(v ...interface{}) {
    println(Error, v...)
}

func PrintlnInfo(v ...interface{}) {
    println(Info, v...)
}
<|endoftext|>


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

// Package logger provides a logger that logs to Google Cloud Logging. It's a thin wrapper around
// golang/cloud/logging package.
package logger

import (
	"context"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/logging"
	"flag"
	"github.com/golang/glog"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"
)

var (
	debugLog     = flag.Bool("debug_log", false, "Whether to output debug logs or not")
	debugLogList = flag.String("debug_logname_regex", "", "Enable debug logs for only for log names that match this regex (e.g. --debug_logname_regex=.*probe1.*")

	// Enable/Disable cloud logging
	disableCloudLogging = flag.Bool("disable_cloud_logging", false, "Disable cloud logging.")

	// LogPrefixEnvVar environment variable is used to determine the stackdriver
	// log name prefix. Default prefix is "cloudprober".
	LogPrefixEnvVar = "CLOUDPROBER_LOG_PREFIX"
)

// EnvVars defines environment variables that can be used to modify the logging
// behavior.
var EnvVars = struct {
	DisableCloudLogging, DebugLog string
}{
	"CLOUDPROBER_DISABLE_CLOUD_LOGGING",
	"CLOUDPROBER_DEBUG_LOG",
}

const (
	// Prefix for the cloudprober stackdriver log names.
	cloudproberPrefix = "cloudprober"
)

const (
	// Regular Expression for all characters that are illegal for log names
	//	Ref: https://cloud.google.com/logging/docs/api/ref_v2beta1/rest/v2beta1/LogEntry
	disapprovedRegExp = "[^A-Za-z0-9_/.-]"

	// MaxLogEntrySize Max size of each log entry (100 KB)
	// This limit helps avoid creating very large log lines in case someone
	// accidentally creates a large EventMetric, which in turn is possible due to
	// unbounded nature of "map" metric where keys are created on demand.
	//
	// TODO(manugarg): We can possibly get rid of this check now as the code that
	// could cause a large map metric has been fixed now. Earlier, cloudprober's
	// HTTP server used to respond to all URLs and used to record access to those
	// URLs as a "map" metric. Now, it responds only to pre-configured URLs.
	MaxLogEntrySize = 102400
)

func enableDebugLog(debugLog bool, debugLogRe string, logName string) bool {
	if !debugLog && debugLogRe == "" {
		return false
	}

	if debugLog && debugLogRe == "" {
		// Enable for all logs, regardless of log names.
		return true
	}

	r, err := regexp.Compile(debugLogRe)
	if err != nil {
		panic(fmt.Sprintf("error while parsing log name regex (%s): %v", debugLogRe, err))
	}

	if r.MatchString(logName) {
		return true
	}

	return false
}

// Logger implements a logger that logs messages to Google Cloud Logging. It
// provides a suite of methods where each method corresponds to a specific
// logging.Level, e.g. Error(paylod interface{}). Each method takes a payload
// that has to either be a JSON-encodable object, a string or a []byte slice
// (all other types of payload will result in error).
//
// It falls back to logging through the traditional logger if:
//
//   * Not running on GCE,
//   * Logging client is uninitialized (e.g. for testing),
//   * Logging to cloud fails for some reason.
//
// Logger{} is a valid object that will log through the traditional logger.
//
type Logger struct {
	name                string
	logc                *logging.Client
	logger              *logging.Logger
	debugLog            bool
	disableCloudLogging bool
	// TODO(manugarg): Logger should eventually embed the probe id and each probe
	// should get a different Logger object (embedding that probe's probe id) but
	// sharing the same logging client. We could then make probe id one of the
	// metadata on all logging messages.
}

// NewCloudproberLog is a convenient wrapper around New that sets context to
// context.Background and attaches cloudprober prefix to log names.
func NewCloudproberLog(component string) (*Logger, error) {
	cpPrefix := cloudproberPrefix

	envLogPrefix := os.Getenv(LogPrefixEnvVar)
	if envLogPrefix != "" {
		cpPrefix = envLogPrefix
	}

	return New(context.Background(), cpPrefix+"."+component)
}

// New returns a new Logger object with cloud logging client initialized if running on GCE.
func New(ctx context.Context, logName string) (*Logger, error) {
	l := &Logger{
		name:                logName,
		debugLog:            enableDebugLog(*debugLog, *debugLogList, logName),
		disableCloudLogging: *disableCloudLogging,
	}

	if !metadata.OnGCE() || l.disableCloudLogging {
		return l, nil
	}

	l.Infof("Running on GCE. Logs for %s will go to Cloud (Stackdriver).", logName)
	if err := l.EnableStackdriverLogging(ctx); err != nil {
		return nil, err
	}
	return l, nil
}

// EnableStackdriverLogging enables logging to stackdriver.
func (l *Logger) EnableStackdriverLogging(ctx context.Context) error {
	if !metadata.OnGCE() {
		return fmt.Errorf("not running on GCE")
	}

	projectID, err := metadata.ProjectID()
	if err != nil {
		return err
	}

	if l.name == "" {
		return fmt.Errorf("logName cannot be empty")
	}
	// Check for illegal characters in the log name
	if match, err := regexp.Match(disapprovedRegExp, []byte(l.name)); err != nil || match {
		if err != nil {
			return fmt.Errorf("unable to parse logName: %v", err)
		}
		return fmt.Errorf("logName of %s contains an invalid character, valid characters are [A-Za-z0-9_/.-]", l.name)
	}
	// Any forward slashes need to be URL encoded, so we query escape to replace them
	logName := url.QueryEscape(l.name)

	l.logc, err = logging.NewClient(ctx, projectID, option.WithTokenSource(google.ComputeTokenSource("")))
	if err != nil {
		return err
	}

	loggerOpts := []logging.LoggerOption{
		// Encourage batching of write requests.
		// Flush logs to remote logging after 1000 entries (default is 10).
		logging.EntryCountThreshold(1000),
		// Maximum amount of time that an item should remain buffered in memory
		// before being flushed to the logging service. Default is 1 second.
		// We want flushing to be mostly driven by the buffer size (configured
		// above), rather than time.
		logging.DelayThreshold(10 * time.Second),
	}

	// Add instance_name to common labels if available.
	instanceName, err := metadata.InstanceName()
	if err != nil {
		l.Infof("Error getting instance name on GCE. Possibly running on GKE: %v", err)
	} else {
		loggerOpts = append(loggerOpts, logging.CommonLabels(map[string]string{"instance_name": instanceName}))
	}

	l.logger = l.logc.Logger(logName, loggerOpts...)
	return nil
}

func payloadToString(payload ...string) string {
	if len(payload) == 1 {
		return payload[0]
	}

	var b strings.Builder
	for _, s := range payload {
		b.WriteString(s)
	}
	return b.String()
}

// log sends payload ([]string) to cloud logging. If cloud logging client is
// not initialized (e.g. if not running on GCE) or cloud logging fails for some
// reason, it writes logs through the traditional logger.
func (l *Logger) log(severity logging.Severity, payload ...string) {
	payloadStr := payloadToString(payload...)

	if len(payloadStr) > MaxLogEntrySize {
		truncateMsg := "... (truncated)"
		truncateMsgLen := len(truncateMsg)
		payloadStr = payloadStr[:MaxLogEntrySize-truncateMsgLen] + truncateMsg
	}

	if l == nil {
		genericLog(severity, "nil", payloadStr)
		return
	}

	if l.logger == nil {
		genericLog(severity, l.name, payloadStr)
		return
	}

	l.logger.Log(logging.Entry{
		Severity: severity,
		Payload:  payloadStr,
	})
}

// Close closes the cloud logging client if it exists. This flushes the buffer
// and should be called before exiting the program to ensure all logs are persisted.
func (l *Logger) Close() error {
	if l != nil && l.logc != nil {
		return l.logc.Close()
	}

	return nil
}

// Debug logs messages with logging level set to "Debug".
func (l *Logger) Debug(payload ...string) {
	if l.debugLog {
		l.log(logging.Debug, payload...)
	}
}

// Info logs messages with logging level set to "Info".
func (l *Logger) Info(payload ...string) {
	l.log(logging.Info, payload...)
}

// Warning logs messages with logging level set to "Warning".
func (l *Logger) Warning(payload ...string) {
	l.log(logging.Warning, payload...)
}

// Error logs messages with logging level set to "Error".
func (l *Logger) Error(payload ...string) {
	l.log(logging.Error, payload...)
}

// Critical logs messages with logging level set to "Critical" and
// exits the process with error status. The buffer is flushed before exiting.
func (l *Logger) Critical(payload ...string) {
	l.log(logging.Critical, payload...)
	if err := l.Close(); err != nil {
		panic(fmt.Sprintf("could not close client: %v", err))
	}
	os.Exit(1)
}

// Debugf logs formatted text messages with logging level "Debug".
func (l *Logger) Debugf(format string, args ...interface{}) {
	if l != nil && l.debugLog {
		l.log(logging.Debug, fmt.Sprintf(format, args...))
	}
}

// Infof logs formatted text messages with logging level "Info".
func (l *Logger) Infof(format string, args ...interface{}) {
	l.log(logging.Info, fmt.Sprintf(format, args...))
}

// Warningf logs formatted text messages with logging level "Warning".
func (l *Logger) Warningf(format string, args ...interface{}) {
	l.log(logging.Warning, fmt.Sprintf(format, args...))
}

// Errorf logs formatted text messages with logging level "Error".
func (l *Logger) Errorf(format string, args ...interface{}) {
	l.log(logging.Error, fmt.Sprintf(format, args...))
}

// Criticalf logs formatted text messages with logging level "Critical" and
// exits the process with error status. The buffer is flushed before exiting.
func (l *Logger) Criticalf(format string, args ...interface{}) {
	l.log(logging.Critical, fmt.Sprintf(format, args...))
	if err := l.Close(); err != nil {
		panic(fmt.Sprintf("could not close client: %v", err))
	}
	os.Exit(1)
}

func genericLog(severity logging.Severity, name string, s string) {
	// Set the caller frame depth to 3 so that can get to the actual caller of
	// the logger. genericLog -> log -> Info* -> actualCaller
	depth := 3

	s = fmt.Sprintf("[%s] %s", name, s)

	switch severity {
	case logging.Debug, logging.Info:
		glog.InfoDepth(depth, s)
	case logging.Warning:
		glog.WarningDepth(depth, s)
	case logging.Error:
		glog.ErrorDepth(depth, s)
	case logging.Critical:
		glog.FatalDepth(depth, s)
	}
}

func envVarSet(key string) bool {
	v, ok := os.LookupEnv(key)
	if ok && strings.ToUpper(v) != "NO" && strings.ToUpper(v) != "FALSE" {
		return true
	}
	return false
}

func init() {
	if envVarSet(EnvVars.DisableCloudLogging) {
		*disableCloudLogging = true
	}

	if envVarSet(EnvVars.DebugLog) {
		*debugLog = true
	}
}
<commit_msg>[logger] Don't add instance name label on Kubernetes.<commit_after>// Copyright 2017-2020 The Cloudprober Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package logger provides a logger that logs to Google Cloud Logging. It's a thin wrapper around
// golang/cloud/logging package.
package logger

import (
	"context"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"flag"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/logging"
	md "github.com/cloudprober/cloudprober/common/metadata"
	"github.com/golang/glog"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/option"
)

var (
	debugLog     = flag.Bool("debug_log", false, "Whether to output debug logs or not")
	debugLogList = flag.String("debug_logname_regex", "", "Enable debug logs for only for log names that match this regex (e.g. --debug_logname_regex=.*probe1.*")

	// Enable/Disable cloud logging
	disableCloudLogging = flag.Bool("disable_cloud_logging", false, "Disable cloud logging.")

	// LogPrefixEnvVar environment variable is used to determine the stackdriver
	// log name prefix. Default prefix is "cloudprober".
	LogPrefixEnvVar = "CLOUDPROBER_LOG_PREFIX"
)

// EnvVars defines environment variables that can be used to modify the logging
// behavior.
var EnvVars = struct {
	DisableCloudLogging, DebugLog string
}{
	"CLOUDPROBER_DISABLE_CLOUD_LOGGING",
	"CLOUDPROBER_DEBUG_LOG",
}

const (
	// Prefix for the cloudprober stackdriver log names.
	cloudproberPrefix = "cloudprober"
)

const (
	// Regular Expression for all characters that are illegal for log names
	//	Ref: https://cloud.google.com/logging/docs/api/ref_v2beta1/rest/v2beta1/LogEntry
	disapprovedRegExp = "[^A-Za-z0-9_/.-]"

	// MaxLogEntrySize Max size of each log entry (100 KB)
	// This limit helps avoid creating very large log lines in case someone
	// accidentally creates a large EventMetric, which in turn is possible due to
	// unbounded nature of "map" metric where keys are created on demand.
	//
	// TODO(manugarg): We can possibly get rid of this check now as the code that
	// could cause a large map metric has been fixed now. Earlier, cloudprober's
	// HTTP server used to respond to all URLs and used to record access to those
	// URLs as a "map" metric. Now, it responds only to pre-configured URLs.
	MaxLogEntrySize = 102400
)

func enableDebugLog(debugLog bool, debugLogRe string, logName string) bool {
	if !debugLog && debugLogRe == "" {
		return false
	}

	if debugLog && debugLogRe == "" {
		// Enable for all logs, regardless of log names.
		return true
	}

	r, err := regexp.Compile(debugLogRe)
	if err != nil {
		panic(fmt.Sprintf("error while parsing log name regex (%s): %v", debugLogRe, err))
	}

	if r.MatchString(logName) {
		return true
	}

	return false
}

// Logger implements a logger that logs messages to Google Cloud Logging. It
// provides a suite of methods where each method corresponds to a specific
// logging.Level, e.g. Error(paylod interface{}). Each method takes a payload
// that has to either be a JSON-encodable object, a string or a []byte slice
// (all other types of payload will result in error).
//
// It falls back to logging through the traditional logger if:
//
//   * Not running on GCE,
//   * Logging client is uninitialized (e.g. for testing),
//   * Logging to cloud fails for some reason.
//
// Logger{} is a valid object that will log through the traditional logger.
//
type Logger struct {
	name                string
	logc                *logging.Client
	logger              *logging.Logger
	debugLog            bool
	disableCloudLogging bool
	// TODO(manugarg): Logger should eventually embed the probe id and each probe
	// should get a different Logger object (embedding that probe's probe id) but
	// sharing the same logging client. We could then make probe id one of the
	// metadata on all logging messages.
}

// NewCloudproberLog is a convenient wrapper around New that sets context to
// context.Background and attaches cloudprober prefix to log names.
func NewCloudproberLog(component string) (*Logger, error) {
	cpPrefix := cloudproberPrefix

	envLogPrefix := os.Getenv(LogPrefixEnvVar)
	if envLogPrefix != "" {
		cpPrefix = envLogPrefix
	}

	return New(context.Background(), cpPrefix+"."+component)
}

// New returns a new Logger object with cloud logging client initialized if running on GCE.
func New(ctx context.Context, logName string) (*Logger, error) {
	l := &Logger{
		name:                logName,
		debugLog:            enableDebugLog(*debugLog, *debugLogList, logName),
		disableCloudLogging: *disableCloudLogging,
	}

	if !metadata.OnGCE() || l.disableCloudLogging {
		return l, nil
	}

	l.Infof("Running on GCE. Logs for %s will go to Cloud (Stackdriver).", logName)
	if err := l.EnableStackdriverLogging(ctx); err != nil {
		return nil, err
	}
	return l, nil
}

// EnableStackdriverLogging enables logging to stackdriver.
func (l *Logger) EnableStackdriverLogging(ctx context.Context) error {
	if !metadata.OnGCE() {
		return fmt.Errorf("not running on GCE")
	}

	projectID, err := metadata.ProjectID()
	if err != nil {
		return err
	}

	if l.name == "" {
		return fmt.Errorf("logName cannot be empty")
	}
	// Check for illegal characters in the log name
	if match, err := regexp.Match(disapprovedRegExp, []byte(l.name)); err != nil || match {
		if err != nil {
			return fmt.Errorf("unable to parse logName: %v", err)
		}
		return fmt.Errorf("logName of %s contains an invalid character, valid characters are [A-Za-z0-9_/.-]", l.name)
	}
	// Any forward slashes need to be URL encoded, so we query escape to replace them
	logName := url.QueryEscape(l.name)

	l.logc, err = logging.NewClient(ctx, projectID, option.WithTokenSource(google.ComputeTokenSource("")))
	if err != nil {
		return err
	}

	loggerOpts := []logging.LoggerOption{
		// Encourage batching of write requests.
		// Flush logs to remote logging after 1000 entries (default is 10).
		logging.EntryCountThreshold(1000),
		// Maximum amount of time that an item should remain buffered in memory
		// before being flushed to the logging service. Default is 1 second.
		// We want flushing to be mostly driven by the buffer size (configured
		// above), rather than time.
		logging.DelayThreshold(10 * time.Second),
	}

	// Add instance_name to common labels if available.
	if !md.IsKubernetes() {
		instanceName, err := metadata.InstanceName()
		if err != nil {
			l.Infof("Error getting instance name on GCE: %v", err)
		} else {
			loggerOpts = append(loggerOpts, logging.CommonLabels(map[string]string{"instance_name": instanceName}))
		}
	}

	l.logger = l.logc.Logger(logName, loggerOpts...)
	return nil
}

func payloadToString(payload ...string) string {
	if len(payload) == 1 {
		return payload[0]
	}

	var b strings.Builder
	for _, s := range payload {
		b.WriteString(s)
	}
	return b.String()
}

// log sends payload ([]string) to cloud logging. If cloud logging client is
// not initialized (e.g. if not running on GCE) or cloud logging fails for some
// reason, it writes logs through the traditional logger.
func (l *Logger) log(severity logging.Severity, payload ...string) {
	payloadStr := payloadToString(payload...)

	if len(payloadStr) > MaxLogEntrySize {
		truncateMsg := "... (truncated)"
		truncateMsgLen := len(truncateMsg)
		payloadStr = payloadStr[:MaxLogEntrySize-truncateMsgLen] + truncateMsg
	}

	if l == nil {
		genericLog(severity, "nil", payloadStr)
		return
	}

	if l.logger == nil {
		genericLog(severity, l.name, payloadStr)
		return
	}

	l.logger.Log(logging.Entry{
		Severity: severity,
		Payload:  payloadStr,
	})
}

// Close closes the cloud logging client if it exists. This flushes the buffer
// and should be called before exiting the program to ensure all logs are persisted.
func (l *Logger) Close() error {
	if l != nil && l.logc != nil {
		return l.logc.Close()
	}

	return nil
}

// Debug logs messages with logging level set to "Debug".
func (l *Logger) Debug(payload ...string) {
	if l.debugLog {
		l.log(logging.Debug, payload...)
	}
}

// Info logs messages with logging level set to "Info".
func (l *Logger) Info(payload ...string) {
	l.log(logging.Info, payload...)
}

// Warning logs messages with logging level set to "Warning".
func (l *Logger) Warning(payload ...string) {
	l.log(logging.Warning, payload...)
}

// Error logs messages with logging level set to "Error".
func (l *Logger) Error(payload ...string) {
	l.log(logging.Error, payload...)
}

// Critical logs messages with logging level set to "Critical" and
// exits the process with error status. The buffer is flushed before exiting.
func (l *Logger) Critical(payload ...string) {
	l.log(logging.Critical, payload...)
	if err := l.Close(); err != nil {
		panic(fmt.Sprintf("could not close client: %v", err))
	}
	os.Exit(1)
}

// Debugf logs formatted text messages with logging level "Debug".
func (l *Logger) Debugf(format string, args ...interface{}) {
	if l != nil && l.debugLog {
		l.log(logging.Debug, fmt.Sprintf(format, args...))
	}
}

// Infof logs formatted text messages with logging level "Info".
func (l *Logger) Infof(format string, args ...interface{}) {
	l.log(logging.Info, fmt.Sprintf(format, args...))
}

// Warningf logs formatted text messages with logging level "Warning".
func (l *Logger) Warningf(format string, args ...interface{}) {
	l.log(logging.Warning, fmt.Sprintf(format, args...))
}

// Errorf logs formatted text messages with logging level "Error".
func (l *Logger) Errorf(format string, args ...interface{}) {
	l.log(logging.Error, fmt.Sprintf(format, args...))
}

// Criticalf logs formatted text messages with logging level "Critical" and
// exits the process with error status. The buffer is flushed before exiting.
func (l *Logger) Criticalf(format string, args ...interface{}) {
	l.log(logging.Critical, fmt.Sprintf(format, args...))
	if err := l.Close(); err != nil {
		panic(fmt.Sprintf("could not close client: %v", err))
	}
	os.Exit(1)
}

func genericLog(severity logging.Severity, name string, s string) {
	// Set the caller frame depth to 3 so that can get to the actual caller of
	// the logger. genericLog -> log -> Info* -> actualCaller
	depth := 3

	s = fmt.Sprintf("[%s] %s", name, s)

	switch severity {
	case logging.Debug, logging.Info:
		glog.InfoDepth(depth, s)
	case logging.Warning:
		glog.WarningDepth(depth, s)
	case logging.Error:
		glog.ErrorDepth(depth, s)
	case logging.Critical:
		glog.FatalDepth(depth, s)
	}
}

func envVarSet(key string) bool {
	v, ok := os.LookupEnv(key)
	if ok && strings.ToUpper(v) != "NO" && strings.ToUpper(v) != "FALSE" {
		return true
	}
	return false
}

func init() {
	if envVarSet(EnvVars.DisableCloudLogging) {
		*disableCloudLogging = true
	}

	if envVarSet(EnvVars.DebugLog) {
		*debugLog = true
	}
}
<|endoftext|>


<commit_before>package logger

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"
)

var (
	Trace   *log.Logger
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
)

func Init(traceHandle io.Writer, infoHandle io.Writer, warningHandle io.Writer, errorHandle io.Writer) {
	Trace = log.New(traceHandle, "TRACE: ", log.Ldate|log.Ltime|log.Lshortfile)
	Info = log.New(infoHandle, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
	Warning = log.New(warningHandle, "WARNING: ", log.Ldate|log.Ltime|log.Lshortfile)
	Error = log.New(errorHandle, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
}

var already_initialized bool = false
func DefaultLoggerInit() {
	if(already_initialized == false ) {
		Init(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)
		already_initialized = true
	}
}

func TrackTime(start time.Time, name string) {
	elapsed := time.Since(start)
	Info.Printf("%s %d\n", name, int(elapsed/1000000))
}
<commit_msg>add the JSONPrintDS<commit_after>package logger

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"
)

var (
	Trace   *log.Logger
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
)

func Init(traceHandle io.Writer, infoHandle io.Writer, warningHandle io.Writer, errorHandle io.Writer) {
	Trace = log.New(traceHandle, "TRACE: ", log.Ldate|log.Ltime|log.Lshortfile)
	Info = log.New(infoHandle, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
	Warning = log.New(warningHandle, "WARNING: ", log.Ldate|log.Ltime|log.Lshortfile)
	Error = log.New(errorHandle, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
}

var already_initialized bool = false

func DefaultLoggerInit() {
	if already_initialized == false {
		Init(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)
		already_initialized = true
	}
}

func TrackTime(start time.Time, name string) {
	elapsed := time.Since(start)
	Info.Printf("%s %d\n", name, int(elapsed/1000000))
}

func JSONPrintDS(input_data interface{}) {
	data, err := json.Marshal(input_data)
	if err == nil {
		fmt.Println(string(data[:]))
	} else {
		log.Printf("got err %v while marshaling", err)
	}
}
<|endoftext|>


<commit_before>package logger

type Logger interface {
	Print(v ...interface{})
	Println(v ...interface{})
}
<commit_msg>[logger] Add godoc<commit_after>// Package logger defines interface for logging and provide a reference
// implementation.
package logger

// Logger is an interface for logging.
type Logger interface {
	Print(v ...interface{})
	Println(v ...interface{})
}
<|endoftext|>


<commit_before>package logger

import (
	"net/http"

	"go.uber.org/zap"
)

var Log = setUpDefaultLogger().Sugar()

func setUpDefaultLogger() *zap.Logger {
	logger, e := zap.NewDevelopment()
	if e != nil {
		panic(e)
	}
	return logger
}

func SetLogger(logger *zap.Logger) {
	Log = logger.Sugar()
}

func From(r *http.Request) *zap.SugaredLogger {
	log := r.Context().Value("logger")
	if log != nil {
		return log.(*zap.SugaredLogger)
	} else {
		return Log
	}
}
<commit_msg>Make zap standard logger available that logs as Error<commit_after>package logger

import (
	"bytes"
	"log"
	"net/http"

	"go.uber.org/zap"
)

var Log = setUpDefaultLogger().Sugar()

func setUpDefaultLogger() *zap.Logger {
	logger, e := zap.NewDevelopment()
	if e != nil {
		panic(e)
	}
	return logger
}

func SetLogger(logger *zap.Logger) {
	Log = logger.Sugar()
}

func From(r *http.Request) *zap.SugaredLogger {
	log := r.Context().Value("logger")
	if log != nil {
		return log.(*zap.SugaredLogger)
	} else {
		return Log
	}
}

const (
	_stdLogDefaultDepth = 2
	_loggerWriterDepth  = 1
)

// Copied from go.uber.org/zap/global.go and changed to use Error instead of Info:
func NewStdLog(l *zap.Logger) *log.Logger {
	return log.New(&loggerWriter{l.WithOptions(
		zap.AddCallerSkip(_stdLogDefaultDepth + _loggerWriterDepth),
	)}, "" /* prefix */, 0 /* flags */)
}

type loggerWriter struct{ logger *zap.Logger }

func (l *loggerWriter) Write(p []byte) (int, error) {
	p = bytes.TrimSpace(p)
	l.logger.Error(string(p))
	return len(p), nil
}
<|endoftext|>


<commit_before>package login

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/dgrijalva/jwt-go"
	"github.com/tarent/loginsrv/logging"
	"github.com/tarent/loginsrv/model"
	"github.com/tarent/loginsrv/oauth2"
)

const contentTypeHTML = "text/html; charset=utf-8"
const contentTypeJWT = "application/jwt"
const contentTypePlain = "text/plain"

// Handler is the mail login handler.
// It serves the login ressource and does the authentication against the backends or oauth provider.
type Handler struct {
	backends []Backend
	oauth    oauthManager
	config   *Config
}

// NewHandler creates a login handler based on the supplied configuration.
func NewHandler(config *Config) (*Handler, error) {
	if len(config.Backends) == 0 && len(config.Oauth) == 0 {
		return nil, errors.New("No login backends or oauth provider configured")
	}

	backends := []Backend{}
	for pName, opts := range config.Backends {
		p, exist := GetProvider(pName)
		if !exist {
			return nil, fmt.Errorf("No such provider: %v", pName)
		}
		b, err := p(opts)
		if err != nil {
			return nil, err
		}
		backends = append(backends, b)
	}

	oauth := oauth2.NewManager()
	for providerName, opts := range config.Oauth {
		err := oauth.AddConfig(providerName, opts)
		if err != nil {
			return nil, err
		}
	}

	return &Handler{
		backends: backends,
		config:   config,
		oauth:    oauth,
	}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !strings.HasPrefix(r.URL.Path, h.config.LoginPath) {
		h.respondNotFound(w, r)
		return
	}

	if (h.shouldRedirect(r)) && (r.Method != "POST") {
		queries, _ := url.ParseQuery(r.URL.RawQuery)
		if queries.Get(h.config.RedirectQueryParameter) != "" {
			cookie := http.Cookie{
				Name:  h.config.RedirectQueryParameter,
				Value: queries.Get(h.config.RedirectQueryParameter),
			}
			spew.Dump(cookie)
			http.SetCookie(w, &cookie)
		}
	}

	_, err := h.oauth.GetConfigFromRequest(r)
	if err == nil {
		h.handleOauth(w, r)
		return
	}

	h.handleLogin(w, r)
	return
}

func (h *Handler) handleOauth(w http.ResponseWriter, r *http.Request) {
	startedFlow, authenticated, userInfo, err := h.oauth.Handle(w, r)

	if startedFlow {
		// the oauth flow started
		return
	}

	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if authenticated {
		logging.Application(r.Header).
			WithField("username", userInfo.Sub).Info("successfully authenticated")
		h.respondAuthenticated(w, r, userInfo)
		return
	}
	logging.Application(r.Header).
		WithField("username", userInfo.Sub).Info("failed authentication")

	h.respondAuthFailure(w, r)
	return
}

func (h *Handler) handleLogin(w http.ResponseWriter, r *http.Request) {
	contentType := r.Header.Get("Content-Type")
	if !(r.Method == "GET" || r.Method == "DELETE" ||
		(r.Method == "POST" &&
			(strings.HasPrefix(contentType, "application/json") ||
				strings.HasPrefix(contentType, "application/x-www-form-urlencoded") ||
				strings.HasPrefix(contentType, "multipart/form-data") ||
				contentType == ""))) {
		h.respondBadRequest(w, r)
		return
	}

	r.ParseForm()
	if r.Method == "DELETE" || r.FormValue("logout") == "true" {
		h.deleteCookie(w, h.config.CookieName)
		if h.config.LogoutURL != "" {
			w.Header().Set("Location", h.config.LogoutURL)
			w.WriteHeader(303)
			return
		}
		writeLoginForm(w,
			loginFormData{
				Config: h.config,
			})
		return
	}

	if r.Method == "GET" {
		userInfo, valid := h.GetToken(r)
		writeLoginForm(w,
			loginFormData{
				Config:        h.config,
				Authenticated: valid,
				UserInfo:      userInfo,
			})
		return
	}

	if r.Method == "POST" {
		username, password, err := getCredentials(r)
		if err != nil {
			h.respondBadRequest(w, r)
			return
		}
		if username != "" {
			// No token found or credentials found, assuming new authentication
			h.handleAuthentication(w, r, username, password)
			return
		}
		userInfo, valid := h.GetToken(r)
		if valid {
			h.handleRefresh(w, r, userInfo)
			return
		}
		h.respondBadRequest(w, r)
		return
	}
}

func (h *Handler) handleAuthentication(w http.ResponseWriter, r *http.Request, username string, password string) {
	authenticated, userInfo, err := h.authenticate(username, password)
	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if authenticated {
		logging.Application(r.Header).
			WithField("username", username).Info("successfully authenticated")
		h.respondAuthenticated(w, r, userInfo)
		return
	}
	logging.Application(r.Header).
		WithField("username", username).Info("failed authentication")

	h.respondAuthFailure(w, r)
}

func (h *Handler) handleRefresh(w http.ResponseWriter, r *http.Request, userInfo model.UserInfo) {
	if userInfo.Refreshes >= h.config.JwtRefreshes {
		h.respondMaxRefreshesReached(w, r)
	} else {
		userInfo.Refreshes++
		h.respondAuthenticated(w, r, userInfo)
		logging.Application(r.Header).WithField("username", userInfo.Sub).Info("refreshed jwt")
	}
}

func (h *Handler) deleteCookie(w http.ResponseWriter, cookieName string) {
	cookie := &http.Cookie{
		Name:     cookieName,
		Value:    "delete",
		HttpOnly: true,
		Expires:  time.Unix(0, 0),
		Path:     "/",
	}
	if h.config.CookieDomain != "" {
		cookie.Domain = h.config.CookieDomain
	}
	http.SetCookie(w, cookie)
}

func (h *Handler) respondAuthenticated(w http.ResponseWriter, r *http.Request, userInfo model.UserInfo) {
	userInfo.Expiry = time.Now().Add(h.config.JwtExpiry).Unix()
	token, err := h.createToken(userInfo)
	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if wantHTML(r) {
		cookie := &http.Cookie{
			Name:     h.config.CookieName,
			Value:    token,
			HttpOnly: h.config.CookieHTTPOnly,
			Path:     "/",
		}
		if h.config.CookieExpiry != 0 {
			cookie.Expires = time.Now().Add(h.config.CookieExpiry)
		}
		if h.config.CookieDomain != "" {
			cookie.Domain = h.config.CookieDomain
		}

		http.SetCookie(w, cookie)

		w.Header().Set("Location", h.redirectURL(r, w))
		_, err := r.Cookie(h.config.RedirectQueryParameter)
		if err == nil {
			h.deleteCookie(w, h.config.RedirectQueryParameter)
		}
		w.WriteHeader(303)
		return
	}

	w.Header().Set("Content-Type", contentTypeJWT)
	w.WriteHeader(200)
	fmt.Fprintf(w, "%s", token)
}

func (h *Handler) createToken(userInfo jwt.Claims) (string, error) {
	token := jwt.NewWithClaims(jwt.SigningMethodHS512, userInfo)
	return token.SignedString([]byte(h.config.JwtSecret))
}

func (h *Handler) GetToken(r *http.Request) (userInfo model.UserInfo, valid bool) {
	c, err := r.Cookie(h.config.CookieName)
	if err != nil {
		return model.UserInfo{}, false
	}

	token, err := jwt.ParseWithClaims(c.Value, &model.UserInfo{}, func(*jwt.Token) (interface{}, error) {
		return []byte(h.config.JwtSecret), nil
	})
	if err != nil {
		return model.UserInfo{}, false
	}

	u, ok := token.Claims.(*model.UserInfo)
	if !ok {
		return model.UserInfo{}, false
	}

	return *u, u.Valid() == nil
}

func (h *Handler) respondError(w http.ResponseWriter, r *http.Request) {
	if wantHTML(r) {
		username, _, _ := getCredentials(r)
		writeLoginForm(w,
			loginFormData{
				Error:    true,
				Config:   h.config,
				UserInfo: model.UserInfo{Sub: username},
			})
		return
	}
	w.Header().Set("Content-Type", contentTypePlain)
	w.WriteHeader(500)
	fmt.Fprintf(w, "Internal Server Error")
}

func (h *Handler) respondBadRequest(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(400)
	fmt.Fprintf(w, "Bad Request: Method or content-type not supported")
}

func (h *Handler) respondNotFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "Not Found: The requested page does not exist")
}

func (h *Handler) respondMaxRefreshesReached(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(403)
	fmt.Fprint(w, "Max JWT refreshes reached")
}

func (h *Handler) respondAuthFailure(w http.ResponseWriter, r *http.Request) {
	if wantHTML(r) {
		w.Header().Set("Content-Type", contentTypeHTML)
		w.WriteHeader(403)
		username, _, _ := getCredentials(r)
		writeLoginForm(w,
			loginFormData{
				Failure:  true,
				Config:   h.config,
				UserInfo: model.UserInfo{Sub: username},
			})
		return
	}

	w.Header().Set("Content-Type", contentTypePlain)
	w.WriteHeader(403)
	fmt.Fprintf(w, "Wrong credentials")
}

func wantHTML(r *http.Request) bool {
	return strings.Contains(r.Header.Get("Accept"), "text/html")
}

func getCredentials(r *http.Request) (string, string, error) {
	if r.Header.Get("Content-Type") == "application/json" {
		m := map[string]string{}
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return "", "", err
		}
		err = json.Unmarshal(body, &m)
		if err != nil {
			return "", "", err
		}
		return m["username"], m["password"], nil
	}
	return r.PostForm.Get("username"), r.PostForm.Get("password"), nil
}

func (h *Handler) authenticate(username, password string) (bool, model.UserInfo, error) {
	for _, b := range h.backends {
		authenticated, userInfo, err := b.Authenticate(username, password)
		if err != nil {
			return false, model.UserInfo{}, err
		}
		if authenticated {
			return authenticated, userInfo, nil
		}
	}
	return false, model.UserInfo{}, nil
}

type oauthManager interface {
	Handle(w http.ResponseWriter, r *http.Request) (
		startedFlow bool,
		authenticated bool,
		userInfo model.UserInfo,
		err error)
	AddConfig(providerName string, opts map[string]string) error
	GetConfigFromRequest(r *http.Request) (oauth2.Config, error)
}
<commit_msg>Remove forgotten spew<commit_after>package login

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/tarent/loginsrv/logging"
	"github.com/tarent/loginsrv/model"
	"github.com/tarent/loginsrv/oauth2"
)

const contentTypeHTML = "text/html; charset=utf-8"
const contentTypeJWT = "application/jwt"
const contentTypePlain = "text/plain"

// Handler is the mail login handler.
// It serves the login ressource and does the authentication against the backends or oauth provider.
type Handler struct {
	backends []Backend
	oauth    oauthManager
	config   *Config
}

// NewHandler creates a login handler based on the supplied configuration.
func NewHandler(config *Config) (*Handler, error) {
	if len(config.Backends) == 0 && len(config.Oauth) == 0 {
		return nil, errors.New("No login backends or oauth provider configured")
	}

	backends := []Backend{}
	for pName, opts := range config.Backends {
		p, exist := GetProvider(pName)
		if !exist {
			return nil, fmt.Errorf("No such provider: %v", pName)
		}
		b, err := p(opts)
		if err != nil {
			return nil, err
		}
		backends = append(backends, b)
	}

	oauth := oauth2.NewManager()
	for providerName, opts := range config.Oauth {
		err := oauth.AddConfig(providerName, opts)
		if err != nil {
			return nil, err
		}
	}

	return &Handler{
		backends: backends,
		config:   config,
		oauth:    oauth,
	}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !strings.HasPrefix(r.URL.Path, h.config.LoginPath) {
		h.respondNotFound(w, r)
		return
	}

	if (h.shouldRedirect(r)) && (r.Method != "POST") {
		queries, _ := url.ParseQuery(r.URL.RawQuery)
		if queries.Get(h.config.RedirectQueryParameter) != "" {
			cookie := http.Cookie{
				Name:  h.config.RedirectQueryParameter,
				Value: queries.Get(h.config.RedirectQueryParameter),
			}
			http.SetCookie(w, &cookie)
		}
	}

	_, err := h.oauth.GetConfigFromRequest(r)
	if err == nil {
		h.handleOauth(w, r)
		return
	}

	h.handleLogin(w, r)
	return
}

func (h *Handler) handleOauth(w http.ResponseWriter, r *http.Request) {
	startedFlow, authenticated, userInfo, err := h.oauth.Handle(w, r)

	if startedFlow {
		// the oauth flow started
		return
	}

	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if authenticated {
		logging.Application(r.Header).
			WithField("username", userInfo.Sub).Info("successfully authenticated")
		h.respondAuthenticated(w, r, userInfo)
		return
	}
	logging.Application(r.Header).
		WithField("username", userInfo.Sub).Info("failed authentication")

	h.respondAuthFailure(w, r)
	return
}

func (h *Handler) handleLogin(w http.ResponseWriter, r *http.Request) {
	contentType := r.Header.Get("Content-Type")
	if !(r.Method == "GET" || r.Method == "DELETE" ||
		(r.Method == "POST" &&
			(strings.HasPrefix(contentType, "application/json") ||
				strings.HasPrefix(contentType, "application/x-www-form-urlencoded") ||
				strings.HasPrefix(contentType, "multipart/form-data") ||
				contentType == ""))) {
		h.respondBadRequest(w, r)
		return
	}

	r.ParseForm()
	if r.Method == "DELETE" || r.FormValue("logout") == "true" {
		h.deleteCookie(w, h.config.CookieName)
		if h.config.LogoutURL != "" {
			w.Header().Set("Location", h.config.LogoutURL)
			w.WriteHeader(303)
			return
		}
		writeLoginForm(w,
			loginFormData{
				Config: h.config,
			})
		return
	}

	if r.Method == "GET" {
		userInfo, valid := h.GetToken(r)
		writeLoginForm(w,
			loginFormData{
				Config:        h.config,
				Authenticated: valid,
				UserInfo:      userInfo,
			})
		return
	}

	if r.Method == "POST" {
		username, password, err := getCredentials(r)
		if err != nil {
			h.respondBadRequest(w, r)
			return
		}
		if username != "" {
			// No token found or credentials found, assuming new authentication
			h.handleAuthentication(w, r, username, password)
			return
		}
		userInfo, valid := h.GetToken(r)
		if valid {
			h.handleRefresh(w, r, userInfo)
			return
		}
		h.respondBadRequest(w, r)
		return
	}
}

func (h *Handler) handleAuthentication(w http.ResponseWriter, r *http.Request, username string, password string) {
	authenticated, userInfo, err := h.authenticate(username, password)
	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if authenticated {
		logging.Application(r.Header).
			WithField("username", username).Info("successfully authenticated")
		h.respondAuthenticated(w, r, userInfo)
		return
	}
	logging.Application(r.Header).
		WithField("username", username).Info("failed authentication")

	h.respondAuthFailure(w, r)
}

func (h *Handler) handleRefresh(w http.ResponseWriter, r *http.Request, userInfo model.UserInfo) {
	if userInfo.Refreshes >= h.config.JwtRefreshes {
		h.respondMaxRefreshesReached(w, r)
	} else {
		userInfo.Refreshes++
		h.respondAuthenticated(w, r, userInfo)
		logging.Application(r.Header).WithField("username", userInfo.Sub).Info("refreshed jwt")
	}
}

func (h *Handler) deleteCookie(w http.ResponseWriter, cookieName string) {
	cookie := &http.Cookie{
		Name:     cookieName,
		Value:    "delete",
		HttpOnly: true,
		Expires:  time.Unix(0, 0),
		Path:     "/",
	}
	if h.config.CookieDomain != "" {
		cookie.Domain = h.config.CookieDomain
	}
	http.SetCookie(w, cookie)
}

func (h *Handler) respondAuthenticated(w http.ResponseWriter, r *http.Request, userInfo model.UserInfo) {
	userInfo.Expiry = time.Now().Add(h.config.JwtExpiry).Unix()
	token, err := h.createToken(userInfo)
	if err != nil {
		logging.Application(r.Header).WithError(err).Error()
		h.respondError(w, r)
		return
	}

	if wantHTML(r) {
		cookie := &http.Cookie{
			Name:     h.config.CookieName,
			Value:    token,
			HttpOnly: h.config.CookieHTTPOnly,
			Path:     "/",
		}
		if h.config.CookieExpiry != 0 {
			cookie.Expires = time.Now().Add(h.config.CookieExpiry)
		}
		if h.config.CookieDomain != "" {
			cookie.Domain = h.config.CookieDomain
		}

		http.SetCookie(w, cookie)

		w.Header().Set("Location", h.redirectURL(r, w))
		_, err := r.Cookie(h.config.RedirectQueryParameter)
		if err == nil {
			h.deleteCookie(w, h.config.RedirectQueryParameter)
		}
		w.WriteHeader(303)
		return
	}

	w.Header().Set("Content-Type", contentTypeJWT)
	w.WriteHeader(200)
	fmt.Fprintf(w, "%s", token)
}

func (h *Handler) createToken(userInfo jwt.Claims) (string, error) {
	token := jwt.NewWithClaims(jwt.SigningMethodHS512, userInfo)
	return token.SignedString([]byte(h.config.JwtSecret))
}

func (h *Handler) GetToken(r *http.Request) (userInfo model.UserInfo, valid bool) {
	c, err := r.Cookie(h.config.CookieName)
	if err != nil {
		return model.UserInfo{}, false
	}

	token, err := jwt.ParseWithClaims(c.Value, &model.UserInfo{}, func(*jwt.Token) (interface{}, error) {
		return []byte(h.config.JwtSecret), nil
	})
	if err != nil {
		return model.UserInfo{}, false
	}

	u, ok := token.Claims.(*model.UserInfo)
	if !ok {
		return model.UserInfo{}, false
	}

	return *u, u.Valid() == nil
}

func (h *Handler) respondError(w http.ResponseWriter, r *http.Request) {
	if wantHTML(r) {
		username, _, _ := getCredentials(r)
		writeLoginForm(w,
			loginFormData{
				Error:    true,
				Config:   h.config,
				UserInfo: model.UserInfo{Sub: username},
			})
		return
	}
	w.Header().Set("Content-Type", contentTypePlain)
	w.WriteHeader(500)
	fmt.Fprintf(w, "Internal Server Error")
}

func (h *Handler) respondBadRequest(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(400)
	fmt.Fprintf(w, "Bad Request: Method or content-type not supported")
}

func (h *Handler) respondNotFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "Not Found: The requested page does not exist")
}

func (h *Handler) respondMaxRefreshesReached(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(403)
	fmt.Fprint(w, "Max JWT refreshes reached")
}

func (h *Handler) respondAuthFailure(w http.ResponseWriter, r *http.Request) {
	if wantHTML(r) {
		w.Header().Set("Content-Type", contentTypeHTML)
		w.WriteHeader(403)
		username, _, _ := getCredentials(r)
		writeLoginForm(w,
			loginFormData{
				Failure:  true,
				Config:   h.config,
				UserInfo: model.UserInfo{Sub: username},
			})
		return
	}

	w.Header().Set("Content-Type", contentTypePlain)
	w.WriteHeader(403)
	fmt.Fprintf(w, "Wrong credentials")
}

func wantHTML(r *http.Request) bool {
	return strings.Contains(r.Header.Get("Accept"), "text/html")
}

func getCredentials(r *http.Request) (string, string, error) {
	if r.Header.Get("Content-Type") == "application/json" {
		m := map[string]string{}
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return "", "", err
		}
		err = json.Unmarshal(body, &m)
		if err != nil {
			return "", "", err
		}
		return m["username"], m["password"], nil
	}
	return r.PostForm.Get("username"), r.PostForm.Get("password"), nil
}

func (h *Handler) authenticate(username, password string) (bool, model.UserInfo, error) {
	for _, b := range h.backends {
		authenticated, userInfo, err := b.Authenticate(username, password)
		if err != nil {
			return false, model.UserInfo{}, err
		}
		if authenticated {
			return authenticated, userInfo, nil
		}
	}
	return false, model.UserInfo{}, nil
}

type oauthManager interface {
	Handle(w http.ResponseWriter, r *http.Request) (
		startedFlow bool,
		authenticated bool,
		userInfo model.UserInfo,
		err error)
	AddConfig(providerName string, opts map[string]string) error
	GetConfigFromRequest(r *http.Request) (oauth2.Config, error)
}
<|endoftext|>