<commit_before>// Package crawler provides helper methods and defines an interface for lauching
// source repository crawlers that retrieve files from a source and forwards
// to a channel for indexing and retrieval.
package crawler

import (
	"context"
	"fmt"
	"log"
	"os"
	"sync"

	"sigs.k8s.io/kustomize/api/internal/crawl/index"

	_ "github.com/gomodule/redigo/redis"

	"sigs.k8s.io/kustomize/api/internal/crawl/doc"
)

var (
	logger = log.New(os.Stdout, "Crawler: ", log.LstdFlags|log.LUTC|log.Llongfile)
)

// Crawler forwards documents from source repositories to index and store them
// for searching. Each crawler is responsible for querying it's source of
// information, and forwarding files that have not been seen before or that need
// updating.
type Crawler interface {
	// Crawl returns when it is done processing. This method does not take
	// ownership of the channel. The channel is write only, and it
	// designates where the crawler should forward the documents.
	Crawl(ctx context.Context, output chan<- CrawledDocument, seen SeenMap) error

	// Get the document data given the FilePath, Repo, and Ref/Tag/Branch.
	FetchDocument(context.Context, *doc.Document) error
	// Write to the document what the created time is.
	SetCreated(context.Context, *doc.Document) error

	Match(*doc.Document) bool
}

type CrawledDocument interface {
	ID() string
	GetDocument() *doc.Document
	// Get all the Documents directly referred in a Document.
	// For a Document representing a non-kustomization file, an empty slice will be returned.
	// For a Document representing a kustomization file:
	// the `includeResources` parameter determines whether the documents referred in the `resources` field are returned or not;
	// the `includeTransformers` parameter determines whether the documents referred in the `transformers` field are returned or not;
	// the `includeGenerators` parameter determines whether the documents referred in the `generators` field are returned or not.
	GetResources(includeResources, includeTransformers, includeGenerators bool) ([]*doc.Document, error)
	WasCached() bool
}

type SeenMap map[string]struct{}

func (seen SeenMap) Seen(item string) bool {
	_, ok := seen[item]
	return ok
}

func (seen SeenMap) Add(item string) {
	seen[item] = struct{}{}
}

func NewSeenMap() SeenMap {
	return make(map[string]struct{})
}

type CrawlSeed []*doc.Document

type IndexFunc func(CrawledDocument, index.Mode) error
type Converter func(*doc.Document) (CrawledDocument, error)

func logIfErr(err error) {
	if err == nil {
		return
	}
	logger.Println("error: ", err)
}

func findMatch(d *doc.Document, crawlers []Crawler) Crawler {
	for _, crawl := range crawlers {
		if crawl.Match(d) {
			return crawl
		}
	}
	return nil
}

func addBranches(cdoc CrawledDocument, match Crawler, indx IndexFunc,
	seen SeenMap, stack *CrawlSeed) {

	seen.Add(cdoc.ID())

	// Insert into index
	if err := indx(cdoc, index.InsertOrUpdate); err != nil {
		logger.Printf("Failed to insert or update %s %s: %v",
			cdoc.GetDocument().RepositoryURL, cdoc.GetDocument().FilePath, err)
		return
	}

	deps, err := cdoc.GetResources(true, false, false)
	if err != nil {
		logger.Println(err)
		return
	}

	for _, dep := range deps {
		if seen.Seen(dep.ID()) {
			continue
		}
		*stack = append(*stack, dep)
	}
}

func doCrawl(ctx context.Context, docsPtr *CrawlSeed, crawlers []Crawler, conv Converter, indx IndexFunc,
	seen SeenMap, stack *CrawlSeed) {

	UpdatedDocCount := 0
	seenDocCount := 0
	cachedDocCount := 0
	findMatchErrCount := 0
	FetchDocumentErrCount := 0
	SetCreatedErrCount := 0
	convErrCount := 0
	deleteDocCount := 0
	crawledDocCount := 0

	// During the execution of the for loop, more Documents may be added into (*docsPtr).
	for len(*docsPtr) > 0 {
		// get the last Document in (*docPtr), which will be crawled in this iteration.
		tail := (*docsPtr)[len(*docsPtr)-1]

		// remove the last Document in (*docPtr)
		*docsPtr = (*docsPtr)[:(len(*docsPtr) - 1)]

		crawledDocCount++
		logger.Printf("Crawling doc %d: %s %s", crawledDocCount, tail.RepositoryURL, tail.FilePath)

		if seen.Seen(tail.ID()) {
			logger.Printf("this doc has been seen before")
			seenDocCount++
			continue
		}

		if tail.WasCached() {
			logger.Printf("%s %s is cached already", tail.RepositoryURL, tail.FilePath)
			cachedDocCount++
			continue
		}

		match := findMatch(tail, crawlers)
		if match == nil {
			logIfErr(fmt.Errorf("%v could not match any crawler", tail))
			findMatchErrCount++
			continue
		}

		// If the Document represents a kustomization root, FetchDcoument will change
		// the `filePath` field of the Document by adding `kustomization.yaml` or
		// `kustomization.yml` or `kustomization` into the the field.
		// Therefore, it is necessary to add the ID of the Document into seen before
		// calling FetchDocument. Otherwise, the binary may enter into an infinite loop
		// if a kustomization file points to its kustmozation root in its `resources` or
		// `bases` field.
		seen.Add(tail.ID())

		if err := match.FetchDocument(ctx, tail); err != nil {
			logger.Printf("FetchDocument failed on %s %s: %v",
				tail.RepositoryURL, tail.FilePath, err)
			FetchDocumentErrCount++
			// delete the document from the index
			cdoc := &doc.KustomizationDocument{
				Document: *tail,
			}
			seen.Add(cdoc.ID())
			if err := indx(cdoc, index.Delete); err != nil {
				logger.Printf("Failed to delete %s %s: %v",
					cdoc.RepositoryURL, cdoc.FilePath, err)
			}
			deleteDocCount++
			continue
		}

		if err := match.SetCreated(ctx, tail); err != nil {
			logger.Printf("SetCreated failed on %s %s: %v",
				tail.RepositoryURL, tail.FilePath, err)
			SetCreatedErrCount++
		}

		cdoc, err := conv(tail)
		// If conv returns an error, cdoc can still be added into the index so that
		// cdoc.Document can be searched.
		if err != nil {
			logger.Printf("conv failed on %s %s: %v",
				tail.RepositoryURL, tail.FilePath, err)
			convErrCount++
		}

		UpdatedDocCount++
		addBranches(cdoc, match, indx, seen, stack)
	}
	logger.Printf("Summary of doCrawl:\n")
	logger.Printf("\t%d documents were updated\n", UpdatedDocCount)
	logger.Printf("\t%d documents were seen by the crawler already and skipped\n", seenDocCount)
	logger.Printf("\t%d documents were cached already and skipped\n", cachedDocCount)
	logger.Printf("\t%d documents didn't have a matching crawler and skipped\n", findMatchErrCount)
	logger.Printf("\t%d documents cannot be fetched, %d out of them are deleted\n",
		FetchDocumentErrCount, deleteDocCount)
	logger.Printf("\t%d documents cannot update its creation time but still were inserted or updated in the index\n", SetCreatedErrCount)
	logger.Printf("\t%d documents cannot be converted but still were inserted or updated in the index\n", convErrCount)
}

// CrawlFromSeed updates all the documents in seed, and crawls all the new
// documents referred in the seed.
func CrawlFromSeed(ctx context.Context, seed CrawlSeed, crawlers []Crawler,
	conv Converter, indx IndexFunc, seen SeenMap) {

	// stack tracks the documents directly referred in other documents.
	stack := make(CrawlSeed, 0)

	// Exploit seed to update bulk of corpus.
	logger.Printf("updating %d documents from seed\n", len(seed))
	// each unique document in seed will be crawled once.
	doCrawl(ctx, &seed, crawlers, conv, indx, seen, &stack)

	// Traverse any new documents added while updating corpus.
	logger.Printf("crawling %d new documents found in the seed\n", len(stack))
	// While crawling each document in stack, the documents directly referred in the document
	// will be added into stack.
	// After this statement is done, stack will become empty.
	doCrawl(ctx, &stack, crawlers, conv, indx, seen, &stack)
}

// CrawlGithubRunner is a blocking function and only returns once all of the
// crawlers are finished with execution.
//
// This function uses the output channel to forward kustomization documents
// from a list of crawlers. The output is to be consumed by a database/search
// indexer for later retrieval.
//
// The return value is an array of errors in which each index represents the
// index of the crawler that emitted the error. Although the errors themselves
// can be nil, the array will always be exactly the size of the crawlers array.
//
// CrawlGithubRunner takes in a seed, which represents the documents stored in an
// index somewhere. The document data is not required to be populated. If there
// are many documents, this is preferable. The order of iteration over the seed
// is not guaranteed, but the CrawlGithub does guarantee that every element
// from the seed will be processed before any other documents from the
// crawlers.
func CrawlGithubRunner(ctx context.Context, output chan<- CrawledDocument,
	crawlers []Crawler, seen SeenMap) []error {

	errs := make([]error, len(crawlers))
	wg := sync.WaitGroup{}

	for i, crawler := range crawlers {
		// Crawler implementations get their own channels to prevent a
		// crawler from closing the main output channel.
		docs := make(chan CrawledDocument)
		wg.Add(2)

		// Forward all of the documents from this crawler's channel to
		// the main output channel.
		go func(docs <-chan CrawledDocument) {
			defer wg.Done()
			for d := range docs {
				output <- d
			}
		}(docs)

		// Run this crawler and capture its returned error.
		go func(idx int, crawler Crawler,
			docs chan<- CrawledDocument) {

			defer func() {
				wg.Done()
				if r := recover(); r != nil {
					errs[idx] = fmt.Errorf(
						"%+v panicked: %v, additional error %v",
						crawler, r, errs[idx],
					)
				}
			}()
			defer close(docs)
			errs[idx] = crawler.Crawl(ctx, docs, seen)
		}(i, crawler, docs) // Copies the index and the crawler
	}

	wg.Wait()
	return errs
}

// CrawlGithub crawls all the kustomization files on Github.
func CrawlGithub(ctx context.Context, crawlers []Crawler, conv Converter,
	indx IndexFunc, seen SeenMap) {
	// stack tracks the documents directly referred in other documents.
	stack := make(CrawlSeed, 0)

	// ch is channel where all the crawlers sends the crawled documents to.
	ch := make(chan CrawledDocument, 1<<10)

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		defer wg.Done()
		docCount := 0
		for cdoc := range ch {
			docCount++
			logger.Printf("Processing doc %d found on Github", docCount)
			if seen.Seen(cdoc.ID()) {
				logger.Printf("the doc has been seen before")
				continue
			}
			match := findMatch(cdoc.GetDocument(), crawlers)
			if match == nil {
				logIfErr(fmt.Errorf(
					"%v could not match any crawler", cdoc))
				continue
			}
			addBranches(cdoc, match, indx, seen, &stack)
		}
	}()

	logger.Println("processing the documents found from crawling github")
	if errs := CrawlGithubRunner(ctx, ch, crawlers, seen); errs != nil {
		for _, err := range errs {
			logIfErr(err)
		}
	}
	close(ch)
	wg.Wait()

	// Handle deps of newly discovered documents.
	logger.Printf("crawling the %d new documents referred by other documents",
		len(stack))
	doCrawl(ctx, &stack, crawlers, conv, indx, seen, &stack)
}
<commit_msg>Use Document.Path instead of its fields<commit_after>// Package crawler provides helper methods and defines an interface for lauching
// source repository crawlers that retrieve files from a source and forwards
// to a channel for indexing and retrieval.
package crawler

import (
	"context"
	"fmt"
	"log"
	"os"
	"sync"

	"sigs.k8s.io/kustomize/api/internal/crawl/index"

	_ "github.com/gomodule/redigo/redis"

	"sigs.k8s.io/kustomize/api/internal/crawl/doc"
)

var (
	logger = log.New(os.Stdout, "Crawler: ", log.LstdFlags|log.LUTC|log.Llongfile)
)

// Crawler forwards documents from source repositories to index and store them
// for searching. Each crawler is responsible for querying it's source of
// information, and forwarding files that have not been seen before or that need
// updating.
type Crawler interface {
	// Crawl returns when it is done processing. This method does not take
	// ownership of the channel. The channel is write only, and it
	// designates where the crawler should forward the documents.
	Crawl(ctx context.Context, output chan<- CrawledDocument, seen SeenMap) error

	// Get the document data given the FilePath, Repo, and Ref/Tag/Branch.
	FetchDocument(context.Context, *doc.Document) error
	// Write to the document what the created time is.
	SetCreated(context.Context, *doc.Document) error

	Match(*doc.Document) bool
}

type CrawledDocument interface {
	ID() string
	GetDocument() *doc.Document
	// Get all the Documents directly referred in a Document.
	// For a Document representing a non-kustomization file, an empty slice will be returned.
	// For a Document representing a kustomization file:
	// the `includeResources` parameter determines whether the documents referred in the `resources` field are returned or not;
	// the `includeTransformers` parameter determines whether the documents referred in the `transformers` field are returned or not;
	// the `includeGenerators` parameter determines whether the documents referred in the `generators` field are returned or not.
	GetResources(includeResources, includeTransformers, includeGenerators bool) ([]*doc.Document, error)
	WasCached() bool
}

type SeenMap map[string]struct{}

func (seen SeenMap) Seen(item string) bool {
	_, ok := seen[item]
	return ok
}

func (seen SeenMap) Add(item string) {
	seen[item] = struct{}{}
}

func NewSeenMap() SeenMap {
	return make(map[string]struct{})
}

type CrawlSeed []*doc.Document

type IndexFunc func(CrawledDocument, index.Mode) error
type Converter func(*doc.Document) (CrawledDocument, error)

func logIfErr(err error) {
	if err == nil {
		return
	}
	logger.Println("error: ", err)
}

func findMatch(d *doc.Document, crawlers []Crawler) Crawler {
	for _, crawl := range crawlers {
		if crawl.Match(d) {
			return crawl
		}
	}
	return nil
}

func addBranches(cdoc CrawledDocument, match Crawler, indx IndexFunc,
	seen SeenMap, stack *CrawlSeed) {

	seen.Add(cdoc.ID())

	// Insert into index
	if err := indx(cdoc, index.InsertOrUpdate); err != nil {
		logger.Printf("Failed to insert or update doc(%s): %v",
			cdoc.GetDocument().Path(), err)
		return
	}

	deps, err := cdoc.GetResources(true, false, false)
	if err != nil {
		logger.Println(err)
		return
	}

	for _, dep := range deps {
		if seen.Seen(dep.ID()) {
			continue
		}
		*stack = append(*stack, dep)
	}
}

func doCrawl(ctx context.Context, docsPtr *CrawlSeed, crawlers []Crawler, conv Converter, indx IndexFunc,
	seen SeenMap, stack *CrawlSeed) {

	UpdatedDocCount := 0
	seenDocCount := 0
	cachedDocCount := 0
	findMatchErrCount := 0
	FetchDocumentErrCount := 0
	SetCreatedErrCount := 0
	convErrCount := 0
	deleteDocCount := 0
	crawledDocCount := 0

	// During the execution of the for loop, more Documents may be added into (*docsPtr).
	for len(*docsPtr) > 0 {
		// get the last Document in (*docPtr), which will be crawled in this iteration.
		tail := (*docsPtr)[len(*docsPtr)-1]

		// remove the last Document in (*docPtr)
		*docsPtr = (*docsPtr)[:(len(*docsPtr) - 1)]

		crawledDocCount++
		logger.Printf("Crawling doc %d: %s", crawledDocCount, tail.Path())

		if seen.Seen(tail.ID()) {
			logger.Printf("this doc has been seen before")
			seenDocCount++
			continue
		}

		if tail.WasCached() {
			logger.Printf("doc(%s) is cached already", tail.Path())
			cachedDocCount++
			continue
		}

		match := findMatch(tail, crawlers)
		if match == nil {
			logIfErr(fmt.Errorf("%v could not match any crawler", tail))
			findMatchErrCount++
			continue
		}

		// If the Document represents a kustomization root, FetchDcoument will change
		// the `filePath` field of the Document by adding `kustomization.yaml` or
		// `kustomization.yml` or `kustomization` into the the field.
		// Therefore, it is necessary to add the ID of the Document into seen before
		// calling FetchDocument. Otherwise, the binary may enter into an infinite loop
		// if a kustomization file points to its kustmozation root in its `resources` or
		// `bases` field.
		seen.Add(tail.ID())

		if err := match.FetchDocument(ctx, tail); err != nil {
			logger.Printf("FetchDocument failed on doc(%s): %v", tail.Path(), err)
			FetchDocumentErrCount++
			// delete the document from the index
			cdoc := &doc.KustomizationDocument{
				Document: *tail,
			}
			seen.Add(cdoc.ID())
			if err := indx(cdoc, index.Delete); err != nil {
				logger.Printf("Failed to delete doc(%s): %v", cdoc.Path(), err)
			}
			deleteDocCount++
			continue
		}

		if err := match.SetCreated(ctx, tail); err != nil {
			logger.Printf("SetCreated failed on doc(%s): %v", tail.Path(), err)
			SetCreatedErrCount++
		}

		cdoc, err := conv(tail)
		// If conv returns an error, cdoc can still be added into the index so that
		// cdoc.Document can be searched.
		if err != nil {
			logger.Printf("conv failed on doc(%s): %v", tail.Path(), err)
			convErrCount++
		}

		UpdatedDocCount++
		addBranches(cdoc, match, indx, seen, stack)
	}
	logger.Printf("Summary of doCrawl:\n")
	logger.Printf("\t%d documents were updated\n", UpdatedDocCount)
	logger.Printf("\t%d documents were seen by the crawler already and skipped\n", seenDocCount)
	logger.Printf("\t%d documents were cached already and skipped\n", cachedDocCount)
	logger.Printf("\t%d documents didn't have a matching crawler and skipped\n", findMatchErrCount)
	logger.Printf("\t%d documents cannot be fetched, %d out of them are deleted\n",
		FetchDocumentErrCount, deleteDocCount)
	logger.Printf("\t%d documents cannot update its creation time but still were inserted or updated in the index\n", SetCreatedErrCount)
	logger.Printf("\t%d documents cannot be converted but still were inserted or updated in the index\n", convErrCount)
}

// CrawlFromSeed updates all the documents in seed, and crawls all the new
// documents referred in the seed.
func CrawlFromSeed(ctx context.Context, seed CrawlSeed, crawlers []Crawler,
	conv Converter, indx IndexFunc, seen SeenMap) {

	// stack tracks the documents directly referred in other documents.
	stack := make(CrawlSeed, 0)

	// Exploit seed to update bulk of corpus.
	logger.Printf("updating %d documents from seed\n", len(seed))
	// each unique document in seed will be crawled once.
	doCrawl(ctx, &seed, crawlers, conv, indx, seen, &stack)

	// Traverse any new documents added while updating corpus.
	logger.Printf("crawling %d new documents found in the seed\n", len(stack))
	// While crawling each document in stack, the documents directly referred in the document
	// will be added into stack.
	// After this statement is done, stack will become empty.
	doCrawl(ctx, &stack, crawlers, conv, indx, seen, &stack)
}

// CrawlGithubRunner is a blocking function and only returns once all of the
// crawlers are finished with execution.
//
// This function uses the output channel to forward kustomization documents
// from a list of crawlers. The output is to be consumed by a database/search
// indexer for later retrieval.
//
// The return value is an array of errors in which each index represents the
// index of the crawler that emitted the error. Although the errors themselves
// can be nil, the array will always be exactly the size of the crawlers array.
//
// CrawlGithubRunner takes in a seed, which represents the documents stored in an
// index somewhere. The document data is not required to be populated. If there
// are many documents, this is preferable. The order of iteration over the seed
// is not guaranteed, but the CrawlGithub does guarantee that every element
// from the seed will be processed before any other documents from the
// crawlers.
func CrawlGithubRunner(ctx context.Context, output chan<- CrawledDocument,
	crawlers []Crawler, seen SeenMap) []error {

	errs := make([]error, len(crawlers))
	wg := sync.WaitGroup{}

	for i, crawler := range crawlers {
		// Crawler implementations get their own channels to prevent a
		// crawler from closing the main output channel.
		docs := make(chan CrawledDocument)
		wg.Add(2)

		// Forward all of the documents from this crawler's channel to
		// the main output channel.
		go func(docs <-chan CrawledDocument) {
			defer wg.Done()
			for d := range docs {
				output <- d
			}
		}(docs)

		// Run this crawler and capture its returned error.
		go func(idx int, crawler Crawler,
			docs chan<- CrawledDocument) {

			defer func() {
				wg.Done()
				if r := recover(); r != nil {
					errs[idx] = fmt.Errorf(
						"%+v panicked: %v, additional error %v",
						crawler, r, errs[idx],
					)
				}
			}()
			defer close(docs)
			errs[idx] = crawler.Crawl(ctx, docs, seen)
		}(i, crawler, docs) // Copies the index and the crawler
	}

	wg.Wait()
	return errs
}

// CrawlGithub crawls all the kustomization files on Github.
func CrawlGithub(ctx context.Context, crawlers []Crawler, conv Converter,
	indx IndexFunc, seen SeenMap) {
	// stack tracks the documents directly referred in other documents.
	stack := make(CrawlSeed, 0)

	// ch is channel where all the crawlers sends the crawled documents to.
	ch := make(chan CrawledDocument, 1<<10)

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		defer wg.Done()
		docCount := 0
		for cdoc := range ch {
			docCount++
			logger.Printf("Processing doc %d found on Github", docCount)
			if seen.Seen(cdoc.ID()) {
				logger.Printf("the doc has been seen before")
				continue
			}
			match := findMatch(cdoc.GetDocument(), crawlers)
			if match == nil {
				logIfErr(fmt.Errorf(
					"%v could not match any crawler", cdoc))
				continue
			}
			addBranches(cdoc, match, indx, seen, &stack)
		}
	}()

	logger.Println("processing the documents found from crawling github")
	if errs := CrawlGithubRunner(ctx, ch, crawlers, seen); errs != nil {
		for _, err := range errs {
			logIfErr(err)
		}
	}
	close(ch)
	wg.Wait()

	// Handle deps of newly discovered documents.
	logger.Printf("crawling the %d new documents referred by other documents",
		len(stack))
	doCrawl(ctx, &stack, crawlers, conv, indx, seen, &stack)
}
<|endoftext|>


<commit_before>// Copyright 2017, OpenCensus 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.

//go:generate protoc -I ../proto --go_out=plugins=grpc:../proto ../proto/helloworld.proto

package main

import (
	"log"
	"math/rand"
	"net"
	"net/http"
	"time"

	"go.opencensus.io/examples/exporter"
	pb "go.opencensus.io/examples/grpc/proto"
	"go.opencensus.io/plugin/ocgrpc"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
	"go.opencensus.io/zpages"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

const port = ":50051"

// server is used to implement helloworld.GreeterServer.
type server struct{}

// SayHello implements helloworld.GreeterServer
func (s *server) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
	ctx, span := trace.StartSpan(ctx, "sleep")
	time.Sleep(time.Duration(rand.Float64() * float64(time.Second)))
	span.End()
	return &pb.HelloReply{Message: "Hello " + in.Name}, nil
}

func main() {
	go func() {
		http.Handle("/debug/", http.StripPrefix("/debug", zpages.Handler))
		log.Fatal(http.ListenAndServe(":8081", nil))
	}()
	// Register stats and trace exporters to export
	// the collected data.
	view.RegisterExporter(&exporter.PrintExporter{})

	// Register the views to collect server request count.
	if err := view.Register(ocgrpc.DefaultServerViews...); err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("Failed to listen: %v", err)
	}

	// Set up a new server with the OpenCensus
	// stats handler to enable stats and tracing.
	s := grpc.NewServer(grpc.StatsHandler(&ocgrpc.ServerHandler{}))
	pb.RegisterGreeterServer(s, &server{})
	// Register reflection service on gRPC server.
	reflection.Register(s)
	if err := s.Serve(lis); err != nil {
		log.Fatalf("Failed to serve: %v", err)
	}
}
<commit_msg>Remove reflection server from gRPC example (#732)<commit_after>// Copyright 2017, OpenCensus 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.

//go:generate protoc -I ../proto --go_out=plugins=grpc:../proto ../proto/helloworld.proto

package main

import (
	"log"
	"math/rand"
	"net"
	"net/http"
	"time"

	"go.opencensus.io/examples/exporter"
	pb "go.opencensus.io/examples/grpc/proto"
	"go.opencensus.io/plugin/ocgrpc"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
	"go.opencensus.io/zpages"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

const port = ":50051"

// server is used to implement helloworld.GreeterServer.
type server struct{}

// SayHello implements helloworld.GreeterServer
func (s *server) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
	ctx, span := trace.StartSpan(ctx, "sleep")
	time.Sleep(time.Duration(rand.Float64() * float64(time.Second)))
	span.End()
	return &pb.HelloReply{Message: "Hello " + in.Name}, nil
}

func main() {
	go func() {
		http.Handle("/debug/", http.StripPrefix("/debug", zpages.Handler))
		log.Fatal(http.ListenAndServe(":8081", nil))
	}()
	// Register stats and trace exporters to export
	// the collected data.
	view.RegisterExporter(&exporter.PrintExporter{})

	// Register the views to collect server request count.
	if err := view.Register(ocgrpc.DefaultServerViews...); err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("Failed to listen: %v", err)
	}

	// Set up a new server with the OpenCensus
	// stats handler to enable stats and tracing.
	s := grpc.NewServer(grpc.StatsHandler(&ocgrpc.ServerHandler{}))
	pb.RegisterGreeterServer(s, &server{})

	if err := s.Serve(lis); err != nil {
		log.Fatalf("Failed to serve: %v", err)
	}
}
<|endoftext|>


<commit_before>// Package auth defines a standard interface for request access controllers.
//
// An access controller has a simple interface with a single `Authorized`
// method which checks that a given request is authorized to perform one or
// more actions on one or more resources. This method should return a non-nil
// error if the request is not authorized.
//
// An implementation registers its access controller by name with a constructor
// which accepts an options map for configuring the access controller.
//
//		options := map[string]interface{}{"sillySecret": "whysosilly?"}
// 		accessController, _ := auth.GetAccessController("silly", options)
//
// This `accessController` can then be used in a request handler like so:
//
// 		func updateOrder(w http.ResponseWriter, r *http.Request) {
//			orderNumber := r.FormValue("orderNumber")
//			resource := auth.Resource{Type: "customerOrder", Name: orderNumber}
// 			access := auth.Access{Resource: resource, Action: "update"}
//
// 			if ctx, err := accessController.Authorized(ctx, access); err != nil {
//				if challenge, ok := err.(auth.Challenge) {
//					// Let the challenge write the response.
//					challenge.ServeHTTP(w, r)
//				} else {
//					// Some other error.
//				}
//			}
// 		}
//
package auth

import (
	"fmt"
	"net/http"

	"github.com/docker/distribution/context"
)

// UserInfo carries information about
// an autenticated/authorized client.
type UserInfo struct {
	Name string
}

// Resource describes a resource by type and name.
type Resource struct {
	Type string
	Name string
}

// Access describes a specific action that is
// requested or allowed for a given resource.
type Access struct {
	Resource
	Action string
}

// Challenge is a special error type which is used for HTTP 401 Unauthorized
// responses and is able to write the response with WWW-Authenticate challenge
// header values based on the error.
type Challenge interface {
	error

	// SetHeaders prepares the request to conduct a challenge response by
	// adding the an HTTP challenge header on the response message. Callers
	// are expected to set the appropriate HTTP status code (e.g. 401)
	// themselves.
	SetHeaders(w http.ResponseWriter)
}

// AccessController controls access to registry resources based on a request
// and required access levels for a request. Implementations can support both
// complete denial and http authorization challenges.
type AccessController interface {
	// Authorized returns a non-nil error if the context is granted access and
	// returns a new authorized context. If one or more Access structs are
	// provided, the requested access will be compared with what is available
	// to the context. The given context will contain a "http.request" key with
	// a `*http.Request` value. If the error is non-nil, access should always
	// be denied. The error may be of type Challenge, in which case the caller
	// may have the Challenge handle the request or choose what action to take
	// based on the Challenge header or response status. The returned context
	// object should have a "auth.user" value set to a UserInfo struct.
	Authorized(ctx context.Context, access ...Access) (context.Context, error)
}

// WithUser returns a context with the authorized user info.
func WithUser(ctx context.Context, user UserInfo) context.Context {
	return userInfoContext{
		Context: ctx,
		user:    user,
	}
}

type userInfoContext struct {
	context.Context
	user UserInfo
}

func (uic userInfoContext) Value(key interface{}) interface{} {
	switch key {
	case "auth.user":
		return uic.user
	case "auth.user.name":
		return uic.user.Name
	}

	return uic.Context.Value(key)
}

// InitFunc is the type of an AccessController factory function and is used
// to register the constructor for different AccesController backends.
type InitFunc func(options map[string]interface{}) (AccessController, error)

var accessControllers map[string]InitFunc

func init() {
	accessControllers = make(map[string]InitFunc)
}

// Register is used to register an InitFunc for
// an AccessController backend with the given name.
func Register(name string, initFunc InitFunc) error {
	if _, exists := accessControllers[name]; exists {
		return fmt.Errorf("name already registered: %s", name)
	}

	accessControllers[name] = initFunc

	return nil
}

// GetAccessController constructs an AccessController
// with the given options using the named backend.
func GetAccessController(name string, options map[string]interface{}) (AccessController, error) {
	if initFunc, exists := accessControllers[name]; exists {
		return initFunc(options)
	}

	return nil, fmt.Errorf("no access controller registered with name: %s", name)
}
<commit_msg>Update "type auth.Challenge" comment example code<commit_after>// Package auth defines a standard interface for request access controllers.
//
// An access controller has a simple interface with a single `Authorized`
// method which checks that a given request is authorized to perform one or
// more actions on one or more resources. This method should return a non-nil
// error if the request is not authorized.
//
// An implementation registers its access controller by name with a constructor
// which accepts an options map for configuring the access controller.
//
//		options := map[string]interface{}{"sillySecret": "whysosilly?"}
// 		accessController, _ := auth.GetAccessController("silly", options)
//
// This `accessController` can then be used in a request handler like so:
//
// 		func updateOrder(w http.ResponseWriter, r *http.Request) {
//			orderNumber := r.FormValue("orderNumber")
//			resource := auth.Resource{Type: "customerOrder", Name: orderNumber}
// 			access := auth.Access{Resource: resource, Action: "update"}
//
// 			if ctx, err := accessController.Authorized(ctx, access); err != nil {
//				if challenge, ok := err.(auth.Challenge) {
//					// Let the challenge write the response.
//					challenge.SetHeaders(w)
//					w.WriteHeader(http.StatusUnauthorized)
//					return
//				} else {
//					// Some other error.
//				}
//			}
// 		}
//
package auth

import (
	"fmt"
	"net/http"

	"github.com/docker/distribution/context"
)

// UserInfo carries information about
// an autenticated/authorized client.
type UserInfo struct {
	Name string
}

// Resource describes a resource by type and name.
type Resource struct {
	Type string
	Name string
}

// Access describes a specific action that is
// requested or allowed for a given resource.
type Access struct {
	Resource
	Action string
}

// Challenge is a special error type which is used for HTTP 401 Unauthorized
// responses and is able to write the response with WWW-Authenticate challenge
// header values based on the error.
type Challenge interface {
	error

	// SetHeaders prepares the request to conduct a challenge response by
	// adding the an HTTP challenge header on the response message. Callers
	// are expected to set the appropriate HTTP status code (e.g. 401)
	// themselves.
	SetHeaders(w http.ResponseWriter)
}

// AccessController controls access to registry resources based on a request
// and required access levels for a request. Implementations can support both
// complete denial and http authorization challenges.
type AccessController interface {
	// Authorized returns a non-nil error if the context is granted access and
	// returns a new authorized context. If one or more Access structs are
	// provided, the requested access will be compared with what is available
	// to the context. The given context will contain a "http.request" key with
	// a `*http.Request` value. If the error is non-nil, access should always
	// be denied. The error may be of type Challenge, in which case the caller
	// may have the Challenge handle the request or choose what action to take
	// based on the Challenge header or response status. The returned context
	// object should have a "auth.user" value set to a UserInfo struct.
	Authorized(ctx context.Context, access ...Access) (context.Context, error)
}

// WithUser returns a context with the authorized user info.
func WithUser(ctx context.Context, user UserInfo) context.Context {
	return userInfoContext{
		Context: ctx,
		user:    user,
	}
}

type userInfoContext struct {
	context.Context
	user UserInfo
}

func (uic userInfoContext) Value(key interface{}) interface{} {
	switch key {
	case "auth.user":
		return uic.user
	case "auth.user.name":
		return uic.user.Name
	}

	return uic.Context.Value(key)
}

// InitFunc is the type of an AccessController factory function and is used
// to register the constructor for different AccesController backends.
type InitFunc func(options map[string]interface{}) (AccessController, error)

var accessControllers map[string]InitFunc

func init() {
	accessControllers = make(map[string]InitFunc)
}

// Register is used to register an InitFunc for
// an AccessController backend with the given name.
func Register(name string, initFunc InitFunc) error {
	if _, exists := accessControllers[name]; exists {
		return fmt.Errorf("name already registered: %s", name)
	}

	accessControllers[name] = initFunc

	return nil
}

// GetAccessController constructs an AccessController
// with the given options using the named backend.
func GetAccessController(name string, options map[string]interface{}) (AccessController, error) {
	if initFunc, exists := accessControllers[name]; exists {
		return initFunc(options)
	}

	return nil, fmt.Errorf("no access controller registered with name: %s", name)
}
<|endoftext|>


<commit_before>package forwarder

import (
	"bytes"
	"gopkg.in/stretchr/testify.v1/assert"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

type testResponseWriter struct {
	Headers    http.Header
	StatusCode int
	Body       *bytes.Buffer
}

func NewTestResponseWriter() *testResponseWriter {
	return &testResponseWriter{
		Headers:    make(http.Header),
		StatusCode: -1,
		Body:       &bytes.Buffer{},
	}
}
func (w *testResponseWriter) Header() http.Header {
	return w.Headers
}
func (w *testResponseWriter) Write(body []byte) (int, error) {
	return w.Body.Write(body)
}
func (w *testResponseWriter) WriteHeader(statusCode int) {
	w.StatusCode = statusCode
}

type testInterceptor func(http.ResponseWriter, *http.Request, io.Reader) bool

func (i testInterceptor) Handle(w http.ResponseWriter, req *http.Request, body io.Reader) bool {
	return i(w, req, body)
}

func TestInterceptor(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatal("Request was forwarded but should not have been")
	}))
	defer ts.Close()
	tsURL, _ := url.Parse(ts.URL)

	bodyText := "This is the request body"

	fw := NewForwarder(tsURL, 10000, testInterceptor(func(w http.ResponseWriter, req *http.Request, body io.Reader) bool {
		buf := make([]byte, 10000)
		if n, err := body.Read(buf); err != nil && err != io.EOF {
			t.Fatalf("Got error reading body: %s", err.Error())
		} else if string(buf[0:n]) != bodyText {
			t.Fatalf("Read body failed: X%vX, expected X%vX", string(buf[0:n]), bodyText)
		}
		w.Header().Set("foo", "bar")
		w.WriteHeader(204)
		w.Write([]byte("This is a response"))
		return true
	}))

	req, _ := http.NewRequest("POST", "http://localhost:9999/bar?crazy=true", strings.NewReader(bodyText))
	w := NewTestResponseWriter()

	fw.HandleRequest(w, req)

	if w.StatusCode != 204 {
		t.Fatalf("Response: StatusCode was %d, expected %d", w.StatusCode, 204)
	}
	if w.Header().Get("Foo") != "bar" {
		t.Fatalf("Response: Header['foo'] not set")
	}
	if w.Body.String() != "This is a response" {
		t.Fatalf("Response: Body is: %s", w.Body.String())
	}
}

func TestForwarding(t *testing.T) {
	requestText := "This is a request"
	responseText := "This is the response"

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/bar" {
			t.Fatalf("Request path incorrect: %s, expected %s", r.URL.Path, "bar")
		}
		if r.URL.RawQuery != "crazy=true" {
			t.Fatalf("Request raw query incorrect: %s, expected %s", r.URL.RawQuery, "crazy=true")
		}
		if r.Header.Get("myheader") != "headervalue" {
			t.Fatalf("Request header 'myheader' is not set")
		}
		if r.Header.Get("X-Forwarded-For") != "foobar" {
			t.Fatalf("Request header 'X-Forwarded-For' is not set")
		}
		if r.Body == nil {
			t.Fatal("Forwarded request has no body")
		}
		defer r.Body.Close()
		if body, err := ioutil.ReadAll(r.Body); err != nil {
			t.Fatal("Unexpected error reading request body:", err)
		} else if string(body) != requestText {
			t.Fatalf("Request body is %s, expected %s", string(body), requestText)
		}
		w.Header().Add("foo", "bar")
		w.WriteHeader(202)
		w.Write([]byte(responseText))
	}))
	defer ts.Close()
	tsURL, _ := url.Parse(ts.URL)

	fw := NewForwarder(tsURL, 10000, nil)

	req, _ := http.NewRequest("POST", "http://localhost:99999/bar?crazy=true", strings.NewReader(requestText))
	req.Header.Set("myheader", "headervalue")
	req.RemoteAddr = "foobar:1234"
	w := NewTestResponseWriter()

	fw.HandleRequest(w, req)

	if w.StatusCode != 202 {
		t.Fatalf("Response: StatusCode was %d, expected %d", w.StatusCode, 202)
	}
	if w.Header().Get("Foo") == "" {
		t.Fatalf("Response: Header['foo'] not set: %v", w.Headers)
	}
	if w.Body.String() != responseText {
		t.Fatalf("Response: Body is: %s", w.Body.String())
	}
}

func TestMutliForwarder(t *testing.T) {
	requestText := "This is a request"

	fw := NewForwarder(nil, 10000, nil)

	// Build two backends
	backend1ResponseText := "backend1"
	backend1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(202)
		w.Write([]byte(backend1ResponseText))
	}))
	defer backend1.Close()
	backend1URL, _ := url.Parse(backend1.URL)

	backend2ResponseText := "backend2"
	backend2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(202)
		w.Write([]byte(backend2ResponseText))
	}))
	defer backend2.Close()
	backend2URL, _ := url.Parse(backend2.URL)

	// Ensure they have different urls
	assert.NotEqual(t, backend1URL, backend2URL)

	// Test a request for backend1
	req1, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	req1.Header.Set("X-Clammit-Backend", backend1URL.String())
	w1 := NewTestResponseWriter()

	fw.HandleRequest(w1, req1)

	assert.Equal(t, w1.StatusCode, 202)
	assert.Equal(t, w1.Body.String(), backend1ResponseText)

	// Test a request for backend2
	req2, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	req2.Header.Set("X-Clammit-Backend", backend2URL.String())
	w2 := NewTestResponseWriter()

	fw.HandleRequest(w2, req2)

	assert.Equal(t, w2.StatusCode, 202)
	assert.Equal(t, w2.Body.String(), backend2ResponseText)

	// Test a request without the backend header
	req3, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	w3 := NewTestResponseWriter()

	fw.HandleRequest(w3, req3)

	assert.Equal(t, w3.StatusCode, 500)
	assert.Equal(t, w3.Body.String(), "Internal Server Error\n")
}
<commit_msg>Test: Host header handling<commit_after>package forwarder

import (
	"bytes"
	"gopkg.in/stretchr/testify.v1/assert"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

type testResponseWriter struct {
	Headers    http.Header
	StatusCode int
	Body       *bytes.Buffer
}

func NewTestResponseWriter() *testResponseWriter {
	return &testResponseWriter{
		Headers:    make(http.Header),
		StatusCode: -1,
		Body:       &bytes.Buffer{},
	}
}
func (w *testResponseWriter) Header() http.Header {
	return w.Headers
}
func (w *testResponseWriter) Write(body []byte) (int, error) {
	return w.Body.Write(body)
}
func (w *testResponseWriter) WriteHeader(statusCode int) {
	w.StatusCode = statusCode
}

type testInterceptor func(http.ResponseWriter, *http.Request, io.Reader) bool

func (i testInterceptor) Handle(w http.ResponseWriter, req *http.Request, body io.Reader) bool {
	return i(w, req, body)
}

func TestInterceptor(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatal("Request was forwarded but should not have been")
	}))
	defer ts.Close()
	tsURL, _ := url.Parse(ts.URL)

	bodyText := "This is the request body"

	fw := NewForwarder(tsURL, 10000, testInterceptor(func(w http.ResponseWriter, req *http.Request, body io.Reader) bool {
		buf := make([]byte, 10000)
		if n, err := body.Read(buf); err != nil && err != io.EOF {
			t.Fatalf("Got error reading body: %s", err.Error())
		} else if string(buf[0:n]) != bodyText {
			t.Fatalf("Read body failed: X%vX, expected X%vX", string(buf[0:n]), bodyText)
		}
		w.Header().Set("foo", "bar")
		w.WriteHeader(204)
		w.Write([]byte("This is a response"))
		return true
	}))

	req, _ := http.NewRequest("POST", "http://localhost:9999/bar?crazy=true", strings.NewReader(bodyText))
	w := NewTestResponseWriter()

	fw.HandleRequest(w, req)

	if w.StatusCode != 204 {
		t.Fatalf("Response: StatusCode was %d, expected %d", w.StatusCode, 204)
	}
	if w.Header().Get("Foo") != "bar" {
		t.Fatalf("Response: Header['foo'] not set")
	}
	if w.Body.String() != "This is a response" {
		t.Fatalf("Response: Body is: %s", w.Body.String())
	}
}

func TestForwarding(t *testing.T) {
	requestText := "This is a request"
	responseText := "This is the response"

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/bar" {
			t.Fatalf("Request path incorrect: %s, expected %s", r.URL.Path, "bar")
		}
		if r.URL.RawQuery != "crazy=true" {
			t.Fatalf("Request raw query incorrect: %s, expected %s", r.URL.RawQuery, "crazy=true")
		}
		if r.Header.Get("myheader") != "headervalue" {
			t.Fatalf("Request header 'myheader' is not set")
		}
		if r.Header.Get("X-Forwarded-For") != "foobar" {
			t.Fatalf("Request header 'X-Forwarded-For' is not set")
		}
		if r.Body == nil {
			t.Fatal("Forwarded request has no body")
		}
		defer r.Body.Close()
		if body, err := ioutil.ReadAll(r.Body); err != nil {
			t.Fatal("Unexpected error reading request body:", err)
		} else if string(body) != requestText {
			t.Fatalf("Request body is %s, expected %s", string(body), requestText)
		}
		w.Header().Add("foo", "bar")
		w.WriteHeader(202)
		w.Write([]byte(responseText))
	}))
	defer ts.Close()
	tsURL, _ := url.Parse(ts.URL)

	fw := NewForwarder(tsURL, 10000, nil)

	req, _ := http.NewRequest("POST", "http://localhost:99999/bar?crazy=true", strings.NewReader(requestText))
	req.Header.Set("myheader", "headervalue")
	req.RemoteAddr = "foobar:1234"
	w := NewTestResponseWriter()

	fw.HandleRequest(w, req)

	if w.StatusCode != 202 {
		t.Fatalf("Response: StatusCode was %d, expected %d", w.StatusCode, 202)
	}
	if w.Header().Get("Foo") == "" {
		t.Fatalf("Response: Header['foo'] not set: %v", w.Headers)
	}
	if w.Body.String() != responseText {
		t.Fatalf("Response: Body is: %s", w.Body.String())
	}
}

func TestHostForwarding(t *testing.T) {
	requestText := "This is a request"
	responseText := "This is the response"

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hostname := strings.Split(r.Host, ":")[0]
		if hostname != "localhost" {
			t.Fatalf("Host field incorrect: %v, expected %v", hostname, "localhost")
		}
		defer r.Body.Close()
		w.WriteHeader(202)
		w.Write([]byte(responseText))
	}))
	defer ts.Close()
	tsURL, _ := url.Parse(ts.URL)

	fw := NewForwarder(tsURL, 10000, nil)

	req, _ := http.NewRequest("POST", "http://localhost:99999/bar?crazy=true", strings.NewReader(requestText))
	req.Header.Set("myheader", "headervalue")
	req.RemoteAddr = "foobar:1234"
	w := NewTestResponseWriter()

	fw.HandleRequest(w, req)

	ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hostname := strings.Split(r.Host, ":")[0]
		if hostname != "testhost" {
			t.Fatalf("Host field incorrect: %v, expected %v", hostname, "testhost")
		}
		defer r.Body.Close()
		w.WriteHeader(202)
		w.Write([]byte(responseText))
	}))
	defer ts.Close()
	tsURL, _ = url.Parse(ts.URL)

	fw = NewForwarder(tsURL, 10000, nil)

	req, _ = http.NewRequest("POST", "http://localhost:99999/bar?crazy=true", strings.NewReader(requestText))
	req.Host = "testhost"
	req.Header.Set("myheader", "headervalue")
	req.RemoteAddr = "foobar:1234"
	w = NewTestResponseWriter()

	fw.HandleRequest(w, req)
}

func TestMutliForwarder(t *testing.T) {
	requestText := "This is a request"

	fw := NewForwarder(nil, 10000, nil)

	// Build two backends
	backend1ResponseText := "backend1"
	backend1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(202)
		w.Write([]byte(backend1ResponseText))
	}))
	defer backend1.Close()
	backend1URL, _ := url.Parse(backend1.URL)

	backend2ResponseText := "backend2"
	backend2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(202)
		w.Write([]byte(backend2ResponseText))
	}))
	defer backend2.Close()
	backend2URL, _ := url.Parse(backend2.URL)

	// Ensure they have different urls
	assert.NotEqual(t, backend1URL, backend2URL)

	// Test a request for backend1
	req1, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	req1.Header.Set("X-Clammit-Backend", backend1URL.String())
	w1 := NewTestResponseWriter()

	fw.HandleRequest(w1, req1)

	assert.Equal(t, w1.StatusCode, 202)
	assert.Equal(t, w1.Body.String(), backend1ResponseText)

	// Test a request for backend2
	req2, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	req2.Header.Set("X-Clammit-Backend", backend2URL.String())
	w2 := NewTestResponseWriter()

	fw.HandleRequest(w2, req2)

	assert.Equal(t, w2.StatusCode, 202)
	assert.Equal(t, w2.Body.String(), backend2ResponseText)

	// Test a request without the backend header
	req3, _ := http.NewRequest("POST", "http://localhost:99999/bar", strings.NewReader(requestText))
	w3 := NewTestResponseWriter()

	fw.HandleRequest(w3, req3)

	assert.Equal(t, w3.StatusCode, 500)
	assert.Equal(t, w3.Body.String(), "Internal Server Error\n")
}
<|endoftext|>


<commit_before>package log

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"testing"
)

func TestLogRouter(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name     string
		fields   Fields
		expected string
	}{
		{"empty", nil, `{}` + "\n"},
		{"nil", Fields{"undefined": nil}, `{"undefined":null}` + "\n"},
		{"string", Fields{"string": "message1"}, `{"string":"message1"}` + "\n"},
		{"int", Fields{"number": 42}, `{"number":42}` + "\n"},
		{"float", Fields{"number": 99.1}, `{"number":99.1}` + "\n"},
		{"bool", Fields{"bool": true}, `{"bool":true}` + "\n"},
		{"object", Fields{"object": Fields{"key1": "value1"}}, `{"object":{"key1":"value1"}}` + "\n"},
		{"array_of_nils", Fields{"x": [2]interface{}{nil, nil}}, `{"x":[null,null]}` + "\n"},
		{"array_of_strings", Fields{"x": [2]interface{}{"msg1", "msg2"}}, `{"x":["msg1","msg2"]}` + "\n"},
		{"array_of_ints", Fields{"x": [2]interface{}{42, 82}}, `{"x":[42,82]}` + "\n"},
		{"array_of_floats", Fields{"x": [2]interface{}{99.1, 33.1}}, `{"x":[99.1,33.1]}` + "\n"},
		{"array_of_bools", Fields{"x": [2]interface{}{true, false}}, `{"x":[true,false]}` + "\n"},
		{"array_of_objects", Fields{"x": [2]interface{}{Fields{"key1": "msg1"}, Fields{"key2": "msg2"}}},
			`{"x":[{"key1":"msg1"},{"key2":"msg2"}]}` + "\n"},
		{"array_of_mixed", Fields{"x": [6]interface{}{nil, "msg", 42, 33.6, true, Fields{"key1": "msg1"}}},
			`{"x":[null,"msg",42,33.6,true,{"key1":"msg1"}]}` + "\n"},
		{"slice_of_nils", Fields{"x": []interface{}{nil, nil}}, `{"x":[null,null]}` + "\n"},
		{"slice_of_strings", Fields{"x": []interface{}{"msg1", "msg2"}}, `{"x":["msg1","msg2"]}` + "\n"},
		{"slice_of_ints", Fields{"x": []interface{}{42, 82}}, `{"x":[42,82]}` + "\n"},
		{"slice_of_floats", Fields{"x": []interface{}{99.1, 33.1}}, `{"x":[99.1,33.1]}` + "\n"},
		{"slice_of_bools", Fields{"x": []interface{}{true, false}}, `{"x":[true,false]}` + "\n"},
		{"slice_of_objects", Fields{"x": []interface{}{Fields{"key1": "msg1"}, Fields{"key2": "msg2"}}},
			`{"x":[{"key1":"msg1"},{"key2":"msg2"}]}` + "\n"},
		{"slice_of_mixed", Fields{"x": []interface{}{nil, "msg", 42, 33.6, true, Fields{"key1": "msg1"}}},
			`{"x":[null,"msg",42,33.6,true,{"key1":"msg1"}]}` + "\n"},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			r := NewRouter()
			buf := &bytes.Buffer{}
			r.Output("output1", buf, nil)
			r.Log(tc.fields)

			actual := buf.String()
			if actual != tc.expected {
				t.Fatalf("expected %q, but got: %q", tc.expected, actual)
			}
		})
	}
}

func TestLoggerFlags(t *testing.T) {
	t.Parallel()

	rfc3339Re := regexp.MustCompile(`^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}[+-][0-9]{2}:[0-9]{2}$`)
	rfc3339UTCRe := regexp.MustCompile(`^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z$`)
	shortfileRe := regexp.MustCompile(`log_test.go:[0-9]+$`)
	longfileRe := regexp.MustCompile(`.+(\\|/)log_test.go:[0-9]+$`)
	unixRe := regexp.MustCompile(`^[0-9]+$`)

	type e struct {
		field   string
		pattern interface{}
	}

	testCases := []struct {
		testName   string
		loggerName string
		flags      int64
		fields     Fields
		expected   []*e
	}{
		{"no flags", "logger", 0, nil, []*e{&e{"time", rfc3339UTCRe}}},
		{"std flags", "logger", FlagStd, nil, []*e{&e{"time", rfc3339UTCRe}}},
		{"rfc3339", "logger", FlagRFC3339, nil, []*e{&e{"time", rfc3339Re}}},
		{"rfc3339 utc", "logger", FlagRFC3339 | FlagUTC, nil, []*e{&e{"time", rfc3339UTCRe}}},
		{"unix", "logger", FlagUnix, nil, []*e{&e{"time", unixRe}}},
		{"unix nano", "logger", FlagUnixNano, nil, []*e{&e{"time", unixRe}}},
		{"logger name", "duck duck", FlagLogger, nil, []*e{&e{"logger", "duck duck"}}},
		{"short file line", "logger", FlagShortfile, nil, []*e{&e{"file", shortfileRe}}},
		{"long file line", "logger", FlagLongfile, nil, []*e{&e{"file", longfileRe}}},
		{"custom time1", "logger", 0, Fields{"time": "now1"}, []*e{&e{"time", "now1"}}},
		{"custom time2", "logger", FlagStd, Fields{"time": "now2"}, []*e{&e{"time", "now2"}}},
		{"custom time3", "logger", FlagRFC3339, Fields{"time": "now3"}, []*e{&e{"time", "now3"}}},
		{"custom time4", "logger", FlagRFC3339 | FlagUTC, Fields{"time": "now4"}, []*e{&e{"time", "now4"}}},
		{"custom time5", "logger", FlagUnix, Fields{"time": "now5"}, []*e{&e{"time", "now5"}}},
		{"custom time6", "logger", FlagUnixNano, Fields{"time": "now6"}, []*e{&e{"time", "now6"}}},
		{"custom logger name", "monkey", FlagLogger, Fields{"logger": "elephant"}, []*e{&e{"logger", "elephant"}}},
		{"custom short file line", "logger", FlagShortfile, Fields{"file": "line1"}, []*e{&e{"file", "line1"}}},
		{"custom long file line", "logger", FlagLongfile, Fields{"file": "line2"}, []*e{&e{"file", "line2"}}},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.testName, func(t *testing.T) {
			t.Parallel()

			spy := &routerSpy{}
			l := NewLogger(tc.loggerName, tc.flags, spy)
			l.Log(tc.fields)

			for _, e := range tc.expected {
				actual, ok := spy.fields[e.field]
				if !ok {
					t.Fatalf("field not found:", e.field)
				}

				if re, ok := e.pattern.(*regexp.Regexp); ok {
					if !re.MatchString(fmt.Sprintf("%v", actual)) {
						t.Fatalf("expected %v, but got %v", re.String(), actual)
					}
				} else {
					if actual != e.pattern {
						t.Fatalf("expected %v, but got %v", e.pattern, actual)
					}
				}
			}
		})
	}
}

// no goroutine safe
type routerSpy struct {
	fields Fields
}

func (r *routerSpy) Output(id string, w io.Writer, filter Filter) {}

func (r *routerSpy) Log(fields Fields) {
	r.fields = fields
}
<commit_msg>fix tests<commit_after>package log

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"testing"
)

func TestLogRouter(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name     string
		fields   Fields
		expected string
	}{
		{"empty", nil, `{}` + "\n"},
		{"nil", Fields{"undefined": nil}, `{"undefined":null}` + "\n"},
		{"string", Fields{"string": "message1"}, `{"string":"message1"}` + "\n"},
		{"int", Fields{"number": 42}, `{"number":42}` + "\n"},
		{"float", Fields{"number": 99.1}, `{"number":99.1}` + "\n"},
		{"bool", Fields{"bool": true}, `{"bool":true}` + "\n"},
		{"object", Fields{"object": Fields{"key1": "value1"}}, `{"object":{"key1":"value1"}}` + "\n"},
		{"array_of_nils", Fields{"x": [2]interface{}{nil, nil}}, `{"x":[null,null]}` + "\n"},
		{"array_of_strings", Fields{"x": [2]interface{}{"msg1", "msg2"}}, `{"x":["msg1","msg2"]}` + "\n"},
		{"array_of_ints", Fields{"x": [2]interface{}{42, 82}}, `{"x":[42,82]}` + "\n"},
		{"array_of_floats", Fields{"x": [2]interface{}{99.1, 33.1}}, `{"x":[99.1,33.1]}` + "\n"},
		{"array_of_bools", Fields{"x": [2]interface{}{true, false}}, `{"x":[true,false]}` + "\n"},
		{"array_of_objects", Fields{"x": [2]interface{}{Fields{"key1": "msg1"}, Fields{"key2": "msg2"}}},
			`{"x":[{"key1":"msg1"},{"key2":"msg2"}]}` + "\n"},
		{"array_of_mixed", Fields{"x": [6]interface{}{nil, "msg", 42, 33.6, true, Fields{"key1": "msg1"}}},
			`{"x":[null,"msg",42,33.6,true,{"key1":"msg1"}]}` + "\n"},
		{"slice_of_nils", Fields{"x": []interface{}{nil, nil}}, `{"x":[null,null]}` + "\n"},
		{"slice_of_strings", Fields{"x": []interface{}{"msg1", "msg2"}}, `{"x":["msg1","msg2"]}` + "\n"},
		{"slice_of_ints", Fields{"x": []interface{}{42, 82}}, `{"x":[42,82]}` + "\n"},
		{"slice_of_floats", Fields{"x": []interface{}{99.1, 33.1}}, `{"x":[99.1,33.1]}` + "\n"},
		{"slice_of_bools", Fields{"x": []interface{}{true, false}}, `{"x":[true,false]}` + "\n"},
		{"slice_of_objects", Fields{"x": []interface{}{Fields{"key1": "msg1"}, Fields{"key2": "msg2"}}},
			`{"x":[{"key1":"msg1"},{"key2":"msg2"}]}` + "\n"},
		{"slice_of_mixed", Fields{"x": []interface{}{nil, "msg", 42, 33.6, true, Fields{"key1": "msg1"}}},
			`{"x":[null,"msg",42,33.6,true,{"key1":"msg1"}]}` + "\n"},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			r := NewRouter()
			buf := &bytes.Buffer{}
			r.Output("output1", buf, nil)
			r.Log(tc.fields)

			actual := buf.String()
			if actual != tc.expected {
				t.Fatalf("expected %q, but got: %q", tc.expected, actual)
			}
		})
	}
}

func TestLoggerFlags(t *testing.T) {
	t.Parallel()

	rfc3339Re := regexp.MustCompile(`^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(Z|[+-][0-9]{2}:[0-9]{2})$`)
	shortfileRe := regexp.MustCompile(`log_test.go:[0-9]+$`)
	longfileRe := regexp.MustCompile(`.+(\\|/)log_test.go:[0-9]+$`)
	unixRe := regexp.MustCompile(`^[0-9]+$`)

	type e struct {
		field   string
		pattern interface{}
	}

	testCases := []struct {
		testName   string
		loggerName string
		flags      int64
		fields     Fields
		expected   []*e
	}{
		{"no flags", "logger", 0, nil, []*e{&e{"time", rfc3339Re}}},
		{"std flags", "logger", FlagStd, nil, []*e{&e{"time", rfc3339Re}}},
		{"rfc3339", "logger", FlagRFC3339, nil, []*e{&e{"time", rfc3339Re}}},
		{"rfc3339 utc", "logger", FlagRFC3339 | FlagUTC, nil, []*e{&e{"time", rfc3339Re}}},
		{"unix", "logger", FlagUnix, nil, []*e{&e{"time", unixRe}}},
		{"unix nano", "logger", FlagUnixNano, nil, []*e{&e{"time", unixRe}}},
		{"logger name", "duck duck", FlagLogger, nil, []*e{&e{"logger", "duck duck"}}},
		{"short file line", "logger", FlagShortfile, nil, []*e{&e{"file", shortfileRe}}},
		{"long file line", "logger", FlagLongfile, nil, []*e{&e{"file", longfileRe}}},
		{"custom time1", "logger", 0, Fields{"time": "now1"}, []*e{&e{"time", "now1"}}},
		{"custom time2", "logger", FlagStd, Fields{"time": "now2"}, []*e{&e{"time", "now2"}}},
		{"custom time3", "logger", FlagRFC3339, Fields{"time": "now3"}, []*e{&e{"time", "now3"}}},
		{"custom time4", "logger", FlagRFC3339 | FlagUTC, Fields{"time": "now4"}, []*e{&e{"time", "now4"}}},
		{"custom time5", "logger", FlagUnix, Fields{"time": "now5"}, []*e{&e{"time", "now5"}}},
		{"custom time6", "logger", FlagUnixNano, Fields{"time": "now6"}, []*e{&e{"time", "now6"}}},
		{"custom logger name", "monkey", FlagLogger, Fields{"logger": "elephant"}, []*e{&e{"logger", "elephant"}}},
		{"custom short file line", "logger", FlagShortfile, Fields{"file": "line1"}, []*e{&e{"file", "line1"}}},
		{"custom long file line", "logger", FlagLongfile, Fields{"file": "line2"}, []*e{&e{"file", "line2"}}},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.testName, func(t *testing.T) {
			t.Parallel()

			spy := &routerSpy{}
			l := NewLogger(tc.loggerName, tc.flags, spy)
			l.Log(tc.fields)

			for _, e := range tc.expected {
				actual, ok := spy.fields[e.field]
				if !ok {
					t.Fatalf("field not found:", e.field)
				}

				if re, ok := e.pattern.(*regexp.Regexp); ok {
					if !re.MatchString(fmt.Sprintf("%v", actual)) {
						t.Fatalf("expected %v, but got %v", re.String(), actual)
					}
				} else {
					if actual != e.pattern {
						t.Fatalf("expected %v, but got %v", e.pattern, actual)
					}
				}
			}
		})
	}
}

// no goroutine safe
type routerSpy struct {
	fields Fields
}

func (r *routerSpy) Output(id string, w io.Writer, filter Filter) {}

func (r *routerSpy) Log(fields Fields) {
	r.fields = fields
}
<|endoftext|>


<commit_before>package object

import (
	"net/http"
	"sync"
)

var once sync.Once

var defaultArchetype = Archetype{
	ContentMarshaler: ContentMarshaler{
		"":                 Json,
		"application/json": Json,
		"application/xml":  Xml,
	},
}

type Archetype struct {
	ContentMarshaler
}

func (a *Archetype) Handler() Handler {
	return Handler{
		Archetype: a,
	}
}

type HTTPHandler struct {
	Getter
	Handler
}

func (h HTTPHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request) {
	h.Handler.ServeObject(h.Getter.Get(rq), rw, rq)
}

type Handler struct {
	*Archetype
	ContentMarshaler
}

func (h Handler) ServeObject(o interface{}, rw http.ResponseWriter, rq *http.Request) {
	var ms []ContentMarshaler
	if h.Archetype != nil {
		ms = []ContentMarshaler{
			h.ContentMarshaler,
			h.Archetype.ContentMarshaler,
		}
	} else {
		ms = []ContentMarshaler{h.ContentMarshaler}
	}

	mf, ct := RequestMarshaler(rq, ms...)
	switch {
	case mf == nil:
		rw.WriteHeader(406)
		if h.ContentMarshaler != nil {
			if mf = h.ContentMarshaler[""]; mf != nil {
				break
			}
		}
		if h.Archetype != nil {
			if mf = h.Archetype.ContentMarshaler[""]; mf != nil {
				break
			}
		}
		//wtf man
		mf = plain
		o = "None of the specified Content-Types supported."
	}

	dt, ctt, err := mf(o, ct.MediaType, ct.Params)
	if err != nil {
		panic(err)
	}
	rw.Header().Add("Content-Type", ctt)
	rw.Write(dt)
}
<commit_msg>Forgot to expose this<commit_after>package object

import (
	"net/http"
	"sync"
)

var once sync.Once

var DefaultArchetype = Archetype{
	ContentMarshaler: ContentMarshaler{
		"":                 Json,
		"application/json": Json,
		"application/xml":  Xml,
	},
}

type Archetype struct {
	ContentMarshaler
}

func (a *Archetype) Handler() Handler {
	return Handler{
		Archetype: a,
	}
}

type HTTPHandler struct {
	Getter
	Handler
}

func (h HTTPHandler) ServeHTTP(rw http.ResponseWriter, rq *http.Request) {
	h.Handler.ServeObject(h.Getter.Get(rq), rw, rq)
}

type Handler struct {
	*Archetype
	ContentMarshaler
}

func (h Handler) ServeObject(o interface{}, rw http.ResponseWriter, rq *http.Request) {
	var ms []ContentMarshaler
	if h.Archetype != nil {
		ms = []ContentMarshaler{
			h.ContentMarshaler,
			h.Archetype.ContentMarshaler,
		}
	} else {
		ms = []ContentMarshaler{h.ContentMarshaler}
	}

	mf, ct := RequestMarshaler(rq, ms...)
	switch {
	case mf == nil:
		rw.WriteHeader(406)
		if h.ContentMarshaler != nil {
			if mf = h.ContentMarshaler[""]; mf != nil {
				break
			}
		}
		if h.Archetype != nil {
			if mf = h.Archetype.ContentMarshaler[""]; mf != nil {
				break
			}
		}
		//wtf man
		mf = plain
		o = "None of the specified Content-Types supported."
	}

	dt, ctt, err := mf(o, ct.MediaType, ct.Params)
	if err != nil {
		panic(err)
	}
	rw.Header().Add("Content-Type", ctt)
	rw.Write(dt)
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 cgroups

import (
	"fmt"
	"path/filepath"
	"strings"
	"sync"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/godbus/dbus"
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

const (
	SystemdDbus  Name = "systemd"
	defaultSlice      = "system.slice"
)

var (
	canDelegate bool
	once sync.Once
)

func Systemd() ([]Subsystem, error) {
	root, err := v1MountPoint()
	if err != nil {
		return nil, err
	}
	defaultSubsystems, err := defaults(root)
	if err != nil {
		return nil, err
	}
	s, err := NewSystemd(root)
	if err != nil {
		return nil, err
	}
	// make sure the systemd controller is added first
	return append([]Subsystem{s}, defaultSubsystems...), nil
}

func Slice(slice, name string) Path {
	if slice == "" {
		slice = defaultSlice
	}
	return func(subsystem Name) (string, error) {
		return filepath.Join(slice, name), nil
	}
}

func NewSystemd(root string) (*SystemdController, error) {
	return &SystemdController{
		root: root,
	}, nil
}

type SystemdController struct {
	mu   sync.Mutex
	root string
}

func (s *SystemdController) Name() Name {
	return SystemdDbus
}

func (s *SystemdController) Create(path string, resources *specs.LinuxResources) error {
	conn, err := systemdDbus.New()
	if err != nil {
		return err
	}
	defer conn.Close()
	slice, name := splitName(path)
	// We need to see if systemd can handle the delegate property
	// Systemd will return an error if it cannot handle delegate regardless
	// of its bool setting.
	checkDelegate := func() {
		canDelegate = true
		dlSlice := newProperty("Delegate", true)
		if _, err := conn.StartTransientUnit(slice, "testdelegate", []systemdDbus.Property{dlSlice}, nil); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				// Starting with systemd v237, Delegate is not even a property of slices anymore,
				// so the D-Bus call fails with "InvalidArgs" error.
				if strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.PropertyReadOnly") || strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.InvalidArgs") {
					canDelegate = false
				}
			}
		}

		conn.StopUnit(slice, "testDelegate", nil)
	}
	once.Do(checkDelegate)
	properties := []systemdDbus.Property{
		systemdDbus.PropDescription(fmt.Sprintf("cgroup %s", name)),
		systemdDbus.PropWants(slice),
		newProperty("DefaultDependencies", false),
		newProperty("MemoryAccounting", true),
		newProperty("CPUAccounting", true),
		newProperty("BlockIOAccounting", true),
	}

	// If we can delegate, we add the property back in
	if canDelegate {
		properties = append(properties, newProperty("Delegate", true))
	}

	ch := make(chan string)
	_, err = conn.StartTransientUnit(name, "replace", properties, ch)
	if err != nil {
		return err
	}
	<-ch
	return nil
}

func (s *SystemdController) Delete(path string) error {
	conn, err := systemdDbus.New()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, name := splitName(path)
	ch := make(chan string)
	_, err = conn.StopUnit(name, "replace", ch)
	if err != nil {
		return err
	}
	<-ch
	return nil
}

func newProperty(name string, units interface{}) systemdDbus.Property {
	return systemdDbus.Property{
		Name:  name,
		Value: dbus.MakeVariant(units),
	}
}

func unitName(name string) string {
	return fmt.Sprintf("%s.slice", name)
}

func splitName(path string) (slice string, unit string) {
	slice, unit = filepath.Split(path)
	return strings.TrimSuffix(slice, "/"), unit
}
<commit_msg>Fix gofmt of systemd.go<commit_after>/*
   Copyright The containerd 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 cgroups

import (
	"fmt"
	"path/filepath"
	"strings"
	"sync"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/godbus/dbus"
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

const (
	SystemdDbus  Name = "systemd"
	defaultSlice      = "system.slice"
)

var (
	canDelegate bool
	once        sync.Once
)

func Systemd() ([]Subsystem, error) {
	root, err := v1MountPoint()
	if err != nil {
		return nil, err
	}
	defaultSubsystems, err := defaults(root)
	if err != nil {
		return nil, err
	}
	s, err := NewSystemd(root)
	if err != nil {
		return nil, err
	}
	// make sure the systemd controller is added first
	return append([]Subsystem{s}, defaultSubsystems...), nil
}

func Slice(slice, name string) Path {
	if slice == "" {
		slice = defaultSlice
	}
	return func(subsystem Name) (string, error) {
		return filepath.Join(slice, name), nil
	}
}

func NewSystemd(root string) (*SystemdController, error) {
	return &SystemdController{
		root: root,
	}, nil
}

type SystemdController struct {
	mu   sync.Mutex
	root string
}

func (s *SystemdController) Name() Name {
	return SystemdDbus
}

func (s *SystemdController) Create(path string, resources *specs.LinuxResources) error {
	conn, err := systemdDbus.New()
	if err != nil {
		return err
	}
	defer conn.Close()
	slice, name := splitName(path)
	// We need to see if systemd can handle the delegate property
	// Systemd will return an error if it cannot handle delegate regardless
	// of its bool setting.
	checkDelegate := func() {
		canDelegate = true
		dlSlice := newProperty("Delegate", true)
		if _, err := conn.StartTransientUnit(slice, "testdelegate", []systemdDbus.Property{dlSlice}, nil); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				// Starting with systemd v237, Delegate is not even a property of slices anymore,
				// so the D-Bus call fails with "InvalidArgs" error.
				if strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.PropertyReadOnly") || strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.InvalidArgs") {
					canDelegate = false
				}
			}
		}

		conn.StopUnit(slice, "testDelegate", nil)
	}
	once.Do(checkDelegate)
	properties := []systemdDbus.Property{
		systemdDbus.PropDescription(fmt.Sprintf("cgroup %s", name)),
		systemdDbus.PropWants(slice),
		newProperty("DefaultDependencies", false),
		newProperty("MemoryAccounting", true),
		newProperty("CPUAccounting", true),
		newProperty("BlockIOAccounting", true),
	}

	// If we can delegate, we add the property back in
	if canDelegate {
		properties = append(properties, newProperty("Delegate", true))
	}

	ch := make(chan string)
	_, err = conn.StartTransientUnit(name, "replace", properties, ch)
	if err != nil {
		return err
	}
	<-ch
	return nil
}

func (s *SystemdController) Delete(path string) error {
	conn, err := systemdDbus.New()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, name := splitName(path)
	ch := make(chan string)
	_, err = conn.StopUnit(name, "replace", ch)
	if err != nil {
		return err
	}
	<-ch
	return nil
}

func newProperty(name string, units interface{}) systemdDbus.Property {
	return systemdDbus.Property{
		Name:  name,
		Value: dbus.MakeVariant(units),
	}
}

func unitName(name string) string {
	return fmt.Sprintf("%s.slice", name)
}

func splitName(path string) (slice string, unit string) {
	slice, unit = filepath.Split(path)
	return strings.TrimSuffix(slice, "/"), unit
}
<|endoftext|>


<commit_before>// +build linux

package btrfs

/*
#include <stdlib.h>
#include <dirent.h>
#include <btrfs/ioctl.h>
#include <btrfs/ctree.h>
*/
import "C"

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"syscall"
	"unsafe"

	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/pkg/idtools"
	"github.com/docker/docker/pkg/mount"
	"github.com/opencontainers/runc/libcontainer/label"
)

func init() {
	graphdriver.Register("btrfs", Init)
}

// Init returns a new BTRFS driver.
// An error is returned if BTRFS is not supported.
func Init(home string, options []string, uidMaps, gidMaps []idtools.IDMap) (graphdriver.Driver, error) {
	rootdir := path.Dir(home)

	var buf syscall.Statfs_t
	if err := syscall.Statfs(rootdir, &buf); err != nil {
		return nil, err
	}

	if graphdriver.FsMagic(buf.Type) != graphdriver.FsMagicBtrfs {
		return nil, graphdriver.ErrPrerequisites
	}

	rootUID, rootGID, err := idtools.GetRootUIDGID(uidMaps, gidMaps)
	if err != nil {
		return nil, err
	}
	if err := idtools.MkdirAllAs(home, 0700, rootUID, rootGID); err != nil {
		return nil, err
	}

	if err := mount.MakePrivate(home); err != nil {
		return nil, err
	}

	driver := &Driver{
		home:    home,
		uidMaps: uidMaps,
		gidMaps: gidMaps,
	}

	return graphdriver.NewNaiveDiffDriver(driver, uidMaps, gidMaps), nil
}

// Driver contains information about the filesystem mounted.
type Driver struct {
	//root of the file system
	home    string
	uidMaps []idtools.IDMap
	gidMaps []idtools.IDMap
}

// String prints the name of the driver (btrfs).
func (d *Driver) String() string {
	return "btrfs"
}

// Status returns current driver information in a two dimensional string array.
// Output contains "Build Version" and "Library Version" of the btrfs libraries used.
// Version information can be used to check compatibility with your kernel.
func (d *Driver) Status() [][2]string {
	status := [][2]string{}
	if bv := btrfsBuildVersion(); bv != "-" {
		status = append(status, [2]string{"Build Version", bv})
	}
	if lv := btrfsLibVersion(); lv != -1 {
		status = append(status, [2]string{"Library Version", fmt.Sprintf("%d", lv)})
	}
	return status
}

// GetMetadata returns empty metadata for this driver.
func (d *Driver) GetMetadata(id string) (map[string]string, error) {
	return nil, nil
}

// Cleanup unmounts the home directory.
func (d *Driver) Cleanup() error {
	return mount.Unmount(d.home)
}

func free(p *C.char) {
	C.free(unsafe.Pointer(p))
}

func openDir(path string) (*C.DIR, error) {
	Cpath := C.CString(path)
	defer free(Cpath)

	dir := C.opendir(Cpath)
	if dir == nil {
		return nil, fmt.Errorf("Can't open dir")
	}
	return dir, nil
}

func closeDir(dir *C.DIR) {
	if dir != nil {
		C.closedir(dir)
	}
}

func getDirFd(dir *C.DIR) uintptr {
	return uintptr(C.dirfd(dir))
}

func subvolCreate(path, name string) error {
	dir, err := openDir(path)
	if err != nil {
		return err
	}
	defer closeDir(dir)

	var args C.struct_btrfs_ioctl_vol_args
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(dir), C.BTRFS_IOC_SUBVOL_CREATE,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to create btrfs subvolume: %v", errno.Error())
	}
	return nil
}

func subvolSnapshot(src, dest, name string) error {
	srcDir, err := openDir(src)
	if err != nil {
		return err
	}
	defer closeDir(srcDir)

	destDir, err := openDir(dest)
	if err != nil {
		return err
	}
	defer closeDir(destDir)

	var args C.struct_btrfs_ioctl_vol_args_v2
	args.fd = C.__s64(getDirFd(srcDir))
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(destDir), C.BTRFS_IOC_SNAP_CREATE_V2,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to create btrfs snapshot: %v", errno.Error())
	}
	return nil
}

func isSubvolume(p string) (bool, error) {
	var bufStat syscall.Stat_t
	if err := syscall.Lstat(p, &bufStat); err != nil {
		return false, err
	}

	// return true if it is a btrfs subvolume
	return bufStat.Ino == C.BTRFS_FIRST_FREE_OBJECTID, nil
}

func subvolDelete(dirpath, name string) error {
	dir, err := openDir(dirpath)
	if err != nil {
		return err
	}
	defer closeDir(dir)
	fullPath := path.Join(dirpath, name)

	var args C.struct_btrfs_ioctl_vol_args

	// walk the btrfs subvolumes
	walkSubvolumes := func(p string, f os.FileInfo, err error) error {
		if err != nil {
			if os.IsNotExist(err) && p != fullPath {
				// missing most likely because the path was a subvolume that got removed in the previous iteration
				// since it's gone anyway, we don't care
				return nil
			}
			return fmt.Errorf("error walking subvolumes: %v", err)
		}
		// we want to check children only so skip itself
		// it will be removed after the filepath walk anyways
		if f.IsDir() && p != fullPath {
			sv, err := isSubvolume(p)
			if err != nil {
				return fmt.Errorf("Failed to test if %s is a btrfs subvolume: %v", p, err)
			}
			if sv {
				if err := subvolDelete(path.Dir(p), f.Name()); err != nil {
					return fmt.Errorf("Failed to destroy btrfs child subvolume (%s) of parent (%s): %v", p, dirpath, err)
				}
			}
		}
		return nil
	}
	if err := filepath.Walk(path.Join(dirpath, name), walkSubvolumes); err != nil {
		return fmt.Errorf("Recursively walking subvolumes for %s failed: %v", dirpath, err)
	}

	// all subvolumes have been removed
	// now remove the one originally passed in
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(dir), C.BTRFS_IOC_SNAP_DESTROY,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to destroy btrfs snapshot %s for %s: %v", dirpath, name, errno.Error())
	}
	return nil
}

func (d *Driver) subvolumesDir() string {
	return path.Join(d.home, "subvolumes")
}

func (d *Driver) subvolumesDirID(id string) string {
	return path.Join(d.subvolumesDir(), id)
}

// Create the filesystem with given id.
func (d *Driver) Create(id, parent, mountLabel string) error {
	subvolumes := path.Join(d.home, "subvolumes")
	rootUID, rootGID, err := idtools.GetRootUIDGID(d.uidMaps, d.gidMaps)
	if err != nil {
		return err
	}
	if err := idtools.MkdirAllAs(subvolumes, 0700, rootUID, rootGID); err != nil {
		return err
	}
	if parent == "" {
		if err := subvolCreate(subvolumes, id); err != nil {
			return err
		}
	} else {
		parentDir, err := d.Get(parent, "")
		if err != nil {
			return err
		}
		if err := subvolSnapshot(parentDir, subvolumes, id); err != nil {
			return err
		}
	}

	// if we have a remapped root (user namespaces enabled), change the created snapshot
	// dir ownership to match
	if rootUID != 0 || rootGID != 0 {
		if err := os.Chown(path.Join(subvolumes, id), rootUID, rootGID); err != nil {
			return err
		}
	}

	return label.Relabel(path.Join(subvolumes, id), mountLabel, false)
}

// Remove the filesystem with given id.
func (d *Driver) Remove(id string) error {
	dir := d.subvolumesDirID(id)
	if _, err := os.Stat(dir); err != nil {
		return err
	}
	if err := subvolDelete(d.subvolumesDir(), id); err != nil {
		return err
	}
	if err := os.RemoveAll(dir); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

// Get the requested filesystem id.
func (d *Driver) Get(id, mountLabel string) (string, error) {
	dir := d.subvolumesDirID(id)
	st, err := os.Stat(dir)
	if err != nil {
		return "", err
	}

	if !st.IsDir() {
		return "", fmt.Errorf("%s: not a directory", dir)
	}

	return dir, nil
}

// Put is not implemented for BTRFS as there is no cleanup required for the id.
func (d *Driver) Put(id string) error {
	// Get() creates no runtime resources (like e.g. mounts)
	// so this doesn't need to do anything.
	return nil
}

// Exists checks if the id exists in the filesystem.
func (d *Driver) Exists(id string) bool {
	dir := d.subvolumesDirID(id)
	_, err := os.Stat(dir)
	return err == nil
}
<commit_msg>Make btrfs call same interface as others<commit_after>// +build linux

package btrfs

/*
#include <stdlib.h>
#include <dirent.h>
#include <btrfs/ioctl.h>
#include <btrfs/ctree.h>
*/
import "C"

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"syscall"
	"unsafe"

	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/pkg/idtools"
	"github.com/docker/docker/pkg/mount"
	"github.com/opencontainers/runc/libcontainer/label"
)

func init() {
	graphdriver.Register("btrfs", Init)
}

// Init returns a new BTRFS driver.
// An error is returned if BTRFS is not supported.
func Init(home string, options []string, uidMaps, gidMaps []idtools.IDMap) (graphdriver.Driver, error) {

	fsMagic, err := graphdriver.GetFSMagic(home)
	if err != nil {
		return nil, err
	}

	if fsMagic != graphdriver.FsMagicBtrfs {
		return nil, graphdriver.ErrPrerequisites
	}

	rootUID, rootGID, err := idtools.GetRootUIDGID(uidMaps, gidMaps)
	if err != nil {
		return nil, err
	}
	if err := idtools.MkdirAllAs(home, 0700, rootUID, rootGID); err != nil {
		return nil, err
	}

	if err := mount.MakePrivate(home); err != nil {
		return nil, err
	}

	driver := &Driver{
		home:    home,
		uidMaps: uidMaps,
		gidMaps: gidMaps,
	}

	return graphdriver.NewNaiveDiffDriver(driver, uidMaps, gidMaps), nil
}

// Driver contains information about the filesystem mounted.
type Driver struct {
	//root of the file system
	home    string
	uidMaps []idtools.IDMap
	gidMaps []idtools.IDMap
}

// String prints the name of the driver (btrfs).
func (d *Driver) String() string {
	return "btrfs"
}

// Status returns current driver information in a two dimensional string array.
// Output contains "Build Version" and "Library Version" of the btrfs libraries used.
// Version information can be used to check compatibility with your kernel.
func (d *Driver) Status() [][2]string {
	status := [][2]string{}
	if bv := btrfsBuildVersion(); bv != "-" {
		status = append(status, [2]string{"Build Version", bv})
	}
	if lv := btrfsLibVersion(); lv != -1 {
		status = append(status, [2]string{"Library Version", fmt.Sprintf("%d", lv)})
	}
	return status
}

// GetMetadata returns empty metadata for this driver.
func (d *Driver) GetMetadata(id string) (map[string]string, error) {
	return nil, nil
}

// Cleanup unmounts the home directory.
func (d *Driver) Cleanup() error {
	return mount.Unmount(d.home)
}

func free(p *C.char) {
	C.free(unsafe.Pointer(p))
}

func openDir(path string) (*C.DIR, error) {
	Cpath := C.CString(path)
	defer free(Cpath)

	dir := C.opendir(Cpath)
	if dir == nil {
		return nil, fmt.Errorf("Can't open dir")
	}
	return dir, nil
}

func closeDir(dir *C.DIR) {
	if dir != nil {
		C.closedir(dir)
	}
}

func getDirFd(dir *C.DIR) uintptr {
	return uintptr(C.dirfd(dir))
}

func subvolCreate(path, name string) error {
	dir, err := openDir(path)
	if err != nil {
		return err
	}
	defer closeDir(dir)

	var args C.struct_btrfs_ioctl_vol_args
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(dir), C.BTRFS_IOC_SUBVOL_CREATE,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to create btrfs subvolume: %v", errno.Error())
	}
	return nil
}

func subvolSnapshot(src, dest, name string) error {
	srcDir, err := openDir(src)
	if err != nil {
		return err
	}
	defer closeDir(srcDir)

	destDir, err := openDir(dest)
	if err != nil {
		return err
	}
	defer closeDir(destDir)

	var args C.struct_btrfs_ioctl_vol_args_v2
	args.fd = C.__s64(getDirFd(srcDir))
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(destDir), C.BTRFS_IOC_SNAP_CREATE_V2,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to create btrfs snapshot: %v", errno.Error())
	}
	return nil
}

func isSubvolume(p string) (bool, error) {
	var bufStat syscall.Stat_t
	if err := syscall.Lstat(p, &bufStat); err != nil {
		return false, err
	}

	// return true if it is a btrfs subvolume
	return bufStat.Ino == C.BTRFS_FIRST_FREE_OBJECTID, nil
}

func subvolDelete(dirpath, name string) error {
	dir, err := openDir(dirpath)
	if err != nil {
		return err
	}
	defer closeDir(dir)
	fullPath := path.Join(dirpath, name)

	var args C.struct_btrfs_ioctl_vol_args

	// walk the btrfs subvolumes
	walkSubvolumes := func(p string, f os.FileInfo, err error) error {
		if err != nil {
			if os.IsNotExist(err) && p != fullPath {
				// missing most likely because the path was a subvolume that got removed in the previous iteration
				// since it's gone anyway, we don't care
				return nil
			}
			return fmt.Errorf("error walking subvolumes: %v", err)
		}
		// we want to check children only so skip itself
		// it will be removed after the filepath walk anyways
		if f.IsDir() && p != fullPath {
			sv, err := isSubvolume(p)
			if err != nil {
				return fmt.Errorf("Failed to test if %s is a btrfs subvolume: %v", p, err)
			}
			if sv {
				if err := subvolDelete(path.Dir(p), f.Name()); err != nil {
					return fmt.Errorf("Failed to destroy btrfs child subvolume (%s) of parent (%s): %v", p, dirpath, err)
				}
			}
		}
		return nil
	}
	if err := filepath.Walk(path.Join(dirpath, name), walkSubvolumes); err != nil {
		return fmt.Errorf("Recursively walking subvolumes for %s failed: %v", dirpath, err)
	}

	// all subvolumes have been removed
	// now remove the one originally passed in
	for i, c := range []byte(name) {
		args.name[i] = C.char(c)
	}
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, getDirFd(dir), C.BTRFS_IOC_SNAP_DESTROY,
		uintptr(unsafe.Pointer(&args)))
	if errno != 0 {
		return fmt.Errorf("Failed to destroy btrfs snapshot %s for %s: %v", dirpath, name, errno.Error())
	}
	return nil
}

func (d *Driver) subvolumesDir() string {
	return path.Join(d.home, "subvolumes")
}

func (d *Driver) subvolumesDirID(id string) string {
	return path.Join(d.subvolumesDir(), id)
}

// Create the filesystem with given id.
func (d *Driver) Create(id, parent, mountLabel string) error {
	subvolumes := path.Join(d.home, "subvolumes")
	rootUID, rootGID, err := idtools.GetRootUIDGID(d.uidMaps, d.gidMaps)
	if err != nil {
		return err
	}
	if err := idtools.MkdirAllAs(subvolumes, 0700, rootUID, rootGID); err != nil {
		return err
	}
	if parent == "" {
		if err := subvolCreate(subvolumes, id); err != nil {
			return err
		}
	} else {
		parentDir, err := d.Get(parent, "")
		if err != nil {
			return err
		}
		if err := subvolSnapshot(parentDir, subvolumes, id); err != nil {
			return err
		}
	}

	// if we have a remapped root (user namespaces enabled), change the created snapshot
	// dir ownership to match
	if rootUID != 0 || rootGID != 0 {
		if err := os.Chown(path.Join(subvolumes, id), rootUID, rootGID); err != nil {
			return err
		}
	}

	return label.Relabel(path.Join(subvolumes, id), mountLabel, false)
}

// Remove the filesystem with given id.
func (d *Driver) Remove(id string) error {
	dir := d.subvolumesDirID(id)
	if _, err := os.Stat(dir); err != nil {
		return err
	}
	if err := subvolDelete(d.subvolumesDir(), id); err != nil {
		return err
	}
	if err := os.RemoveAll(dir); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

// Get the requested filesystem id.
func (d *Driver) Get(id, mountLabel string) (string, error) {
	dir := d.subvolumesDirID(id)
	st, err := os.Stat(dir)
	if err != nil {
		return "", err
	}

	if !st.IsDir() {
		return "", fmt.Errorf("%s: not a directory", dir)
	}

	return dir, nil
}

// Put is not implemented for BTRFS as there is no cleanup required for the id.
func (d *Driver) Put(id string) error {
	// Get() creates no runtime resources (like e.g. mounts)
	// so this doesn't need to do anything.
	return nil
}

// Exists checks if the id exists in the filesystem.
func (d *Driver) Exists(id string) bool {
	dir := d.subvolumesDirID(id)
	_, err := os.Stat(dir)
	return err == nil
}
<|endoftext|>


<commit_before>package log

import (
	"bufio"
	"io"
	"os"
	"path"
	"sync"
	"time"
)

var (
	HourlySuffix = ".20060102-15"
	DailySuffix  = ".20060102"
)

type Appender interface {
	Output(level Level, t time.Time, data []byte)
}

type console struct {
	io.Writer
	mu sync.Mutex
}

func NewConsoleAppender() Appender {
	return &console{Writer: os.Stdout}
}

func (c *console) Output(level Level, t time.Time, data []byte) {
	c.mu.Lock()
	c.Write(data)
	c.mu.Unlock()
}

type RotateAppender struct {
	mu       sync.Mutex
	rt       time.Time
	filename string
	rtfn     func(time.Time) (time.Time, string)
	w        io.Writer
	file     *os.File
}

func hourly() time.Time {
	n := time.Now()
	y, m, d := n.Date()
	h, _, _ := n.Clock()
	return time.Date(y, m, d, h+1, 0, 0, 0, time.Local)
}

func daily() time.Time {
	y, m, d := time.Now().Date()
	return time.Date(y, m, d+1, 0, 0, 0, 0, time.Local)
}

func NewHourlyRotateAppender(filename string) (*RotateAppender, error) {
	return NewHourlyRotateBufAppender(filename, 0)
}

func NewHourlyRotateBufAppender(filename string, bufsize int) (*RotateAppender, error) {
	a := &RotateAppender{
		filename: path.Clean(filename),
		rt:       hourly(),
	}

	a.rtfn = func(t time.Time) (time.Time, string) {
		return hourly(), t.Format(HourlySuffix)
	}

	return a.open(bufsize)
}

func NewDailyRotateAppender(filename string) (*RotateAppender, error) {
	return NewDailyRotateBufAppender(filename, 0)
}

func NewDailyRotateBufAppender(filename string, bufsize int) (*RotateAppender, error) {
	a := &RotateAppender{
		filename: path.Clean(filename),
		rt:       daily(),
	}

	a.rtfn = func(t time.Time) (time.Time, string) {
		return daily(), t.Format(DailySuffix)
	}

	return a.open(bufsize)
}

func (a *RotateAppender) open(bufsize int) (*RotateAppender, error) {
	err := os.MkdirAll(path.Dir(a.filename), 0755)
	if err != nil && !os.IsExist(err) {
		return nil, err
	}
	a.file, err = os.OpenFile(a.filename,
		os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if bufsize > 0 {
		a.w = bufio.NewWriterSize(a.file, bufsize)
	} else {
		a.w = a.file
	}
	return a, err
}

func (a *RotateAppender) Close() error {
	a.mu.Lock()
	e := a.close()
	a.mu.Unlock()
	return e
}

func (a *RotateAppender) close() error {
	var e1, e2 error
	if bw, ok := a.w.(*bufio.Writer); ok {
		if e1 = bw.Flush(); e1 != nil {
			println("appender close bufio flush error: ", e1)
		}
	}

	if e2 = a.file.Close(); e2 != nil {
		println("appender close filename: ", a.filename, "error: ", e2)
	} else {
		a.file = nil
	}

	if e1 != nil {
		return e1
	} else if e2 != nil {
		return e2
	}
	return nil
}

func (a *RotateAppender) reset(file *os.File) {
	if bw, ok := a.w.(*bufio.Writer); ok {
		bw.Reset(file)
	} else {
		a.w = file
	}
}

func (a *RotateAppender) Output(_ Level, t time.Time, data []byte) {
	a.mu.Lock()
	if t.After(a.rt) {
		var suffix string
		a.rt, suffix = a.rtfn(a.rt)
		filename := a.filename + suffix
		err := a.close()
		if err != nil {
			println("appender close ", a.filename, "error: ", err)
		}
		if err = os.Rename(a.filename, filename); err != nil {
			println("appender rename ", filename, "error: ", err)
		}

		a.file, err = os.OpenFile(a.filename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			println("appender open ", a.filename, "error: ", err)
		}
		a.reset(a.file)
	}
	if a.file == nil {
		a.mu.Unlock()
		return
	}
	a.w.Write(data)
	a.mu.Unlock()
}
<commit_msg>add internal detail error message<commit_after>package log

import (
	"bufio"
	"io"
	"os"
	"path"
	"sync"
	"time"
)

var (
	HourlySuffix = ".20060102-15"
	DailySuffix  = ".20060102"
)

type Appender interface {
	Output(level Level, t time.Time, data []byte)
}

type console struct {
	io.Writer
	mu sync.Mutex
}

func NewConsoleAppender() Appender {
	return &console{Writer: os.Stdout}
}

func (c *console) Output(level Level, t time.Time, data []byte) {
	c.mu.Lock()
	c.Write(data)
	c.mu.Unlock()
}

type RotateAppender struct {
	mu       sync.Mutex
	rt       time.Time
	filename string
	rtfn     func(time.Time) (time.Time, string)
	w        io.Writer
	file     *os.File
}

func hourly() time.Time {
	n := time.Now()
	y, m, d := n.Date()
	h, _, _ := n.Clock()
	return time.Date(y, m, d, h+1, 0, 0, 0, time.Local)
}

func daily() time.Time {
	y, m, d := time.Now().Date()
	return time.Date(y, m, d+1, 0, 0, 0, 0, time.Local)
}

func NewHourlyRotateAppender(filename string) (*RotateAppender, error) {
	return NewHourlyRotateBufAppender(filename, 0)
}

func NewHourlyRotateBufAppender(filename string, bufsize int) (*RotateAppender, error) {
	a := &RotateAppender{
		filename: path.Clean(filename),
		rt:       hourly(),
	}

	a.rtfn = func(t time.Time) (time.Time, string) {
		return hourly(), t.Format(HourlySuffix)
	}

	return a.open(bufsize)
}

func NewDailyRotateAppender(filename string) (*RotateAppender, error) {
	return NewDailyRotateBufAppender(filename, 0)
}

func NewDailyRotateBufAppender(filename string, bufsize int) (*RotateAppender, error) {
	a := &RotateAppender{
		filename: path.Clean(filename),
		rt:       daily(),
	}

	a.rtfn = func(t time.Time) (time.Time, string) {
		return daily(), t.Format(DailySuffix)
	}

	return a.open(bufsize)
}

func (a *RotateAppender) open(bufsize int) (*RotateAppender, error) {
	err := os.MkdirAll(path.Dir(a.filename), 0755)
	if err != nil && !os.IsExist(err) {
		return nil, err
	}
	a.file, err = os.OpenFile(a.filename,
		os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if bufsize > 0 {
		a.w = bufio.NewWriterSize(a.file, bufsize)
	} else {
		a.w = a.file
	}
	return a, err
}

func (a *RotateAppender) Close() error {
	a.mu.Lock()
	e := a.close()
	a.mu.Unlock()
	return e
}

func (a *RotateAppender) close() error {
	var e1, e2 error
	if bw, ok := a.w.(*bufio.Writer); ok {
		if e1 = bw.Flush(); e1 != nil {
			println("appender close bufio flush error: ", e1.Error())
		}
	}

	if e2 = a.file.Close(); e2 != nil {
		println("appender close filename: ", a.filename, "error: ", e2.Error())
	} else {
		a.file = nil
	}

	if e1 != nil {
		return e1
	} else if e2 != nil {
		return e2
	}
	return nil
}

func (a *RotateAppender) reset(file *os.File) {
	if bw, ok := a.w.(*bufio.Writer); ok {
		bw.Reset(file)
	} else {
		a.w = file
	}
}

func (a *RotateAppender) Output(_ Level, t time.Time, data []byte) {
	a.mu.Lock()
	if t.After(a.rt) {
		var suffix string
		a.rt, suffix = a.rtfn(a.rt)
		filename := a.filename + suffix
		err := a.close()
		if err != nil {
			println("appender close ", a.filename, "error: ", err.Error())
		}
		if err = os.Rename(a.filename, filename); err != nil {
			println("appender rename ", filename, "error: ", err.Error())
		}

		a.file, err = os.OpenFile(a.filename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			println("appender open ", a.filename, "error: ", err.Error())
		}
		a.reset(a.file)
	}
	if a.file == nil {
		a.mu.Unlock()
		return
	}
	a.w.Write(data)
	a.mu.Unlock()
}
<|endoftext|>